Implement resumable jobs

Here, implement "resumable" jobs, which are jobs that can checkpoint
their progress so that in case they have to stop early, they're picked
up from a point that lets them skip work that's already been done. This
is especially useful for long running jobs that are at risk of being
interrupted from something like a deploy.

Here's roughly the shape of the API, with the same normal `Work`
function that all jobs implement, and with a series of `ResumableStep`
calls within, each of which take a name for the step and function
representing it:

    func (w *ResumableWorker) Work(ctx context.Context, job *river.Job[ResumableArgs]) error {
        river.ResumableStep(ctx, "step1", func(ctx context.Context) error {
            fmt.Println("Step 1")
            return nil
        })

        river.ResumableStep(ctx, "step2", func(ctx context.Context) error {
            fmt.Println("Step 2")
            return nil
        })

        river.ResumableStep(ctx, "step3", func(ctx context.Context) error {
            fmt.Println("Step 3")
            return nil
        })

        return nil
    }

We also provide a cursor API for more granularity. This lets a step set
an arbitrary cursor value periodically as it's doing something like
looping over records in a set:

    river.ResumableStepCursor(ctx, "process_ids", func(ctx context.Context, cursor ResumableCursor) error {
        for _, id := range job.Args.IDs {
            if id <= cursor.LastProcessedID {
                continue
            }

            fmt.Printf("Processed %d\n", id)

            if err := river.ResumableSetCursor(ctx, ResumableCursor{LastProcessedID: id}); err != nil {
                return err
            }
        }

        return nil
    })

The function is `ResumableStepCursor[TCursor any]` where `TCursor` can
be defined arbitrarily by the user. This could be a simple scalar value
representing an ID, or a more complex `struct` value containing multiple
IDs, enabling nested loops that set inner and outer IDs at the same time.

`ResumableStep` and `ResumableStepCursor` steps can be freely
intermingled, and multiple `ResumableStepCursor` steps with different
cursor types are supported. Cursors must be JSON marshable because
they're stored to a job's metadata.

Lastly, we provide `ResumableSetStepTx` and `ResumableSetStepCursorTx`
for cases where a transaction guarantee is necessary. Normally,
resumable step and cursor are set as a job's being completed, but
there's a chance this is never called in case of sudden failure.
`ResumableSetStepTx` (and its cursor version) is available to durably
persist a step at the cost of an extra database operation similar to how
`JobCompleteTx` does the same for job completion.

One neat aspect the implementation here is that I was able to make it
entirely middleware-only. So all the resumable job logic goes in an
internal `resumableMiddleware` that's included in all clients by
default. This is kind of nice because it keeps its code highly modular
and will hopefully act as a template for future features.

Author: brandur

client.go

@@ -780,7 +780,8 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 	// the more abstract config.Middleware for middleware are set, but not both,
 	// so in practice we never append all three of these to each other.
 	{
-		middleware := config.Middleware
+		middleware := defaultMiddleware()
+		middleware = append(middleware, config.Middleware...)
 		for _, jobInsertMiddleware := range config.JobInsertMiddleware {
 			middleware = append(middleware, jobInsertMiddleware)
 		}
@@ -1002,6 +1003,10 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 	return client, nil
 }
 
+func defaultMiddleware() []rivertype.Middleware {
+	return []rivertype.Middleware{&resumableMiddleware{}}
+}
+
 // Start starts the client's job fetching and working loops. Once this is called,
 // the client will run in a background goroutine until stopped. All jobs are
 // run with a context inheriting from the provided context, but with a timeout

client_test.go

@@ -82,6 +82,62 @@ func (w *periodicJobWorker) Work(ctx context.Context, job *Job[periodicJobArgs])
 	return nil
 }
 
