Fix graceful shutdown example flakiness

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},
 		},

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},
 		},

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()))
 				}

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())
+	})
+}