From 76c48d67353a25f922a4df10b491bf6003baa14d Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Wed, 14 Oct 2026 08:54:07 -0700 Subject: [PATCH] avoid scheduler error logs during shutdown When shutdown cancels a pgx query during a socket write, the driver can return an I/O timeout instead of wrapping `context.Canceled`. The scheduler logs this expected cancellation as an error, and the shutdown examples' stdout loggers make their output checks flaky. Check the scheduler's service context before reporting failed passes and exit when it is canceled. Route both shutdown examples' diagnostic logs to stderr so asynchronous diagnostics do not alter expected output. Cover write timeout errors after parent cancellation and explicit stop, and preserve logging for database errors with an active context. --- ..._graceful_shutdown_stop_and_cancel_test.go | 2 +- example_graceful_shutdown_test.go | 2 +- internal/maintenance/job_scheduler.go | 7 +- internal/maintenance/job_scheduler_test.go | 120 ++++++++++++++++++ 4 files changed, 128 insertions(+), 3 deletions(-) diff --git a/example_graceful_shutdown_stop_and_cancel_test.go b/example_graceful_shutdown_stop_and_cancel_test.go index 30217c2d5..23eb387c3 100644 --- a/example_graceful_shutdown_stop_and_cancel_test.go +++ b/example_graceful_shutdown_stop_and_cancel_test.go @@ -39,7 +39,7 @@ func Example_gracefulShutdownStopAndCancel() { river.AddWorker(workers, &WaitsForCancelOnlyWorker{jobStarted: jobStarted}) riverClient, err := river.NewClient(riverpgxv5.New(dbPool), initTestConfig(ctx, dbPool, &river.Config{ - Logger: slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError, ReplaceAttr: slogutil.NoLevelTimeJobID})), + Logger: slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError, ReplaceAttr: slogutil.NoLevelTimeJobID})), Queues: map[string]river.QueueConfig{ river.QueueDefault: {MaxWorkers: 100}, }, diff --git a/example_graceful_shutdown_test.go b/example_graceful_shutdown_test.go index bf99b8ac6..7a32dc6a5 100644 --- a/example_graceful_shutdown_test.go +++ b/example_graceful_shutdown_test.go @@ -60,7 +60,7 @@ func Example_gracefulShutdown() { river.AddWorker(workers, &WaitsForCancelOnlyWorker{jobStarted: jobStarted}) riverClient, err := river.NewClient(riverpgxv5.New(dbPool), initTestConfig(ctx, dbPool, &river.Config{ - Logger: slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError, ReplaceAttr: slogutil.NoLevelTimeJobID})), + Logger: slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError, ReplaceAttr: slogutil.NoLevelTimeJobID})), Queues: map[string]river.QueueConfig{ river.QueueDefault: {MaxWorkers: 100}, }, diff --git a/internal/maintenance/job_scheduler.go b/internal/maintenance/job_scheduler.go index a2d289a24..6e532298a 100644 --- a/internal/maintenance/job_scheduler.go +++ b/internal/maintenance/job_scheduler.go @@ -107,7 +107,7 @@ func NewJobScheduler(archetype *baseservice.Archetype, config *JobSchedulerConfi }) } -func (s *JobScheduler) Start(ctx context.Context) error { //nolint:dupl +func (s *JobScheduler) Start(ctx context.Context) error { ctx, shouldStart, started, stopped := s.StartInit(ctx) if !shouldStart { return nil @@ -132,6 +132,11 @@ func (s *JobScheduler) Start(ctx context.Context) error { //nolint:dupl res, err := s.runOnce(ctx) if err != nil { + // Cancellation can surface as a driver-specific error, like a + // socket write timeout, instead of wrapping context.Canceled. + if ctx.Err() != nil { + return + } if !errors.Is(err, context.Canceled) { s.Logger.ErrorContext(ctx, s.Name+": Error scheduling jobs", slog.String("error", err.Error())) } diff --git a/internal/maintenance/job_scheduler_test.go b/internal/maintenance/job_scheduler_test.go index e960a9302..5e33cfcd7 100644 --- a/internal/maintenance/job_scheduler_test.go +++ b/internal/maintenance/job_scheduler_test.go @@ -1,7 +1,11 @@ package maintenance import ( + "bytes" "context" + "fmt" + "log/slog" + "os" "sort" "testing" "time" @@ -16,10 +20,41 @@ import ( "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/startstoptest" "github.com/riverqueue/river/rivershared/testfactory" + "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivershared/uniquestates" "github.com/riverqueue/river/rivertype" ) +type jobSchedulerExecutorMock struct { + riverdriver.ExecutorTx + + TestSignals jobSchedulerExecutorMockTestSignals + jobScheduleFunc func(context.Context, *riverdriver.JobScheduleParams) ([]*riverdriver.JobScheduleResult, error) +} + +func (e *jobSchedulerExecutorMock) Begin(ctx context.Context) (riverdriver.ExecutorTx, error) { + return e, nil +} + +func (e *jobSchedulerExecutorMock) JobSchedule(ctx context.Context, params *riverdriver.JobScheduleParams) ([]*riverdriver.JobScheduleResult, error) { + e.TestSignals.JobScheduleStarted.Signal(struct{}{}) + return e.jobScheduleFunc(ctx, params) +} + +func (e *jobSchedulerExecutorMock) Rollback(ctx context.Context) error { + return nil +} + +type jobSchedulerExecutorMockTestSignals struct { + JobScheduleStarted testsignal.TestSignal[struct{}] +} + +func (ts *jobSchedulerExecutorMockTestSignals) Init(t *testing.T) { + t.Helper() + + ts.JobScheduleStarted.Init(t) +} + func TestJobScheduler(t *testing.T) { t.Parallel() @@ -440,3 +475,88 @@ func TestJobScheduler(t *testing.T) { } }) } + +func TestJobScheduler_ErrorLogging(t *testing.T) { + t.Parallel() + + type testBundle struct { + exec *jobSchedulerExecutorMock + logOutput bytes.Buffer + scheduler *JobScheduler + } + + // pgx can return a write timeout instead of context.Canceled when shutdown + // interrupts a query write. Preserve that error shape in the regression. + writeErr := fmt.Errorf("write failed: %w", os.ErrDeadlineExceeded) + require.NotErrorIs(t, writeErr, context.Canceled) + + setup := func(t *testing.T) *testBundle { + t.Helper() + + bundle := &testBundle{exec: &jobSchedulerExecutorMock{}} + bundle.exec.TestSignals.Init(t) + bundle.exec.jobScheduleFunc = func(ctx context.Context, params *riverdriver.JobScheduleParams) ([]*riverdriver.JobScheduleResult, error) { + <-ctx.Done() + return nil, writeErr + } + + archetype := riversharedtest.BaseServiceArchetype(t) + archetype.Logger = slog.New(slog.NewTextHandler(&bundle.logOutput, &slog.HandlerOptions{Level: slog.LevelError})) + bundle.scheduler = NewJobScheduler(archetype, &JobSchedulerConfig{Interval: time.Millisecond}, bundle.exec) + bundle.scheduler.StaggerStartupDisable(true) + t.Cleanup(bundle.scheduler.Stop) + + return bundle + } + + t.Run("ContextCancellationWithoutLogging", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + + ctx, cancelFunc := context.WithCancel(context.Background()) + defer cancelFunc() + require.NoError(t, bundle.scheduler.Start(ctx)) + bundle.exec.TestSignals.JobScheduleStarted.WaitOrTimeout() + + cancelFunc() + bundle.scheduler.Stop() + require.Empty(t, bundle.logOutput.String()) + }) + + t.Run("LogsDatabaseErrors", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + + var numCalls int + bundle.exec.jobScheduleFunc = func(ctx context.Context, params *riverdriver.JobScheduleParams) ([]*riverdriver.JobScheduleResult, error) { + numCalls++ + if numCalls > 1 { + <-ctx.Done() + return nil, ctx.Err() + } + return nil, writeErr + } + + require.NoError(t, bundle.scheduler.Start(context.Background())) + bundle.exec.TestSignals.JobScheduleStarted.WaitOrTimeout() + // The next query starts only after the first error has been logged. + bundle.exec.TestSignals.JobScheduleStarted.WaitOrTimeout() + bundle.scheduler.Stop() + require.Contains(t, bundle.logOutput.String(), "Error scheduling jobs") + require.Contains(t, bundle.logOutput.String(), writeErr.Error()) + }) + + t.Run("StopsWithoutLogging", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + + require.NoError(t, bundle.scheduler.Start(context.Background())) + bundle.exec.TestSignals.JobScheduleStarted.WaitOrTimeout() + + bundle.scheduler.Stop() + require.Empty(t, bundle.logOutput.String()) + }) +}