fix(metrics): order aggregator shutdown and reset the singleton

Shutdown closed the Redis client while the publisher goroutine could still use it, so the instance key sometimes lingered in the cluster set, and the goroutine could read a package global after shutdown returned. Shutdown now removes the instance key, joins the publisher goroutine, then closes the client, and clears the singleton so a later initialisation creates a fresh instance.
2026-10-11 05:27:47 +00:00 · 2026-08-14 11:10:22 +01:00
parent 2911dd9b1d
commit e6a43356c0
2 changed files with 176 additions and 9 deletions
@@ -25,6 +25,10 @@ type MetricsAggregator struct {
 	ctx          context.Context
 	cancel       context.CancelFunc
 	mu           sync.RWMutex
+	// wg tracks the publisher goroutine started by startPublishingAsync.
+	// Shutdown waits on it so it never returns while that goroutine is
+	// still running (see Shutdown's doc comment).
+	wg sync.WaitGroup
 }

 // InstanceMetrics represents metrics for a single proxy instance
@@ -141,7 +145,7 @@ func InitializeMetricsAggregator(redisURL, redisPassword string, redisDB int, lo
 	metricsAggregator = aggregator

 	// Start publishing metrics
-	go aggregator.startPublishing()
+	aggregator.startPublishingAsync()

 	if logger != nil {
 		logger.Info(&libpack_logger.LogMessage{
@@ -164,8 +168,19 @@ func GetMetricsAggregator() *MetricsAggregator {
 	return metricsAggregator
 }

+// startPublishingAsync launches the publisher goroutine and registers it
+// with ma.wg before starting it, so Shutdown can join on it via wg.Wait().
+// wg.Add must happen-before the goroutine's wg.Done (WaitGroup's own
+// contract), so Add is called here, synchronously in the caller, rather
+// than as the first line inside startPublishing.
+func (ma *MetricsAggregator) startPublishingAsync() {
+	ma.wg.Add(1)
+	go ma.startPublishing()
+}
+
 // startPublishing periodically publishes metrics to Redis
 func (ma *MetricsAggregator) startPublishing() {
+	defer ma.wg.Done()
 	defer ma.publishTimer.Stop()

 	// Publish immediately on start
@@ -174,8 +189,12 @@ func (ma *MetricsAggregator) startPublishing() {
 	for {
 		select {
 		case <-ma.ctx.Done():
-			// Clean up our metrics on shutdown
-			ma.removeInstanceMetrics()
+			// Shutdown() removes this instance's Redis keys synchronously,
+			// before it closes the Redis client, so this goroutine must not
+			// also call removeInstanceMetrics() here: doing so raced
+			// Shutdown()'s Close() call (the two used the same client
+			// concurrently with no ordering guarantee) and left the
+			// instance's key/set-membership behind until TTL expiry.
 			return
 		case <-ma.publishTimer.C:
 			ma.publishMetrics()
@@ -778,15 +797,42 @@ func (ma *MetricsAggregator) aggregateStats(instances []InstanceMetrics) map[str
 	return result
 }

-// Shutdown stops the metrics aggregator
+// Shutdown stops the metrics aggregator.
+//
+// It cancels the publisher goroutine, synchronously removes this instance's
+// Redis keys via removeInstanceMetrics (so the removal always completes
+// before the Redis client closes, instead of racing it), waits for the
+// publisher goroutine (started by startPublishingAsync) to fully exit, and
+// clears the package-level singleton if this instance is currently
+// registered as it, so a later InitializeMetricsAggregator call creates a
+// fresh instance rather than silently no-oping against a closed client.
+//
+// ma.mu is held only across cancel/removeInstanceMetrics/Close, not across
+// wg.Wait() or the aggregatorMutex update below. publishMetrics is the only
+// other user of ma.mu (via RLock), and it is only ever invoked by the
+// publisher goroutine. Holding ma.mu across wg.Wait() would deadlock: a
+// publishMetrics call that raced the ctx-cancellation select (picked the
+// timer case just as cancel() ran) blocks on ma.mu.RLock() until Shutdown
+// unlocks, so waiting for that same goroutine to exit while still holding
+// the write lock would wait forever. Unlocking before wg.Wait() avoids that.
+// The straggler is still safe: it can only resume after Unlock(), by which
+// point Close() has already run (still inside the lock, unchanged from
+// before), so it hits a closed client and fails fast instead of
+// resurrecting the instance key removeInstanceMetrics just deleted.
+// wg.Wait() then blocks Shutdown from returning until the goroutine has
+// truly exited, closing the window where it reads the package-level cfg
+// after a caller has already reset it.
 func (ma *MetricsAggregator) Shutdown() {
 	ma.mu.Lock()
-	defer ma.mu.Unlock()

 	if ma.cancel != nil {
 		ma.cancel()
 	}

+	// Remove this instance's key and set membership before closing the
+	// client so the cleanup cannot race Close() (see startPublishing).
+	ma.removeInstanceMetrics()
+
 	if ma.redisClient != nil {
 		_ = ma.redisClient.Close() // Best-effort cleanup
 	}
@@ -796,6 +842,20 @@ func (ma *MetricsAggregator) Shutdown() {
 			Message: "Metrics aggregator shut down",
 		})
 	}
+
+	ma.mu.Unlock()
+
+	// Join the publisher goroutine. Must happen unlocked - see doc comment.
+	ma.wg.Wait()
+
+	// Clear the singleton only if this instance is still registered as it;
+	// a test-constructed or already-replaced instance must not clobber a
+	// different, currently-active aggregator.
+	aggregatorMutex.Lock()
+	if metricsAggregator == ma {
+		metricsAggregator = nil
+	}
+	aggregatorMutex.Unlock()
 }

 // GetInstanceID returns the current instance ID
@@ -119,12 +119,17 @@ func TestMetricsAggregator_InitializeMetricsAggregator_BadURL(t *testing.T) {
 	metricsAggregator = nil
 	aggregatorMutex.Unlock()
 	t.Cleanup(func() {
+		// Shutdown() itself acquires aggregatorMutex (to clear the singleton
+		// if it owns it), so it must not be called while this goroutine
+		// already holds that lock - grab-and-release first, then call
+		// Shutdown() unlocked.
 		aggregatorMutex.Lock()
-		if metricsAggregator != nil {
-			metricsAggregator.Shutdown()
-		}
+		current := metricsAggregator
 		metricsAggregator = old
 		aggregatorMutex.Unlock()
+		if current != nil {
+			current.Shutdown()
+		}
 	})

 	// An unreachable address should cause Ping to fail and return an error.
@@ -214,7 +219,7 @@ func TestMetricsAggregator_StartPublishing_PublishesOnStart(t *testing.T) {
 	ma, _ := newTestAggregator(t)

 	// Run startPublishing in background; it calls publishMetrics immediately.
-	go ma.startPublishing()
+	ma.startPublishingAsync()

 	// Give the initial synchronous publish time to complete, then cancel.
 	time.Sleep(80 * time.Millisecond)
@@ -628,3 +633,105 @@ func TestMetricsAggregator_Shutdown_Idempotent(t *testing.T) {
 		ma.Shutdown()
 	})
 }
+
+// TestMetricsAggregator_Shutdown_RemovesInstanceMetricsBeforeClose is the L3
+// regression test. It runs the real publisher goroutine (as
+// InitializeMetricsAggregator does) with a tight publish interval to provoke
+// the shutdown-ordering race, then asserts the instance's key and set
+// membership are gone the moment Shutdown() returns - not left lingering
+// until Redis TTL expiry.
+func TestMetricsAggregator_Shutdown_RemovesInstanceMetricsBeforeClose(t *testing.T) {
+	minimalCfg(t)
+
+	mr, err := miniredis.Run()
+	require.NoError(t, err)
+	t.Cleanup(func() { mr.Close() })
+
+	client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
+	ctx, cancel := context.WithCancel(context.Background())
+
+	ma := &MetricsAggregator{
+		redisClient:  client,
+		logger:       libpack_logger.New(),
+		instanceID:   "shutdown-race-test",
+		publishKey:   "graphql-proxy:metrics:instances",
+		ttl:          30 * time.Second,
+		publishTimer: time.NewTicker(5 * time.Millisecond), // tight interval to provoke the race
+		ctx:          ctx,
+		cancel:       cancel,
+	}
+
+	// Run the publisher goroutine exactly like InitializeMetricsAggregator does.
+	ma.startPublishingAsync()
+
+	// Let it publish at least once so the key/set-membership exist before shutdown.
+	time.Sleep(30 * time.Millisecond)
+
+	ma.Shutdown()
+
+	// Give the (now-exiting) publisher goroutine a moment to observe
+	// ctx.Done(); the race detector will flag it if it still touches Redis.
+	time.Sleep(30 * time.Millisecond)
+
+	key := fmt.Sprintf("%s:%s", ma.publishKey, ma.instanceID)
+
+	// Query miniredis directly - ma.redisClient is closed by now.
+	assert.False(t, mr.Exists(key), "instance key must be removed before Shutdown returns, not left lingering until TTL")
+
+	// SREM-ing the only member deletes the set key itself (matches real
+	// Redis semantics for empty collections), so the set key must be gone.
+	assert.False(t, mr.Exists(ma.publishKey), "instance must be removed from the cluster set before Shutdown returns")
+
+	// Confirm the client really is closed - proves the cleanup ran before
+	// Close(), not that Close() never happened.
+	closedErr := ma.redisClient.Ping(context.Background()).Err()
+	assert.Error(t, closedErr, "redis client must be closed after Shutdown")
+}
+
+// TestMetricsAggregator_Shutdown_ClearsSingleton_AllowsReinit is the L6
+// regression test: Shutdown() must clear the package-level singleton so a
+// later InitializeMetricsAggregator call produces a fresh, usable instance
+// instead of silently no-oping against the closed one.
+func TestMetricsAggregator_Shutdown_ClearsSingleton_AllowsReinit(t *testing.T) {
+	// Register mr.Close() before the singleton-restore cleanup below: t.Cleanup
+	// runs LIFO, and the restore cleanup calls Shutdown() on the live instance,
+	// which needs a still-running Redis server. Registering mr.Close() first
+	// means it runs last, after that Shutdown() call has completed.
+	mr, err := miniredis.Run()
+	require.NoError(t, err)
+	t.Cleanup(func() { mr.Close() })
+
+	aggregatorMutex.Lock()
+	old := metricsAggregator
+	metricsAggregator = nil
+	aggregatorMutex.Unlock()
+	t.Cleanup(func() {
+		// Grab-and-release before calling Shutdown(): Shutdown() itself
+		// locks aggregatorMutex, so it must not be invoked while this
+		// goroutine still holds it.
+		aggregatorMutex.Lock()
+		current := metricsAggregator
+		metricsAggregator = old
+		aggregatorMutex.Unlock()
+		if current != nil {
+			current.Shutdown()
+		}
+	})
+
+	require.NoError(t, InitializeMetricsAggregator(mr.Addr(), "", 0, libpack_logger.New()))
+	first := GetMetricsAggregator()
+	require.NotNil(t, first, "first Initialize must produce a usable aggregator")
+
+	first.Shutdown()
+
+	assert.Nil(t, GetMetricsAggregator(), "singleton must be cleared after Shutdown so a closed instance is never returned again")
+
+	require.NoError(t, InitializeMetricsAggregator(mr.Addr(), "", 0, libpack_logger.New()), "re-Initialize after Shutdown must succeed, not silently no-op")
+	second := GetMetricsAggregator()
+	require.NotNil(t, second)
+	assert.NotSame(t, first, second, "re-Initialize must produce a fresh instance, not the closed one")
+
+	// The new instance's Redis client must actually work (not closed).
+	_, err = second.GetAggregatedMetrics()
+	assert.NoError(t, err, "the re-initialized aggregator's Redis client must be usable")
+}