+type resumableClientTestArgs struct{}
+
+func (resumableClientTestArgs) Kind() string { return "resumable_client_test" }
+
+type resumableClientTestWorker struct {
+	WorkerDefaults[resumableClientTestArgs]
+
+	calls      []string
+	callsMu    sync.Mutex
+	failedOnce atomic.Bool
+}
+
+func (w *resumableClientTestWorker) Calls() []string {
+	w.callsMu.Lock()
+	defer w.callsMu.Unlock()
+
+	return append([]string(nil), w.calls...)
+}
+
+func (w *resumableClientTestWorker) Work(ctx context.Context, job *Job[resumableClientTestArgs]) error {
+	appendCall := func(call string) {
+		w.callsMu.Lock()
+		defer w.callsMu.Unlock()
+
+		w.calls = append(w.calls, call)
+	}
+
+	ResumableStep(ctx, "step1", func(ctx context.Context) error {
+		appendCall("step1")
+		return nil
+	})
+
+	ResumableStepCursor(ctx, "step2", func(ctx context.Context, cursor int) error {
+		appendCall("step2:" + strconv.Itoa(cursor))
+
+		for itemID := cursor + 1; itemID <= 2; itemID++ {
+			appendCall("item:" + strconv.Itoa(itemID))
+			if err := ResumableSetCursor(ctx, itemID); err != nil {
+				return err
+			}
+			if !w.failedOnce.Swap(true) {
+				return errors.New("retry me")
+			}
+		}
+
+		return nil
+	})
+
+	ResumableStep(ctx, "step3", func(ctx context.Context) error {
+		appendCall("step3")
+		return nil
+	})
+
+	return nil
+}
+
 func makeAwaitWorker[T JobArgs](startedCh chan<- int64, doneCh chan struct{}) Worker[T] {
 	return WorkFunc(func(ctx context.Context, job *Job[T]) error {
 		client := ClientFromContext[pgx.Tx](ctx)
@@ -6936,6 +6992,60 @@ func Test_Client_JobCompletion(t *testing.T) {
 		require.Nil(t, reloadedJob.FinalizedAt)
 	})
 
+	t.Run("ResumableJobRetriesAndResumes", func(t *testing.T) {
+		t.Parallel()
+
+		config := newTestConfig(t, "")
+		config.RetryPolicy = &retrypolicytest.RetryPolicyNoJitter{}
+
+		worker := &resumableClientTestWorker{}
+		AddWorker(config.Workers, worker)
+
+		client, bundle := setup(t, config)
+
+		insertRes, err := client.Insert(ctx, resumableClientTestArgs{}, nil)
+		require.NoError(t, err)
+
+		// Wait for the first attempt to fail after step2 checkpoints cursor
+		// progress and intentionally returns "retry me", leaving the job retryable.
+		eventFailed := riversharedtest.WaitOrTimeout(t, bundle.subscribeChan)
+		require.Equal(t, EventKindJobFailed, eventFailed.Kind)
+		require.Equal(t, insertRes.Job.ID, eventFailed.Job.ID)
+
+		reloadedRetryableJob, err := client.JobGet(ctx, insertRes.Job.ID)
+		require.NoError(t, err)
+		require.Equal(t, rivertype.JobStateRetryable, reloadedRetryableJob.State)
+
+		var retryableMetadata map[string]any
+		require.NoError(t, json.Unmarshal(reloadedRetryableJob.Metadata, &retryableMetadata))
+		require.Equal(t, "step1", retryableMetadata["river:resumable_step"])
+		require.Equal(t, map[string]any{"step2": float64(1)}, retryableMetadata["river:resumable_cursor"])
+
+		// Wait for the retried attempt to resume and then complete successfully.
+		eventCompleted := riversharedtest.WaitOrTimeout(t, bundle.subscribeChan)
+		require.Equal(t, EventKindJobCompleted, eventCompleted.Kind)
+		require.Equal(t, insertRes.Job.ID, eventCompleted.Job.ID)
+
+		reloadedJob, err := client.JobGet(ctx, insertRes.Job.ID)
+		require.NoError(t, err)
+		require.Equal(t, rivertype.JobStateCompleted, reloadedJob.State)
+		require.Len(t, reloadedJob.Errors, 1)
+
+		var metadata map[string]any
+		require.NoError(t, json.Unmarshal(reloadedJob.Metadata, &metadata))
+		require.Equal(t, "step1", metadata["river:resumable_step"])
+		require.Equal(t, map[string]any{"step2": float64(1)}, metadata["river:resumable_cursor"])
+
+		require.Equal(t, []string{
+			"step1",
+			"step2:0",
+			"item:1",
+			"step2:1",
+			"item:2",
+			"step3",
+		}, worker.Calls())
+	})
+
 	t.Run("JobThatReturnsJobCancelErrorIsImmediatelyCancelled", func(t *testing.T) {
 		t.Parallel()
 
@@ -7602,7 +7712,7 @@ func Test_NewClient_Validations(t *testing.T) {
 			},
 			validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
 				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 1)
-				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 1)
+				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2)
 			},
 		},
 		{
@@ -7613,7 +7723,7 @@ func Test_NewClient_Validations(t *testing.T) {
 			},
 			validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
 				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 2)
