apiserver/apic: migrate metrics lifecycle to context

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: mmetc

pkg/apiserver/apic.go

@@ -63,7 +63,8 @@ type apic struct {
 	mu            sync.Mutex
 	pushTomb      tomb.Tomb
 	pullTomb      tomb.Tomb
-	metricsTomb   tomb.Tomb
+	metricsCancel context.CancelFunc
+	metricsDone   chan struct{}
 	startup       bool
 	consoleConfig *csconfig.ConsoleConfig
 	isPulling     chan bool
@@ -195,7 +196,6 @@ func NewAPIC(ctx context.Context, config *csconfig.OnlineApiClientCfg, dbClient
 		startup:                   true,
 		pullTomb:                  tomb.Tomb{},
 		pushTomb:                  tomb.Tomb{},
-		metricsTomb:               tomb.Tomb{},
 		consoleConfig:             consoleConfig,
 		pullInterval:              pullIntervalDefault,
 		pullIntervalFirst:         randomDuration(pullIntervalDefault, pullIntervalDelta),
@@ -295,7 +295,7 @@ func (a *apic) Push(ctx context.Context) error {
 		select {
 		case <-a.pushTomb.Dying(): // if one apic routine is dying, do we kill the others?
 			a.pullTomb.Kill(nil)
-			a.metricsTomb.Kill(nil)
+			a.StopMetrics()
 			log.Infof("push tomb is dying, sending cache (%d elements) before exiting", len(cache))
 
 			if len(cache) == 0 {
@@ -1073,18 +1073,70 @@ func (a *apic) Pull(ctx context.Context) error {
 				continue
 			}
 		case <-a.pullTomb.Dying(): // if one apic routine is dying, do we kill the others?
-			a.metricsTomb.Kill(nil)
+			a.StopMetrics()
 			a.pushTomb.Kill(nil)
 
 			return nil
 		}
 	}
 }
 
+func (a *apic) StartMetrics(ctx context.Context, sendUsageMetrics bool) {
+	a.mu.Lock()
+	if a.metricsCancel != nil {
+		a.mu.Unlock()
+		return
+	}
+
+	metricsCtx, cancel := context.WithCancel(ctx)
+	done := make(chan struct{})
+	a.metricsCancel = cancel
+	a.metricsDone = done
+	a.mu.Unlock()
+
+	go func() {
+		var wg sync.WaitGroup
+
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			a.SendMetrics(metricsCtx, make(chan bool))
+		}()
+
+		if sendUsageMetrics {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				a.SendUsageMetrics(metricsCtx)
+			}()
+		}
+
+		wg.Wait()
+		close(done)
+	}()
+}
+
+func (a *apic) StopMetrics() {
+	a.mu.Lock()
+	cancel := a.metricsCancel
+	done := a.metricsDone
+	a.metricsCancel = nil
+	a.metricsDone = nil
+	a.mu.Unlock()
+
+	if cancel != nil {
+		cancel()
+	}
+
+	if done != nil {
+		<-done
+	}
+}
+
 func (a *apic) Shutdown() {
 	a.pushTomb.Kill(nil)
 	a.pullTomb.Kill(nil)
-	a.metricsTomb.Kill(nil)
+	a.StopMetrics()
 }
 
 func makeAddAndDeleteCounters() (map[string]map[string]int, map[string]map[string]int) {

pkg/apiserver/apic_metrics.go

@@ -291,13 +291,16 @@ func (a *apic) SendMetrics(ctx context.Context, stop chan bool) {
 
 	checkTicker := time.NewTicker(checkInt)
 	metTicker := time.NewTicker(nextMetInt())
+	defer checkTicker.Stop()
+	defer metTicker.Stop()
 
 	for {
 		select {
 		case <-stop:
-			checkTicker.Stop()
-			metTicker.Stop()
-
+			return
+		case <-ctx.Done():
+			a.pullTomb.Kill(nil)
+			a.pushTomb.Kill(nil)
 			return
 		case <-checkTicker.C:
 			oldIDs := machineIDs
@@ -326,13 +329,6 @@ func (a *apic) SendMetrics(ctx context.Context, stop chan bool) {
 			}
 
 			metTicker.Reset(nextMetInt())
-		case <-a.metricsTomb.Dying(): // if one apic routine is dying, do we kill the others?
-			checkTicker.Stop()
-			metTicker.Stop()
-			a.pullTomb.Kill(nil)
-			a.pushTomb.Kill(nil)
-
-			return
 		}
 	}
 }
@@ -345,8 +341,7 @@ func (a *apic) SendUsageMetrics(ctx context.Context) {
 
 	for {
 		select {
-		case <-a.metricsTomb.Dying():
-			// The normal metrics routine also kills push/pull tombs, does that make sense ?
+		case <-ctx.Done():
 			ticker.Stop()
 			return
 		case <-ticker.C:

pkg/apiserver/apic_test.go

@@ -57,12 +57,11 @@ func getAPIC(t *testing.T, ctx context.Context) *apic {
 	return &apic{
 		AlertsAddChan: make(chan []*models.Alert),
 		// DecisionDeleteChan: make(chan []*models.Decision),
-		dbClient:    dbClient,
-		mu:          sync.Mutex{},
-		startup:     true,
-		pullTomb:    tomb.Tomb{},
-		pushTomb:    tomb.Tomb{},
-		metricsTomb: tomb.Tomb{},
+		dbClient: dbClient,
+		mu:       sync.Mutex{},
+		startup:  true,
+		pullTomb: tomb.Tomb{},
+		pushTomb: tomb.Tomb{},
 		consoleConfig: &csconfig.ConsoleConfig{
 			ShareManualDecisions:  ptr.Of(false),
 			ShareTaintedScenarios: ptr.Of(false),

pkg/apiserver/apiserver.go

@@ -302,19 +302,7 @@ func (s *APIServer) initAPIC(ctx context.Context) {
 		}
 	}
 
-	s.apic.metricsTomb.Go(func() error {
-		defer trace.ReportPanic()
-		s.apic.SendMetrics(ctx, make(chan bool))
-		return nil
-	})
-
-	if !s.cfg.DisableUsageMetricsExport {
-		s.apic.metricsTomb.Go(func() error {
-			defer trace.ReportPanic()
-			s.apic.SendUsageMetrics(ctx)
-			return nil
-		})
-	}
+	s.apic.StartMetrics(ctx, !s.cfg.DisableUsageMetricsExport)
 }
 
 func (s *APIServer) Run(ctx context.Context, apiReady chan bool) error {