From 60e13cc7881e9ae39bc3b5a129fd0277e3e2fc6d Mon Sep 17 00:00:00 2001 From: Jonathon Leight Date: Sun, 5 Jul 2026 19:24:18 -0400 Subject: [PATCH] Connection pool bounds and timeouts --- internal/store/store.go | 50 +++++++++++++++- internal/store/store_test.go | 110 +++++++++++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 3 deletions(-) create mode 100644 internal/store/store_test.go diff --git a/internal/store/store.go b/internal/store/store.go index c98fabd..8bf9594 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -3,8 +3,11 @@ package store import ( "context" + "database/sql" "embed" "fmt" + "strings" + "time" "github.com/jackc/pgx/v5/pgxpool" "github.com/jackc/pgx/v5/stdlib" @@ -14,14 +17,42 @@ import ( //go:embed migrations/*.sql var migrationsFS embed.FS +// defaultStatementTimeout caps any single statement on the app pool so one slow +// query can't hold a pooled connection indefinitely and starve other requests. +// A backstop, not an SLA — normal queries finish in milliseconds. Migrations are +// exempt (see migrationDB). Overridable via a statement_timeout in the DSN. +const defaultStatementTimeout = "30s" + // Store wraps a pgx connection pool and the data-access methods built on it. type Store struct { pool *pgxpool.Pool } // New opens a connection pool against the given DSN and verifies connectivity. +// The pool is bounded and connections are recycled, and each statement is capped +// (defaultStatementTimeout) unless the DSN sets its own statement_timeout. func New(ctx context.Context, dsn string) (*Store, error) { - pool, err := pgxpool.New(ctx, dsn) + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + return nil, fmt.Errorf("parse config: %w", err) + } + // Bound the pool so a burst of slow requests can't exhaust connections, and + // recycle connections so none lives forever (jitter avoids synchronized churn). + // Respect a pool_max_conns set in the DSN (operator tuning; the test harness + // pins it low), applying our default only when unspecified. + if !strings.Contains(dsn, "pool_max_conns") { + cfg.MaxConns = 16 + } + cfg.MaxConnLifetime = time.Hour + cfg.MaxConnLifetimeJitter = 5 * time.Minute + cfg.MaxConnIdleTime = 30 * time.Minute + if cfg.ConnConfig.RuntimeParams == nil { + cfg.ConnConfig.RuntimeParams = map[string]string{} + } + if _, ok := cfg.ConnConfig.RuntimeParams["statement_timeout"]; !ok { + cfg.ConnConfig.RuntimeParams["statement_timeout"] = defaultStatementTimeout + } + pool, err := pgxpool.NewWithConfig(ctx, cfg) if err != nil { return nil, fmt.Errorf("open pool: %w", err) } @@ -45,11 +76,24 @@ func (s *Store) Migrate(ctx context.Context) error { if err := goose.SetDialect("postgres"); err != nil { return fmt.Errorf("set dialect: %w", err) } - // goose needs a database/sql handle; derive one from the pool config. - db := stdlib.OpenDBFromPool(s.pool) + db := s.migrationDB() defer func() { _ = db.Close() }() if err := goose.UpContext(ctx, db, "migrations"); err != nil { return fmt.Errorf("goose up: %w", err) } return nil } + +// migrationDB opens a dedicated database/sql handle for migrations with no +// statement_timeout. Migrations legitimately run long (index builds, backfills), +// so they must not inherit the app pool's per-statement cap (see New). It derives +// from the pool's already-parsed connection config (so pgxpool-only params like +// pool_max_conns are stripped) and overrides just the timeout. +func (s *Store) migrationDB() *sql.DB { + cc := s.pool.Config().ConnConfig.Copy() + if cc.RuntimeParams == nil { + cc.RuntimeParams = map[string]string{} + } + cc.RuntimeParams["statement_timeout"] = "0" // uncapped for migrations + return stdlib.OpenDB(*cc) +} diff --git a/internal/store/store_test.go b/internal/store/store_test.go new file mode 100644 index 0000000..fd0f111 --- /dev/null +++ b/internal/store/store_test.go @@ -0,0 +1,110 @@ +package store + +import ( + "context" + "net/url" + "strings" + "testing" + "time" + + "github.com/jleight/meshtender/internal/testdb" +) + +// freshDSN returns an isolated test database DSN, optionally with extra query +// params merged in (e.g. a short statement_timeout). +func freshDSN(t *testing.T, params map[string]string) string { + t.Helper() + dsn := testdb.Fresh(t, migrateTemplate) + if len(params) == 0 { + return dsn + } + u, err := url.Parse(dsn) + if err != nil { + t.Fatalf("parse dsn: %v", err) + } + q := u.Query() + for k, v := range params { + q.Set(k, v) + } + u.RawQuery = q.Encode() + return u.String() +} + +// TestNewAppliesPoolConfig: New bounds and configures the pool, respects a +// pool_max_conns from the DSN, and caps statements with the default timeout. +func TestNewAppliesPoolConfig(t *testing.T) { + t.Parallel() + ctx := context.Background() + st, err := New(ctx, freshDSN(t, nil)) // testdb DSN pins pool_max_conns=4 + if err != nil { + t.Fatalf("New: %v", err) + } + defer st.Close() + + cfg := st.pool.Config() + if cfg.MaxConns != 4 { + t.Errorf("MaxConns = %d, want 4 (DSN pool_max_conns respected)", cfg.MaxConns) + } + if cfg.MaxConnLifetime != time.Hour { + t.Errorf("MaxConnLifetime = %v, want 1h", cfg.MaxConnLifetime) + } + if cfg.MaxConnIdleTime != 30*time.Minute { + t.Errorf("MaxConnIdleTime = %v, want 30m", cfg.MaxConnIdleTime) + } + + var timeout string + if err := st.pool.QueryRow(ctx, `SHOW statement_timeout`).Scan(&timeout); err != nil { + t.Fatalf("show statement_timeout: %v", err) + } + if timeout != "30s" { + t.Errorf("statement_timeout = %q, want 30s", timeout) + } +} + +// TestStatementTimeoutEnforced: a query exceeding the pool's statement_timeout is +// aborted, so one slow query can't hold a connection indefinitely. +func TestStatementTimeoutEnforced(t *testing.T) { + t.Parallel() + ctx := context.Background() + st, err := New(ctx, freshDSN(t, map[string]string{"statement_timeout": "200"})) // 200ms + if err != nil { + t.Fatalf("New: %v", err) + } + defer st.Close() + + _, err = st.pool.Exec(ctx, `SELECT pg_sleep(2)`) + if err == nil { + t.Fatal("pg_sleep(2) succeeded under a 200ms statement_timeout, want it aborted") + } + if !strings.Contains(err.Error(), "statement timeout") { + t.Fatalf("err = %v, want a statement timeout cancellation", err) + } +} + +// TestMigrationDBUncapped: migrations run on a connection with no +// statement_timeout even when the app pool caps it short, so long DDL isn't +// killed. This is the safety property behind running migrations off the pool. +func TestMigrationDBUncapped(t *testing.T) { + t.Parallel() + ctx := context.Background() + st, err := New(ctx, freshDSN(t, map[string]string{"statement_timeout": "200"})) // pool capped at 200ms + if err != nil { + t.Fatalf("New: %v", err) + } + defer st.Close() + + db := st.migrationDB() + defer func() { _ = db.Close() }() + + var setting string + if err := db.QueryRowContext(ctx, `SELECT current_setting('statement_timeout')`).Scan(&setting); err != nil { + t.Fatalf("read statement_timeout: %v", err) + } + if setting != "0" { + t.Errorf("migration statement_timeout = %q, want 0 (uncapped)", setting) + } + // And a statement longer than the pool's cap must not be killed here. + if _, err := db.ExecContext(ctx, `SELECT pg_sleep(0.5)`); err != nil { + t.Fatalf("pg_sleep(0.5) on migration conn was killed: %v", err) + } +}