diff --git a/metrics_aggregator.go b/metrics_aggregator.go index 27de1fa..e076e51 100644 --- a/metrics_aggregator.go +++ b/metrics_aggregator.go @@ -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 diff --git a/metrics_aggregator_test.go b/metrics_aggregator_test.go index f838a09..571af86 100644 --- a/metrics_aggregator_test.go +++ b/metrics_aggregator_test.go @@ -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") +}