Skip to content

PR_DESCRIPTION_685.md

Latest commit

 

History

History
232 lines (175 loc) · 12.6 KB

PR_DESCRIPTION_685.md

File metadata and controls

232 lines (175 loc) · 12.6 KB
Raw

PR: Enable Graceful Pool Draining on SIGTERM Before Pod Exit

Closes: #685
Branch: feat/graceful-pool-drain
Author: Buffy (Freebuff AI)
Type: Feature / Reliability

Summary

On SIGTERM, the server now stops accepting new HTTP requests, waits for in-flight requests to complete (bounded by a configurable grace period), drains and closes the pgxpool.Pool cleanly, and records a shutdown_duration_seconds histogram — all before the process exits. This eliminates half-open database connections that previously survived pod restarts and could block the next replica from starting.

Problem

Before this PR, the application handled SIGTERM by performing a graceful HTTP server shutdown (srv.Shutdown()), but the *pgxpool.Pool was never explicitly drained or closed. This meant:

  1. Half-open connections — Acquired connections with in-flight queries were left dangling when the process exited without waiting for them to finish or rollback.
  2. Startup blocking — The next replica (in a rolling deployment) could be blocked from acquiring connections if the previous pod's connections were still visible to the database as "idle in transaction" or "active".
  3. No observability — No metric captured how long shutdown took, making it hard to tune terminationGracePeriodSeconds in Kubernetes.
  4. Silent cleanup failures — If cleanup failed, there was no log output to alert operators.

Solution

High-Level Flow

SIGTERM received
    │
    ▼
┌─────────────────────────────────┐
│  HTTP server stops accepting     │  srv.Shutdown(shutdownCtx)
│  new connections                 │  bounded by GracefulShutdownTimeout
│  Drains in-flight HTTP requests  │  (default 30s, env: GRACEFUL_SHUTDOWN_TIMEOUT)
└──────────────┬──────────────────┘
               │
               ▼
┌─────────────────────────────────┐
│  DB pool drain                   │  DrainPool(ctx, pool)
│  - Calls pool.Close() in         │  bounded by fresh GracefulShutdownTimeout
│    goroutine with timeout ctx    │  (always runs even if HTTP shutdown timed out)
│  - Logs errors to stderr         │
│  - Records ShutdownDuration      │
│    histogram                     │
└──────────────┬──────────────────┘
               │
               ▼
           Process exits

Key Design Decisions

  1. Sequential draining — HTTP shutdown completes first, then the DB pool. This ensures no new queries are submitted while connections are being drained.
  2. Cleanup always runs — A defer in runHTTPServer guarantees the pool drain fires even when srv.Shutdown() times out or a second signal forces an immediate close. If we didn't do this, a stuck HTTP handler would prevent the pool from ever closing.
  3. Fresh timeout context for drain — The pool drain gets its own context.WithTimeout (matching GracefulShutdownTimeout) rather than reusing the potentially-expired HTTP shutdown context.
  4. Errors are logged, not swallowed — Drain failures are written to stderr via log.Printf so operators can detect when connections were forcefully terminated rather than drained cleanly.
  5. Nil-safe — DrainPool is a no-op when the pool is nil (e.g., local dev without DATABASE_URL), so the server still shuts down gracefully in all environments.

Files Changed (10 files, +321 / −25 lines)

Core Implementation

File Change
internal/db/pool.go Added DrainPool(ctx, pool) — runs pool.Close() in a goroutine with a bounded context deadline. Returns context.DeadlineExceeded on timeout. Nil-safe.
cmd/server/providers.go Added ProvideDBPool(cfg) — creates *pgxpool.Pool from validated config with a connect-timeout context. Returns (nil, nil) when cfg.DBConn is empty.
cmd/server/wire.go Added ProvideDBPool to AppProviders. Changed InitializeServer() return type from (*http.Server, error) → (*pgxpool.Pool, *http.Server, error).
cmd/server/wire_gen.go Updated generated DI code to match wire.go. InitializeServer now calls ProvideDBPool and returns the pool alongside the server.
cmd/server/main.go main(): receives pool from InitializeServer, creates cleanup closure that calls DrainPool and records ShutdownDuration histogram. Reads GRACEFUL_SHUTDOWN_TIMEOUT via shutdownTimeoutSecs() helper (default 30s, validated 1–600s). runHTTPServer(): cleanup is now deferred and always runs (even on HTTP shutdown timeout). Errors are logged rather than silently discarded.

