Add Config.SoftStopTimeout for a cleaner way to shut a client down (#1239)

A part of River code that's always bothered me is the graceful shutdown
example [1]. It's really, really complicated, and I think that
complicated orchestration is very likely a sign that our code isn't
quite right.

I initially tried putting the graceful shutdown code into a helper like
`river.GracefulShutdown`, but I found that also problematic because it's
not super clean to call, and even minor customizations would require a
user to copy out the entire function.

Playing with an LLM a bit, I came up with this alternative: build an
understanding of soft/hard stop into the client's default `Stop`
function by adding `Config.SoftStopTimeout`. By default, leaving this
value unconfigured keeps the existing behavior of River today where a
call to `Stop` waits unbounded time for jobs to finish working. Adding a
`SoftStopTimeout` value brings in functionality similar to the graceful
shutdown example: stopping proceeds normally waiting for jobs to finish,
but in case they don't inside of `SoftStopTimeout`, they're cancelled as
if `StopAndCancel` had been called.

The soft timeout is also respected for a cancelled `Start` context,
which makes the basic use of River for soft/hard stop very nice (you
don't even need to call `Stop` anymore):

    // Use signal.NotifyContext to cancel the start context on SIGINT/SIGTERM.
    // When the signal fires, the client initiates a soft stop. If running jobs
    // don't finish within SoftStopTimeout, their contexts are automatically
    // cancelled.
    ctx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
    defer stop()

    if err := riverClient.Start(ctx); err != nil {
        panic(err)
    }

[1] https://github.com/riverqueue/river/blob/master/example_graceful_shutdown_test.go

Author: brandur

client.go

@@ -329,6 +329,32 @@ type Config struct {
 	// setting of Postgres `search_path`.
 	Schema string
 
+	// SoftStopTimeout is the maximum amount of time that the client will wait
+	// for running jobs to finish during a stop before their contexts are
+	// cancelled. After the timeout elapses, the client escalates to a hard stop
+	// by cancelling the context of all running jobs. This applies regardless of
+	// how stop is initiated — whether by calling Stop, StopAndCancel, or by
+	// cancelling the context passed to Start.
+	//
+	// In combination with signal.NotifyContext on the context passed to Start,
+	// this can simplify graceful stop to:
+	//
+	//	ctx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
+	//	defer stop()
+	//
+	//	if err := client.Start(ctx); err != nil { ... }
+	//	<-client.Stopped()
+	//
+	// The signal cancels the Start context, which initiates a soft stop. If
+	// running jobs haven't finished after SoftStopTimeout, their contexts are
+	// automatically cancelled to trigger a hard stop.
+	//
+	// StopAndCancel bypasses the timeout entirely and cancels job contexts
+	// immediately.
+	//
+	// Defaults to no timeout (wait indefinitely for jobs to finish).
+	SoftStopTimeout time.Duration
+
 	// SkipJobKindValidation causes the job kind format validation check to be
 	// skipped. This is available as an interim stopgap for users that have
 	// invalid job kind names, but would rather disable the check rather than
@@ -458,6 +484,7 @@ func (c *Config) WithDefaults() *Config {
 		RescueStuckJobsAfter:        cmp.Or(c.RescueStuckJobsAfter, rescueAfter),
 		RetryPolicy:                 retryPolicy,
 		Schema:                      c.Schema,
+		SoftStopTimeout:             c.SoftStopTimeout,
 		SkipJobKindValidation:       c.SkipJobKindValidation,
 		SkipUnknownJobCheck:         c.SkipUnknownJobCheck,
 		Test:                        c.Test,
@@ -1084,10 +1111,19 @@ func (c *Client[TTx]) Start(ctx context.Context) error {
 			return err
 		}
 
-		// We use separate contexts for fetching and working to allow for a graceful
-		// stop. Both inherit from the provided context, so if it's cancelled, a
-		// more aggressive stop will be initiated.
-		workCtx, workCancel := context.WithCancelCause(ctx)
+		// We use separate contexts for fetching and working to allow for a
+		// graceful stop. When SoftStopTimeout is configured, the work context
+		// is detached from the start context so that cancelling the start
+		// context initiates a soft stop (with timeout escalation) rather than
+		// an immediate hard stop. When SoftStopTimeout is not configured, the
+		// work context inherits from the start context to preserve the
+		// existing behavior where cancelling the start context is equivalent
+		// to StopAndCancel.
+		workParentCtx := ctx
+		if c.config.SoftStopTimeout > 0 {
+			workParentCtx = context.WithoutCancel(ctx)
+		}
+		workCtx, workCancel := context.WithCancelCause(workParentCtx)
 
 		// Client available to executors and to various service hooks.
 		fetchCtx := withClient(fetchCtx, c)
@@ -1151,6 +1187,18 @@ func (c *Client[TTx]) Start(ctx context.Context) error {
 		c.queues.startStopMu.Lock()
 		defer c.queues.startStopMu.Unlock()
 
+		// If SoftStopTimeout is configured, start a timer that will cancel
+		// the work context (escalating to a hard stop) if producers don't
+		// finish in time. StopAndCancel also calls workCancel, in which case
+		// this timer is a harmless no-op because the context is already done.
+		if c.config.SoftStopTimeout > 0 {
+			softStopTimer := time.AfterFunc(c.config.SoftStopTimeout, func() {
+				c.baseService.Logger.WarnContext(ctx, c.baseService.Name+": Soft stop timeout; cancelling remaining job contexts", slog.Duration("soft_stop_timeout", c.config.SoftStopTimeout))
+				c.workCancel(rivercommon.ErrStop)
+			})
+			defer softStopTimer.Stop()
+		}
+
 		// On stop, have the producers stop fetching first of all.
 		c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": Stopping producers")
 		startstop.StopAllParallel(producersAsServices()...)
@@ -1183,6 +1231,10 @@ func (c *Client[TTx]) Start(ctx context.Context) error {
 // complete before exiting. If the provided context is done before shutdown has
 // completed, Stop will return immediately with the context's error.
 //
+// If SoftStopTimeout is configured, running job contexts will be automatically
+// cancelled after the timeout elapses, escalating to a hard stop. This also
+// applies when stop is initiated by cancelling the context passed to Start.
+//
 // There's no need to call this method if a hard stop has already been initiated
 // by cancelling the context passed to Start or by calling StopAndCancel.
 func (c *Client[TTx]) Stop(ctx context.Context) error {
@@ -1211,6 +1263,12 @@ func (c *Client[TTx]) Stop(ctx context.Context) error {
 // This can also be initiated by cancelling the context passed to Start. There is
 // no need to call this method if the context passed to Start is cancelled
 // instead.
+//
+// In most cases, using Stop with SoftStopTimeout configured is preferable to
+// calling StopAndCancel directly. SoftStopTimeout gives running jobs a chance
+// to finish before automatically escalating to context cancellation, providing
+// graceful stop semantics without requiring manual orchestration of Stop and
+// StopAndCancel.
 func (c *Client[TTx]) StopAndCancel(ctx context.Context) error {
 	c.baseService.Logger.InfoContext(ctx, c.baseService.Name+": Hard stop started; cancelling all work")
 	c.workCancel(rivercommon.ErrStop)

client_test.go

@@ -2361,6 +2361,122 @@ func Test_Client_StopAndCancel(t *testing.T) {
 	})
 }
 
+func Test_Client_SoftStopTimeout(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	type JobArgs struct {
+		testutil.JobArgsReflectKind[JobArgs]
+	}
+
+	t.Run("EscalatesToHardStopAfterTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		config := newTestConfig(t, "")
+		config.SoftStopTimeout = 100 * time.Millisecond
+
+		jobDoneChan := make(chan struct{})
+		jobStartedChan := make(chan struct{})
+		AddWorker(config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error {
+			close(jobStartedChan)
+			<-ctx.Done() // only finishes when context is cancelled
+			close(jobDoneChan)
+			return nil
+		}))
+
+		client := runNewTestClient(ctx, t, config)
+
+		_, err := client.Insert(ctx, JobArgs{}, nil)
+		require.NoError(t, err)
+
+		riversharedtest.WaitOrTimeout(t, jobStartedChan)
+
+		// Stop initiates a soft stop. The job won't finish on its own, but
+		// SoftStopTimeout should escalate to a hard stop after 100ms.
+		require.NoError(t, client.Stop(ctx))
+
+		// Verify the job's context was indeed cancelled.
+		select {
+		case <-jobDoneChan:
+		default:
+			t.Fatal("expected job to have been cancelled by soft stop timeout")
+		}
+	})
+
+	t.Run("SoftStopSucceedsBeforeTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		config := newTestConfig(t, "")
+		config.SoftStopTimeout = 5 * time.Second
+
+		jobStartedChan := make(chan struct{})
+		AddWorker(config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error {
+			close(jobStartedChan)
+			return nil // finishes immediately
+		}))
+
+		client := runNewTestClient(ctx, t, config)
+
+		_, err := client.Insert(ctx, JobArgs{}, nil)
+		require.NoError(t, err)
+
+		riversharedtest.WaitOrTimeout(t, jobStartedChan)
+
+		// Stop should succeed quickly since the job finishes on its own.
+		// The 5s timeout should not fire.
+		require.NoError(t, client.Stop(ctx))
+	})
+
+	t.Run("ContextCancellationEscalatesAfterTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		config := newTestConfig(t, "")
+		config.SoftStopTimeout = 100 * time.Millisecond
+
+		jobDoneChan := make(chan struct{})
+		jobStartedChan := make(chan struct{})
+		AddWorker(config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error {
+			close(jobStartedChan)
+			<-ctx.Done()
+			close(jobDoneChan)
+			return nil
+		}))
+
+		var (
+			dbPool = riversharedtest.DBPool(ctx, t)
+			driver = riverpgxv5.New(dbPool)
+			schema = riverdbtest.TestSchema(ctx, t, driver, nil)
+		)
+		config.Schema = schema
+
+		client, err := NewClient(driver, config)
+		require.NoError(t, err)
+
+		startCtx, startCtxCancel := context.WithCancel(ctx)
+		defer startCtxCancel()
+
+		require.NoError(t, client.Start(startCtx))
+
+		_, err = client.Insert(ctx, JobArgs{}, nil)
+		require.NoError(t, err)
+
+		riversharedtest.WaitOrTimeout(t, jobStartedChan)
+
+		// Cancel the start context. This should initiate a soft stop, then
+		// escalate to hard stop after SoftStopTimeout.
+		startCtxCancel()
+
+		riversharedtest.WaitOrTimeout(t, client.Stopped())
+
+		select {
+		case <-jobDoneChan:
+		default:
+			t.Fatal("expected job to have been cancelled by soft stop timeout")
+		}
+	})
+}
+
 type callbackWithCustomTimeoutArgs struct {
 	TimeoutValue time.Duration `json:"timeout"`
 }