-				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2)
+				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 3)
 			},
 		},
 		{
@@ -7625,7 +7735,7 @@ func Test_NewClient_Validations(t *testing.T) {
 			},
 			validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper
 				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindJobInsert), 1)
-				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 1)
+				require.Len(t, client.middlewareLookupGlobal.ByMiddlewareKind(middlewarelookup.MiddlewareKindWorker), 2)
 			},
 		},
 		{

example_resumable_cursor_job_test.go

@@ -0,0 +1,105 @@
+package river_test
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+
+	"github.com/riverqueue/river"
+	"github.com/riverqueue/river/riverdbtest"
+	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+	"github.com/riverqueue/river/rivershared/util/slogutil"
+	"github.com/riverqueue/river/rivershared/util/testutil"
+)
+
+type ResumableCursorArgs struct {
+	IDs []int `json:"ids"`
+}
+
+func (ResumableCursorArgs) Kind() string { return "resumable_cursor" }
+
+type ResumableCursor struct {
+	LastProcessedID int `json:"last_processed_id"`
+}
+
+type ResumableCursorWorker struct {
+	river.WorkerDefaults[ResumableCursorArgs]
+}
+
+func (w *ResumableCursorWorker) Work(ctx context.Context, job *river.Job[ResumableCursorArgs]) error {
+	river.ResumableStepCursor(ctx, "process_ids", func(ctx context.Context, cursor ResumableCursor) error {
+		for _, id := range job.Args.IDs {
+			if id <= cursor.LastProcessedID {
+				continue
+			}
+
+			fmt.Printf("Processed %d\n", id)
+
+			if err := river.ResumableSetCursor(ctx, ResumableCursor{LastProcessedID: id}); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	})
+
+	return nil
+}
+
+// Example_resumableCursor demonstrates the use of a resumable cursor step, a
+// step that can store arbitrary JSON state to resume a partially completed loop.
+func Example_resumableCursor() { //nolint:dupl
+	ctx := context.Background()
+
+	dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL())
+	if err != nil {
+		panic(err)
+	}
+	defer dbPool.Close()
+
+	workers := river.NewWorkers()
+	river.AddWorker(workers, &ResumableCursorWorker{})
+
+	riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{
+		Logger: slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelWarn, ReplaceAttr: slogutil.NoLevelTime})),
+		Queues: map[string]river.QueueConfig{
+			river.QueueDefault: {MaxWorkers: 100},
+		},
+		Schema:   riverdbtest.TestSchema(ctx, testutil.PanicTB(), riverpgxv5.New(dbPool), nil), // only necessary for the example test
+		TestOnly: true,                                                                         // suitable only for use in tests; remove for live environments
+		Workers:  workers,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	// Out of example scope, but used to wait until a job is worked.
+	subscribeChan, subscribeCancel := riverClient.Subscribe(river.EventKindJobCompleted)
+	defer subscribeCancel()
+
+	if err := riverClient.Start(ctx); err != nil {
+		panic(err)
+	}
+
+	if _, err = riverClient.Insert(ctx, ResumableCursorArgs{
+		IDs: []int{1, 2, 3},
+	}, nil); err != nil {
+		panic(err)
+	}
+
+	// Wait for jobs to complete. Only needed for purposes of the example test.
+	riversharedtest.WaitOrTimeoutN(testutil.PanicTB(), subscribeChan, 1)
+
+	if err := riverClient.Stop(ctx); err != nil {
+		panic(err)
+	}
+
+	// Output:
+	// Processed 1
+	// Processed 2
+	// Processed 3
+}

