database: add failover-safe pool defaults and pgxpool health checks for Postgres/AlloyDB

database/database.go

@@ -36,11 +36,15 @@ func New(ctx context.Context, logger log.Logger, config DatabaseConfig) (*sql.DB
 		}
 		return db, nil
 	} else if config.Postgres != nil {
+		// Pool settings are applied to pgxpool inside postgresConnection.
+		// Do not call ApplyConnectionsConfig on the returned *sql.DB:
+		// OpenDBFromPool requires MaxIdleConns=0, and sql.DB setters do not
+		// configure the underlying pgxpool.
 		db, err := postgresConnection(ctx, logger, *config.Postgres, config.DatabaseName)
 		if err != nil {
 			return nil, fmt.Errorf("connecting to postgres: %w", err)
 		}
-		return ApplyConnectionsConfig(db, &config.Postgres.Connections, logger), nil
+		return db, nil
 	}
 
 	return nil, ErrMissingConfig
@@ -87,6 +91,19 @@ func DeadlockFound(err error) bool {
 	return MySQLDeadlockFound(err) || PostgresDeadlockFound(err)
 }
 
+// ApplyPostgresConnectionsConfig applies connection pool settings onto a *sql.DB.
+//
+// Deprecated: Postgres connections from New use pgxpool under the hood. Pool
+// settings are applied via ApplyPostgresPoolConfig inside postgresConnection.
+// Calling this on a Postgres *sql.DB from New is incorrect: SetMaxIdleConns
+// with a non-zero value breaks OpenDBFromPool, and the other setters do not
+// configure the underlying pgxpool. Prefer ConnectionsConfig on PostgresConfig
+// (applied automatically) or ApplyPostgresPoolConfig when building a pool.
+func ApplyPostgresConnectionsConfig(db *sql.DB, connections *ConnectionsConfig, logger log.Logger) *sql.DB {
+	applied := ResolvePostgresConnectionsConfig(*connections)
+	return ApplyConnectionsConfig(db, &applied, logger)
+}
+
 func ApplyConnectionsConfig(db *sql.DB, connections *ConnectionsConfig, logger log.Logger) *sql.DB {
 	if connections.MaxOpen > 0 {
 		logger.Logf("setting SQL max open connections to %d", connections.MaxOpen)

database/model_config.go

@@ -104,3 +104,25 @@ type RetryConfig struct {
 	MinDuration time.Duration
 	MaxDuration time.Duration
 }
+
+// DefaultPostgresConnectionsConfig returns connection pool defaults tuned for
+// database failover recovery (e.g. AlloyDB maintenance switchovers).
+//
+// These are applied by ResolvePostgresConnectionsConfig / ApplyPostgresPoolConfig
+// whenever a field on ConnectionsConfig is zero. pgxpool always uses a finite
+// MaxConns (unlike database/sql, where MaxOpen=0 means unlimited), so leaving
+// MaxOpen unset would otherwise silently fall back to max(4, NumCPU()).
+// Explicit defaults keep pool size and eviction policy predictable across
+// services that never set Connections.
+//
+// Short MaxLifetime / MaxIdleTime help the background reaper drop stale
+// connections after a primary change; acquire-time liveness is separate
+// (pgxpool ShouldPing / PingTimeout).
+func DefaultPostgresConnectionsConfig() ConnectionsConfig {
+	return ConnectionsConfig{
+		MaxOpen:     25,
+		MaxIdle:     5,
+		MaxLifetime: 5 * time.Minute,
+		MaxIdleTime: 2 * time.Minute,
+	}
+}

database/postgres.go

@@ -3,14 +3,17 @@ package database
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"errors"
 	"fmt"
 	"net"
 	"strings"
+	"sync"
+	"time"
 
 	"cloud.google.com/go/alloydbconn"
-	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/jackc/pgx/v5/stdlib"
 	"github.com/moov-io/base/log"
 )
@@ -20,38 +23,176 @@ const (
 	// https://www.postgresql.org/docs/current/errcodes-appendix.html
 	postgresErrUniqueViolation = "23505"
 	postgresErrDeadlockFound   = "40P01"
+
+	// Bound ShouldPing wait so acquire retries during failover stay inside
+	// typical request budgets. AlloyDB disconnects usually fail fast (TCP RST);
+	// this caps hung/TIME_WAIT peers.
+	defaultPostgresPingTimeout = time.Second
 )
 
 func postgresConnection(ctx context.Context, logger log.Logger, config PostgresConfig, databaseName string) (*sql.DB, error) {
-	var connStr string
-	if config.Alloy != nil {
-		c, err := getAlloyDBConnectorConnStr(ctx, config, databaseName)
-		if err != nil {
-			return nil, logger.LogErrorf("creating alloydb connection: %w", err).Err()
-		}
-		connStr = c
-	} else {
-		c, err := getPostgresConnStr(config, databaseName)
-		if err != nil {
-			return nil, logger.LogErrorf("creating postgres connection: %w", err).Err()
-		}
-		connStr = c
+	poolConfig, dialer, err := buildPgxPoolConfig(ctx, config, databaseName)
+	if err != nil {
+		return nil, logger.LogErrorf("building pgx pool config: %w", err).Err()
+	}
+
+	// Apply connection limits to pgxpool (not database/sql). OpenDBFromPool
+	// requires sql.DB MaxIdleConns=0; sql.DB setters do not configure the
+	// underlying pool and SetMaxIdleConns(n>0) actively breaks it.
+	ApplyPostgresPoolConfig(logger, poolConfig, config.Connections)
+
+	// Ping connections that have been idle for more than 200ms before handing
+	// them to the caller. This catches dead connections left by an AlloyDB
+	// switchover before a query is attempted, without adding overhead on
+	// hot connections used moments ago.
+	// HealthCheckPeriod (the background reaper) does NOT ping — it only evicts
+	// connections that have exceeded their age thresholds. ShouldPing is the
+	// mechanism that actually tests liveness at acquire time.
+	poolConfig.ShouldPing = func(_ context.Context, p pgxpool.ShouldPingParams) bool {
+		return p.IdleDuration > 200*time.Millisecond
 	}
 
-	db, err := sql.Open("pgx", connStr)
+	if poolConfig.PingTimeout <= 0 {
+		poolConfig.PingTimeout = defaultPostgresPingTimeout
+	}
+
+	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
 	if err != nil {
-		return nil, logger.LogErrorf("opening database: %w", err).Err()
+		_ = closeAlloyDialer(dialer)
+		return nil, logger.LogErrorf("creating pgx pool: %w", err).Err()
 	}
 
-	err = db.Ping()
+	err = pool.Ping(ctx)
 	if err != nil {
-		_ = db.Close()
+		pool.Close()
+		_ = closeAlloyDialer(dialer)
 		return nil, logger.LogErrorf("connecting to database: %w", err).Err()
 	}
 
+	// OpenDBFromPool does not close the pool when *sql.DB is closed. Wrap the
+	// connector so db.Close() shuts down the pool (and AlloyDB dialer).
+	db := openDBFromPool(pool, dialer)
+
 	return db, nil
 }
 
+// ApplyPostgresPoolConfig fills zero-valued fields in connections with
+// DefaultPostgresConnectionsConfig, then maps them onto poolConfig.
+//
+// Unlike database/sql (where MaxOpen=0 means unlimited), pgxpool always has a
+// finite MaxConns. Leaving MaxOpen unset previously fell through to pgxpool's
+// default of max(4, NumCPU()), which silently shrinks pools for services that
+// never configured Connections. We instead apply explicit library defaults so
+// behavior is predictable and logged.
+//
+// MaxIdle has no pgxpool "max idle" equivalent. When set (or defaulted), it is
+// applied as MinIdleConns (warm floor), capped by MaxConns, so the field still
+// influences pool shape rather than being dropped on the floor.
+func ApplyPostgresPoolConfig(logger log.Logger, poolConfig *pgxpool.Config, connections ConnectionsConfig) {
+	if poolConfig == nil {
+		return
+	}
+
+	applied := ResolvePostgresConnectionsConfig(connections)
+
+	logger.Logf("setting pgx pool MaxConns to %d", applied.MaxOpen)
+	poolConfig.MaxConns = int32(applied.MaxOpen)
+
+	minIdle := applied.MaxIdle
+	if minIdle > applied.MaxOpen {
+		minIdle = applied.MaxOpen
+	}
+	if minIdle < 0 {
+		minIdle = 0
+	}
+	logger.Logf("setting pgx pool MinIdleConns to %d (from ConnectionsConfig.MaxIdle)", minIdle)
+	poolConfig.MinIdleConns = int32(minIdle)
+
+	logger.Logf("setting pgx pool MaxConnIdleTime to %v", applied.MaxIdleTime)
+	poolConfig.MaxConnIdleTime = applied.MaxIdleTime
+
+	logger.Logf("setting pgx pool MaxConnLifetime to %v", applied.MaxLifetime)
+	poolConfig.MaxConnLifetime = applied.MaxLifetime
+}
+
+// ResolvePostgresConnectionsConfig returns connections with zero-valued fields
+// replaced by DefaultPostgresConnectionsConfig.
+func ResolvePostgresConnectionsConfig(connections ConnectionsConfig) ConnectionsConfig {
+	defaults := DefaultPostgresConnectionsConfig()
+	if connections.MaxOpen <= 0 {
+		connections.MaxOpen = defaults.MaxOpen
+	}
+	if connections.MaxIdle <= 0 {
+		connections.MaxIdle = defaults.MaxIdle
+	}
+	if connections.MaxLifetime <= 0 {
+		connections.MaxLifetime = defaults.MaxLifetime
+	}
+	if connections.MaxIdleTime <= 0 {
+		connections.MaxIdleTime = defaults.MaxIdleTime
+	}
+	return connections
+}
+
+// openDBFromPool wraps pgxpool in a *sql.DB whose Close also closes the pool
+// and optional AlloyDB dialer. stdlib.OpenDBFromPool alone leaks both.
+func openDBFromPool(pool *pgxpool.Pool, dialer *alloydbconn.Dialer) *sql.DB {
+	c := &poolConnector{
+		Connector: stdlib.GetPoolConnector(pool),
+		pool:      pool,
+		dialer:    dialer,
+	}
+	db := sql.OpenDB(c)
+	// Required when using a pgxpool-backed connector: non-zero idle conns on
+	// sql.DB prevent connections from being released back to the pool.
+	db.SetMaxIdleConns(0)
+	return db
+}
+
+// poolConnector delegates to pgx stdlib's pool connector and implements
+// io.Closer so database/sql.DB.Close shuts down the underlying pgxpool.
+type poolConnector struct {
+	driver.Connector
+	pool   *pgxpool.Pool
+	dialer *alloydbconn.Dialer
+
+	closeOnce sync.Once
+	closeErr  error
+}
+
+func (c *poolConnector) Close() error {
+	c.closeOnce.Do(func() {
+		if c.pool != nil {
+			c.pool.Close()
+		}
+		c.closeErr = closeAlloyDialer(c.dialer)
+	})
+	return c.closeErr
+}
+
+func closeAlloyDialer(dialer *alloydbconn.Dialer) error {
+	if dialer == nil {
+		return nil
+	}
+	return dialer.Close()
+}
+
+func buildPgxPoolConfig(ctx context.Context, config PostgresConfig, databaseName string) (*pgxpool.Config, *alloydbconn.Dialer, error) {
+	if config.Alloy != nil {
+		return buildAlloyDBPoolConfig(ctx, config, databaseName)
+	}
+
+	connStr, err := getPostgresConnStr(config, databaseName)
+	if err != nil {
+		return nil, nil, err
+	}
+	poolConfig, err := pgxpool.ParseConfig(connStr)
+	if err != nil {
+		return nil, nil, err
+	}
+	return poolConfig, nil, nil
+}
+
 func getPostgresConnStr(config PostgresConfig, databaseName string) (string, error) {
 	url := fmt.Sprintf("postgres://%s:%s@%s/%s", config.User, config.Password, config.Address, databaseName)
 
@@ -81,9 +222,9 @@ func getPostgresConnStr(config PostgresConfig, databaseName string) (string, err
 	return connStr, nil
 }
 
-func getAlloyDBConnectorConnStr(ctx context.Context, config PostgresConfig, databaseName string) (string, error) {
+func buildAlloyDBPoolConfig(ctx context.Context, config PostgresConfig, databaseName string) (*pgxpool.Config, *alloydbconn.Dialer, error) {
 	if config.Alloy == nil {
-		return "", fmt.Errorf("missing alloy config")
+		return nil, nil, fmt.Errorf("missing alloy config")
 	}
 
 	var dialer *alloydbconn.Dialer
@@ -92,7 +233,7 @@ func getAlloyDBConnectorConnStr(ctx context.Context, config PostgresConfig, data
 	if config.Alloy.UseIAM {
 		d, err := alloydbconn.NewDialer(ctx, alloydbconn.WithIAMAuthN())
 		if err != nil {
-			return "", fmt.Errorf("creating alloydb dialer: %v", err)
+			return nil, nil, fmt.Errorf("creating alloydb dialer: %w", err)
 		}
 		dialer = d
 		dsn = fmt.Sprintf(
@@ -104,7 +245,7 @@ func getAlloyDBConnectorConnStr(ctx context.Context, config PostgresConfig, data
 	} else {
 		d, err := alloydbconn.NewDialer(ctx)
 		if err != nil {
-			return "", fmt.Errorf("creating alloydb dialer: %v", err)
+			return nil, nil, fmt.Errorf("creating alloydb dialer: %w", err)
 		}
 		dialer = d
 		dsn = fmt.Sprintf(
@@ -114,25 +255,22 @@ func getAlloyDBConnectorConnStr(ctx context.Context, config PostgresConfig, data
 		)
 	}
 
-	// TODO
-	//cleanup := func() error { return d.Close() }
-
-	connConfig, err := pgx.ParseConfig(dsn)
+	poolConfig, err := pgxpool.ParseConfig(dsn)
 	if err != nil {
-		return "", fmt.Errorf("failed to parse pgx config: %v", err)
+		_ = closeAlloyDialer(dialer)
+		return nil, nil, fmt.Errorf("failed to parse pgx pool config: %w", err)
 	}
 
 	var connOptions []alloydbconn.DialOption
 	if config.Alloy.UsePSC {
 		connOptions = append(connOptions, alloydbconn.WithPSC())
 	}
 
-	connConfig.DialFunc = func(ctx context.Context, _ string, _ string) (net.Conn, error) {
+	poolConfig.ConnConfig.DialFunc = func(ctx context.Context, _ string, _ string) (net.Conn, error) {
 		return dialer.Dial(ctx, config.Alloy.InstanceURI, connOptions...)
 	}
 
-	connStr := stdlib.RegisterConnConfig(connConfig)
-	return connStr, nil
+	return poolConfig, dialer, nil
 }
 
 // PostgresUniqueViolation returns true when the provided error matches the Postgres code