From 863cf76ce77f70f808ddf5d72680ccdb8f4d27c6 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Fri, 23 Oct 2026 08:47:17 -0500 Subject: [PATCH] respect maintenance recovery budgets in tests Maintenance recovery uses the generic ten-second CI wait even though resignation attempts and backoff can take about thirteen seconds, and periodic insertion permits a thirty-second transaction. Slow database operations make the test abort a valid recovery and cancel its commit during cleanup. Give leadership recovery thirty seconds and allow periodic insertion its transaction timeout plus the normal event margin. Preserve assertions for a fresh leadership term, completed work, and exactly four startup attempts. Exercise exhausted resignation attempts and a slow periodic commit with controlled driver faults. Capture Postgres activity and blocking PIDs before shutdown when recovery fails so database delays remain observable. --- .../client_maintenance_test.go | 157 +++++++++++++++++- 1 file changed, 155 insertions(+), 2 deletions(-) diff --git a/riverdriver/riverdrivertest/client_maintenance_test.go b/riverdriver/riverdrivertest/client_maintenance_test.go index 77ef4e00a..028c47258 100644 --- a/riverdriver/riverdrivertest/client_maintenance_test.go +++ b/riverdriver/riverdrivertest/client_maintenance_test.go @@ -7,15 +7,126 @@ import ( "testing" "time" + "github.com/jackc/pgx/v5" "github.com/stretchr/testify/require" "github.com/riverqueue/river" + "github.com/riverqueue/river/riverdbtest" "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/riverdriver/riverpgxv5" + "github.com/riverqueue/river/rivershared/riversharedmaintenance" "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivertype" ) +type maintenanceRecoveryDriver struct { + riverdriver.Driver[pgx.Tx] + + commitDelay time.Duration + commitDelays atomic.Int32 + resignAttempts atomic.Int32 + resignTimeouts int32 +} + +func (d *maintenanceRecoveryDriver) GetExecutor() riverdriver.Executor { + return &maintenanceRecoveryExecutor{ + Executor: d.Driver.GetExecutor(), + driver: d, + } +} + +type maintenanceRecoveryExecutor struct { + riverdriver.Executor + + driver *maintenanceRecoveryDriver +} + +func (e *maintenanceRecoveryExecutor) Begin(ctx context.Context) (riverdriver.ExecutorTx, error) { + tx, err := e.Executor.Begin(ctx) + if err != nil { + return nil, err + } + return &maintenanceRecoveryExecutorTx{ExecutorTx: tx, driver: e.driver}, nil +} + +func (e *maintenanceRecoveryExecutor) LeaderResign(ctx context.Context, params *riverdriver.LeaderResignParams) (bool, error) { + if e.driver.resignAttempts.Add(1) <= e.driver.resignTimeouts { + <-ctx.Done() + return false, ctx.Err() + } + return e.Executor.LeaderResign(ctx, params) +} + +type maintenanceRecoveryExecutorTx struct { + riverdriver.ExecutorTx + + driver *maintenanceRecoveryDriver + insertNotified bool +} + +func (e *maintenanceRecoveryExecutorTx) Commit(ctx context.Context) error { + if e.insertNotified && e.driver.commitDelay > 0 { + e.driver.commitDelays.Add(1) + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(e.driver.commitDelay): + } + } + return e.ExecutorTx.Commit(ctx) +} + +func (e *maintenanceRecoveryExecutorTx) NotifyMany(ctx context.Context, params *riverdriver.NotifyManyParams) error { + if params.Topic == "river_insert" { + e.insertNotified = true + } + return e.ExecutorTx.NotifyMany(ctx, params) +} + +func TestMaintenanceStartRecovery(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + driver *maintenanceRecoveryDriver + schema string + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + driver := &maintenanceRecoveryDriver{Driver: riverpgxv5.New(riversharedtest.DBPool(ctx, t))} + return &testBundle{ + driver: driver, + schema: riverdbtest.TestSchema(ctx, t, driver, nil), + } + } + + t.Run("ResignTimeouts", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.driver.resignTimeouts = 3 + + exerciseClientMaintenanceStartRecovery(ctx, t, bundle.driver, bundle.schema, false) + require.GreaterOrEqual(t, bundle.driver.resignAttempts.Load(), int32(3)) + }) + + t.Run("SlowPeriodicCommit", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + // Exceed the ordinary 10-second CI event wait while remaining within + // the enqueuer's 30-second transaction timeout. + bundle.driver.commitDelay = 11 * time.Second + + exerciseClientMaintenanceStartRecovery(ctx, t, bundle.driver, bundle.schema, false) + require.Equal(t, int32(1), bundle.driver.commitDelays.Load()) + }) +} + func exerciseClientMaintenanceStartRecovery[TTx any](ctx context.Context, t *testing.T, driver riverdriver.Driver[TTx], schema string, pollOnly bool) { t.Helper() @@ -51,6 +162,33 @@ func exerciseClientMaintenanceStartRecovery[TTx any](ctx context.Context, t *tes events := subscribe(t, client) startClient(ctx, t, client) + t.Cleanup(func() { + if !t.Failed() || driver.DatabaseName() != riverdriver.DatabaseNamePostgres { + return + } + + // Register after startClient so diagnostics run before it stops the + // client and cancels the operations that explain a recovery timeout. + diagnosticCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + + var activity string + err := riversharedtest.DBPool(diagnosticCtx, t).QueryRow(diagnosticCtx, ` +SELECT coalesce(json_agg(activity)::text, '[]') +FROM ( + SELECT pid, state, wait_event_type, wait_event, + pg_blocking_pids(pid) AS blocking_pids, + clock_timestamp() - query_start AS query_age, + left(query, 2000) AS query + FROM pg_stat_activity + WHERE datname = current_database() AND pid <> pg_backend_pid() AND state <> 'idle' +) AS activity`).Scan(&activity) + if err != nil { + t.Logf("Unable to capture maintenance recovery database activity: %v", err) + return + } + t.Logf("Maintenance recovery database activity before shutdown: %s", activity) + }) first := attemptLeaders.WaitOrTimeout() require.NotNil(t, first) @@ -64,14 +202,29 @@ func exerciseClientMaintenanceStartRecovery[TTx any](ctx context.Context, t *tes // Exhausting startup retries must end the term. With only one client, it // wins again and retries maintenance in a fresh term without a notification. - recovered := attemptLeaders.WaitOrTimeout() + // Three resignation attempts can take 1+2+3 seconds, plus about 1+2+4 + // seconds of backoff. If all fail, the 15-second lease must expire before + // a new term can start. Both paths can exceed the usual test wait. + var recovered *riverdriver.Leader + select { + case recovered = <-attemptLeaders.WaitC(): + case <-time.After(30 * time.Second): + t.Fatal("maintenance did not recover in a new leadership term within 30s") + } require.NotNil(t, recovered) require.Equal(t, client.ID(), recovered.LeaderID) require.True(t, recovered.ElectedAt.After(first.ElectedAt)) // Prove maintenance actually recovered, rather than just receiving a // resignation request: its run-on-start periodic job must be worked. - event := riversharedtest.WaitOrTimeout(t, events) + // A run-on-start insertion has its own maintenance transaction timeout. + // Allow that operation to finish before applying the normal event margin. + var event *river.Event + select { + case event = <-events: + case <-time.After(riversharedmaintenance.TimeoutDefault + riversharedtest.WaitTimeout()): + t.Fatal("maintenance recovery periodic job did not finish within its transaction timeout and event margin") + } require.Equal(t, river.EventKindJobCompleted, event.Kind) require.Equal(t, (noOpArgs{}).Kind(), event.Job.Kind) require.Equal(t, int32(4), startAttempts.Load())