Tests

File Change
internal/db/pool_test.go 5 new tests: TestDrainPool_NilPool (nil safety), TestDrainPool_AlreadyClosedPool, TestDrainPool_ContextCancelled (pre-cancelled context), TestDrainPool_ClosesPoolCleanly (happy path), TestDrainPool_TimeoutWithAcquiredConnection (edge case: long-running query via pg_sleep(2) with 200ms drain timeout → verifies DeadlineExceeded). Test helper newTestPool(t) skips gracefully when DATABASE_URL is not set or in short mode.
cmd/server/main_test.go Updated TestRunHTTPServer_SecondSignalForcesClose — fixed assertion (was checking for non-existent "forced shutdown after second signal" string), now verifies cleanup is called after forced close. New TestRunHTTPServer_CleanupCalledOnTimeout — verifies cleanup runs even when HTTP shutdown times out (the key invariant).

Kubernetes Manifests

File Change
deploy/kustomize/base/deployment.yaml Added terminationGracePeriodSeconds: 65 with inline documentation explaining the formula: 2 × GRACEFUL_SHUTDOWN_TIMEOUT + 5s buffer.
deploy/helm/stellabill/templates/deployment.yaml Added conditional terminationGracePeriodSeconds block that reads from values.yaml.
deploy/helm/stellabill/values.yaml Added terminationGracePeriodSeconds: 65 to both api and worker sections with documentation.

Timing & Kubernetes Coordination

Application startup
    │
    ├─ config.Load()  ── reads GRACEFUL_SHUTDOWN_TIMEOUT (default 30s)
    ├─ ProvideDBPool() ── creates pgxpool.Pool (bounded by DBPoolConnectTimeout)
    └─ ListenAndServe()
            │
    SIGTERM from kubelet
            │
    ┌───────┴─────────────────────────────────────────────┐
    │  Phase 1: HTTP Shutdown  (≤ 30s)                     │
    │  srv.Shutdown(shutdownCtx) blocks until all handlers  │
    │  return or deadline fires                             │
    └───────┬─────────────────────────────────────────────┘
            │
    ┌───────┴─────────────────────────────────────────────┐
    │  Phase 2: DB Pool Drain  (≤ 30s)                     │
    │  DrainPool(cleanupCtx, pool)                          │
    │    → pool.Close() in goroutine                        │
    │    → waits for acquired connections                    │
    │    → returns on completion or DeadlineExceeded         │
    │  ShutdownDuration.Observe(elapsed)                     │
    └───────┬─────────────────────────────────────────────┘
            │
        exit(0)

Total worst-case application drain: ≤ 60s (2 × 30s)
Kubernetes terminationGracePeriodSeconds: 65s (60s + 5s buffer)

Important: If you change GRACEFUL_SHUTDOWN_TIMEOUT from the default, update terminationGracePeriodSeconds accordingly. The formula is:

terminationGracePeriodSeconds ≥ 2 × GRACEFUL_SHUTDOWN_TIMEOUT + 5

Edge Cases Covered

