fix(shutdown): sequence components into phases so backends close after drain

The manager ran every component concurrently, so on SIGTERM the cache
and connection backends closed while the proxy was still draining
in-flight requests, which then failed with closed clients.

Add an ordered phase to each component. Run phases in ascending order,
components within a phase concurrently, under one timeout budget. Drain
the proxy in an earlier phase than the backend closers.

main.go

@@ -47,6 +47,30 @@ var (
 	shutdownManager *ShutdownManager
 )
 
+// Shutdown phases for components registered on shutdownManager.
+//
+// shutdownPhaseDrainTraffic runs first: it owns the request-serving
+// lifecycle (accepting connections, draining in-flight requests) and must
+// finish before anything it depends on is torn down.
+//
+// shutdownPhaseCloseBackends runs after the drain completes. Every
+// component here (connection pool, backend health manager, metrics
+// aggregator, cache, tracer) is a dependency that in-flight requests read
+// from during the drain - closing them earlier reproduces the redis
+// "client is closed" failures the phased shutdown exists to prevent.
+const (
+	// shutdownPhaseDrainTraffic is negative so it always sorts strictly
+	// before DefaultShutdownPhase (0, see shutdown.go). A plain
+	// RegisterComponent call must never land in the same phase as traffic
+	// draining, or it would be torn down concurrently with in-flight
+	// requests instead of after they finish draining.
+	shutdownPhaseDrainTraffic = -1
+	// shutdownPhaseCloseBackends stays after DefaultShutdownPhase too, so an
+	// unphased component (DefaultShutdownPhase) is guaranteed to shut down
+	// before the backends it may depend on.
+	shutdownPhaseCloseBackends = 1
+)
+
 // getDetailsFromEnv retrieves the value from the environment or returns the default.
 // It first checks for a prefixed environment variable (GMP_KEY), then falls back to the unprefixed version.
 func getDetailsFromEnv[T any](key string, defaultValue T) T {
@@ -604,24 +628,30 @@ func main() {
 		}
 	})
 
