Day 107 — Multi-Stage Graceful Shutdown

Key Concepts

3-stage shutdown:

1. Stage 1 (stop accepting new): Close listener, stop accepting new connections. Budget: 0s (immediate).
2. Stage 2 (drain in-flight): Wait for active requests to complete. Budget: 30s.
3. Stage 3 (force-close): Force-close any remaining connections. Budget: 5s.

Exercise — 3-Stage Shutdown

package ops

import (
    "context"
    "log/slog"
    "net/http"
    "os"
    "os/signal"
    "sync"
    "sync/atomic"
    "syscall"
    "time"
)

// ShutdownManager orchestrates the 3-stage shutdown sequence
type ShutdownManager struct {
    server      *http.Server
    inFlight    atomic.Int64
    drainBudget time.Duration
    forceBudget time.Duration
    logger      *slog.Logger
    hooks       []func() // called in order during shutdown (e.g. close DB connections)
}

func NewShutdownManager(server *http.Server, logger *slog.Logger) *ShutdownManager {
    return &ShutdownManager{
        server:      server,
        drainBudget: 30 * time.Second,
        forceBudget: 5 * time.Second,
        logger:      logger,
    }
}

func (sm *ShutdownManager) AddHook(fn func()) {
    sm.hooks = append(sm.hooks, fn)
}

// TrackRequest wraps a handler, tracking in-flight request count
func (sm *ShutdownManager) TrackRequest(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        sm.inFlight.Add(1)
        defer sm.inFlight.Add(-1)
        next.ServeHTTP(w, r)
    })
}

// WaitForSignal blocks until SIGTERM or SIGINT, then executes shutdown sequence
func (sm *ShutdownManager) WaitForSignal() {
    sigCh := make(chan os.Signal, 1)
    signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)

    sig := <-sigCh
    sm.logger.Info("shutdown signal received", "signal", sig.String())

    sm.shutdown()
}

func (sm *ShutdownManager) shutdown() {
    start := time.Now()

    // Stage 1: Stop accepting new connections (immediate)
    sm.logger.Info("shutdown stage 1: stopping acceptance of new connections")
    drainCtx, drainCancel := context.WithTimeout(context.Background(), sm.drainBudget)
    defer drainCancel()

    // http.Server.Shutdown stops accepting and waits for in-flight to complete
    shutdownDone := make(chan error, 1)
    go func() {
        shutdownDone <- sm.server.Shutdown(drainCtx)
    }()

    // Stage 2: Wait for in-flight requests to drain
    sm.logger.Info("shutdown stage 2: draining in-flight requests",
        "in_flight", sm.inFlight.Load(),
        "budget_s", sm.drainBudget.Seconds(),
    )

    ticker := time.NewTicker(1 * time.Second)
    defer ticker.Stop()

drainLoop:
    for {
        select {
        case <-drainCtx.Done():
            sm.logger.Warn("shutdown stage 2: drain budget exhausted",
                "remaining_in_flight", sm.inFlight.Load(),
            )
            break drainLoop
        case err := <-shutdownDone:
            if err != nil {
                sm.logger.Error("shutdown stage 2: server shutdown error", "error", err)
            } else {
                sm.logger.Info("shutdown stage 2: all requests drained",
                    "elapsed_ms", time.Since(start).Milliseconds(),
                )
            }
            break drainLoop
        case <-ticker.C:
            remaining := sm.inFlight.Load()
            if remaining == 0 {
                break drainLoop
            }
            sm.logger.Info("shutdown stage 2: waiting", "in_flight", remaining)
        }
    }

    // Stage 3: Force-close remaining connections
    sm.logger.Info("shutdown stage 3: force closing remaining connections")
    forceCtx, forceCancel := context.WithTimeout(context.Background(), sm.forceBudget)
    defer forceCancel()
    _ = forceCtx // In production: use this to force-close lingering connections

    // Run shutdown hooks (close DB pools, flush metrics, etc.)
    var wg sync.WaitGroup
    for i, hook := range sm.hooks {
        wg.Add(1)
        i, hook := i, hook
        go func() {
            defer wg.Done()
            sm.logger.Info("shutdown hook running", "index", i)
            hook()
        }()
    }

    hookDone := make(chan struct{})
    go func() {
        wg.Wait()
        close(hookDone)
    }()

    select {
    case <-hookDone:
        sm.logger.Info("shutdown complete", "total_elapsed_ms", time.Since(start).Milliseconds())
    case <-time.After(sm.forceBudget):
        sm.logger.Warn("shutdown hooks timed out", "budget_s", sm.forceBudget.Seconds())
    }
}