hoststats: cancel host collector in base deps Close

agent/setup.go

@@ -4,6 +4,7 @@
 package agent
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net"
@@ -59,6 +60,7 @@ type BaseDeps struct {
 	WatchedFiles  []string
 
 	deregisterBalancer, deregisterResolver func()
+	stopHostCollector                      func()
 }
 
 type ConfigLoader func(source config.Source) (config.LoadResult, error)
@@ -106,8 +108,10 @@ func NewBaseDeps(configLoader ConfigLoader, logOut io.Writer, providedLogger hcl
 	if err != nil {
 		return d, fmt.Errorf("failed to initialize telemetry: %w", err)
 	}
-	if d.MetricsConfig != nil && !cfg.Telemetry.DisableHostMetrics {
-		hoststats.NewCollector(d.MetricsConfig.Context(), d.Logger, cfg.DataDir)
+	if !cfg.Telemetry.Disable && !cfg.Telemetry.DisableHostMetrics {
+		ctx, cancel := context.WithCancel(context.Background())
+		hoststats.NewCollector(ctx, d.Logger, cfg.DataDir)
+		d.stopHostCollector = cancel
 	}
 
 	d.TLSConfigurator, err = tlsutil.NewConfigurator(cfg.TLS, d.Logger)
@@ -212,11 +216,10 @@ func (bd BaseDeps) Close() {
 	bd.AutoConfig.Stop()
 	bd.MetricsConfig.Cancel()
 
-	if fn := bd.deregisterBalancer; fn != nil {
-		fn()
-	}
-	if fn := bd.deregisterResolver; fn != nil {
-		fn()
+	for _, fn := range []func(){bd.deregisterBalancer, bd.deregisterResolver, bd.stopHostCollector} {
+		if fn != nil {
+			fn()
+		}
 	}
 }
 

lib/telemetry.go

@@ -226,16 +226,9 @@ type MetricsHandler interface {
 type MetricsConfig struct {
 	Handler  MetricsHandler
 	mu       sync.Mutex
-	ctx      context.Context
 	cancelFn context.CancelFunc
 }
 
-func (cfg *MetricsConfig) Context() context.Context {
-	cfg.mu.Lock()
-	defer cfg.mu.Unlock()
-	return cfg.ctx
-}
-
 func (cfg *MetricsConfig) Cancel() {
 	cfg.mu.Lock()
 	defer cfg.mu.Unlock()
@@ -415,7 +408,6 @@ func InitTelemetry(cfg TelemetryConfig, logger hclog.Logger) (*MetricsConfig, er
 			ctx, cancel = context.WithCancel(context.Background())
 
 			metricsConfig.mu.Lock()
-			metricsConfig.ctx = ctx
 			metricsConfig.cancelFn = cancel
 			metricsConfig.mu.Unlock()
 			go retryWithBackoff()