Add Client.JobWaitFinalized

CHANGELOG.md

@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `Config.FetchOnlyKnownKinds` to restrict job fetching to registered worker kinds, including aliases. Clients with different workers can share a queue while leaving unknown jobs available without consuming attempts. Disabled by default; leader election and stuck-job rescue behavior are unchanged. [PR #1396](https://github.com/riverqueue/river/pull/1396).
 - Added `Config.LeaderElectionDisabled` to let a client work jobs without participating in leader election or running maintenance services. Other eligible clients in the same database and schema continue handling scheduling, retries, periodic enqueueing, rescue, and cleanup. [PR #1382](https://github.com/riverqueue/river/pull/1382).
 - Added support for YugabyteDB. When `LISTEN/NOTIFY` is unavailable or disabled, clients automatically poll for running job cancellations and queue pause, resume, and metadata changes, and skip unsupported notification broadcasts. This works with the default `PollOnly: false`. Native notifications require YugabyteDB 2025.2.3 or later with `ysql_yb_enable_listen_notify=true` on both Masters and TServers. [PR #1347](https://github.com/riverqueue/river/pull/1347).
+- Added `Client.JobWaitFinalized` to wait for a job to reach a cancelled, completed, or discarded state using shared, batched polling. It works across clients and with clients that don't run workers, and supports context cancellation and a configurable poll interval through `JobWaitFinalizedOpts`. [PR #1393](https://github.com/riverqueue/river/pull/1393).
 
 ### Changed
 

client.go

@@ -749,6 +749,7 @@ type Client[TTx any] struct {
 	pluginLookupByJob     *pluginlookup.JobPluginLookup
 	pluginLookupGlobal    *pluginlookup.PluginLookup
 	insertNotifyLimiter   *notifylimiter.Limiter
+	jobWaiter             *jobWaiter
 	notifier              *notifier.Notifier // may be nil in poll-only mode
 	periodicJobs          *PeriodicJobBundle
 	pilot                 riverpilot.Pilot
@@ -881,6 +882,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 		},
 		config:               config,
 		driver:               driver,
+		jobWaiter:            newJobWaiter(driver.GetExecutor, config.Schema),
 		pluginLookupByJob:    pluginLookupByJob,
 		pluginLookupGlobal:   pluginLookupGlobal,
 		producersByQueueName: make(map[string]*producer),
@@ -1754,6 +1756,40 @@ func (c *Client[TTx]) jobUpdate(ctx context.Context, exec riverdriver.Executor,
 	})
 }
 
+// JobWaitFinalized waits until the job with the given ID is observed in a
+// finalized state (cancelled, completed, or discarded), and returns its persisted
+// row. A cancelled or discarded job is returned with a nil error; callers should
+// inspect State to determine the job's outcome. Retried and snoozed jobs continue
+// waiting until they finalize.
+//
+// The job is checked immediately, then polled approximately every 250 milliseconds
+// by default. Pass opts to customize the poll interval, or nil to use the defaults.
+// Concurrent calls on the same client share a single poll loop that batches
+// database queries for all jobs being waited on, using the shortest poll interval
+// requested by an active call. The client does not need to be started or configured
+// with workers, and the job may be executed by another client. Stop and StopAndCancel
+// do not cancel waits.
+//
+// Cancelling ctx stops only this wait, without cancelling the job. Database
+// errors are returned to the caller. ErrNotFound is returned if the job does not
+// exist or is deleted before its finalized state is observed. Commit any
+// transaction that inserts the job before calling JobWaitFinalized.
+//
+// This method observes current state, not completion history. If a finalized job
+// is manually retried before being observed, the wait continues for that execution.
+func (c *Client[TTx]) JobWaitFinalized(ctx context.Context, id int64, opts *JobWaitFinalizedOpts) (*rivertype.JobRow, error) {
+	if !c.driver.PoolIsSet() {
+		return nil, errNoDriverDBPool
+	}
+	if opts == nil {
+		opts = &JobWaitFinalizedOpts{}
+	}
+	if opts.PollInterval < 0 {
+		return nil, errors.New("PollInterval cannot be less than zero")
+	}
+	return c.jobWaiter.wait(ctx, id, cmp.Or(opts.PollInterval, jobWaitPollIntervalDefault))
+}
+
 // ID returns the unique ID of this client as set in its config or
 // auto-generated if not specified.
 func (c *Client[TTx]) ID() string {

client_test.go

@@ -14,6 +14,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/jackc/pgerrcode"
@@ -5955,6 +5956,231 @@ func Test_Client_JobUpdateTx(t *testing.T) {
 	})
 }
 