docs/README.md

@@ -87,17 +87,39 @@ that inserted job kinds have a worker that can run them.
 
 ### Stopping
 
-The client should also be stopped on program shutdown:
+The client should also be stopped on program shutdown. There's a number of ways
+to go about this (see [graceful shutdown]), but the shortest is to cancel the
+context send to `Start` when the program's ready to stop. For example, to stop
+on `SIGINT`/`SIGTERM`:
 
 ```go
-// Stop fetching new work and wait for active jobs to finish.
-if err := riverClient.Stop(ctx); err != nil {
+riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{
+    SoftStopTimeout: 10 * time.Second,
+    ...
+})
+if err != nil {
+    panic(err)
+}
+
+signalCtx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
+defer stop()
+
+// Stop fetching new work and wait for active jobs to finish. Cancel jobs after
+// SoftStopTimeout elapses.
+if err := riverClient.Start(signalCtx); err != nil {
     panic(err)
 }
+
+<-riverClient.Stopped()
 ```
 
-There are some complexities around ensuring clients stop cleanly, but also in a
-timely manner. See [graceful shutdown] for more details on River's stop modes.
+Alternatively, use an explicit call to `Stop`:
+
+```go
+if err := riverClient.Stop(ctx); err != nil {
+    panic(err)
+}
+```
 
 [Insert-only clients](/docs/insert-only-clients) will insert jobs, but not work
 them, and don't need to be started or stopped.

