From 1acdebf395eda4d88b15a50de626bc5ce452d091 Mon Sep 17 00:00:00 2001 From: Bethuel Mmbaga Date: Wed, 28 Oct 2026 17:02:11 +0300 Subject: [PATCH] [management] Add configurable statement timeout to postgres store (#8381) --- .../accesslogs/manager/repository.go | 37 +++-- .../accesslogs/manager/repository_test.go | 70 +++++++++ management/internals/shared/db/conn.go | 12 +- management/internals/shared/db/open.go | 27 +++- management/internals/shared/db/timeout.go | 52 +++++++ .../internals/shared/db/timeout_test.go | 146 ++++++++++++++++++ util/duration.go | 16 ++ util/duration_test.go | 25 +++ 8 files changed, 366 insertions(+), 19 deletions(-) create mode 100644 management/internals/shared/db/timeout.go create mode 100644 management/internals/shared/db/timeout_test.go create mode 100644 util/duration_test.go diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/repository.go b/management/internals/modules/reverseproxy/accesslogs/manager/repository.go index 3967442d8..1b2cd4976 100644 --- a/management/internals/modules/reverseproxy/accesslogs/manager/repository.go +++ b/management/internals/modules/reverseproxy/accesslogs/manager/repository.go @@ -14,18 +14,21 @@ import ( "github.com/netbirdio/netbird/shared/management/status" ) +const deleteBatchSize = 10000 + type sqlRepository struct { - conn *db.Conn - db *gorm.DB + conn *db.Conn + db *gorm.DB + deleteBatchSize int } // NewRepository returns the access log repository backed by conn. func NewRepository(conn *db.Conn) accesslogs.Repository { - return &sqlRepository{conn: conn, db: conn.DB(nil)} + return &sqlRepository{conn: conn, db: conn.DB(nil), deleteBatchSize: deleteBatchSize} } func (r *sqlRepository) WithTx(tx *db.Tx) accesslogs.Repository { - return &sqlRepository{conn: r.conn, db: r.conn.DB(tx)} + return &sqlRepository{conn: r.conn, db: r.conn.DB(tx), deleteBatchSize: r.deleteBatchSize} } func (r *sqlRepository) Create(ctx context.Context, entry *accesslogs.AccessLogEntry) error { @@ -73,12 +76,28 @@ func (r *sqlRepository) ListByAccount(ctx context.Context, lockStrength db.Locki } func (r *sqlRepository) DeleteOlderThan(ctx context.Context, olderThan time.Time) (int64, error) { - result := r.db.Where("timestamp < ?", olderThan).Delete(&accesslogs.AccessLogEntry{}) - if result.Error != nil { - log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error) - return 0, status.Errorf(status.Internal, "failed to delete old access logs") + var deleted int64 + for { + if err := ctx.Err(); err != nil { + return deleted, err + } + + batch := r.db.Model(&accesslogs.AccessLogEntry{}). + Select("id"). + Where("timestamp < ?", olderThan). + Order("timestamp"). + Limit(r.deleteBatchSize) + result := r.db.Where("id IN (?)", batch).Delete(&accesslogs.AccessLogEntry{}) + if result.Error != nil { + log.WithContext(ctx).Errorf("failed to delete old access logs: %v", result.Error) + return deleted, status.Errorf(status.Internal, "failed to delete old access logs") + } + + deleted += result.RowsAffected + if result.RowsAffected < int64(r.deleteBatchSize) { + return deleted, nil + } } - return result.RowsAffected, nil } func applyFilters(query *gorm.DB, filter accesslogs.AccessLogFilter) *gorm.DB { diff --git a/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go b/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go index 7fea12a85..6cbbd3c60 100644 --- a/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go +++ b/management/internals/modules/reverseproxy/accesslogs/manager/repository_test.go @@ -3,6 +3,8 @@ package manager import ( "context" "errors" + "fmt" + "os" "testing" "time" @@ -12,6 +14,7 @@ import ( "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/accesslogs" "github.com/netbirdio/netbird/management/internals/shared/db" "github.com/netbirdio/netbird/management/internals/shared/db/dbtest" + "github.com/netbirdio/netbird/management/server/testutil" ) func newTestRepository(t *testing.T) (accesslogs.Repository, *db.Conn) { @@ -157,3 +160,70 @@ func TestSqlRepository_ListByAccount_Search(t *testing.T) { assert.Equal(t, []string{"req-api", "req-web"}, search("example.com"), "search should list every match newest first") assert.Empty(t, search("nothing-matches")) } + +func TestSqlRepository_DeleteOlderThan_DeletesInBatches(t *testing.T) { + cases := map[string]struct{ old, recent int }{ + "partial last batch": {old: 5, recent: 2}, + "exact multiple of batch": {old: 4, recent: 1}, + "nothing to delete": {old: 0, recent: 3}, + } + for name, tc := range cases { + t.Run(name, func(t *testing.T) { + repo, _ := newTestRepository(t) + assertBatchedDelete(t, repo, tc.old, tc.recent) + }) + } +} + +func TestSqlRepository_DeleteOlderThan_DeletesInBatchesOnPostgres(t *testing.T) { + if os.Getenv("NETBIRD_STORE_ENGINE") != string(db.PostgresStoreEngine) { + t.Skip("requires NETBIRD_STORE_ENGINE=postgres") + } + cleanup, dsn, err := testutil.CreatePostgresTestContainer() + require.NoError(t, err) + t.Cleanup(cleanup) + conn, err := db.OpenPostgres(context.Background(), dsn, db.DefaultPoolConfig) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + require.NoError(t, conn.AutoMigrate(&accesslogs.AccessLogEntry{})) + require.NoError(t, conn.DB(nil).Exec("DELETE FROM access_log_entries").Error) + + assertBatchedDelete(t, NewRepository(conn), 5, 2) +} + +func TestSqlRepository_DeleteOlderThan_StopsWhenContextDone(t *testing.T) { + repo, _ := newTestRepository(t) + require.NoError(t, repo.Create(context.Background(), newEntry("old", "acc", "GET", 48*time.Hour))) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour)) + + require.ErrorIs(t, err, context.Canceled) + assert.Zero(t, deleted) + assert.EqualValues(t, 1, countEntries(t, repo), "nothing should be deleted once the context is done") +} + +func assertBatchedDelete(t *testing.T, repo accesslogs.Repository, old, recent int) { + t.Helper() + repo.(*sqlRepository).deleteBatchSize = 2 + ctx := context.Background() + for i := 0; i < old; i++ { + require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("old-%d", i), "acc", "GET", 48*time.Hour))) + } + for i := 0; i < recent; i++ { + require.NoError(t, repo.Create(ctx, newEntry(fmt.Sprintf("new-%d", i), "acc", "GET", time.Hour))) + } + + deleted, err := repo.DeleteOlderThan(ctx, time.Now().Add(-24*time.Hour)) + require.NoError(t, err) + assert.EqualValues(t, old, deleted, "every old entry should be deleted across batches") + assert.EqualValues(t, recent, countEntries(t, repo), "recent entries should be kept") +} + +func countEntries(t *testing.T, repo accesslogs.Repository) int64 { + t.Helper() + _, total, err := repo.ListByAccount(context.Background(), db.LockingStrengthNone, "acc", accesslogs.AccessLogFilter{Page: 1, PageSize: 10}) + require.NoError(t, err) + return total +} diff --git a/management/internals/shared/db/conn.go b/management/internals/shared/db/conn.go index b21b066e2..e12d8e8c8 100644 --- a/management/internals/shared/db/conn.go +++ b/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")) diff --git a/management/internals/shared/db/open.go b/management/internals/shared/db/open.go index f813b0ce8..a64a4003f 100644 --- a/management/internals/shared/db/open.go +++ b/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 diff --git a/management/internals/shared/db/timeout.go b/management/internals/shared/db/timeout.go new file mode 100644 index 000000000..d6bc70296 --- /dev/null +++ b/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) +} diff --git a/management/internals/shared/db/timeout_test.go b/management/internals/shared/db/timeout_test.go new file mode 100644 index 000000000..b9678082d --- /dev/null +++ b/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 +} diff --git a/util/duration.go b/util/duration.go index b657a582d..db4fb3eec 100644 --- a/util/duration.go +++ b/util/duration.go @@ -3,7 +3,10 @@ package util import ( "encoding/json" "errors" + "os" "time" + + log "github.com/sirupsen/logrus" ) // Duration is used strictly for JSON requests/responses due to duration marshalling issues @@ -35,3 +38,16 @@ func (d *Duration) UnmarshalJSON(b []byte) error { return errors.New("invalid duration") } } + +func DurationFromEnv(key string, fallback time.Duration) time.Duration { + val := os.Getenv(key) + if val == "" { + return fallback + } + parsed, err := time.ParseDuration(val) + if err != nil || parsed < 0 { + log.Warnf("ignoring invalid %s value %q, using %v", key, val, fallback) + return fallback + } + return parsed +} diff --git a/util/duration_test.go b/util/duration_test.go new file mode 100644 index 000000000..784bfcf2f --- /dev/null +++ b/util/duration_test.go @@ -0,0 +1,25 @@ +package util + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestDurationFromEnv(t *testing.T) { + const key = "NB_TEST_DURATION" + cases := map[string]time.Duration{ + "": time.Minute, + "30s": 30 * time.Second, + "0": 0, + "soon": time.Minute, + "-1s": time.Minute, + } + for value, want := range cases { + t.Run(value, func(t *testing.T) { + t.Setenv(key, value) + assert.Equal(t, want, DurationFromEnv(key, time.Minute), "duration for %q", value) + }) + } +}