Connection pool bounds and timeouts

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)
+}

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)
+	}
+}