diff --git a/main.go b/main.go index 0c15d19..0876a50 100644 --- a/main.go +++ b/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 - shutdownManager.RegisterComponent("http-connection-pool", func(ctx context.Context) error { + // Register connection pool for cleanup. It wraps the same fasthttp + // 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 - shutdownManager.RegisterComponent("backend-health-manager", func(ctx context.Context) error { + // Register backend health manager for cleanup. performProxyRequestWithRetries + // 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 - shutdownManager.RegisterComponent("metrics-aggregator", func(ctx context.Context) error { + // Register metrics aggregator for cleanup. Shutdown() closes its Redis + // 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. - shutdownManager.RegisterComponent("http-proxy", shutdownHTTPProxy) + // instead of being cut off when the process exits. This runs in the + // 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). - shutdownManager.RegisterComponent("cache", func(context.Context) error { + // caches are no-ops). Runs after the traffic drain so in-flight + // 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) }) } diff --git a/shutdown.go b/shutdown.go index da84e39..bdba246 100644 --- a/shutdown.go +++ b/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) - go func(c ShutdownComponent) { - defer shutdownWg.Done() - cfgMutex.RLock() - logger := cfg.Logger - cfgMutex.RUnlock() - if logger != nil { - logger.Info(&libpack_logging.LogMessage{ - Message: "Shutting down component", - Pairs: map[string]any{"component": c.Name}, - }) - } - if err := c.Shutdown(shutdownCtx); err != nil { + if _, seen := phaseGroups[comp.Phase]; !seen { + phases = append(phases, comp.Phase) + } + phaseGroups[comp.Phase] = append(phaseGroups[comp.Phase], comp) + } + sort.Ints(phases) + + // allWg tracks every component across every phase, so the final wait + // below reports success only once nothing is left running anywhere. + var allWg sync.WaitGroup + for _, phase := range phases { + var phaseWg sync.WaitGroup + 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{ - Message: "Error shutting down component", - Pairs: map[string]any{ - "component": c.Name, - "error": err.Error(), - }, + logger.Info(&libpack_logging.LogMessage{ + Message: "Shutting down component", + Pairs: map[string]any{"component": c.Name, "phase": c.Phase}, }) } + 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) }() diff --git a/shutdown_test.go b/shutdown_test.go index 5b39377..4e9db2e 100644 --- a/shutdown_test.go +++ b/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())