example_graceful_shutdown_stop_and_cancel_test.go

@@ -0,0 +1,102 @@
+package river_test
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log/slog"
+	"os"
+	"os/signal"
+	"syscall"
+	"time"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+
+	"github.com/riverqueue/river"
+	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+	"github.com/riverqueue/river/rivershared/util/slogutil"
+)
+
+// Example_gracefulShutdownStopCancel demonstrates graceful stop with explicit
+// fallback to StopAndCancel. When a SIGINT/SIGTERM arrives, Stop initiates a
+// soft stop. If running jobs don't finish before the soft stop context expires,
+// StopAndCancel cancels their contexts (hard stop). This example is intended to
+// demonstrate advanced use of StopAndCancel. Generally, prefer the method shown
+// in Example_gracefulShutdown over the one here.
+func Example_gracefulShutdownStopAndCancel() {
+	ctx := context.Background()
+
+	jobStarted := make(chan struct{})
+
+	dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL())
+	if err != nil {
+		panic(err)
+	}
+	defer dbPool.Close()
+
+	workers := river.NewWorkers()
+	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})),
+		Queues: map[string]river.QueueConfig{
+			river.QueueDefault: {MaxWorkers: 100},
+		},
+		Workers: workers,
+	}))
+	if err != nil {
+		panic(err)
+	}
+
+	_, err = riverClient.Insert(ctx, WaitsForCancelOnlyArgs{}, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	if err := riverClient.Start(ctx); err != nil {
+		panic(err)
+	}
+
+	// Use signal.NotifyContext to detect SIGINT/SIGTERM, but don't pass the
+	// signal context to Start. Cancelling the Start context cancels running job
+	// contexts immediately, which is equivalent to StopAndCancel.
+	signalCtx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
+
+	// Make sure our job starts being worked before doing anything else.
+	<-jobStarted
+
+	// Cheat by sending ourselves a SIGTERM for the purpose of this example
+	// (normally this is sent by user or supervisory process).
+	selfProcess, _ := os.FindProcess(os.Getpid())
+	_ = selfProcess.Signal(syscall.SIGTERM)
+
+	// Wait for the first signal handler to consume the SIGTERM and cancel the
+	// start context before tearing it down.
+	<-signalCtx.Done()
+	stop()
+
+	fmt.Printf("Received SIGINT/SIGTERM; initiating soft stop (waiting for jobs to finish)\n")
+
+	softStopCtx, softStopCancel := context.WithTimeout(ctx, 100*time.Millisecond)
+	defer softStopCancel()
+
+	if err := riverClient.Stop(softStopCtx); err != nil {
+		if !errors.Is(err, context.DeadlineExceeded) {
+			panic(err)
+		}
+
+		fmt.Printf("Soft stop timeout; cancelling remaining jobs\n")
+
+		if err := riverClient.StopAndCancel(ctx); err != nil {
+			panic(err)
+		}
+	}
+
+	// Output:
+	// Working job that doesn't finish until cancelled
+	// Received SIGINT/SIGTERM; initiating soft stop (waiting for jobs to finish)
+	// Soft stop timeout; cancelling remaining jobs
+	// Job cancelled
+}