+func Test_Client_JobWaitFinalized(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+
+	type testBundle struct {
+		client *Client[pgx.Tx]
+		driver *riverpgxv5.Driver
+		exec   riverdriver.Executor
+		schema string
+	}
+
+	setup := func(t *testing.T) *testBundle {
+		t.Helper()
+
+		var (
+			dbPool = riversharedtest.DBPool(ctx, t)
+			driver = riverpgxv5.New(dbPool)
+			schema = riverdbtest.TestSchema(ctx, t, driver, nil)
+		)
+		client, err := NewClient(driver, &Config{Logger: riversharedtest.Logger(t), Schema: schema})
+		require.NoError(t, err)
+
+		return &testBundle{
+			client: client,
+			driver: driver,
+			exec:   driver.GetExecutor(),
+			schema: schema,
+		}
+	}
+
+	t.Run("AlreadyFinalized", func(t *testing.T) {
+		t.Parallel()
+
+		for _, state := range []rivertype.JobState{rivertype.JobStateCancelled, rivertype.JobStateCompleted, rivertype.JobStateDiscarded} {
+			t.Run(string(state), func(t *testing.T) {
+				t.Parallel()
+
+				bundle := setup(t)
+
+				job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{
+					Metadata: []byte(`{"output":{"answer":42}}`),
+					Schema:   bundle.schema,
+					State:    &state,
+				})
+
+				waitCtx, cancel := context.WithTimeout(ctx, riversharedtest.WaitTimeout())
+				defer cancel()
+				finalized, err := bundle.client.JobWaitFinalized(waitCtx, job.ID, nil)
+				require.NoError(t, err)
+				require.Equal(t, job.ID, finalized.ID)
+				require.Equal(t, state, finalized.State)
+				require.NotNil(t, finalized.FinalizedAt)
+				require.JSONEq(t, `{"answer":42}`, string(finalized.Output()))
+			})
+		}
+	})
+
+	t.Run("CancellationDoesNotCancelJob", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{Schema: bundle.schema})
+
+		waitCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
+		defer cancel()
+		finalized, err := bundle.client.JobWaitFinalized(waitCtx, job.ID, nil)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.Nil(t, finalized)
+
+		job, err = bundle.client.JobGet(ctx, job.ID)
+		require.NoError(t, err)
+		require.Equal(t, rivertype.JobStateAvailable, job.State)
+	})
+
+	t.Run("CompletedByAnotherClient", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		config := newTestConfig(t, bundle.schema)
+		config.PollOnly = true
+		workerClient, err := NewClient(bundle.driver, config)
+		require.NoError(t, err)
+		startClient(ctx, t, workerClient)
+
+		insertRes, err := bundle.client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+		waitCtx, cancel := context.WithTimeout(ctx, riversharedtest.WaitTimeout())
+		defer cancel()
+		finalized, err := bundle.client.JobWaitFinalized(waitCtx, insertRes.Job.ID, nil)
+		require.NoError(t, err)
+		require.Equal(t, insertRes.Job.ID, finalized.ID)
+		require.Equal(t, rivertype.JobStateCompleted, finalized.State)
+		require.Equal(t, 1, finalized.Attempt)
+	})
+
+	t.Run("FinalizationMustCommit", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{Schema: bundle.schema, State: new(rivertype.JobStateRunning)})
+
+		finalizeInTxFunc := func() riverdriver.ExecutorTx {
+			execTx, err := bundle.exec.Begin(ctx)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = execTx.Rollback(ctx) })
+			_, err = execTx.JobSetStateIfRunningMany(ctx, &riverdriver.JobSetStateIfRunningManyParams{
+				ID:              []int64{job.ID},
+				Attempt:         []*int{nil},
+				ErrData:         [][]byte{nil},
+				FinalizedAt:     []*time.Time{new(time.Now())},
+				MetadataDoMerge: []bool{false},
+				MetadataUpdates: [][]byte{nil},
+				ScheduledAt:     []*time.Time{nil},
+				Schema:          bundle.schema,
+				State:           []rivertype.JobState{rivertype.JobStateCompleted},
+			})
+			require.NoError(t, err)
+			return execTx
+		}
+
+		execTx := finalizeInTxFunc()
+		waitCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
+		defer cancel()
+		finalized, err := bundle.client.JobWaitFinalized(waitCtx, job.ID, nil)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.Nil(t, finalized)
+		require.NoError(t, execTx.Rollback(ctx))
+
+		persisted, err := bundle.client.JobGet(ctx, job.ID)
+		require.NoError(t, err)
+		require.Equal(t, rivertype.JobStateRunning, persisted.State)
+
+		execTx = finalizeInTxFunc()
+		require.NoError(t, execTx.Commit(ctx))
+		waitCtx, cancel = context.WithTimeout(ctx, riversharedtest.WaitTimeout())
+		defer cancel()
+		finalized, err = bundle.client.JobWaitFinalized(waitCtx, job.ID, nil)
+		require.NoError(t, err)
+		require.Equal(t, rivertype.JobStateCompleted, finalized.State)
+	})
+
+	t.Run("NegativePollInterval", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		job, err := bundle.client.JobWaitFinalized(ctx, 0, &JobWaitFinalizedOpts{PollInterval: -time.Second})
+		require.EqualError(t, err, "PollInterval cannot be less than zero")
+		require.Nil(t, job)
+	})
+
+	t.Run("NotFound", func(t *testing.T) {
+		t.Parallel()
+
+		bundle := setup(t)
+
+		waitCtx, cancel := context.WithTimeout(ctx, riversharedtest.WaitTimeout())
+		defer cancel()
+		job, err := bundle.client.JobWaitFinalized(waitCtx, 0, nil)
+		require.ErrorIs(t, err, rivertype.ErrNotFound)
+		require.Nil(t, job)
+	})
+
+	t.Run("PollInterval", func(t *testing.T) {
+		t.Parallel()
+
+		for _, testCase := range []struct {
+			name     string
+			interval time.Duration
+			opts     *JobWaitFinalizedOpts
+		}{
+			{name: "CustomLonger", interval: time.Second, opts: &JobWaitFinalizedOpts{PollInterval: time.Second}},
+			{name: "CustomShorter", interval: 10 * time.Millisecond, opts: &JobWaitFinalizedOpts{PollInterval: 10 * time.Millisecond}},
+			{name: "Nil", interval: 250 * time.Millisecond},
+			{name: "Zero", interval: 250 * time.Millisecond, opts: &JobWaitFinalizedOpts{}},
+		} {
+			t.Run(testCase.name, func(t *testing.T) {
+				t.Parallel()
+
+				synctest.Test(t, func(t *testing.T) {
+					// Only PoolIsSet is used; the executor stub keeps all work in
+					// the bubble so polling intervals use deterministic fake time.
+					client, err := NewClient(riverpgxv5.New(&pgxpool.Pool{}), &Config{})
+					require.NoError(t, err)
+					exec := &jobWaiterExecutorStub{}
+					exec.testSignals.Init(t)
+					client.jobWaiter.getExecutorFunc = func() riverdriver.Executor { return exec }
+
+					var numCalls atomic.Int32
+					exec.getByIDManyFunc = func(context.Context, *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) {
+						state := rivertype.JobStateAvailable
+						if numCalls.Add(1) == 2 {
+							state = rivertype.JobStateCompleted
+						}
+						return []*rivertype.JobRow{{ID: 1, State: state}}, nil
+					}
+
+					waitCtx, cancel := context.WithCancel(t.Context())
+					defer cancel()
+					go func() {
+						job, err := client.JobWaitFinalized(waitCtx, 1, testCase.opts)
+						exec.testSignals.WaitFinished.Signal(jobWaiterResult{err: err, job: job})
+					}()
+					synctest.Wait()
+					require.EqualValues(t, 1, numCalls.Load())
+
+					time.Sleep(testCase.interval - time.Nanosecond)
+					synctest.Wait()
+					require.EqualValues(t, 1, numCalls.Load())
+					time.Sleep(time.Nanosecond)
+					synctest.Wait()
+					require.EqualValues(t, 2, numCalls.Load())
+					result := exec.testSignals.WaitFinished.WaitOrTimeout()
+					require.NoError(t, result.err)
+					require.Equal(t, rivertype.JobStateCompleted, result.job.State)
+				})
+			})
+		}
+	})
+}
+
 func Test_Client_ErrorHandler(t *testing.T) {
 	t.Parallel()