Add RollbackWithoutCancel helper to have rollbacks supersede context cancellation (#1062)

Here, add a `dbutil.RollbackWithoutCancel` helper. The purpose of this
is in situations where an operation is in a transaction but ends up
cancelling due to a context timeout. From there, there tends to be a
rollback in a `defer`, but the rollback ends up not running because its
context is already cancelled.

The new helper removes a cancelled context, then adds a new 5 second
context timeout to run the rollback command. This is mainly used for
`JobSchedule` in the hopes of coming up with a fix for #1059, but we use
the new function everywhere that we previously had a manual rollback,
and it also gets add to the `dbutil.WithTx*` helpers so that it gets
used anywhere those do (and these should probably be preferred most of
the time because it makes rollbacks and commits harder to forget).

Fixes #1059.

Author: brandur

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Fix snoozed events emitted from `rivertest.Worker` when snooze duration is zero seconds. [PR #1057](https://github.com/riverqueue/river/pull/1057).
+- Rollbacks now use an uncancelled context so as to not leave transactions in an ambiguous state if a transaction in them fails due to context cancellation. [PR #1062](https://github.com/riverqueue/river/pull/1062).
 
 ## [0.26.0] - 2025-10-07
 

client.go

@@ -2403,7 +2403,7 @@ func (c *Client[TTx]) QueuePause(ctx context.Context, name string, opts *QueuePa
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback(ctx)
+	defer dbutil.RollbackWithoutCancel(ctx, tx)
 
 	if err := tx.QueuePause(ctx, &riverdriver.QueuePauseParams{
 		Name:   name,
@@ -2473,7 +2473,7 @@ func (c *Client[TTx]) QueueResume(ctx context.Context, name string, opts *QueueP
 	if err != nil {
 		return err
 	}
-	defer tx.Rollback(ctx)
+	defer dbutil.RollbackWithoutCancel(ctx, tx)
 
 	if err := tx.QueueResume(ctx, &riverdriver.QueueResumeParams{
 		Name:   name,
@@ -2541,7 +2541,7 @@ func (c *Client[TTx]) QueueUpdate(ctx context.Context, name string, params *Queu
 	if err != nil {
 		return nil, err
 	}
-	defer tx.Rollback(ctx)
+	defer dbutil.RollbackWithoutCancel(ctx, tx)
 
 	queue, controlEvent, err := c.queueUpdate(ctx, tx, name, params)
 	if err != nil {

internal/maintenance/job_scheduler.go

@@ -14,6 +14,7 @@ import (
 	"github.com/riverqueue/river/rivershared/riversharedmaintenance"
 	"github.com/riverqueue/river/rivershared/startstop"
 	"github.com/riverqueue/river/rivershared/testsignal"
+	"github.com/riverqueue/river/rivershared/util/dbutil"
 	"github.com/riverqueue/river/rivershared/util/randutil"
 	"github.com/riverqueue/river/rivershared/util/serviceutil"
 	"github.com/riverqueue/river/rivershared/util/testutil"
@@ -37,7 +38,7 @@ func (ts *JobSchedulerTestSignals) Init(tb testutil.TestingTB) {
 
 // NotifyInsert is a function to call to emit notifications for queues where
 // jobs were scheduled.
-type NotifyInsertFunc func(ctx context.Context, tx riverdriver.ExecutorTx, queues []string) error
+type NotifyInsertFunc func(ctx context.Context, execTx riverdriver.ExecutorTx, queues []string) error
 
 type JobSchedulerConfig struct {
 	riversharedmaintenance.BatchSizes
@@ -167,16 +168,16 @@ func (s *JobScheduler) runOnce(ctx context.Context) (*schedulerRunOnceResult, er
 			ctx, cancelFunc := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault)
 			defer cancelFunc()
 
-			tx, err := s.exec.Begin(ctx)
+			execTx, err := s.exec.Begin(ctx)
 			if err != nil {
 				return 0, fmt.Errorf("error starting transaction: %w", err)
 			}
-			defer tx.Rollback(ctx)
+			defer dbutil.RollbackWithoutCancel(ctx, execTx)
 
 			now := s.Time.NowUTC()
 			nowWithLookAhead := now.Add(s.config.Interval)
 
-			scheduledJobResults, err := tx.JobSchedule(ctx, &riverdriver.JobScheduleParams{
+			scheduledJobResults, err := execTx.JobSchedule(ctx, &riverdriver.JobScheduleParams{
 				Max:    s.batchSize(),
 				Now:    &nowWithLookAhead,
 				Schema: s.config.Schema,
@@ -204,13 +205,13 @@ func (s *JobScheduler) runOnce(ctx context.Context) (*schedulerRunOnceResult, er
 			}
 
 			if len(queues) > 0 {
-				if err := s.config.NotifyInsert(ctx, tx, queues); err != nil {
+				if err := s.config.NotifyInsert(ctx, execTx, queues); err != nil {
 					return 0, fmt.Errorf("error notifying insert: %w", err)
 				}
 				s.TestSignals.NotifiedQueues.Signal(queues)
 			}
 
-			return len(scheduledJobResults), tx.Commit(ctx)
+			return len(scheduledJobResults), execTx.Commit(ctx)
 		}()
 		if err != nil {
 			if errors.Is(err, context.DeadlineExceeded) {

riverdriver/riversqlite/river_sqlite_driver.go

@@ -243,7 +243,7 @@ func (e *Executor) JobCancel(ctx context.Context, params *riverdriver.JobCancelP
 	// exists and is not running, only one database operation is needed, but if
 	// the initial update comes back empty, it does one more fetch to return the
 	// most appropriate error.
-	return dbutil.WithTxV(ctx, e, func(ctx context.Context, execTx riverdriver.ExecutorTx) (*rivertype.JobRow, error) {
+	return dbutil.WithTxV(ctx, e, func(ctx context.Context, execTx riverdriver.ExecutorTx) (*rivertype.JobRow, error) { // TODO
 		dbtx := templateReplaceWrapper{dbtx: e.driver.UnwrapTx(execTx), replacer: &e.driver.replacer}
 
 		cancelledAt, err := params.CancelAttemptedAt.UTC().MarshalJSON()
@@ -320,7 +320,7 @@ func (e *Executor) JobDelete(ctx context.Context, params *riverdriver.JobDeleteP
 	// exists and is not running, only one database operation is needed, but if
 	// the initial delete comes back empty, it does one more fetch to return the
 	// most appropriate error.
-	return dbutil.WithTxV(ctx, e, func(ctx context.Context, execTx riverdriver.ExecutorTx) (*rivertype.JobRow, error) {
+	return dbutil.WithTxV(ctx, e, func(ctx context.Context, execTx riverdriver.ExecutorTx) (*rivertype.JobRow, error) { // TODO
 		dbtx := templateReplaceWrapper{dbtx: e.driver.UnwrapTx(execTx), replacer: &e.driver.replacer}
 
 		job, err := dbsqlc.New().JobDelete(schemaTemplateParam(ctx, params.Schema), dbtx, params.ID)
@@ -495,7 +495,7 @@ func (e *Executor) JobInsertFastMany(ctx context.Context, params *riverdriver.Jo
 		uniqueNonce = randutil.Hex(8)
 	)
 
-	if err := dbutil.WithTx(ctx, e, func(ctx context.Context, execTx riverdriver.ExecutorTx) error {
+	if err := dbutil.WithTx(ctx, e, func(ctx context.Context, execTx riverdriver.ExecutorTx) error { // TODO
 		ctx = schemaTemplateParam(ctx, params.Schema)
 		dbtx := templateReplaceWrapper{dbtx: e.driver.UnwrapTx(execTx), replacer: &e.driver.replacer}
 

rivermigrate/river_migrate_test.go

@@ -27,10 +27,9 @@ import (
 
 const (
 	// The name of an actual migration line embedded in our test data below.
-	migrationLineAlternate                = "alternate"
-	migrationLineAlternateMaxVersion      = 6
-	migrationLineCommitRequired           = "commit_required"
-	migrationLineCommitRequiredMaxVersion = 3
+	migrationLineAlternate           = "alternate"
+	migrationLineAlternateMaxVersion = 6
+	migrationLineCommitRequired      = "commit_required"
 )
 
 //go:embed migration/*/*.sql

rivershared/util/dbutil/db_util.go

@@ -3,12 +3,44 @@ package dbutil
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/riverqueue/river/riverdriver"
 )
 
+// RollbackWithoutCancel initiates a rollback, but one in which context is
+// overridden with context.WithoutCancel so that the rollback can proceed even
+// if a previous operation was cancelled. This decreases the chance that a
+// transaction is accidentally left open and in an ambiguous state.
+//
+// A rollback error is returned the same way most driver rollbacks return an
+// error, but given this is normally expected to be used in a defer statement,
+// it's unusual for the error to be handled.
+func RollbackWithoutCancel[TExec riverdriver.ExecutorTx](ctx context.Context, execTx TExec) error {
+	ctxWithoutCancel := context.WithoutCancel(ctx)
+
+	ctx, cancel := context.WithTimeout(ctxWithoutCancel, 5*time.Second)
+	defer cancel()
+
+	// It might not be the worst idea to log an unexpected error on rollback
+	// here instead of returning it. I had this in place initially, but there's
+	// a number of common errors that need to be ignored like "conn closed",
+	// `pgx.ErrTxClosed`, or `sql.ErrTxDone`. These all turn out to be
+	// driver-specific when this function is meant to be driver agnostic.
+	//
+	// It'd still be possible to make it happen, but we'd have to have a driver
+	// function like `ShouldIgnoreRollbackError` that'd need a lot of plumbing
+	// and it becomes questionable as to whether it's all worth it as Rollback
+	// producing a non-standard error would be quite unusual.
+	return execTx.Rollback(ctx)
+}
+
 // WithTx starts and commits a transaction on a driver executor around
 // the given function, allowing the return of a generic value.
+//
+// Rollbacks use RollbackWithoutCancel to maximize the chance of a successful
+// rollback even where an operation within the transaction timed out due to
+// context timeout.
 func WithTx[TExec riverdriver.Executor](ctx context.Context, exec TExec, innerFunc func(ctx context.Context, execTx riverdriver.ExecutorTx) error) error {
 	_, err := WithTxV(ctx, exec, func(ctx context.Context, tx riverdriver.ExecutorTx) (struct{}, error) {
 		return struct{}{}, innerFunc(ctx, tx)
@@ -18,14 +50,18 @@ func WithTx[TExec riverdriver.Executor](ctx context.Context, exec TExec, innerFu
 
 // WithTxV starts and commits a transaction on a driver executor around
 // the given function, allowing the return of a generic value.
+//
+// Rollbacks use RollbackWithoutCancel to maximize the chance of a successful
+// rollback even where an operation within the transaction timed out due to
+// context timeout.
 func WithTxV[TExec riverdriver.Executor, T any](ctx context.Context, exec TExec, innerFunc func(ctx context.Context, execTx riverdriver.ExecutorTx) (T, error)) (T, error) {
 	var defaultRes T
 
 	tx, err := exec.Begin(ctx)
 	if err != nil {
 		return defaultRes, fmt.Errorf("error beginning transaction: %w", err)
 	}
-	defer tx.Rollback(ctx)
+	defer RollbackWithoutCancel(ctx, tx) //nolint:errcheck
 
 	res, err := innerFunc(ctx, tx)
 	if err != nil {

rivershared/util/dbutil/db_util_test.go

@@ -2,8 +2,10 @@ package dbutil_test
 
 import (
 	"context"
+	"errors"
 	"testing"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/stretchr/testify/require"
 
 	"github.com/riverqueue/river/riverdbtest"
@@ -12,12 +14,74 @@ import (
 	"github.com/riverqueue/river/rivershared/util/dbutil"
 )
 
-func TestWithTx(t *testing.T) {
+func TestRollbackCancelOverride(t *testing.T) {
 	t.Parallel()
 
 	ctx := context.Background()
-	tx := riverdbtest.TestTxPgx(ctx, t)
-	driver := riverpgxv5.New(nil)
+
+	type testBundle struct {
+		driver *riverpgxv5.Driver
+		tx     pgx.Tx
+	}
+
+	setup := func(t *testing.T) *testBundle {
+		t.Helper()
+
+		return &testBundle{
+			driver: riverpgxv5.New(nil),
+			tx:     riverdbtest.TestTxPgx(ctx, t),
+		}
+	}
+
+	t.Run("Success", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		dbutil.RollbackWithoutCancel(ctx, bundle.driver.UnwrapExecutor(bundle.tx))
+	})
+
+	t.Run("WithCancelledContext", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		ctx, cancel := context.WithCancel(ctx)
+		cancel()
+
+		dbutil.RollbackWithoutCancel(ctx, bundle.driver.UnwrapExecutor(bundle.tx))
+	})
+
+	t.Run("RollbackError", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		execTx := &executorTxWithRollbackError{
+			ExecutorTx: bundle.driver.UnwrapExecutor(bundle.tx),
+		}
+
+		err := dbutil.RollbackWithoutCancel(ctx, execTx)
+		require.EqualError(t, err, "rollback error")
+	})
+}
+
+type executorTxWithRollbackError struct {
+	riverdriver.ExecutorTx
+}
+
+func (e *executorTxWithRollbackError) Rollback(ctx context.Context) error {
+	return errors.New("rollback error")
+}
+
+func TestWithTx(t *testing.T) {
+	t.Parallel()
+
+	var (
+		ctx    = context.Background()
+		tx     = riverdbtest.TestTxPgx(ctx, t)
+		driver = riverpgxv5.New(nil)
+	)
 
 	err := dbutil.WithTx(ctx, driver.UnwrapExecutor(tx), func(ctx context.Context, execTx riverdriver.ExecutorTx) error {
 		require.NoError(t, execTx.Exec(ctx, "SELECT 1"))
@@ -29,9 +93,11 @@ func TestWithTx(t *testing.T) {
 func TestWithTxV(t *testing.T) {
 	t.Parallel()
 
-	ctx := context.Background()
-	tx := riverdbtest.TestTxPgx(ctx, t)
-	driver := riverpgxv5.New(nil)
+	var (
+		ctx    = context.Background()
+		tx     = riverdbtest.TestTxPgx(ctx, t)
+		driver = riverpgxv5.New(nil)
+	)
 
 	ret, err := dbutil.WithTxV(ctx, driver.UnwrapExecutor(tx), func(ctx context.Context, execTx riverdriver.ExecutorTx) (int, error) {
 		require.NoError(t, execTx.Exec(ctx, "SELECT 1"))