Fix resumable job progress across runtimes

.github/workflows/js.yaml

@@ -245,6 +245,14 @@ jobs:
         with:
           node-version: ${{ matrix.node-version }}
 
+      # The resumable integration scenarios read Go-generated fixtures.
+      - uses: actions/setup-go@v7
+        with:
+          go-version-file: go.work
+
+      - run: make generate/fixtures
+        working-directory: .
+
       - run: pnpm run build:all
 
       - run: node cli/dist/bin.js migrate-up --database-url "$TEST_DATABASE_URL"

CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed resumable transactional checkpoints being saved again after rollback or a failed write, and first-step cursors being lost on error, snooze, or cooperative shutdown. After a successful checkpoint write, the caller's transaction owns that progress; rollback may replay earlier unpersisted steps. The Rust, JavaScript, and Java runtimes now share the applicable behavior, including preserving progress on snooze.
+
 - Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440).
 - Fixed the notifier retrying without backoff and flooding logs when the context passed to `Client.Start` reaches its deadline. [PR #1495](https://github.com/riverqueue/river/pull/1495).
 

Makefile

@@ -219,6 +219,7 @@ test/js/conformance: generate/fixtures
 .PHONY: test/js/integration
 test/js/integration: ## Run JavaScript integration tests against Postgres
 test/js/integration: build/js
+test/js/integration: generate/fixtures
 	pnpm -C js run test:integration
 
 .PHONY: test/rust
@@ -238,7 +239,7 @@ test/rust: generate/fixtures
 .PHONY: test/rust/conformance
 test/rust/conformance: ## Run Rust tests that check Go-generated conformance fixtures
 test/rust/conformance: generate/fixtures
-	cd rust && cargo test -p riverqueue --features chrono-tz --lib --test protocol_fixtures --locked
+	cd rust && cargo test -p riverqueue --features chrono-tz,sqlite --lib --test protocol_fixtures --test resumable_conformance --locked
 
 .PHONY: test/rust/postgres
 test/rust/postgres: ## Run all Rust tests, including Postgres integration tests (requires RIVER_RUST_DATABASE_URL)

conformance/cmd/generatefixtures/main.go

@@ -49,12 +49,18 @@ func run() error {
 		return err
 	}
 
+	resumable, err := makeResumableCases()
+	if err != nil {
+		return err
+	}
+
 	for _, fixture := range []struct {
 		name  string
 		value any
 	}{
 		{"cron_schedules.json", cronSchedules},
 		{"protocol_values.json", protocolValues},
+		{"resumable.json", resumable},
 		{"snooze_counters.json", makeSnoozeCounters()},
 		{"unique_keys.json", uniqueKeys},
 	} {

conformance/cmd/generatefixtures/resumable.go

@@ -0,0 +1,112 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"maps"
+
+	"github.com/riverqueue/river"
+	"github.com/riverqueue/river/internal/jobexecutor"
+	"github.com/riverqueue/river/internal/rivercommon"
+	"github.com/riverqueue/river/internal/riverplugin"
+	"github.com/riverqueue/river/riverdriver"
+	"github.com/riverqueue/river/rivertype"
+)
+
+// fixtureCheckpointDriver only supplies the transactional update boundary. The
+// staged row is committed or discarded by each scenario, like a caller's tx.
+type fixtureCheckpointDriver struct {
+	riverdriver.Driver[riverdriver.ExecutorTx]
+}
+
+func (*fixtureCheckpointDriver) UnwrapExecutor(tx riverdriver.ExecutorTx) riverdriver.ExecutorTx {
+	return tx
+}
+
+type fixtureCheckpointExecutor struct {
+	riverdriver.ExecutorTx
+
+	fail bool
+	row  *rivertype.JobRow
+}
+
+func (e *fixtureCheckpointExecutor) JobUpdate(_ context.Context, params *riverdriver.JobUpdateParams) (*rivertype.JobRow, error) {
+	if e.fail {
+		return nil, errors.New("checkpoint write failed")
+	}
+	e.row.Metadata = params.Metadata
+	return e.row, nil
+}
+
+type fixtureResumableArgs struct{}
+
+func (fixtureResumableArgs) Kind() string { return "resumable_fixture" }
+
+type resumableCase struct {
+	ExpectedMetadata map[string]any `json:"expected_metadata"`
+	Name             string         `json:"name"`
+}
+
+type resumableCases struct {
+	Cases   []resumableCase `json:"cases"`
+	Comment string          `json:"$comment"`
+}
+
+func makeResumableCases() (resumableCases, error) {
+	fixture := resumableCases{Comment: generatedComment("River resumable steps and transactional checkpoints")}
+	for _, name := range []string{
+		"checkpoint_commit_error", "checkpoint_rollback_error", "checkpoint_then_cursor_error", "checkpoint_write_failure",
+		"first_cursor_error", "step_snooze",
+	} {
+		updates := make(map[string]any)
+		ctx := context.WithValue(context.Background(), jobexecutor.ContextKeyMetadataUpdates, updates)
+		client, err := river.NewClient(&fixtureCheckpointDriver{}, &river.Config{})
+		if err != nil {
+			return fixture, err
+		}
+		ctx = context.WithValue(ctx, rivercommon.ContextKeyClient{}, client)
+		row := &rivertype.JobRow{ID: 1, EncodedArgs: []byte(`{}`), Metadata: []byte(`{}`), State: rivertype.JobStateRunning}
+		staged := *row
+		exec := &fixtureCheckpointExecutor{fail: name == "checkpoint_write_failure", row: &staged}
+		attemptErr := (&riverplugin.ResumableMiddleware{}).Work(ctx, row, func(ctx context.Context) error {
+			if name != "first_cursor_error" {
+				river.ResumableStep(ctx, "first", nil, func(context.Context) error { return nil })
+			}
+			if name == "step_snooze" {
+				return river.JobSnooze(60_000_000_000)
+			}
+			river.ResumableStepCursor(ctx, "loop", nil, func(ctx context.Context, _ int) error {
+				if name == "first_cursor_error" {
+					if err := river.ResumableSetCursor(ctx, 7); err != nil {
+						return err
+					}
+				} else {
+					if _, err := river.ResumableSetStepCursorTx[*fixtureCheckpointDriver](ctx, riverdriver.ExecutorTx(exec), &river.Job[fixtureResumableArgs]{JobRow: row}, 1); err != nil {
+						return err
+					}
+				}
+				if name == "checkpoint_then_cursor_error" {
+					if err := river.ResumableSetCursor(ctx, 2); err != nil {
+						return err
+					}
+				}
+				return errors.New("retry")
+			})
+			return nil
+		})
+		if attemptErr == nil {
+			return fixture, fmt.Errorf("resumable scenario %s unexpectedly succeeded", name)
+		}
+		metadata := make(map[string]any)
+		if name == "checkpoint_commit_error" || name == "checkpoint_then_cursor_error" {
+			if err := json.Unmarshal(staged.Metadata, &metadata); err != nil {
+				return fixture, err
+			}
+		}
+		maps.Copy(metadata, updates)
+		fixture.Cases = append(fixture.Cases, resumableCase{ExpectedMetadata: metadata, Name: name})
+	}
+	return fixture, nil
+}

conformance/cmd/generatefixtures/resumable_test.go

@@ -0,0 +1,33 @@
+package main
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMakeResumableCases(t *testing.T) {
+	t.Parallel()
+
+	fixture, err := makeResumableCases()
+	require.NoError(t, err)
+	want := map[string]string{
+		"checkpoint_commit_error":      `{"river:resumable_step":"loop","river:resumable_cursor":{"loop":1}}`,
+		"checkpoint_then_cursor_error": `{"river:resumable_step":"loop","river:resumable_cursor":{"loop":2}}`,
+		"checkpoint_rollback_error":    `{}`,
+		"checkpoint_write_failure":     `{"river:resumable_step":"first"}`,
+		"first_cursor_error":           `{"river:resumable_cursor":{"loop":7}}`,
+		"step_snooze":                  `{"river:resumable_step":"first"}`,
+	}
+	require.Len(t, fixture.Cases, len(want))
+	for _, testCase := range fixture.Cases {
+		t.Run(testCase.Name, func(t *testing.T) {
+			t.Parallel()
+
+			metadata, err := json.Marshal(testCase.ExpectedMetadata)
+			require.NoError(t, err)
+			require.JSONEq(t, want[testCase.Name], string(metadata))
+		})
+	}
+}

conformance/go.mod

@@ -13,9 +13,12 @@ require (
 	github.com/riverqueue/river/rivershared v0.49.0
 	github.com/riverqueue/river/rivertype v0.49.0
 	github.com/robfig/cron/v3 v3.0.1
+	github.com/stretchr/testify v1.12.1
 	golang.org/x/mod v0.42.0
 )
 
+require go.yaml.in/yaml/v3 v3.0.5 // indirect
+
 require (
 	github.com/jackc/puddle/v2 v2.2.3 // indirect
 	github.com/tidwall/gjson v1.20.0 // indirect

conformance/go.sum

@@ -1,3 +1,5 @@
+github.com/jackc/pgerrcode v0.0.0-20240316143900-6e2875d9b438 h1:Dj0L5fhJ9F82ZJyVOmBx6msDp/kfd1t9GRfny/mfJA0=
+github.com/jackc/pgerrcode v0.0.0-20240316143900-6e2875d9b438/go.mod h1:a/s9Lp5W7n/DD0VrVoyJ00FbP2ytTPDVOivvn2bMlds=
 github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
 github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
 github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=

internal/riverplugin/plugin.go

@@ -44,12 +44,10 @@ func (*ResumableMiddleware) Work(ctx context.Context, job *rivertype.JobRow, doI
 		state.ResumeMatched = false
 	}
 
-	hadCursors := false
 	if cursorJSON := gjson.GetBytes(job.Metadata, rivercommon.MetadataKeyResumableCursor); cursorJSON.Exists() && cursorJSON.Type == gjson.JSON {
 		if err := json.Unmarshal([]byte(cursorJSON.Raw), &state.Cursors); err != nil {
 			return fmt.Errorf("river: unmarshal resumable cursors: %w", err)
 		}
-		hadCursors = len(state.Cursors) > 0
 	}
 
 	ctx = context.WithValue(ctx, ResumableContextKey{}, state)
@@ -64,16 +62,18 @@ func (*ResumableMiddleware) Work(ctx context.Context, job *rivertype.JobRow, doI
 		}
 	}
 
-	if err != nil && state.CompletedStep != "" {
-		if len(state.Cursors) > 0 {
+	if err != nil {
+		if state.CursorsDirty && len(state.Cursors) > 0 {
 			metadataUpdates[rivercommon.MetadataKeyResumableCursor] = state.Cursors
-		} else if hadCursors {
+		} else if state.CursorsDirty {
 			// All cursors were consumed (their steps completed). Write null
 			// to clear the stale cursor metadata, since the metadata merge
 			// is additive and wouldn't remove the old key on its own.
 			metadataUpdates[rivercommon.MetadataKeyResumableCursor] = nil
 		}
-		metadataUpdates[rivercommon.MetadataKeyResumableStep] = state.CompletedStep
+		if state.CompletedStep != "" {
+			metadataUpdates[rivercommon.MetadataKeyResumableStep] = state.CompletedStep
+		}
 	}
 
 	return err
@@ -85,6 +85,7 @@ type ResumableState struct {
 	AllStepNames  map[string]struct{}
 	CompletedStep string
 	Cursors       map[string]json.RawMessage
+	CursorsDirty  bool
 	Err           error
 	ResumeMatched bool
 	ResumeStep    string

java/river/src/main/java/com/riverqueue/WorkContext.java

@@ -149,8 +149,8 @@ synchronized Exception finish(Exception failure) {
     if (stepFailure != null) failure = stepFailure;
     if (failure == null && !resumeMatched)
       failure = new IllegalArgumentException("Saved resumable step was not found");
-    if (failure != null && completedStep != null) {
-      metadata(Protocol.METADATA_RESUMABLE_STEP, completedStep);
+    if (failure != null) {
+      if (completedStep != null) metadata(Protocol.METADATA_RESUMABLE_STEP, completedStep);
       if (!cursors.isEmpty() || job.metadata().has(Protocol.METADATA_RESUMABLE_CURSOR))
         metadata(Protocol.METADATA_RESUMABLE_CURSOR, cursors.isEmpty() ? null : cursors);
     }

java/river/src/test/java/com/riverqueue/ProtocolTest.java

@@ -300,6 +300,65 @@ public double nextDouble() {
     return tests.stream();
   }
 
+  @TestFactory
+  Stream<DynamicTest> resumableConformance() throws Exception {
+    var tests = new ArrayList<DynamicTest>();
+    for (var fixture : Conformance.fixture("resumable.json").required("cases")) {
+      String name = fixture.required("name").asString();
+      // Java intentionally has no transactional checkpoint helper.
+      if (name.startsWith("checkpoint_")) continue;
+      tests.add(
+          DynamicTest.dynamicTest(
+              name,
+              () -> {
+                var context = context(42);
+                Exception failure =
+                    assertThrows(
+                        Exception.class,
+                        () -> {
+                          if (name.equals("step_snooze")) {
+                            context.step("first", () -> {});
+                            context.snooze(java.time.Duration.ofSeconds(60));
+                          } else {
+                            context.stepWithCursor(
+                                "loop",
+                                cursor -> {
+                                  context.cursor(7);
+                                  throw new Exception("retry");
+                                });
+                          }
+                        });
+                context.finish(failure);
+                assertEquals(fixture.required("expected_metadata"), context.metadataUpdates());
+              }));
+    }
+    return tests.stream();
+  }
+
+  @Test
+  void resumableFirstStepCursor() throws Exception {
+    var context = context(42, Json.object());
+    var failure = new Exception("retry");
+    assertThrows(
+        Exception.class,
+        () ->
+            context.stepWithCursor(
+                "loop",
+                cursor -> {
+                  context.cursor(Map.of("n", 7));
+                  throw failure;
+                }));
+    assertSame(failure, context.finish(failure));
+    assertEquals(
+        Json.tree(Map.of(Protocol.METADATA_RESUMABLE_CURSOR, Map.of("loop", Map.of("n", 7)))),
+        context.metadataUpdates());
+    var retry = context(42, context.metadataUpdates());
+    retry.stepWithCursor("loop", cursor -> assertEquals(Json.tree(Map.of("n", 7)), cursor));
+    retry.finish(new Exception("later failure"));
+    assertEquals("loop", retry.metadataUpdates().path(Protocol.METADATA_RESUMABLE_STEP).asString());
+    assertTrue(retry.metadataUpdates().path(Protocol.METADATA_RESUMABLE_CURSOR).isNull());
+  }
+
   @Test
   void resumableMetadata() throws Exception {
     var keys = Conformance.fixture("protocol_values.json").required("metadata_keys");

java/river/src/test/java/com/riverqueue/WorkersTest.java

@@ -403,6 +403,44 @@ void oversizedDurationsFailDuringConfiguration(String option) {
         });
   }
 
+  @ParameterizedTest
+  @ValueSource(strings = {"error", "shutdown", "snooze"})
+  void firstStepCursorSurvivesIncompleteAttempt(String outcome) throws Exception {
+    var client = client();
+    var entered = new CountDownLatch(1);
+    var finished = new CompletableFuture<Workers.Event>();
+    try (var workers =
+            workers(client)
+                .add(
+                    TYPE,
+                    context -> {
+                      context.stepWithCursor(
+                          "loop",
+                          cursor -> {
+                            context.cursor(7);
+                            entered.countDown();
+                            if (outcome.equals("snooze")) context.snooze(Duration.ofSeconds(60));
+                            if (outcome.equals("shutdown")) {
+                              context.awaitCancellation();
+                              context.checkCancelled();
+                            }
+                            throw new Exception("retry");
+                          });
+                    })
+                .start();
+        var subscription = workers.subscribe(finished::complete)) {
+      var job = client.insert(TYPE, "cursor", InsertOptions.builder().maxAttempts(1).build()).job();
+      assertTrue(entered.await(5, TimeUnit.SECONDS));
+      if (outcome.equals("shutdown")) workers.stopAndCancel();
+      var event = finished.get(5, TimeUnit.SECONDS);
+      assertEquals(job.id(), event.job().id());
+      assertFalse(event.job().metadata().has(Protocol.METADATA_RESUMABLE_STEP));
+      assertEquals(
+          Json.tree(java.util.Map.of("loop", 7)),
+          event.job().metadata().path(Protocol.METADATA_RESUMABLE_CURSOR));
+    }
+  }
+
   @Test
   void oversizedSnoozeFailsTheAttemptInsteadOfStrandingIt() throws Exception {
     var client = client();