[management] Add configurable statement timeout to postgres store (#8381)

management/internals/shared/db/conn.go

@@ -11,6 +11,8 @@ import (
 	"github.com/jackc/pgx/v5/pgxpool"
 	log "github.com/sirupsen/logrus"
 	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/util"
 )
 
 const (
@@ -64,13 +66,11 @@ func NewConn(ctx context.Context, gormDB *gorm.DB, engine Engine, pool *pgxpool.
 		opt(&options)
 	}
 
-	txTimeout := defaultTransactionTimeout
-	if v := os.Getenv("NB_STORE_TRANSACTION_TIMEOUT"); v != "" {
-		if parsed, err := time.ParseDuration(v); err == nil {
-			txTimeout = parsed
-		}
-	}
+	txTimeout := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout)
 	log.WithContext(ctx).Infof("Setting transaction timeout to %v", txTimeout)
+	if d := util.DurationFromEnv(statementTimeoutEnv, 0); engine == PostgresStoreEngine && d > 0 {
+		log.WithContext(ctx).Infof("Setting statement timeout to %v", d)
+	}
 
 	conns := runtime.NumCPU()
 	configuredConns, err := strconv.Atoi(os.Getenv("NB_SQL_MAX_OPEN_CONNS"))

management/internals/shared/db/open.go

@@ -10,7 +10,9 @@ import (
 	"strings"
 	"time"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/jackc/pgx/v5/stdlib"
 	"gorm.io/driver/postgres"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
@@ -100,9 +102,10 @@ func OpenSqliteFile(ctx context.Context, dataDir, storeFile string) (*Conn, erro
 	return conn, nil
 }
 
-// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by pool.
+// OpenPostgres opens a Postgres database through gorm and a pgx pool sized by
+// pool, applying the configured timeouts to every session.
 func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Option) (*Conn, error) {
-	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	gormDB, err := openGormPostgres(ctx, dsn)
 	if err != nil {
 		return nil, err
 	}
@@ -123,7 +126,7 @@ func OpenPostgres(ctx context.Context, dsn string, pool PoolConfig, opts ...Opti
 // OpenPostgresWithoutPool opens a Postgres database through gorm alone, for
 // repositories that never take the pgx read paths.
 func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*Conn, error) {
-	gormDB, err := gorm.Open(postgres.Open(dsn), GormConfig())
+	gormDB, err := openGormPostgres(ctx, dsn)
 	if err != nil {
 		return nil, err
 	}
@@ -135,13 +138,29 @@ func OpenPostgresWithoutPool(ctx context.Context, dsn string, opts ...Option) (*
 	return conn, nil
 }
 
-// OpenPgxPool opens a pgx pool to dsn sized by cfg and verifies it with a ping.
+func openGormPostgres(ctx context.Context, dsn string) (*gorm.DB, error) {
+	config, err := pgx.ParseConfig(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("parse database config: %w", err)
+	}
+	sqlDB := stdlib.OpenDB(*config, stdlib.OptionAfterConnect(sessionTimeouts()))
+	gormDB, err := gorm.Open(postgres.New(postgres.Config{Conn: sqlDB}), GormConfig())
+	if err != nil {
+		_ = sqlDB.Close()
+		return nil, err
+	}
+	return gormDB, nil
+}
+
+// OpenPgxPool opens a pgx pool to dsn sized by cfg, applying the configured
+// timeouts to every session, and verifies it with a ping.
 func OpenPgxPool(ctx context.Context, dsn string, cfg PoolConfig) (*pgxpool.Pool, error) {
 	config, err := pgxpool.ParseConfig(dsn)
 	if err != nil {
 		return nil, fmt.Errorf("unable to parse database config: %w", err)
 	}
 
+	config.AfterConnect = sessionTimeouts()
 	config.MaxConns = cfg.MaxConns
 	config.MinConns = cfg.MinConns
 	config.MaxConnLifetime = cfg.MaxConnLifetime

management/internals/shared/db/timeout.go

@@ -0,0 +1,52 @@
+package db
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+
+	"github.com/netbirdio/netbird/util"
+)
+
+const (
+	transactionTimeoutEnv = "NB_STORE_TRANSACTION_TIMEOUT"
+	statementTimeoutEnv   = "NB_STORE_STATEMENT_TIMEOUT"
+)
+
+// sessionTimeouts sets the configured timeouts on each new Postgres session with
+// SET, since poolers such as PgBouncer reject them as startup parameters.
+func sessionTimeouts() func(context.Context, *pgx.Conn) error {
+	var settings []string
+	if d := util.DurationFromEnv(statementTimeoutEnv, 0); d > 0 {
+		settings = append(settings, "SET statement_timeout = "+postgresMillis(d))
+	}
+	if d := util.DurationFromEnv(transactionTimeoutEnv, defaultTransactionTimeout); d > 0 {
+		settings = append(settings, "SET idle_in_transaction_session_timeout = "+postgresMillis(d))
+	}
+	query := strings.Join(settings, "; ")
+
+	return func(ctx context.Context, conn *pgx.Conn) error {
+		if query == "" {
+			return nil
+		}
+		if _, err := conn.Exec(ctx, query); err != nil {
+			return fmt.Errorf("set session timeouts: %w", err)
+		}
+		return nil
+	}
+}
+
+// postgresMillis converts d to milliseconds for a Postgres timeout, rounding up
+// so a positive d never becomes 0 (disabled) and capping at the server maximum.
+func postgresMillis(d time.Duration) string {
+	ms := d.Milliseconds()
+	if d%time.Millisecond != 0 {
+		ms++
+	}
+	return strconv.FormatInt(min(ms, math.MaxInt32), 10)
+}

management/internals/shared/db/timeout_test.go

@@ -0,0 +1,146 @@
+package db
+
+import (
+	"context"
+	"errors"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+
+	"github.com/netbirdio/netbird/management/server/testutil"
+)
+
+// queryCanceledCode is the SQLSTATE Postgres reports when statement_timeout fires.
+const queryCanceledCode = "57014"
+
+func TestPostgresMillis(t *testing.T) {
+	cases := map[time.Duration]string{
+		time.Minute:             "60000",
+		1500 * time.Microsecond: "2",
+		500 * time.Microsecond:  "1",
+		720 * time.Hour:         "2147483647",
+	}
+	for d, want := range cases {
+		assert.Equal(t, want, postgresMillis(d), "milliseconds for %v", d)
+	}
+}
+
+func TestSessionTimeouts_NothingToApply(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "0")
+	t.Setenv(transactionTimeoutEnv, "0")
+	ctx := context.Background()
+
+	assert.NoError(t, sessionTimeouts()(ctx, nil), "the hook should not touch the session when both timeouts are disabled")
+}
+
+func TestOpenPostgres_BoundsEverySession(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	t.Setenv(transactionTimeoutEnv, "3m")
+	conn := openPostgresTestConn(t)
+	ctx := context.Background()
+
+	requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+
+	_, err := conn.Pool(nil).Exec(ctx, "SELECT pg_sleep(1)")
+	requireStatementTimeout(t, err)
+
+	var inTx string
+	require.NoError(t, conn.RunInTx(ctx, func(tx *Tx) error {
+		return conn.DB(tx).Raw("SHOW statement_timeout").Scan(&inTx).Error
+	}))
+	assert.Equal(t, "1min", inTx, "transactions should keep their own 1 minute bound")
+
+	assert.Equal(t, "3min", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+		"the idle transaction timeout should follow the transaction timeout")
+}
+
+func TestOpenPostgres_NoStatementTimeoutKeepsServerSetting(t *testing.T) {
+	for name, value := range map[string]string{"unset": "", "zero": "0"} {
+		t.Run(name, func(t *testing.T) {
+			t.Setenv(statementTimeoutEnv, value)
+			if value == "" {
+				require.NoError(t, os.Unsetenv(statementTimeoutEnv))
+			}
+			ctx := context.Background()
+			dsn := postgresTestDSN(t)
+			admin, err := pgx.Connect(ctx, dsn)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = admin.Close(ctx) })
+			_, err = admin.Exec(ctx, "ALTER ROLE CURRENT_USER SET statement_timeout = '7s'")
+			require.NoError(t, err)
+
+			conn, err := OpenPostgres(ctx, dsn, DefaultPoolConfig)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = conn.Close() })
+
+			assert.Equal(t, "7s", showSetting(t, conn.DB(nil), "statement_timeout"), "sessions should keep the server setting")
+		})
+	}
+}
+
+func TestOpenPostgres_CapsTimeoutsAtServerMaximum(t *testing.T) {
+	t.Setenv(transactionTimeoutEnv, "720h")
+	conn := openPostgresTestConn(t)
+
+	assert.Equal(t, "2147483647ms", showSetting(t, conn.DB(nil), "idle_in_transaction_session_timeout"),
+		"a timeout above the server maximum should be capped instead of failing the connection")
+}
+
+func TestOpenPostgresWithoutPool_BoundsSessions(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	conn, err := OpenPostgresWithoutPool(context.Background(), postgresTestDSN(t))
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	requireStatementTimeout(t, conn.DB(nil).Exec("SELECT pg_sleep(1)").Error)
+}
+
+func TestOpenPgxPool_BoundsSessions(t *testing.T) {
+	t.Setenv(statementTimeoutEnv, "500ms")
+	ctx := context.Background()
+	pool, err := OpenPgxPool(ctx, postgresTestDSN(t), DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(pool.Close)
+
+	_, err = pool.Exec(ctx, "SELECT pg_sleep(1)")
+	requireStatementTimeout(t, err)
+}
+
+func postgresTestDSN(t *testing.T) string {
+	t.Helper()
+	if os.Getenv("NETBIRD_STORE_ENGINE") != string(PostgresStoreEngine) {
+		t.Skip("requires NETBIRD_STORE_ENGINE=postgres")
+	}
+	cleanup, dsn, err := testutil.CreatePostgresTestContainer()
+	require.NoError(t, err)
+	t.Cleanup(cleanup)
+	return dsn
+}
+
+func openPostgresTestConn(t *testing.T) *Conn {
+	t.Helper()
+	conn, err := OpenPostgres(context.Background(), postgresTestDSN(t), DefaultPoolConfig)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	return conn
+}
+
+func requireStatementTimeout(t *testing.T, err error) {
+	t.Helper()
+	var pgErr *pgconn.PgError
+	require.True(t, errors.As(err, &pgErr), "expected a Postgres error, got %v", err)
+	assert.Equal(t, queryCanceledCode, pgErr.Code, "the statement should be cancelled by statement_timeout")
+}
+
+func showSetting(t *testing.T, db *gorm.DB, name string) string {
+	t.Helper()
+	var value string
+	require.NoError(t, db.Raw("SHOW "+name).Scan(&value).Error)
+	return value
+}