example_resumable_job_test.go

@@ -0,0 +1,96 @@
+package river_test
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+
+	"github.com/riverqueue/river"
+	"github.com/riverqueue/river/riverdbtest"
+	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivershared/riversharedtest"
+	"github.com/riverqueue/river/rivershared/util/slogutil"
+	"github.com/riverqueue/river/rivershared/util/testutil"
+)
+
+type ResumableArgs struct{}
+
+func (ResumableArgs) Kind() string { return "resumable" }
+
+type ResumableWorker struct {
+	river.WorkerDefaults[ResumableArgs]
+}
+
+func (w *ResumableWorker) Work(ctx context.Context, job *river.Job[ResumableArgs]) error {
+	river.ResumableStep(ctx, "step1", func(ctx context.Context) error {
+		fmt.Println("Step 1")
+		return nil
+	})
+
+	river.ResumableStep(ctx, "step2", func(ctx context.Context) error {
+		fmt.Println("Step 2")
+		return nil
+	})
+
+	river.ResumableStep(ctx, "step3", func(ctx context.Context) error {
+		fmt.Println("Step 3")
+		return nil
+	})
+
+	return nil
+}
+
+// Example_resumable demonstrates the use of a "resumable job", a job that has
+// multiple steps, and which can be resumed after each one.
+func Example_resumable() { //nolint:dupl
+	ctx := context.Background()
+
+	dbPool, err := pgxpool.New(ctx, riversharedtest.TestDatabaseURL())
+	if err != nil {
+		panic(err)
+	}
+	defer dbPool.Close()
+
+	workers := river.NewWorkers()
+	river.AddWorker(workers, &ResumableWorker{})
+
+	riverClient, err := river.NewClient(riverpgxv5.New(dbPool), &river.Config{
+		Logger: slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelWarn, ReplaceAttr: slogutil.NoLevelTime})),
+		Queues: map[string]river.QueueConfig{
+			river.QueueDefault: {MaxWorkers: 100},
+		},
+		Schema:   riverdbtest.TestSchema(ctx, testutil.PanicTB(), riverpgxv5.New(dbPool), nil), // only necessary for the example test
+		TestOnly: true,                                                                         // suitable only for use in tests; remove for live environments
+		Workers:  workers,
+	})
+	if err != nil {
+		panic(err)
+	}
+
+	// Out of example scope, but used to wait until a job is worked.
+	subscribeChan, subscribeCancel := riverClient.Subscribe(river.EventKindJobCompleted)
+	defer subscribeCancel()
+
+	if err := riverClient.Start(ctx); err != nil {
+		panic(err)
+	}
+
+	if _, err = riverClient.Insert(ctx, ResumableArgs{}, nil); err != nil {
+		panic(err)
+	}
+
+	// Wait for jobs to complete. Only needed for purposes of the example test.
+	riversharedtest.WaitOrTimeoutN(testutil.PanicTB(), subscribeChan, 1)
+
+	if err := riverClient.Stop(ctx); err != nil {
+		panic(err)
+	}
+
+	// Output:
+	// Step 1
+	// Step 2
+	// Step 3
+}