-	// Register connection pool for cleanup
+	// Register connection pool for cleanup. It wraps the same fasthttp
-	shutdownManager.RegisterComponent("http-connection-pool", func(ctx context.Context) error {
+	// client the proxy uses per-request (cfg.Client.FastProxyClient), so it
+	// must outlive the traffic drain below.
+	shutdownManager.RegisterComponentWithPhase("http-connection-pool", shutdownPhaseCloseBackends, func(ctx context.Context) error {
 		if connectionPoolManager != nil {
 			return connectionPoolManager.Shutdown()
 		}
 		return nil
 	})
 
-	// Register backend health manager for cleanup
+	// Register backend health manager for cleanup. performProxyRequestWithRetries
-	shutdownManager.RegisterComponent("backend-health-manager", func(ctx context.Context) error {
+	// reads healthMgr.IsHealthy() on every proxied request, so it must
+	// outlive the traffic drain below.
+	shutdownManager.RegisterComponentWithPhase("backend-health-manager", shutdownPhaseCloseBackends, func(ctx context.Context) error {
 		if healthMgr := GetBackendHealthManager(); healthMgr != nil {
 			healthMgr.Shutdown()
 		}
 		return nil
 	})
 
-	// Register metrics aggregator for cleanup
+	// Register metrics aggregator for cleanup. Shutdown() closes its Redis
-	shutdownManager.RegisterComponent("metrics-aggregator", func(ctx context.Context) error {
+	// client, so it belongs with the other backend closers, after the
+	// traffic drain below.
+	shutdownManager.RegisterComponentWithPhase("metrics-aggregator", shutdownPhaseCloseBackends, func(ctx context.Context) error {
 		if aggregator := GetMetricsAggregator(); aggregator != nil {
 			aggregator.Shutdown()
 		}
@@ -629,12 +659,14 @@ func main() {
 	})
 
 	// Register HTTP proxy so in-flight requests drain gracefully on shutdown
-	// instead of being cut off when the process exits.
+	// instead of being cut off when the process exits. This runs in the
-	shutdownManager.RegisterComponent("http-proxy", shutdownHTTPProxy)
+	// traffic-drain phase, before any of the backends above are closed.
+	shutdownManager.RegisterComponentWithPhase("http-proxy", shutdownPhaseDrainTraffic, shutdownHTTPProxy)
 
 	// Close the Redis cache client's connection pool at shutdown (in-memory
-	// caches are no-ops).
+	// caches are no-ops). Runs after the traffic drain so in-flight
-	shutdownManager.RegisterComponent("cache", func(context.Context) error {
+	// CacheLookup/CacheStore calls do not hit a closed Redis client.
+	shutdownManager.RegisterComponentWithPhase("cache", shutdownPhaseCloseBackends, func(context.Context) error {
 		libpack_cache.Shutdown()
 		return nil
 	})
@@ -734,9 +766,11 @@ func main() {
 		Message: "Shutting down services...",
 	})
 
-	// Register tracer shutdown
+	// Register tracer shutdown. proxy.go calls tracer.StartSpan /
+	// ExtractSpanContext on every proxied request, so this closes after the
+	// traffic drain along with the other backend dependencies.
 	if tracer != nil {
-		shutdownManager.RegisterComponent("tracer", func(ctx context.Context) error {
+		shutdownManager.RegisterComponentWithPhase("tracer", shutdownPhaseCloseBackends, func(ctx context.Context) error {
 			return tracer.Shutdown(ctx)
 		})
 	}

shutdown.go

@@ -2,12 +2,26 @@ package main
 
 import (
 	"context"
+	"sort"
 	"sync"
 	"time"
 
 	libpack_logging "github.com/lukaszraczylo/graphql-monitoring-proxy/logging"
 )
 
+// DefaultShutdownPhase is the phase assigned to components registered via
+// RegisterComponent. Components sharing a phase shut down concurrently;
+// phases run sequentially in ascending order. A caller that never assigns an
+// explicit phase gets every component in DefaultShutdownPhase, which
+// reproduces the pre-phase behavior of shutting down everything at once.
+//
+// Invariant: the traffic-drain phase (shutdownPhaseDrainTraffic in main.go)
+// must be strictly earlier than DefaultShutdownPhase. Otherwise an unphased
+// component would land in the same phase as the drain, or later ones, and
+// could be torn down concurrently with in-flight requests instead of after
+// they drain - the exact hazard phased shutdown exists to prevent.
+const DefaultShutdownPhase = 0
+
 // ShutdownManager manages graceful shutdown for all components
 type ShutdownManager struct {
 	ctx          context.Context
@@ -18,10 +32,15 @@ type ShutdownManager struct {
 	mu           sync.Mutex
 }
 
-// ShutdownComponent represents a component that needs graceful shutdown
+// ShutdownComponent represents a component that needs graceful shutdown.
+// Phase controls ordering relative to other components: lower phases shut
+// down first, and shutdown only moves to the next phase once every
+// component in the current phase has returned or the overall shutdown
+// deadline has passed.
 type ShutdownComponent struct {
 	Shutdown func(context.Context) error
 	Name     string
+	Phase    int
 }
 
 // NewShutdownManager creates a new shutdown manager
@@ -33,12 +52,23 @@ func NewShutdownManager(ctx context.Context) *ShutdownManager {
 	}
 }
 
-// RegisterComponent registers a component for graceful shutdown
+// RegisterComponent registers a component for graceful shutdown in
+// DefaultShutdownPhase. Use RegisterComponentWithPhase to sequence a
+// component before or after others (for example, draining traffic before
+// closing the backends that in-flight requests depend on).
 func (sm *ShutdownManager) RegisterComponent(name string, shutdown func(context.Context) error) {
+	sm.RegisterComponentWithPhase(name, DefaultShutdownPhase, shutdown)
+}
+
+// RegisterComponentWithPhase registers a component for graceful shutdown in
+// the given phase. Shutdown runs phases sequentially in ascending order;
+// components within the same phase shut down concurrently.
+func (sm *ShutdownManager) RegisterComponentWithPhase(name string, phase int, shutdown func(context.Context) error) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 	sm.components = append(sm.components, ShutdownComponent{
 		Name:     name,
+		Phase:    phase,
 		Shutdown: shutdown,
 	})
 }
@@ -97,47 +127,103 @@ func (sm *ShutdownManager) doShutdown(timeout time.Duration) error {
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), timeout)
 	defer shutdownCancel()
 
-	// Shutdown all registered components
+	// Shutdown all registered components, grouped by phase. Phases run
+	// sequentially in ascending order against the single shutdownCtx budget
+	// above (no per-phase timeout reset), so a component in an earlier phase
+	// only shuts down once the whole shutdown has enough of that shared
+	// budget remaining. This budget bounds the phases loop and the
+	// componentsDone wait below - it does not bound all of Shutdown(): the
+	// separate goroutinesDone wait further down (for RunGoroutine-managed
+	// goroutines) uses its own fresh time.After(timeout) budget, so
+	// worst-case total shutdown time is up to ~2x timeout. That is
+	// pre-existing, intentional behavior, unchanged here.
 	sm.mu.Lock()
 	components := make([]ShutdownComponent, len(sm.components))
 	copy(components, sm.components)
 	sm.mu.Unlock()
 
-	var shutdownWg sync.WaitGroup
+	phaseGroups := make(map[int][]ShutdownComponent, len(components))
+	phases := make([]int, 0, len(components))
 	for _, comp := range components {
-		shutdownWg.Add(1)
+		if _, seen := phaseGroups[comp.Phase]; !seen {
-		go func(c ShutdownComponent) {
+			phases = append(phases, comp.Phase)
-			defer shutdownWg.Done()
+		}
-			cfgMutex.RLock()
+		phaseGroups[comp.Phase] = append(phaseGroups[comp.Phase], comp)
-			logger := cfg.Logger
+	}
-			cfgMutex.RUnlock()
+	sort.Ints(phases)
-			if logger != nil {
+
-				logger.Info(&libpack_logging.LogMessage{
+	// allWg tracks every component across every phase, so the final wait
-					Message: "Shutting down component",
+	// below reports success only once nothing is left running anywhere.
-					Pairs:   map[string]any{"component": c.Name},
+	var allWg sync.WaitGroup
-				})
+	for _, phase := range phases {
-			}
+		var phaseWg sync.WaitGroup
-			if err := c.Shutdown(shutdownCtx); err != nil {
+		for _, comp := range phaseGroups[phase] {
+			phaseWg.Add(1)
+			allWg.Add(1)
+			go func(c ShutdownComponent) {
+				defer phaseWg.Done()
+				defer allWg.Done()
 				cfgMutex.RLock()
 				logger := cfg.Logger
 				cfgMutex.RUnlock()
 				if logger != nil {
-					logger.Error(&libpack_logging.LogMessage{
+					logger.Info(&libpack_logging.LogMessage{
-						Message: "Error shutting down component",
+						Message: "Shutting down component",
-						Pairs: map[string]any{
+						Pairs:   map[string]any{"component": c.Name, "phase": c.Phase},
-							"component": c.Name,
-							"error":     err.Error(),
-						},
 					})
 				}
+				if err := c.Shutdown(shutdownCtx); err != nil {
+					cfgMutex.RLock()
+					logger := cfg.Logger
+					cfgMutex.RUnlock()
+					if logger != nil {
+						logger.Error(&libpack_logging.LogMessage{
+							Message: "Error shutting down component",
+							Pairs: map[string]any{
+								"component": c.Name,
+								"phase":     c.Phase,
+								"error":     err.Error(),
+							},
+						})
+					}
+				}
+			}(comp)
+		}
+
+		phaseDone := make(chan struct{})
+		go func() {
+			phaseWg.Wait()
+			close(phaseDone)
+		}()
+
+		// Bound the wait for this phase by the overall shutdown budget, not
+		// a fresh per-phase timeout. If the deadline fires while this phase
+		// is still running, its stragglers keep running in the background
+		// (they are never killed, only ctx-cancelled) but we move on so the
+		// next phase still gets attempted instead of deadlocking here. If it
+		// is specifically the drain phase that times out, later phases go
+		// on to close backends while requests still draining in the
+		// background depend on them - an unavoidable consequence of
+		// enforcing a fixed shutdown budget.
+		select {
+		case <-phaseDone:
+		case <-shutdownCtx.Done():
+			cfgMutex.RLock()
+			logger := cfg.Logger
+			cfgMutex.RUnlock()
+			if logger != nil {
+				logger.Warning(&libpack_logging.LogMessage{
+					Message: "Shutdown phase timed out, proceeding to next phase",
+					Pairs:   map[string]any{"phase": phase},
+				})
 			}
-		}(comp)
+		}
 	}
 
-	// Wait for all components to shutdown
+	// Wait for all components, across all phases, to finish shutting down.
 	componentsDone := make(chan struct{})
 	go func() {
-		shutdownWg.Wait()
+		allWg.Wait()
 		close(componentsDone)
 	}()
 

shutdown_test.go

@@ -297,6 +297,183 @@ func (suite *ShutdownTestSuite) TestContextCancellation() {
 	}
 }
 
+func (suite *ShutdownTestSuite) TestRegisterComponentDefaultPhase() {
+	sm := NewShutdownManager(context.Background())
+
+	sm.RegisterComponent("component1", func(ctx context.Context) error {
+		return nil
+	})
+
+	assert.Len(suite.T(), sm.components, 1)
+	assert.Equal(suite.T(), DefaultShutdownPhase, sm.components[0].Phase)
+}
+
+func (suite *ShutdownTestSuite) TestRegisterComponentWithPhase() {
+	sm := NewShutdownManager(context.Background())
+
+	sm.RegisterComponentWithPhase("early", 1, func(ctx context.Context) error {
+		return nil
+	})
+	sm.RegisterComponentWithPhase("late", 5, func(ctx context.Context) error {
+		return nil
+	})
+
+	assert.Len(suite.T(), sm.components, 2)
+	assert.Equal(suite.T(), 1, sm.components[0].Phase)
+	assert.Equal(suite.T(), 5, sm.components[1].Phase)
+}
+
+// TestShutdownPhasesRunSequentially registers components in two phases and
+// records each component's start/end wall-clock time. It asserts every
+// phase-2 component starts only after every phase-1 component has finished,
+// proving phases execute sequentially rather than all at once.
+func (suite *ShutdownTestSuite) TestShutdownPhasesRunSequentially() {
+	sm := NewShutdownManager(context.Background())
+
+	type span struct {
+		start time.Time
+		end   time.Time
+	}
+	var mu sync.Mutex
+	spans := make(map[string]span)
+
+	recordSpan := func(name string, work func()) {
+		mu.Lock()
+		spans[name] = span{start: time.Now()}
+		mu.Unlock()
+
+		work()
+
+		mu.Lock()
+		s := spans[name]
+		s.end = time.Now()
+		spans[name] = s
+		mu.Unlock()
+	}
+
+	sm.RegisterComponentWithPhase("phase1-a", 1, func(ctx context.Context) error {
+		recordSpan("phase1-a", func() { time.Sleep(30 * time.Millisecond) })
+		return nil
+	})
+	sm.RegisterComponentWithPhase("phase1-b", 1, func(ctx context.Context) error {
+		recordSpan("phase1-b", func() { time.Sleep(30 * time.Millisecond) })
+		return nil
+	})
+	sm.RegisterComponentWithPhase("phase2-a", 2, func(ctx context.Context) error {
+		recordSpan("phase2-a", func() {})
+		return nil
+	})
+	sm.RegisterComponentWithPhase("phase2-b", 2, func(ctx context.Context) error {
+		recordSpan("phase2-b", func() {})
+		return nil
+	})
+
+	err := sm.Shutdown(2 * time.Second)
+	assert.NoError(suite.T(), err)
+
+	mu.Lock()
+	defer mu.Unlock()
+
+	assert.Len(suite.T(), spans, 4, "all four components must have run")
+
+	phase1End := spans["phase1-a"].end
+	if spans["phase1-b"].end.After(phase1End) {
+		phase1End = spans["phase1-b"].end
+	}
+
+	assert.False(suite.T(), spans["phase2-a"].start.Before(phase1End),
+		"phase2-a started before every phase-1 component finished")
+	assert.False(suite.T(), spans["phase2-b"].start.Before(phase1End),
+		"phase2-b started before every phase-1 component finished")
+}
+
+// TestShutdownSamePhaseComponentsRunConcurrently proves components sharing a
+// phase still run concurrently. Every registered component blocks on a
+// shared release gate that only opens once all of them have signaled they
+// started; if the components ran one at a time, the first would block
+// forever on the gate and the test would time out.
+func (suite *ShutdownTestSuite) TestShutdownSamePhaseComponentsRunConcurrently() {
+	sm := NewShutdownManager(context.Background())
+
+	const numComponents = 3
+	started := make(chan struct{}, numComponents)
+	release := make(chan struct{})
+
+	for i := 0; i < numComponents; i++ {
+		sm.RegisterComponentWithPhase("component", 7, func(ctx context.Context) error {
+			started <- struct{}{}
+			<-release
+			return nil
+		})
+	}
+
+	shutdownDone := make(chan struct{})
+	go func() {
+		err := sm.Shutdown(2 * time.Second)
+		assert.NoError(suite.T(), err)
+		close(shutdownDone)
+	}()
+
+	for i := 0; i < numComponents; i++ {
+		select {
+		case <-started:
+		case <-time.After(1 * time.Second):
+			suite.T().Fatal("same-phase components did not all start concurrently")
+		}
+	}
+	close(release)
+
+	select {
+	case <-shutdownDone:
+	case <-time.After(1 * time.Second):
+		suite.T().Fatal("shutdown did not complete after releasing same-phase components")
+	}
+}
+
+// TestShutdownPhaseTimeoutStillAttemptsNextPhase proves a phase-1 component
+// that ignores the ctx deadline and outlives the shutdown budget does not
+// prevent phase 2 from being attempted, and that Shutdown still returns
+// promptly instead of blocking on the hung component.
+func (suite *ShutdownTestSuite) TestShutdownPhaseTimeoutStillAttemptsNextPhase() {
+	sm := NewShutdownManager(context.Background())
+
+	phase1Started := make(chan struct{})
+	sm.RegisterComponentWithPhase("phase1-hung", 1, func(ctx context.Context) error {
+		close(phase1Started)
+		// Ignore ctx entirely to simulate a component that does not
+		// respect the shutdown deadline.
+		time.Sleep(5 * time.Second)
+		return nil
+	})
+
+	phase2Attempted := make(chan struct{})
+	sm.RegisterComponentWithPhase("phase2-after-hang", 2, func(ctx context.Context) error {
+		close(phase2Attempted)
+		return nil
+	})
+
+	start := time.Now()
+	err := sm.Shutdown(100 * time.Millisecond)
+	elapsed := time.Since(start)
+
+	assert.NoError(suite.T(), err)
+	assert.Less(suite.T(), elapsed, 1*time.Second,
+		"shutdown must not block waiting on a hung phase-1 component")
+
+	select {
+	case <-phase1Started:
+	default:
+		suite.T().Fatal("phase-1 component was never started")
+	}
+
+	select {
+	case <-phase2Attempted:
+		// Good: phase 2 was attempted despite phase 1 exceeding the deadline.
+	case <-time.After(500 * time.Millisecond):
+		suite.T().Fatal("phase-2 component was never attempted after phase-1 exceeded the shutdown deadline")
+	}
+}
+
 // Benchmark tests
 func BenchmarkRegisterComponent(b *testing.B) {
 	sm := NewShutdownManager(context.Background())