Edge Case Behavior Test Coverage
No database configured (DATABASE_URL empty) ProvideDBPool returns (nil, nil). DrainPool is a no-op on nil pool. Server shuts down normally. TestNewPool_EmptyDBConnReturnsNilNil + TestDrainPool_NilPool
HTTP shutdown times out (stuck handler) Cleanup defer fires with fresh context, pool drain still runs. Error logged. TestRunHTTPServer_CleanupCalledOnTimeout
Second SIGTERM forces immediate close srv.Close() called. Cleanup defer still runs (pool drain). Error logged. TestRunHTTPServer_SecondSignalForcesClose
Long-running query holds connection pool.Close() blocks. Drain deadline fires → returns DeadlineExceeded. PG sends cancel/terminate → rollback. TestDrainPool_TimeoutWithAcquiredConnection (uses pg_sleep(2))
Pool already closed (e.g., prior drain attempt) pool.Close() is idempotent. Drain returns nil. TestDrainPool_AlreadyClosedPool
Context cancelled before drain Drain observes ctx.Done() and returns context.Canceled. TestDrainPool_ContextCancelled
Pool closes cleanly (no in-flight queries) Drain returns nil. Subsequent Ping fails (pool is closed). TestDrainPool_ClosesPoolCleanly

Metrics

shutdown_duration_seconds (histogram)

Already existed in internal/metrics/metrics.go. Now actually recorded during graceful shutdown.

  • Buckets: [0.1, 0.25, 0.5, 1, 2.5, 5, 10, 25, 30]
  • Labels: none (global metric)
  • Observed: In the cleanup closure, timing the full DrainPool call

Alert ideas:

  • shutdown_duration_seconds approaching GRACEFUL_SHUTDOWN_TIMEOUT → pool drain is slow, investigate query performance or increase timeout
  • shutdown_duration_seconds exceeding GRACEFUL_SHUTDOWN_TIMEOUT → drain timed out, connections may have been force-terminated

Configuration

Env Var Default Valid Range Description
GRACEFUL_SHUTDOWN_TIMEOUT 30 1–600 seconds Maximum time for HTTP shutdown AND pool drain (each gets its own timeout). Already existed in config.go, now also consumed by main.go.
DB_POOL_CONNECT_TIMEOUT 5 1–300 seconds Per-dial timeout when creating the pool at startup. No change.

Migration / Rollback

  • No database schema changes.
  • No configuration changes required. Existing GRACEFUL_SHUTDOWN_TIMEOUT continues to work.
  • Backward compatible. If the pool is nil (e.g., old deployment without DATABASE_URL), shutdown proceeds normally.
  • If reverting: Simply revert this commit. The terminationGracePeriodSeconds: 65 in Kubernetes is safe to leave (it's a maximum, not a minimum).

Verification

Before merging

# Run unit tests (short mode, no DB required)
go test -short -count=1 ./cmd/server/... ./internal/db/...

# Run integration tests (requires DATABASE_URL)
DATABASE_URL="postgres://user:pass@localhost:5432/test?sslmode=disable" \
  go test -count=1 -run 'TestDrainPool' ./internal/db/...

# Lint
golangci-lint run ./cmd/server/... ./internal/db/...

# Regenerate wire (if providers change)
go generate ./cmd/server/...

In staging

  1. Deploy with GRACEFUL_SHUTDOWN_TIMEOUT=30
  2. Trigger a rolling restart: kubectl rollout restart deployment/stellabill-api
  3. Watch logs for cleanup error during shutdown (should not appear in normal operation): 
    kubectl logs -l app=stellabill --tail=50 | grep "cleanup error"
  4. Check Prometheus for shutdown_duration_seconds: 
    shutdown_duration_seconds_bucket{le="30"} - shutdown_duration_seconds_bucket{le="5"}
  5. Verify no pg_stat_activity entries remain from terminated pods: 
    SELECT pid, state, query_start, query
FROM pg_stat_activity
WHERE backend_type = 'client backend' AND state != 'idle';

Post-deploy

  1. Monitor shutdown_duration_seconds for 24h
  2. If p99 approaches 25s during peak traffic, increase GRACEFUL_SHUTDOWN_TIMEOUT and adjust terminationGracePeriodSeconds

Future Work

  • Wire the real *pgxpool.Pool into routes.Register() so handlers use the drained pool (currently uses mock repos)
  • Add a shutdown_connections_leaked gauge that reports pg_stat_activity count after drain
  • Consider draining read-replica pools in internal/db/router.go when a read/write split is active
  • Benchmark drain latency under load (100+ concurrent long-running queries) to validate timeout choices