update discovery to reconfigure sharding

Signed-off-by: Walther Lee <[email protected]>

Author: waltherlee

internal/discovery/discovery.go

@@ -203,10 +203,6 @@ func (r *CRDiscoverer) PollForCacheUpdates(
 ) {
 	// The interval at which we will check the cache for updates.
 	t := time.NewTicker(Interval)
-	// Track previous context to allow refreshing cache.
-	olderContext, olderCancel := context.WithCancel(ctx)
-	// Prevent context leak (kill the last metric handler instance).
-	defer olderCancel()
 	generateMetrics := func() {
 		// Get families for discovered factories.
 		customFactories, err := factoryGenerator()
@@ -239,21 +235,8 @@ func (r *CRDiscoverer) PollForCacheUpdates(
 		r.SafeWrite(func() {
 			r.WasUpdated = false
 		})
-		// Run the metrics handler with updated configs.
-		olderContext, olderCancel = context.WithCancel(ctx)
-		go func() {
-			// Blocks indefinitely until the unbuffered context is cancelled to serve metrics for that duration.
-			err = m.Run(olderContext)
-			if err != nil {
-				// Check if context was cancelled.
-				select {
-				case <-olderContext.Done():
-					// Context cancelled, don't really need to log this though.
-				default:
-					klog.ErrorS(err, "failed to run metrics handler")
-				}
-			}
-		}()
+		// Update metric handler with the new configs.
+		m.ReconfigureSharding(ctx)
 	}
 	go func() {
 		for range t.C {
@@ -269,7 +252,6 @@ func (r *CRDiscoverer) PollForCacheUpdates(
 					shouldGenerateMetrics = r.WasUpdated
 				})
 				if shouldGenerateMetrics {
-					olderCancel()
 					generateMetrics()
 					klog.InfoS("discovery finished, cache updated")
 				}

pkg/app/server.go

@@ -297,14 +297,12 @@ func RunKubeStateMetrics(ctx context.Context, opts *options.Options) error {
 		opts.EnableGZIPEncoding,
 	)
 	// Run MetricsHandler
-	if config == nil {
-		ctxMetricsHandler, cancel := context.WithCancel(ctx)
-		g.Add(func() error {
-			return m.Run(ctxMetricsHandler)
-		}, func(error) {
-			cancel()
-		})
-	}
+	ctxMetricsHandler, cancel := context.WithCancel(ctx)
+	g.Add(func() error {
+		return m.Run(ctxMetricsHandler)
+	}, func(error) {
+		cancel()
+	})
 
 	tlsConfig := opts.TLSConfig
 

pkg/metricshandler/metrics_handler.go

@@ -69,7 +69,16 @@ func New(opts *options.Options, kubeClient kubernetes.Interface, storeBuilder ks
 	}
 }
 
-// ConfigureSharding (re-)configures sharding. Re-configuration can be done
+// ReconfigureSharding reconfigures sharding with the current shard and totalShards, and
+// it's a no-op if both values are 0.
+func (m *MetricsHandler) ReconfigureSharding(ctx context.Context) {
+	if m.curShard == 0 && m.curTotalShards == 0 {
+		return
+	}
+	m.ConfigureSharding(ctx, m.curShard, m.curTotalShards)
+}
+
+// ConfigureSharding configures sharding. Configuration can be used mutlitple times and
 // concurrently.
 func (m *MetricsHandler) ConfigureSharding(ctx context.Context, shard int32, totalShards int) {
 	m.mtx.Lock()
@@ -132,6 +141,14 @@ func (m *MetricsHandler) Run(ctx context.Context) error {
 				return
 			}
 
+			m.mtx.RLock()
+			shardingUnchanged := m.curShard == shard && m.curTotalShards == totalShards
+			m.mtx.RUnlock()
+
+			if shardingUnchanged {
+				return
+			}
+
 			m.ConfigureSharding(ctx, shard, totalShards)
 		},
 		UpdateFunc: func(oldo, curo interface{}) {