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).
 - Fixed `JobSnooze` accepting negative durations despite documenting that they panic. [PR #1529](https://github.com/riverqueue/river/pull/1529).

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/DIFFERENCES.md

@@ -39,18 +39,12 @@
   reproducing Go's nonterminating date loop.
 - SQL is kept in dialect-specific resource catalogs. The Go sqlc driver seam is
   not part of the Java API.
-- Resumable cursor updates are saved even when the first step fails before any
-  step completes. Go currently drops cursor updates in that case. Updating a
-  cursor never marks a step completed; a failed plain step still runs on retry.
 - Resumable steps cannot be nested in Java. Each attempt has one active step and
   checkpoint scope, so rejecting nesting avoids ambiguous resume order and
   checkpoint ownership. Go currently permits nested steps.
 - Transactional checkpoints require the job to still be running with the same
   `attempt` and `attempted_at`. Go updates by job ID alone. Java rejects a stale
   attempt's checkpoint to prevent it from changing a replacement attempt.
-- A transactional checkpoint with no cursors stores `river:resumable_cursor` as
-  JSON null, clearing any older cursor map. Go may omit that key; readers in both
-  languages accept null as an empty cursor map.
 - SQLite follows Go's JSONB storage and millisecond timestamp representation;
   Postgres uses microsecond timestamps. These are shared storage contracts,
   not configurable Java serialization choices.

java/README.md

@@ -610,6 +610,10 @@ disabled. Let transaction failures propagate out of the step so a rollback leave
 the last committed checkpoint authoritative. Neither overload commits the caller's
 transaction.
 
+Commit the transaction before the step callback returns successfully, and propagate
+transaction and commit errors from that callback. A successful callback marks the
+step complete even if an enclosing transaction later rolls back.
+
 ### Scheduled jobs
 
 [Schedule a job](https://riverqueue.com/docs/scheduled-jobs) by supplying an

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

@@ -130,12 +130,22 @@ void checkRuntimeCancellation() {
    * Persists the current resumable step in the caller's transaction. Call inside a step, after its
    * application writes. Let transaction failures escape the step so retries use committed progress.
    * The connection must have auto-commit disabled.
+   *
+   * <p>Commit the transaction before the step callback returns successfully, and propagate
+   * transaction and commit errors from that callback. A successful callback marks the step complete
+   * even if an enclosing transaction later rolls back.
    */
   public Job<A> checkpoint(Connection transaction) {
     return persistCheckpoint(transaction, cursors.deepCopy());
   }
 
-  /** Persists the current resumable step and JSON cursor alongside application writes. */
+  /**
+   * Persists the current resumable step and JSON cursor alongside application writes.
+   *
+   * <p>Commit the transaction before the step callback returns successfully, and propagate
+   * transaction and commit errors from that callback. A successful callback marks the step complete
+   * even if an enclosing transaction later rolls back.
+   */
   public Job<A> checkpoint(Connection transaction, Object cursor) {
     var updates = cursors.deepCopy();
     if (currentStep == null)

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

@@ -120,11 +120,11 @@ public boolean clean(Client client, Instant now) {
                 .retryPolicy(
                     job -> {
                       if (!postgres) {
-                        // Retry selection runs after the read transaction closes, so another
-                        // writer can commit without invalidating rescue's update transaction.
+                        // Retry selection runs after the read transaction closes. Keep the normal
+                        // busy timeout for unrelated worker writes; a transaction held by this
+                        // callback would stay locked until the timeout expires.
                         try (var connection = database.connection();
                             var statement = connection.createStatement()) {
-                          statement.execute("PRAGMA busy_timeout=0");
                           statement.executeUpdate(
                               "UPDATE river_job SET priority = 2 WHERE id = " + job.id());
                           writeFailure.complete(null);
@@ -140,7 +140,7 @@ public boolean clean(Client client, Instant now) {
             workers.subscribe(completed::complete, Workers.EventKind.JOB_COMPLETED)) {
       if (!postgres)
         assertNull(
-            writeFailure.get(5, TimeUnit.SECONDS), "Rescue callbacks must not hold the writer");
+            writeFailure.get(10, TimeUnit.SECONDS), "Rescue callbacks must not hold the writer");
       maintained.get(5, TimeUnit.SECONDS);
       var rescued = client.get(inserted.id());
       assertEquals(Job.State.RETRYABLE, rescued.state());