[aggregator] Propagate cancellation through tick

src/aggregator/aggregator/aggregator.go

@@ -23,6 +23,7 @@ package aggregator
 import (
 	"context"
 	"errors"
+	"fmt"
 	"math"
 	"strconv"
 	"sync"
@@ -93,6 +94,11 @@ type Aggregator interface {
 	Close() error
 }
 
+type tickShardFn func(
+	shard *aggregatorShard,
+	perShardTickDuration time.Duration,
+) tickResult
+
 // aggregator stores aggregations of different types of metrics (e.g., counter,
 // timer, gauges) and periodically flushes them out.
 type aggregator struct {
@@ -101,6 +107,7 @@ type aggregator struct {
 	opts              Options
 	nowFn             clock.NowFn
 	shardFn           sharding.ShardFn
+	tickShardFn       tickShardFn
 	checkInterval     time.Duration
 	placementManager  PlacementManager
 	flushTimesManager FlushTimesManager
@@ -133,7 +140,7 @@ func NewAggregator(opts Options) Aggregator {
 	timerOpts := iOpts.TimerOptions()
 	logger := iOpts.Logger()
 
-	return &aggregator{
+	agg := &aggregator{
 		opts:              opts,
 		nowFn:             opts.ClockOptions().NowFn(),
 		shardFn:           opts.ShardFn(),
@@ -151,6 +158,9 @@ func NewAggregator(opts Options) Aggregator {
 		metrics:           newAggregatorMetrics(scope, timerOpts, opts.MaxAllowedForwardingDelayFn()),
 		logger:            logger,
 	}
+
+	agg.tickShardFn = agg.tickShard
+	return agg
 }
 
 func (agg *aggregator) Open() error {
@@ -383,6 +393,20 @@ func (agg *aggregator) Close() error {
 	}
 	agg.state = aggregatorClosed
 
+	var (
+		lastOpCompleted = time.Now()
+		currTime        time.Time
+		closeLogger     = agg.logger.With(zap.String("closing", "aggregator"))
+
+		logCloseOperation = func(op string) {
+			currTime = time.Now()
+			closeLogger.Info(fmt.Sprintf("closed %s", op),
+				zap.String("took", currTime.Sub(lastOpCompleted).String()))
+			lastOpCompleted = currTime
+		}
+	)
+
+	closeLogger.Info("signaling aggregator done")
 	close(agg.doneCh)
 
 	// Waiting for the ticking goroutines to return.
@@ -391,17 +415,31 @@ func (agg *aggregator) Close() error {
 	agg.wg.Wait()
 	agg.Lock()
 
+	logCloseOperation("ticking wait groups")
 	for _, shardID := range agg.shardIDs {
 		agg.shards[shardID].Close()
 	}
+
+	logCloseOperation("aggregator shards")
 	if agg.shardSetOpen {
 		agg.closeShardSetWithLock()
 	}
+
+	logCloseOperation("flush shard sets")
+
 	agg.flushHandler.Close()
+	logCloseOperation("flush handler")
+
 	agg.passthroughWriter.Close()
+	logCloseOperation("passthrough writer")
+
 	if agg.adminClient != nil {
+		closeLogger.Info("closing admin client")
 		agg.adminClient.Close()
+		logCloseOperation("admin client")
 	}
+
+	closeLogger.Info("done")
 	return nil
 }
 
@@ -723,6 +761,13 @@ func (agg *aggregator) tick() {
 	}
 }
 
+func (agg *aggregator) tickShard(
+	shard *aggregatorShard,
+	perShardTickDuration time.Duration,
+) tickResult {
+	return shard.Tick(perShardTickDuration, agg.doneCh)
+}
+
 func (agg *aggregator) tickInternal() {
 	ownedShards, closingShards := agg.ownedShards()
 	agg.closeShardsAsync(closingShards)
@@ -740,8 +785,14 @@ func (agg *aggregator) tickInternal() {
 		tickResult           tickResult
 	)
 	for _, shard := range ownedShards {
-		shardTickResult := shard.Tick(perShardTickDuration)
-		tickResult = tickResult.merge(shardTickResult)
+		select {
+		case <-agg.doneCh:
+			agg.logger.Info("recevied interrupt on tick; aborting")
+			return
+		default:
+			shardTickResult := agg.tickShardFn(shard, perShardTickDuration)
+			tickResult = tickResult.merge(shardTickResult)
+		}
 	}
 	tickDuration := agg.nowFn().Sub(start)
 	agg.metrics.tick.Report(tickResult, tickDuration)

src/aggregator/aggregator/aggregator_test.go

@@ -841,6 +841,44 @@ func TestAggregatorTick(t *testing.T) {
 	require.NoError(t, agg.Close())
 }
 
+func TestAggregatorTickCancelled(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	flushTimesManager := NewMockFlushTimesManager(ctrl)
+	flushTimesManager.EXPECT().Reset().Return(nil).AnyTimes()
+	flushTimesManager.EXPECT().Open(gomock.Any()).Return(nil).AnyTimes()
+	flushTimesManager.EXPECT().Get().Return(nil, nil).AnyTimes()
+	flushTimesManager.EXPECT().Close().Return(nil).AnyTimes()
+
+	agg, _ := testAggregator(t, ctrl)
+	agg.flushTimesManager = flushTimesManager
+	require.NoError(t, agg.Open())
+
+	var (
+		tickedCh       = make(chan struct{})
+		numTicked      = 0
+		doneAfterTicks = 2
+	)
+
+	agg.tickShardFn = func(*aggregatorShard, time.Duration) tickResult {
+		numTicked++
+		if doneAfterTicks == 2 {
+			close(tickedCh)
+		}
+
+		time.Sleep(time.Millisecond * 50)
+		return tickResult{}
+	}
+
+	go func() {
+		<-tickedCh
+		require.NoError(t, agg.Close())
+	}()
+
+	require.Equal(t, 2, doneAfterTicks)
+}
+
 func TestAggregatorShardSetNotOpenNilInstance(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()

src/aggregator/aggregator/map.go

@@ -218,8 +218,11 @@ func (m *metricMap) AddForwarded(
 	return err
 }
 
-func (m *metricMap) Tick(target time.Duration) tickResult {
-	mapTickRes := m.tick(target)
+func (m *metricMap) Tick(
+	target time.Duration,
+	doneCh chan struct{},
+) tickResult {
+	mapTickRes := m.tick(target, doneCh)
 	listsTickRes := m.metricLists.Tick()
 	mapTickRes.standard.activeElems = listsTickRes.standard
 	mapTickRes.forwarded.activeElems = listsTickRes.forwarded
@@ -320,7 +323,10 @@ func (m *metricMap) lookupEntryWithLock(key entryKey) (*Entry, bool) {
 // tick performs two operations:
 // 1. Delete entries that have expired, and report the number of expired entries.
 // 2. Report number of standard entries and forwarded entries that are active.
-func (m *metricMap) tick(target time.Duration) tickResult {
+func (m *metricMap) tick(
+	target time.Duration,
+	doneCh chan struct{},
+) tickResult {
 	// Determine batch size.
 	m.RLock()
 	numEntries := m.entryList.Len()
@@ -340,8 +346,27 @@ func (m *metricMap) tick(target time.Duration) tickResult {
 		numTimedActive       int
 		numTimedExpired      int
 		entryIdx             int
+
+		done bool
 	)
+
+	// NB: if no doneChan provided, do not interrupt the tick.
+	if doneCh == nil {
+		doneCh = make(chan struct{})
+	}
+
 	m.forEachEntry(func(entry hashedEntry) {
+		if done {
+			return
+		}
+
+		select {
+		case <-doneCh:
+			done = true
+			return
+		default:
+		}
+
 		now := m.nowFn()
 		if entryIdx > 0 && entryIdx%defaultSoftDeadlineCheckEvery == 0 {
 			targetDeadline := start.Add(time.Duration(entryIdx) * perEntrySoftDeadline)

src/aggregator/aggregator/map_test.go

@@ -513,7 +513,7 @@ func TestMetricMapDeleteExpired(t *testing.T) {
 	}
 
 	// Delete expired entries.
-	m.tick(opts.EntryCheckInterval())
+	m.tick(opts.EntryCheckInterval(), nil)
 
 	// Assert there should be only half of the entries left.
 	require.Equal(t, numEntries/2, len(m.entries))
@@ -525,3 +525,46 @@ func TestMetricMapDeleteExpired(t *testing.T) {
 		require.NotNil(t, e.entry)
 	}
 }
+
+func TestMetricMapTickCancellation(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	opts := testOptions(ctrl)
+	m := newMetricMap(testShard, opts)
+
+	numBatchesProcessed := 0
+	tickedCh := make(chan struct{})
+
+	m.sleepFn = func(d time.Duration) {
+		numBatchesProcessed++
+		if numBatchesProcessed == 10 {
+			close(tickedCh)
+		}
+
+		time.Sleep(d)
+	}
+
+	doneCh := make(chan struct{})
+	go func() {
+		<-tickedCh
+		close(doneCh)
+	}()
+
+	// NB: wait/early exit on every defaultSoftDeadlineCheckEvery
+	numEntries := defaultSoftDeadlineCheckEvery * 60
+	for i := 0; i < numEntries; i++ {
+		key := entryKey{
+			metricType: metricType(metric.CounterType),
+			idHash:     hash.Murmur3Hash128([]byte(fmt.Sprintf("%d", i))),
+		}
+
+		m.entries[key] = m.entryList.PushBack(hashedEntry{
+			key:   key,
+			entry: NewEntry(m.metricLists, runtime.NewOptions(), opts),
+		})
+	}
+
+	m.Tick(time.Minute, doneCh)
+	require.Equal(t, 10, numBatchesProcessed)
+}

src/aggregator/aggregator/shard.go

@@ -260,8 +260,11 @@ func (s *aggregatorShard) AddForwarded(
 	return nil
 }
 
-func (s *aggregatorShard) Tick(target time.Duration) tickResult {
-	return s.metricMap.Tick(target)
+func (s *aggregatorShard) Tick(
+	target time.Duration,
+	doneCh chan struct{},
+) tickResult {
+	return s.metricMap.Tick(target, doneCh)
 }
 
 func (s *aggregatorShard) Close() {

src/cmd/services/m3aggregator/serve/serve.go

@@ -22,6 +22,7 @@ package serve
 
 import (
 	"fmt"
+	"time"
 
 	"github.com/m3db/m3/src/aggregator/aggregator"
 	httpserver "github.com/m3db/m3/src/aggregator/server/http"
@@ -38,9 +39,23 @@ func Serve(
 	doneCh chan struct{},
 	opts Options,
 ) error {
-	iOpts := opts.InstrumentOpts()
-	log := iOpts.Logger()
-	defer aggregator.Close()
+	var (
+		iOpts       = opts.InstrumentOpts()
+		log         = iOpts.Logger()
+		closeLogger = log.With(zap.String("closing", "aggregator_server"))
+	)
+
+	defer func() {
+		start := time.Now()
+		closeLogger.Info("closing aggregator")
+		err := aggregator.Close()
+		fields := []zap.Field{zap.String("took", time.Since(start).String())}
+		if err != nil {
+			closeLogger.Warn("closed aggregator with error", append(fields, zap.Error(err))...)
+		} else {
+			closeLogger.Info("closed aggregator", fields...)
+		}
+	}()
 
 	if m3msgAddr := opts.M3MsgAddr(); m3msgAddr != "" {
 		serverOpts := opts.M3MsgServerOpts()
@@ -51,7 +66,14 @@ func Serve(
 		if err := m3msgServer.ListenAndServe(); err != nil {
 			return fmt.Errorf("could not start m3msg server at: addr=%s, err=%v", m3msgAddr, err)
 		}
-		defer m3msgServer.Close()
+
+		defer func() {
+			start := time.Now()
+			closeLogger.Info("closing m3msg server")
+			m3msgServer.Close()
+			closeLogger.Info("m3msg server closed", zap.String("took", time.Since(start).String()))
+		}()
+
 		log.Info("m3msg server listening", zap.String("addr", m3msgAddr))
 	}
 
@@ -61,7 +83,14 @@ func Serve(
 		if err := rawTCPServer.ListenAndServe(); err != nil {
 			return fmt.Errorf("could not start raw TCP server at: addr=%s, err=%v", rawTCPAddr, err)
 		}
-		defer rawTCPServer.Close()
+
+		defer func() {
+			start := time.Now()
+			closeLogger.Info("closing raw TCPServer")
+			rawTCPServer.Close()
+			closeLogger.Info("closed raw TCPServer", zap.String("took", time.Since(start).String()))
+		}()
+
 		log.Info("raw TCP server listening", zap.String("addr", rawTCPAddr))
 	}
 
@@ -72,12 +101,20 @@ func Serve(
 		if err := httpServer.ListenAndServe(); err != nil {
 			return fmt.Errorf("could not start http server at: addr=%s, err=%v", httpAddr, err)
 		}
-		defer httpServer.Close()
+
+		defer func() {
+			start := time.Now()
+			closeLogger.Info("closing http server")
+			httpServer.Close()
+			closeLogger.Info("closed http server", zap.String("took", time.Since(start).String()))
+		}()
+
 		log.Info("http server listening", zap.String("addr", httpAddr))
 	}
 
 	// Wait for exit signal.
 	<-doneCh
+	closeLogger.Info("server signaled on doneCh")
 
 	return nil
 }