diff --git a/CHANGELOG.md b/CHANGELOG.md index 73864e81b..7097814fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- A `JobCancel()` sent while a client's notifier is disconnected/reconnecting is no longer lost. Previously, the cancellation's `NOTIFY` could commit while nothing was `LISTEN`ing (Postgres `NOTIFY` is fire-and-forget), leaving the running job unaware of the cancellation until `JobRescuer`'s much longer stuck-job sweep. Each client now reconciles its currently running jobs against `cancel_attempted_at` by primary key every time its notifier (re)establishes healthy listening, closing the gap with no added steady-state query load. [PR #1361](https://github.com/riverqueue/river/pull/1361), as discussed in [#1358](https://github.com/riverqueue/river/issues/1358). + ## [0.50.0] - 2026-10-26 ### Changed diff --git a/client.go b/client.go index e06af4b0c..09d41d182 100644 --- a/client.go +++ b/client.go @@ -1231,6 +1231,16 @@ func (c *Client[TTx]) Start(ctx context.Context) error { fetchCtx := withClient(fetchCtx, c) workCtx = withClient(workCtx, c) + // Keep cancellation notifications and reconnects live until producers + // finish draining. The notifier is stopped with the other services below. + if c.notifier != nil { + if err := c.notifier.StartWithConnectContext(context.WithoutCancel(fetchCtx), fetchCtx); err != nil { + workCancel(err) + stopServicesOnError() + return err + } + } + if err := startstop.StartAll(fetchCtx, c.services...); err != nil { workCancel(err) stopServicesOnError() @@ -1503,8 +1513,8 @@ func (c *Client[TTx]) Driver() riverdriver.Driver[TTx] { // // If the job is currently running, it is not immediately cancelled, but is // instead marked for cancellation. The client running the job will also be -// notified to cancel the running job's context. When running without a notifier, -// clients poll for cancellation requests every two seconds. Although the job's +// notified to cancel the running job's context. Clients reconcile missed requests +// after reconnecting, or poll every two seconds without a notifier. Although the job's // context will be cancelled, since Go does not provide a mechanism to // interrupt a running goroutine the job will continue running until it returns. // As always, it is important for workers to respect context cancellation and @@ -1533,11 +1543,9 @@ func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.Jo return nil, err } - c.notifyProducerWithoutListenerQueueControlEvent(job.Queue, &controlEventPayload{ - Action: controlActionCancel, - JobID: job.ID, - Queue: job.Queue, - }) + if job.State == rivertype.JobStateRunning { + c.notifyProducerOfCancelWithoutNotifier(ctx, job) + } return job, nil } @@ -1560,8 +1568,8 @@ func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.Jo // // If the job is currently running, it is not immediately cancelled, but is // instead marked for cancellation. The client running the job will also be -// notified to cancel the running job's context. When running without a notifier, -// clients poll for cancellation requests every two seconds. Although the job's +// notified to cancel the running job's context. Clients reconcile missed requests +// after reconnecting, or poll every two seconds without a notifier. Although the job's // context will be cancelled, since Go does not provide a mechanism to // interrupt a running goroutine the job will continue running until it returns. // As always, it is important for workers to respect context cancellation and @@ -2406,14 +2414,15 @@ func (c *Client[TTx]) producerAdd(queueName string, queueConfig QueueConfig) (*p } func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) error { - c.producersMu.Lock() - defer c.producersMu.Unlock() - + c.producersMu.RLock() producer, ok := c.producersByQueueName[queueName] + c.producersMu.RUnlock() if !ok { return &QueueNotFoundError{Name: queueName} } + // QueueBundle serializes additions and removals. Keep the registry readable + // while draining so workers can still receive local cancellation requests. shouldStop, stopped, finalizeStop := producer.StopInit() if shouldStop { select { @@ -2425,7 +2434,9 @@ func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) erro } } + c.producersMu.Lock() delete(c.producersByQueueName, queueName) + c.producersMu.Unlock() return nil } @@ -2813,7 +2824,7 @@ func (c *Client[TTx]) QueuePause(ctx context.Context, name string, opts *QueuePa return err } - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) return nil } @@ -2883,7 +2894,7 @@ func (c *Client[TTx]) QueueResume(ctx context.Context, name string, opts *QueueP return err } - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) return nil } @@ -2943,7 +2954,7 @@ func (c *Client[TTx]) QueueUpdate(ctx context.Context, name string, params *Queu return nil, err } - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) return queue, nil } @@ -2965,21 +2976,65 @@ func (c *Client[TTx]) QueueUpdateTx(ctx context.Context, tx TTx, name string, pa // Should only ever be invoked *outside* a transaction. If invoked within a // transaction, the producer wouldn't yet be able to access the state that // triggered the notification because it's not committed yet. -func (c *Client[TTx]) notifyProducerWithoutListenerQueueControlEvent(queue string, controlEvent *controlEventPayload) { +func (c *Client[TTx]) notifyProducerWithoutListenerQueueControlEvent(ctx context.Context, queue string, controlEvent *controlEventPayload) { if c.driver.SupportsListener() { return } c.producersMu.RLock() - defer c.producersMu.RUnlock() + producer := c.producersByQueueName[queue] + c.producersMu.RUnlock() - if len(c.producersByQueueName) < 1 { + if producer != nil { + producer.TriggerQueueControlEvent(ctx, controlEvent) + } +} + +// Deliver cancellations locally when no notifier is configured, including PollOnly. +// Only call after commit so the producer can observe the cancellation metadata. +func (c *Client[TTx]) notifyProducerOfCancelWithoutNotifier(ctx context.Context, job *rivertype.JobRow) { + if c.notifier != nil { + return + } + + c.producersMu.RLock() + producer := c.producersByQueueName[job.Queue] + c.producersMu.RUnlock() + if producer == nil { + return + } + + state := producer.runState.Load() + if state == nil { + return + } + select { + case <-state.done: return + default: + } + + producer.activeJobsMu.Lock() + executor := producer.activeJobs[job.ID] + if executor != nil { + sameAttemptedAt := executor.JobRow.AttemptedAt == job.AttemptedAt || + (executor.JobRow.AttemptedAt != nil && job.AttemptedAt != nil && executor.JobRow.AttemptedAt.Equal(*job.AttemptedAt)) + if executor.JobRow.Attempt != job.Attempt || !sameAttemptedAt { + executor = nil + } } + producer.activeJobsMu.Unlock() - if producer, ok := c.producersByQueueName[queue]; ok { - producer.TriggerQueueControlEvent(controlEvent) + if executor != nil { + executor.Cancel(ctx) + } else { + producer.TriggerCancelCheck() } + producer.testSignals.QueueControlEventTriggered.Signal(&controlEventPayload{ + Action: controlActionCancel, + JobID: job.ID, + Queue: job.Queue, + }) } func (c *Client[TTx]) queueUpdate(ctx context.Context, executorTx riverdriver.ExecutorTx, name string, params *QueueUpdateParams) (*rivertype.Queue, *controlEventPayload, error) { diff --git a/client_test.go b/client_test.go index 472bac298..d8c7fb48a 100644 --- a/client_test.go +++ b/client_test.go @@ -1465,6 +1465,270 @@ func Test_Client_Common(t *testing.T) { require.WithinDuration(t, time.Now(), *event.Job.FinalizedAt, 2*time.Second) }) + // Explicit PollOnly also bypasses notifications on listener-capable drivers. + t.Run("CancelRunningJobExplicitPollOnly", func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + config.PollOnly = true + config.queuePollInterval = time.Hour + + client, err := NewClient(bundle.driver, config) + require.NoError(t, err) + client.producersByQueueName[QueueDefault].testSignals.QueueControlEventTriggered.Init(t) + + jobStartedChan := make(chan int64) + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + AddWorker(client.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + jobStartedChan <- job.ID + <-ctx.Done() + return ctx.Err() + })) + + subscribeChan := subscribe(t, client) + startClient(ctx, t, client) + riversharedtest.WaitOrTimeout(t, client.baseStartStop.Started()) + + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + + startedJobID := riversharedtest.WaitOrTimeout(t, jobStartedChan) + require.Equal(t, insertRes.Job.ID, startedJobID) + + updatedJob, err := client.JobCancel(ctx, insertRes.Job.ID) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, updatedJob.State) + controlEvent := client.producersByQueueName[QueueDefault].testSignals.QueueControlEventTriggered.WaitOrTimeout() + require.NotNil(t, controlEvent) + require.Equal(t, insertRes.Job.ID, controlEvent.JobID) + + event := riversharedtest.WaitOrTimeout(t, subscribeChan) + require.Equal(t, EventKindJobCancelled, event.Kind) + require.Equal(t, rivertype.JobStateCancelled, event.Job.State) + require.WithinDuration(t, time.Now(), *event.Job.FinalizedAt, 2*time.Second) + }) + + t.Run("CancelRunningJobExplicitPollOnlyDuringGracefulStop", func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + config.PollOnly = true + config.queuePollInterval = time.Hour + + client, err := NewClient(bundle.driver, config) + require.NoError(t, err) + client.producersByQueueName[QueueDefault].testSignals.QueueControlEventTriggered.Init(t) + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + var signals clientJobCancelTestSignals + signals.Init(t) + AddWorker(client.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + signals.JobStarted.Signal(job.ID) + <-ctx.Done() + return ctx.Err() + })) + + subscribeChan := subscribe(t, client) + require.NoError(t, client.Start(ctx)) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + require.NoError(t, client.StopAndCancel(stopCtx)) + }) + riversharedtest.WaitOrTimeout(t, client.baseStartStop.Started()) + + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + require.Equal(t, insertRes.Job.ID, signals.JobStarted.WaitOrTimeout()) + + fetchDone := client.queues.fetchCtx.Done() + var stopped testsignal.TestSignal[error] + stopped.Init(t) + go func() { + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + stopped.Signal(client.Stop(stopCtx)) + }() + riversharedtest.WaitOrTimeout(t, fetchDone) + + cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + _, err = client.JobCancel(cancelCtx, insertRes.Job.ID) + require.NoError(t, err) + controlEvent := client.producersByQueueName[QueueDefault].testSignals.QueueControlEventTriggered.WaitOrTimeout() + require.NotNil(t, controlEvent) + require.Equal(t, insertRes.Job.ID, controlEvent.JobID) + + event := riversharedtest.WaitOrTimeout(t, subscribeChan) + require.Equal(t, EventKindJobCancelled, event.Kind) + require.Equal(t, insertRes.Job.ID, event.Job.ID) + require.NoError(t, stopped.WaitOrTimeout()) + }) + + t.Run("CancelRunningJobExplicitPollOnlyDuringQueueRemoval", func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + config.PollOnly = true + config.queuePollInterval = time.Hour + + client, err := NewClient(bundle.driver, config) + require.NoError(t, err) + producer := client.producersByQueueName[QueueDefault] + producer.testSignals.ExecutorShutdownStarted.Init(t) + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + var signals clientJobCancelTestSignals + signals.Init(t) + var cancelReturned, queueRemoved testsignal.TestSignal[error] + cancelReturned.Init(t) + queueRemoved.Init(t) + AddWorker(client.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + signals.JobStarted.Signal(job.ID) + select { + case <-ctx.Done(): + return ctx.Err() + case <-signals.ContinueWork.WaitC(): + } + _, err := client.JobCancel(ctx, job.ID) + cancelReturned.Signal(err) + return JobCancel(err) + })) + + subscribeChan := subscribe(t, client) + require.NoError(t, client.Start(ctx)) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + require.NoError(t, client.StopAndCancel(stopCtx)) + }) + riversharedtest.WaitOrTimeout(t, client.baseStartStop.Started()) + + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + require.Equal(t, insertRes.Job.ID, signals.JobStarted.WaitOrTimeout()) + + removeCtx, cancelRemove := context.WithCancel(ctx) + t.Cleanup(cancelRemove) + go func() { + queueRemoved.Signal(client.Queues().Remove(removeCtx, QueueDefault)) + }() + producer.testSignals.ExecutorShutdownStarted.WaitOrTimeout() + signals.ContinueWork.Signal(struct{}{}) + + require.NoError(t, cancelReturned.WaitOrTimeout()) + event := riversharedtest.WaitOrTimeout(t, subscribeChan) + require.Equal(t, EventKindJobCancelled, event.Kind) + require.Equal(t, insertRes.Job.ID, event.Job.ID) + require.NoError(t, queueRemoved.WaitOrTimeout()) + }) + + t.Run("CancelManyNonRunningJobsExplicitPollOnlyDoesNotBlock", func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + config.PollOnly = true + + client, err := NewClient(bundle.driver, config) + require.NoError(t, err) + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + AddWorker(client.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + return nil + })) + + const numJobs = 101 // more than the producer's 100-slot control-event buffer + + jobIDs := make([]int64, numJobs) + for i := range jobIDs { + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + jobIDs[i] = insertRes.Job.ID + } + + // Producers exist before Start, but have no consumer yet. + cancelDone := make(chan error, 1) + go func() { + for _, jobID := range jobIDs { + if _, err := client.JobCancel(ctx, jobID); err != nil { + cancelDone <- err + return + } + } + cancelDone <- nil + }() + + select { + case err := <-cancelDone: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("cancelling non-running jobs blocked, presumably on the producer's control-event channel") + } + }) + + for _, stopped := range []bool{false, true} { + name := "CancelManyRunningJobsExplicitPollOnlyBeforeStart" + if stopped { + name = "CancelManyRunningJobsExplicitPollOnlyAfterStop" + } + t.Run(name, func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + config.PollOnly = true + + client, err := NewClient(bundle.driver, config) + require.NoError(t, err) + + if stopped { + require.NoError(t, client.Start(ctx)) + riversharedtest.WaitOrTimeout(t, client.baseStartStop.Started()) + require.NoError(t, client.Stop(ctx)) + } + + producer := client.producersByQueueName[QueueDefault] + jobIDs := make([]int64, cap(producer.queueControlCh)+1) + for i := range jobIDs { + jobIDs[i] = testfactory.Job(ctx, t, bundle.driver.GetExecutor(), &testfactory.JobOpts{ + Schema: bundle.schema, + State: new(rivertype.JobStateRunning), + }).ID + } + + cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + var cancelled testsignal.TestSignal[error] + cancelled.Init(t) + go func() { + for _, jobID := range jobIDs { + if _, err := client.JobCancel(cancelCtx, jobID); err != nil { + cancelled.Signal(err) + return + } + } + cancelled.Signal(nil) + }() + + require.NoError(t, cancelled.WaitOrTimeout()) + require.Empty(t, producer.cancelPollCh) + require.Empty(t, producer.queueControlCh) + }) + } + t.Run("CancelScheduledJob", func(t *testing.T) { t.Parallel() @@ -1530,15 +1794,25 @@ func Test_Client_Common(t *testing.T) { testutil.JobArgsReflectKind[JobArgs] } + jobStartedChan := make(chan int64) + continueChan := make(chan struct{}) + AddWorker(client.config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + jobStartedChan <- job.ID + <-continueChan return nil })) startClient(ctx, t, client) + t.Cleanup(func() { close(continueChan) }) insertRes, err := client.Insert(ctx, &JobArgs{}, nil) require.NoError(t, err) + // Wait for a running attempt before testing local cancellation delivery. + startedJobID := riversharedtest.WaitOrTimeout(t, jobStartedChan) + require.Equal(t, insertRes.Job.ID, startedJobID) + _, err = client.JobCancel(ctx, insertRes.Job.ID) require.NoError(t, err) @@ -2552,6 +2826,95 @@ func Test_Client_FetchOnlyKnownKinds(t *testing.T) { require.Equal(t, unknown, unknownAfter) } +func Test_Client_JobCancelSurvivesNotifierReconnectDuringGracefulStop(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dbPool := riversharedtest.DBPool(ctx, t) + driver := riverpgxv5.New(dbPool) + schema := riverdbtest.TestSchema(ctx, t, driver, nil) + listener := driver.GetListener(&riverdriver.GetListenenerParams{Schema: schema}) + listenerMock := newProducerTestListenerMock(listener) + + var signals clientJobCancelTestSignals + signals.Init(t) + var reconnectStarted, releaseReconnect testsignal.TestSignal[struct{}] + reconnectStarted.Init(t) + releaseReconnect.Init(t) + releaseReconnectFunc := sync.OnceFunc(func() { releaseReconnect.Signal(struct{}{}) }) + var disconnect, reconnect atomic.Bool + listenerMock.waitForNotificationFunc = func(ctx context.Context) (*riverdriver.Notification, error) { + if disconnect.CompareAndSwap(true, false) { + reconnect.Store(true) + return nil, errors.New("simulated connection loss") + } + return listener.WaitForNotification(ctx) + } + listenerMock.connectFunc = func(ctx context.Context) error { + if reconnect.Load() { + reconnectStarted.Signal(struct{}{}) + select { + case <-ctx.Done(): + return ctx.Err() + case <-releaseReconnect.WaitC(): + } + reconnect.Store(false) + } + return listener.Connect(ctx) + } + + config := newTestConfig(t, schema) + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + AddWorker(config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + signals.JobStarted.Signal(job.ID) + <-ctx.Done() + return ctx.Err() + })) + client, err := NewClient(&DriverWithListener{Driver: driver, listener: listenerMock}, config) + require.NoError(t, err) + producer := client.producersByQueueName[QueueDefault] + producer.testSignals.ExecutorShutdownStarted.Init(t) + subscribeChan := subscribe(t, client) + require.NoError(t, client.Start(ctx)) + t.Cleanup(func() { + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + require.NoError(t, client.StopAndCancel(stopCtx)) + }) + t.Cleanup(releaseReconnectFunc) + + insertRes, err := client.Insert(ctx, &JobArgs{}, nil) + require.NoError(t, err) + require.Equal(t, insertRes.Job.ID, signals.JobStarted.WaitOrTimeout()) + + disconnect.Store(true) + require.NoError(t, driver.GetExecutor().NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Payload: []string{`{"queue":"unrelated"}`}, + Schema: schema, + Topic: string(notifier.NotificationTopicInsert), + })) + reconnectStarted.WaitOrTimeout() + _, err = client.JobCancel(ctx, insertRes.Job.ID) + require.NoError(t, err) + + var stopped testsignal.TestSignal[error] + stopped.Init(t) + go func() { + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + stopped.Signal(client.Stop(stopCtx)) + }() + producer.testSignals.ExecutorShutdownStarted.WaitOrTimeout() + releaseReconnectFunc() + + event := riversharedtest.WaitOrTimeout(t, subscribeChan) + require.Equal(t, EventKindJobCancelled, event.Kind) + require.Equal(t, insertRes.Job.ID, event.Job.ID) + require.NoError(t, stopped.WaitOrTimeout()) +} + func Test_Client_LeaderElectionDisabled(t *testing.T) { t.Parallel() @@ -2646,6 +3009,225 @@ func Test_Client_LeaderElectionDisabled(t *testing.T) { }) } +func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { + t.Parallel() + + type testBundle struct { + client *Client[pgx.Tx] + producer *producer + stopFunc context.CancelFunc + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + prod := &producer{cancelPollCh: make(chan struct{}, 1)} + prod.runState.Store(&producerRunState{done: ctx.Done()}) + return &testBundle{ + client: &Client[pgx.Tx]{producersByQueueName: map[string]*producer{QueueDefault: prod}}, + producer: prod, + stopFunc: cancel, + } + } + + t.Run("FullWakeBuffer", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.producer.TriggerCancelCheck() + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{Queue: QueueDefault}) + returned.Signal(struct{}{}) + }() + returned.WaitOrTimeout() + require.Len(t, bundle.producer.cancelPollCh, 1) + }) + + for _, stopped := range []bool{false, true} { + t.Run(fmt.Sprintf("InactiveProducer_Stopped_%t", stopped), func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + if stopped { + bundle.stopFunc() + } else { + bundle.producer.runState.Store(nil) + } + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{Queue: QueueDefault}) + require.Empty(t, bundle.producer.cancelPollCh) + }) + } + + for _, testCase := range []struct { + name string + replaced bool + reuseAttempt bool + }{ + {name: "MatchingAttempt"}, + {name: "NewAttempt", replaced: true}, + {name: "NewAttemptedAt", replaced: true, reuseAttempt: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + archetype := riversharedtest.BaseServiceArchetype(t) + originalCtx, cancelOriginal := context.WithCancelCause(context.Background()) + defer cancelOriginal(nil) + replacementCtx, cancelReplacement := context.WithCancelCause(context.Background()) + defer cancelReplacement(nil) + attemptedAt := time.Now().UTC() + job := &rivertype.JobRow{ID: 456, Attempt: 1, AttemptedAt: &attemptedAt, Queue: QueueDefault} + original := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelOriginal, JobRow: job}) + replacementJob := *job + replacementJob.AttemptedAt = new(attemptedAt.Add(time.Second)) + if !testCase.reuseAttempt { + replacementJob.Attempt++ + } + replacement := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelReplacement, JobRow: &replacementJob}) + bundle.producer.activeJobs = map[int64]*jobexecutor.JobExecutor{job.ID: original} + if testCase.replaced { + bundle.producer.activeJobs[job.ID] = replacement + } + + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), job) + + require.NoError(t, replacementCtx.Err()) + if testCase.replaced { + require.NoError(t, originalCtx.Err()) + require.Len(t, bundle.producer.cancelPollCh, 1) + } else { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + require.Empty(t, bundle.producer.cancelPollCh) + } + }) + } +} + +func Test_Client_NotifyProducerWithoutListenerQueueControlEvent(t *testing.T) { + t.Parallel() + + type testBundle struct { + client *Client[pgx.Tx] + fetchStopped <-chan struct{} + producer *producer + stopFunc func() + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + runCtx, cancelRun := context.WithCancel(context.Background()) + targetProducer := &producer{ + cancelPollCh: make(chan struct{}, 1), + queueControlCh: make(chan *controlEventPayload, 1), + } + ctx, _, started, stopped := targetProducer.StartInit(ctx) + targetProducer.runState.Store(&producerRunState{done: runCtx.Done(), fetchDone: ctx.Done()}) + started() + targetProducer.queueControlCh <- &controlEventPayload{Action: controlActionPause} + stopFunc := sync.OnceFunc(func() { + cancel() + cancelRun() + stopped() + }) + t.Cleanup(stopFunc) + + return &testBundle{ + client: &Client[pgx.Tx]{ + driver: &DriverPollOnly{}, + producersByQueueName: map[string]*producer{QueueDefault: targetProducer}, + }, + fetchStopped: ctx.Done(), + producer: targetProducer, + stopFunc: stopFunc, + } + } + + t.Run("FullBufferContextCancelled", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerWithoutListenerQueueControlEvent(ctx, QueueDefault, &controlEventPayload{}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.queueControlCh, 1) + }) + + t.Run("ProducerNotStarted", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.producer.runState.Store(nil) + + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, &controlEventPayload{}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.queueControlCh, 1) + }) + + t.Run("QueueRemoval", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + var controlsReturned testsignal.TestSignal[struct{}] + var queueRemoved testsignal.TestSignal[error] + controlsReturned.Init(t) + queueRemoved.Init(t) + go func() { + queueRemoved.Signal(bundle.client.producerRemove(context.Background(), QueueDefault)) + }() + riversharedtest.WaitOrTimeout(t, bundle.fetchStopped) + + // The queue is still registered while draining, but queue controls must + // not wait on its full buffer or hold the registry lock needed by removal. + go func() { + for _, action := range []controlAction{controlActionPause, controlActionResume, controlActionMetadataChanged} { + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, &controlEventPayload{Action: action}) + } + controlsReturned.Signal(struct{}{}) + }() + controlsReturned.WaitOrTimeout() + require.Len(t, bundle.producer.queueControlCh, 1) + + // Cancellation delivery remains live until the running attempts drain. + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{ID: 123, Queue: QueueDefault}) + require.Len(t, bundle.producer.cancelPollCh, 1) + bundle.stopFunc() + require.NoError(t, queueRemoved.WaitOrTimeout()) + require.Empty(t, bundle.client.producersByQueueName) + }) + + t.Run("RunningProducer", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + <-bundle.producer.queueControlCh + controlEvent := &controlEventPayload{Action: controlActionResume, Queue: QueueDefault} + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, controlEvent) + + require.Same(t, controlEvent, <-bundle.producer.queueControlCh) + }) +} + func Test_Client_Stop_Common(t *testing.T) { t.Parallel() @@ -8613,6 +9195,47 @@ func Test_Client_Start_Error(t *testing.T) { }) } +func Test_Client_Start_NotifierCancellation(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dbPool := riversharedtest.DBPool(ctx, t) + driver := riverpgxv5.New(dbPool) + schema := riverdbtest.TestSchema(ctx, t, driver, nil) + listener := newProducerTestListenerMock(driver.GetListener(&riverdriver.GetListenenerParams{Schema: schema})) + var connecting, releaseConnect testsignal.TestSignal[struct{}] + var startReturned testsignal.TestSignal[error] + connecting.Init(t) + releaseConnect.Init(t) + startReturned.Init(t) + listener.connectFunc = func(ctx context.Context) error { + connecting.Signal(struct{}{}) + select { + case <-ctx.Done(): + return ctx.Err() + case <-releaseConnect.WaitC(): + return errors.New("test released blocked connection") + } + } + client, err := NewClient(&DriverWithListener{Driver: driver, listener: listener}, newTestConfig(t, schema)) + require.NoError(t, err) + + startCtx, cancel := context.WithCancel(ctx) + defer cancel() + go func() { startReturned.Signal(client.Start(startCtx)) }() + t.Cleanup(func() { + releaseConnect.Signal(struct{}{}) + stopCtx, stopCancel := context.WithTimeout(ctx, 5*time.Second) + defer stopCancel() + require.NoError(t, client.Stop(stopCtx)) + }) + connecting.WaitOrTimeout() + cancel() + // The connection must honor startup cancellation, even though the notifier + // stays alive after a successful start until workers finish draining. + require.NoError(t, startReturned.WaitOrTimeout()) +} + func Test_Client_YugabyteQueueControl(t *testing.T) { t.Parallel() @@ -10504,6 +11127,17 @@ func NewDriverPollOnly(dbPool *pgxpool.Pool) *DriverPollOnly { func (d *DriverPollOnly) SupportsListener() bool { return false } +// DriverWithListener lets tests inject faults into a real listener. +type DriverWithListener struct { + *riverpgxv5.Driver + + listener riverdriver.Listener +} + +func (d *DriverWithListener) GetListener(_ *riverdriver.GetListenenerParams) riverdriver.Listener { + return d.listener +} + // DriverWithoutListenNotify simulates a driver without any listen/notify // support at all. An example of this is SQLite through `riversqlite` where // SQLite doesn't support listen/notify or anything like it. We could use the diff --git a/conformance/cmd/generatefixtures/protocol_values.go b/conformance/cmd/generatefixtures/protocol_values.go index 1fa2f71dd..7b4a9e40b 100644 --- a/conformance/cmd/generatefixtures/protocol_values.go +++ b/conformance/cmd/generatefixtures/protocol_values.go @@ -166,6 +166,8 @@ func makeProtocolNotifications(root string) ([]protocolNotification, error) { switch action { case "cancel": values["job_id"] = exampleJobID + values["attempt"] = 1 + values["attempted_at"] = "2026-01-02T03:04:05.678Z" case "metadata_changed": values["metadata"] = map[string]any{"owner": "candidate"} } diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index e860cb17b..2937d5f9d 100644 --- a/internal/notifier/notifier.go +++ b/internal/notifier/notifier.go @@ -48,6 +48,22 @@ var NotificationTopicLongest = func() NotificationTopic { //nolint:gochecknoglob type NotifyFunc func(topic NotificationTopic, payload string) +// RecoveryFunc requests reconciliation after listening starts or notifications are +// dropped. It must not block; calls may overlap with notification delivery. +type RecoveryFunc func() + +// RecoverySubscription unregisters a recovery callback. +type RecoverySubscription struct { + notifier *Notifier +} + +// Unregister stops future callbacks; an invocation already in progress may finish. +func (s *RecoverySubscription) Unregister() { + s.notifier.mu.Lock() + defer s.notifier.mu.Unlock() + delete(s.notifier.recoveryFuncs, s) +} + type Subscription struct { notifyFunc NotifyFunc notifier *Notifier @@ -92,6 +108,7 @@ type Notifier struct { isConnected bool isStarted bool isWaiting bool + recoveryFuncs map[*RecoverySubscription]RecoveryFunc subscriptions map[NotificationTopic][]*Subscription waitCancel context.CancelFunc } @@ -102,12 +119,29 @@ func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notif notificationBuf: make(chan *riverdriver.Notification, 1000), waitInterruptChan: make(chan func(), 10), + recoveryFuncs: make(map[*RecoverySubscription]RecoveryFunc), subscriptions: make(map[NotificationTopic][]*Subscription), }) return notifier } +// RegisterRecoveryFunc registers a callback for connection recovery and buffer overflow. +func (n *Notifier) RegisterRecoveryFunc(recoveryFunc RecoveryFunc) *RecoverySubscription { + n.mu.Lock() + defer n.mu.Unlock() + + sub := &RecoverySubscription{notifier: n} + n.recoveryFuncs[sub] = recoveryFunc + return sub +} + func (n *Notifier) Start(ctx context.Context) error { + return n.StartWithConnectContext(ctx, ctx) +} + +// StartWithConnectContext starts the notifier with a separate context for its +// initial connection, allowing startup to be canceled independently of draining. +func (n *Notifier) StartWithConnectContext(ctx, connectCtx context.Context) error { ctx, shouldStart, started, stopped := n.StartInit(ctx) if !shouldStart { return nil @@ -115,7 +149,12 @@ func (n *Notifier) Start(ctx context.Context) error { // The loop below will connect/close on every iteration, but do one initial // connect so the notifier fails fast in case of an obvious problem. - if err := n.listenerConnect(ctx, false); err != nil { + connectCtx, cancelConnect := context.WithCancel(connectCtx) + stopCancelConnect := context.AfterFunc(ctx, cancelConnect) + err := n.listenerConnect(connectCtx, false) + stopCancelConnect() + cancelConnect() + if err != nil { stopped() if errors.Is(err, context.Canceled) { return nil @@ -214,6 +253,8 @@ func (n *Notifier) listenAndWait(ctx context.Context) error { n.testSignals.ListeningBegin.Signal(struct{}{}) defer n.testSignals.ListeningEnd.Signal(struct{}{}) + n.requestRecovery() + drainInterrupts := func() { for { select { @@ -274,6 +315,18 @@ func (n *Notifier) listenerClose(ctx context.Context, skipLock bool) { n.isConnected = false } +func (n *Notifier) requestRecovery() { + recoveryFuncs := func() []RecoveryFunc { + n.mu.RLock() + defer n.mu.RUnlock() + + return maputil.Values(n.recoveryFuncs) + }() + for _, recoveryFunc := range recoveryFuncs { + recoveryFunc() + } +} + // shouldIgnoreListenerError returns true if the error is a certain type of // common error that we see when trying to close a listener. // @@ -397,6 +450,7 @@ func (n *Notifier) waitOnce(ctx context.Context) error { case n.notificationBuf <- notification: default: n.Logger.WarnContext(ctx, n.Name+": Dropping notification due to full buffer", "payload", notification.Payload) + n.requestRecovery() } } }() diff --git a/internal/notifier/notifier_test.go b/internal/notifier/notifier_test.go index 76ad0b820..504b720e8 100644 --- a/internal/notifier/notifier_test.go +++ b/internal/notifier/notifier_test.go @@ -19,6 +19,7 @@ import ( "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/startstoptest" + "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivershared/util/maputil" "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" @@ -281,6 +282,57 @@ func TestNotifier(t *testing.T) { require.Equal(t, TopicAndPayload{testTopic1, "msg1"}, riversharedtest.WaitOrTimeout(t, notifyChan)) }) + t.Run("RegisterRecoveryFunc", func(t *testing.T) { + t.Parallel() + + notifier, _ := setup(t, nil) + notifier.testDisableSleep = true + + var errorNum int + + listenerMock := NewListenerMock(notifier.listener) + listenerMock.waitForNotificationFunc = func(ctx context.Context) (*riverdriver.Notification, error) { + // Returns an error the first time, but then works after. + errorNum++ + switch errorNum { + case 1: + return nil, errors.New("error during wait") + default: + return listenerMock.Listener.WaitForNotification(ctx) + } + } + notifier.listener = listenerMock + + var ready testsignal.TestSignal[struct{}] + ready.Init(t) + sub := notifier.RegisterRecoveryFunc(func() { ready.Signal(struct{}{}) }) + t.Cleanup(sub.Unregister) + + start(t, notifier) + + // Fires once on the initial connect. + ready.WaitOrTimeout() + + // And again once the notifier reconnects after the injected failure. + require.EqualError(t, notifier.testSignals.BackoffError.WaitOrTimeout(), "error during wait") + ready.WaitOrTimeout() + + notifier.Stop() + sub.Unregister() + + // WaitForNotification is entered after listener-ready callbacks finish. + var waiting testsignal.TestSignal[struct{}] + waiting.Init(t) + listenerMock.waitForNotificationFunc = func(ctx context.Context) (*riverdriver.Notification, error) { + waiting.Signal(struct{}{}) + return listenerMock.Listener.WaitForNotification(ctx) + } + + require.NoError(t, notifier.Start(ctx)) + waiting.WaitOrTimeout() + ready.RequireEmpty() + }) + t.Run("SingleTopicMultipleSubscribers", func(t *testing.T) { t.Parallel() @@ -625,6 +677,36 @@ func TestNotifier(t *testing.T) { }) } +func TestNotifier_OverflowRequestsRecovery(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var recovered testsignal.TestSignal[struct{}] + recovered.Init(t) + delivered := false + notif := New(riversharedtest.BaseServiceArchetype(t), &ListenerMock{ + waitForNotificationFunc: func(ctx context.Context) (*riverdriver.Notification, error) { + if !delivered { + delivered = true + return &riverdriver.Notification{Topic: string(NotificationTopicControl), Payload: `{"action":"cancel","job_id":1}`}, nil + } + <-ctx.Done() + return nil, ctx.Err() + }, + }) + notif.notificationBuf = make(chan *riverdriver.Notification, 1) + notif.notificationBuf <- &riverdriver.Notification{} + sub := notif.RegisterRecoveryFunc(func() { recovered.Signal(struct{}{}) }) + defer sub.Unregister() + var done testsignal.TestSignal[error] + done.Init(t) + go func() { done.Signal(notif.waitOnce(ctx)) }() + recovered.WaitOrTimeout() + cancel() + require.ErrorIs(t, done.WaitOrTimeout(), context.Canceled) +} + func TestNotifier_Start(t *testing.T) { t.Parallel() diff --git a/java/river/src/main/java/com/riverqueue/Client.java b/java/river/src/main/java/com/riverqueue/Client.java index 752f5a9d8..9b930dd6f 100644 --- a/java/river/src/main/java/com/riverqueue/Client.java +++ b/java/river/src/main/java/com/riverqueue/Client.java @@ -656,7 +656,7 @@ public Job cancel(Connection connection, long id) { if (!rows.next()) return get(c, id); var job = read(rows); pilot.afterStateChange(c, driver(), job); - notify(c, Protocol.cancel(id, job.queue())); + notify(c, Protocol.cancel(job)); return job; } }); @@ -1182,11 +1182,7 @@ public Job retry(Connection connection, long id) { var before = locked(c, id); try (var statement = Sql.prepare( - c, - Sql.query(database, "retry"), - database.timestamp(clock.instant()), - id, - database.timestamp(clock.instant())); + c, Sql.query(database, "retry"), database.timestamp(clock.instant()), id); var rows = statement.executeQuery()) { boolean updated = rows.next(); var job = updated ? read(rows) : before; diff --git a/java/river/src/main/java/com/riverqueue/Protocol.java b/java/river/src/main/java/com/riverqueue/Protocol.java index 1e0a1da43..6a3f57912 100644 --- a/java/river/src/main/java/com/riverqueue/Protocol.java +++ b/java/river/src/main/java/com/riverqueue/Protocol.java @@ -1,5 +1,6 @@ package com.riverqueue; +import java.time.Instant; import java.util.Map; import java.util.function.Consumer; import java.util.regex.Pattern; @@ -22,6 +23,28 @@ final class Protocol { private Protocol() {} + static Notification cancel(Job job) { + if (job.attemptedAt() == null) return cancel(job.id(), job.queue()); + return cancel(job.id(), job.queue(), job.attempt(), job.attemptedAt()); + } + + static Notification cancel(long id, String queue, int attempt, Instant attemptedAt) { + return new Notification( + TOPIC_CONTROL, + Json.encode( + Map.of( + "action", + "cancel", + "attempt", + attempt, + "attempted_at", + attemptedAt, + "job_id", + id, + "queue", + queue))); + } + static Notification cancel(long id, String queue) { return new Notification( TOPIC_CONTROL, Json.encode(Map.of("action", "cancel", "job_id", id, "queue", queue))); @@ -52,13 +75,14 @@ static Notification resigned(String id) { TOPIC_LEADERSHIP, Json.encode(Map.of("action", "resigned", "leader_id", id))); } + record CancellationNotice(long id, int attempt, Instant attemptedAt) {} + record Notification(String topic, String payload) {} /** Shared by both database listeners; dispatching itself requires no database I/O. */ record Dispatcher( String clientId, - Map> active, - Map pendingCancellation, + Consumer cancellation, Consumer queues, Consumer leadership) { void dispatch(String topic, String payload) { @@ -70,14 +94,12 @@ void dispatch(String topic, String payload) { switch (topic) { case TOPIC_CONTROL -> { if (action.equals("cancel") && value.path("job_id").isIntegralNumber()) { - long jobId = value.path("job_id").asLong(); - // Publish pending cancellation first so an attempt registered concurrently sees it. - pendingCancellation.put(jobId, System.nanoTime()); - var context = active.get(jobId); - if (context != null) { - context.requestCancellation(WorkContext.Cancellation.REMOTE); - pendingCancellation.remove(jobId); - } + var attemptedAt = value.path("attempted_at"); + cancellation.accept( + new CancellationNotice( + value.path("job_id").asLong(), + value.path("attempt").asInt(0), + attemptedAt.isString() ? Instant.parse(attemptedAt.asString()) : null)); } else if ((action.equals("metadata_changed") || action.equals("pause") || action.equals("resume")) diff --git a/java/river/src/main/java/com/riverqueue/RunHandle.java b/java/river/src/main/java/com/riverqueue/RunHandle.java index 8346562da..0242dbd99 100644 --- a/java/river/src/main/java/com/riverqueue/RunHandle.java +++ b/java/river/src/main/java/com/riverqueue/RunHandle.java @@ -19,6 +19,7 @@ import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; import org.postgresql.PGConnection; @@ -33,7 +34,6 @@ public final class RunHandle implements AutoCloseable { private final Map> active = new ConcurrentHashMap<>(); private final Map activeThreads = new ConcurrentHashMap<>(); - private final Map pendingCancellation = new ConcurrentHashMap<>(); private final AtomicBoolean maintenanceBusy = new AtomicBoolean(); private volatile Instant nextReindex = Instant.now() @@ -42,6 +42,9 @@ public final class RunHandle implements AutoCloseable { .plusDays(1) .atStartOfDay() .toInstant(java.time.ZoneOffset.UTC); + private final AtomicLong cancellationGeneration = new AtomicLong(); + private volatile boolean cancellationPollOnly; + private final Semaphore cancellationWake = new Semaphore(0); private final Client.Builder config; private final RetryPolicy defaultRetryPolicy = RetryPolicy.defaults(); private final java.util.concurrent.BlockingQueue completions = @@ -79,11 +82,12 @@ public final class RunHandle implements AutoCloseable { RunHandle(Client river, Client.Builder config) { this.config = config; + cancellationPollOnly = config.pollOnly; this.river = river; id = river.id(); notifications = new Protocol.Dispatcher( - id, active, pendingCancellation, this::queueNotification, this::leadershipNotification); + id, this::cancelNotification, this::queueNotification, this::leadershipNotification); queues.putAll(config.queues); periodic = new PeriodicJobBundle( @@ -164,6 +168,7 @@ void launch() { unsubscribeInsert = river.onInsertCommit(wake::release); if (!config.pollOnly) executor.submit(this::listen); executor.submit(this::completeBatches); + executor.submit(this::pollCancellations); executor.submit(this::supervise); executor.submit(this::dispatch); dispatcherStarted = true; @@ -267,7 +272,6 @@ public void stopAndCancel() { private void dispatch() { long nextPoll = 0; long nextMaintenance = 0; - long nextCancelPoll = 0; while (running.get()) { try { long now = System.nanoTime(); @@ -295,25 +299,6 @@ private void dispatch() { queues.keySet().forEach(this::heartbeat); nextMaintenance = now + config.serviceInterval.toNanos(); } - if (now >= nextCancelPoll) { - river.transaction( - c -> { - try (var statement = - Sql.prepare(c, Sql.query(river.database(), "cancel_requested")); - var rows = statement.executeQuery()) { - while (rows.next()) { - var context = active.get(rows.getLong(1)); - if (context != null) - context.requestCancellation(WorkContext.Cancellation.REMOTE); - } - } - return null; - }); - pendingCancellation - .entrySet() - .removeIf(entry -> now - entry.getValue() > Duration.ofMinutes(1).toNanos()); - nextCancelPoll = now + TimeUnit.MILLISECONDS.toNanos(50); - } if (wake.tryAcquire(10, TimeUnit.MILLISECONDS)) wake.release(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); @@ -331,6 +316,7 @@ private void dispatch() { } } dispatcherStopped.countDown(); + cancellationWake.release(); } private record FetchWakeup(Instant scheduledAt, String notifyQueue) @@ -367,7 +353,17 @@ private boolean wakeDueFetches() { return !ready.isEmpty(); } + private void cancelNotification(Protocol.CancellationNotice notice) { + var context = active.get(notice.id()); + if (context != null && notice.attempt() > 0 && notice.attemptedAt() != null) { + if (context.job().attempt() == notice.attempt() + && notice.attemptedAt().equals(context.job().attemptedAt())) + context.requestCancellation(WorkContext.Cancellation.REMOTE); + } else requestCancellationCheck(); + } + private void claim(String queue, int count) { + long generation = cancellationGeneration.get(); List jobs = river.transaction( c -> @@ -413,6 +409,7 @@ private void claim(String queue, int count) { return result; })); for (var decoded : jobs) launch(decoded, config.handlers.get(decoded.job().kind())); + if (!jobs.isEmpty() && cancellationGeneration.get() != generation) cancellationWake.release(); } private void launch(Client.Decoded decoded, Registration registration) { @@ -466,8 +463,6 @@ private void launch(Job row, WorkContext context, WorkContext.Step active.put(row.id(), context); if (row.metadata().has("cancel_attempted_at")) context.requestCancellation(WorkContext.Cancellation.REMOTE); - if (pendingCancellation.remove(row.id()) != null) - context.requestCancellation(WorkContext.Cancellation.REMOTE); if (cancelOnStop) context.requestCancellation(WorkContext.Cancellation.SHUTDOWN); executor.submit( () -> { @@ -512,6 +507,7 @@ private void launch(Job row, WorkContext context, WorkContext.Step queueLock.lock(); try { active.remove(row.id(), context); + if (!running.get() && active.isEmpty()) cancellationWake.release(); queueChanged.signalAll(); } finally { queueLock.unlock(); @@ -533,6 +529,7 @@ private void workMiddleware(WorkContext context, int index, WorkContext.Step List> claimPeers( WorkContext parent, JobType type, Client.Transaction> operation) { parent.checkRuntimeCancellation(); + long generation = cancellationGeneration.get(); var rows = river.transaction( connection -> { @@ -551,6 +548,7 @@ List> claimPeers( // Publish before checking the parent so a concurrent cancellation cannot miss both. context.attach(this, parent); } + if (!rows.isEmpty() && cancellationGeneration.get() != generation) cancellationWake.release(); var contexts = new ArrayList>(); for (var decoded : rows) { if (decoded.failure() != null) throw decoded.failure(); @@ -580,6 +578,7 @@ void completePeer(WorkContext parent, WorkContext context, Exception failu queueLock.lock(); try { active.remove(context.job().id(), context); + if (!running.get() && active.isEmpty()) cancellationWake.release(); queueChanged.signalAll(); } finally { queueLock.unlock(); @@ -747,6 +746,10 @@ private void completeBatches() { } } + private boolean hasActiveWork() { + return running.get() || dispatcherStopped.getCount() > 0 || !active.isEmpty(); + } + private void heartbeat(String queue) { river.transaction( c -> { @@ -780,9 +783,13 @@ private void listen() { listenSqlite(); return; } - while (running.get()) { + while (hasActiveWork()) { try (var connection = river.database().connection()) { - if (!river.database().supportsNotifications(connection)) return; + if (!river.database().supportsNotifications(connection)) { + cancellationPollOnly = true; + cancellationWake.release(); + return; + } listener = connection; String schema = river.database().schema(); if (schema.isEmpty()) @@ -797,10 +804,11 @@ private void listen() { statement.execute("LISTEN \"" + schema.replace("\"", "\"\"") + "." + topic + "\""); } wake.release(); + requestCancellationCheck(); observe("listen_ready"); var postgres = connection.unwrap(PGConnection.class); try { - while (running.get()) { + while (hasActiveWork()) { var notifications = postgres.getNotifications(25); if (notifications != null) for (var notification : notifications) @@ -815,7 +823,7 @@ private void listen() { } } } catch (Exception e) { - if (running.get()) { + if (hasActiveWork()) { report(e); try { TimeUnit.MILLISECONDS.sleep(100); @@ -842,7 +850,7 @@ private void deliver(String topic, String payload) { private void listenSqlite() { long after = 0; boolean initialized = false; - while (running.get()) { + while (hasActiveWork()) { try { if (!initialized) { after = @@ -859,6 +867,7 @@ private void listenSqlite() { // The initial cursor skips earlier notifications, including inserts since the first // fetch. Poll again once listening starts, as the Postgres listener does. wake.release(); + requestCancellationCheck(); observe("listen_ready"); } final long cursor = after; @@ -1417,6 +1426,52 @@ private void observe(String event) { } } + private void pollCancellations() { + // Fetches may block, and graceful stop ends dispatch before active work. + // Keep cancellation checks independent until every claimed attempt drains. + boolean retry = false; + while (hasActiveWork()) { + try { + if (cancellationPollOnly || retry) + cancellationWake.tryAcquire(config.pollInterval.toNanos(), TimeUnit.NANOSECONDS); + else cancellationWake.acquire(); + cancellationWake.drainPermits(); + if (!hasActiveWork()) return; + retry = false; + var attempts = Map.copyOf(active); + if (!attempts.isEmpty()) { + var ids = new ArrayList<>(attempts.keySet()); + for (int start = 0; start < ids.size(); start += 1000) { + var batch = ids.subList(start, Math.min(start + 1000, ids.size())); + river.transaction( + c -> { + try (var statement = + Sql.prepare( + c, + Sql.query(river.database(), "cancel_requested"), + Json.encode(batch)); + var rows = statement.executeQuery()) { + while (rows.next()) { + // A late query result must not cancel a subsequent attempt. + var context = attempts.get(rows.getLong(1)); + if (context != null) + context.requestCancellation(WorkContext.Cancellation.REMOTE); + } + } + return null; + }); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } catch (Exception e) { + retry = true; + report(e); + } + } + } + private void publish(Event event) { for (var subscriber : river.subscribers()) try { @@ -1437,6 +1492,11 @@ private void report(Throwable error) { } } + private void requestCancellationCheck() { + cancellationGeneration.incrementAndGet(); + cancellationWake.release(); + } + private Duration retryDelay(Job job, Instant now) { try { var delay = Objects.requireNonNull(config.retryPolicy.delay(job), "Retry delay"); diff --git a/java/river/src/main/resources/com/riverqueue/postgres.sql b/java/river/src/main/resources/com/riverqueue/postgres.sql index d5a92e3ca..d6151487b 100644 --- a/java/river/src/main/resources/com/riverqueue/postgres.sql +++ b/java/river/src/main/resources/com/riverqueue/postgres.sql @@ -56,9 +56,11 @@ DELETE FROM {schema}river_job WHERE id=? AND state != 'running' RETURNING {colum -- name: delete_filtered DELETE FROM {schema}river_job WHERE {where} AND id=? AND state != 'running' RETURNING {columns} -- name: retry -UPDATE {schema}river_job SET state='available', finalized_at=NULL, -max_attempts=CASE WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=? -WHERE id=? AND state != 'running' AND NOT (state='available' AND scheduled_at < ?) RETURNING {columns} +WITH retry_params AS (SELECT CAST(? AS timestamptz) AS now) +UPDATE {schema}river_job SET state='available', finalized_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, +metadata=CASE WHEN jsonb_typeof(metadata)='object' THEN metadata - 'cancel_attempted_at' ELSE metadata END, +max_attempts=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END +WHERE id=? AND state != 'running' AND ((jsonb_typeof(metadata)='object' AND jsonb_exists(metadata, 'cancel_attempted_at')) OR NOT (state='available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {columns} -- name: reset_jobs DELETE FROM {schema}river_job -- name: reset_queues @@ -179,7 +181,7 @@ SELECT {columns} FROM {schema}river_job WHERE state='running' AND attempted_at>'cancel_attempted_at' IS NOT NULL +SELECT id FROM {schema}river_job WHERE id IN (SELECT value::bigint FROM jsonb_array_elements_text(?::jsonb)) AND state='running' AND jsonb_exists(metadata, 'cancel_attempted_at') -- name: notification_cursor SELECT coalesce(max(id),0) FROM {schema}river_notification diff --git a/java/river/src/main/resources/com/riverqueue/sqlite.sql b/java/river/src/main/resources/com/riverqueue/sqlite.sql index 75e2d53b0..bcea6e2e2 100644 --- a/java/river/src/main/resources/com/riverqueue/sqlite.sql +++ b/java/river/src/main/resources/com/riverqueue/sqlite.sql @@ -55,9 +55,11 @@ DELETE FROM {schema}river_job WHERE id=? AND state != 'running' RETURNING {colum -- name: delete_filtered DELETE FROM {schema}river_job WHERE {where} AND id=? AND state != 'running' RETURNING {columns} -- name: retry -UPDATE {schema}river_job SET state='available', finalized_at=NULL, -max_attempts=CASE WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=? -WHERE id=? AND state != 'running' AND NOT (state='available' AND scheduled_at < ?) RETURNING {columns} +WITH retry_params AS (SELECT CAST(? AS text) AS now) +UPDATE {schema}river_job SET state='available', finalized_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, +metadata=CASE WHEN typeof(metadata)='text' AND NOT json_valid(metadata) THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, +max_attempts=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END +WHERE id=? AND state != 'running' AND ((CASE WHEN typeof(metadata)='text' AND NOT json_valid(metadata) THEN NULL ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL OR NOT (state='available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {columns} -- name: reset_jobs DELETE FROM {schema}river_job -- name: reset_queues @@ -196,7 +198,7 @@ AND queue NOT IN (SELECT value FROM json_each(?)) AND (? OR queue IN (SELECT val SELECT {columns} FROM {schema}river_job WHERE state='running' AND attempted_at? ORDER BY id LIMIT 1000 -- name: cancel_requested -SELECT id FROM river_job WHERE state='running' AND (CASE WHEN typeof(metadata)<>'text' OR json_valid(metadata) THEN json_extract(metadata,'$.cancel_attempted_at') END) IS NOT NULL +SELECT id FROM {schema}river_job WHERE id IN (SELECT CAST(value AS integer) FROM json_each(?)) AND state='running' AND (CASE WHEN typeof(metadata)<>'text' OR json_valid(metadata) THEN json_type(metadata,'$.cancel_attempted_at') END) IS NOT NULL -- name: notification_cursor SELECT coalesce(max(id),0) FROM {schema}river_notification diff --git a/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java b/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java new file mode 100644 index 000000000..87bd4498d --- /dev/null +++ b/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java @@ -0,0 +1,238 @@ +package com.riverqueue; + +import static org.junit.jupiter.api.Assertions.*; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.SQLException; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import javax.sql.DataSource; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +class CancellationCheckTest { + @RegisterExtension final TestDatabase databases = new TestDatabase(); + @TempDir Path directory; + + @Test + void cancellationQueryOnlyReadsRequestedJobs() throws Exception { + var database = databases.open(directory.resolve("river.db")); + var client = Client.builder(database).build(); + var type = JobType.of("cancellation_query", String.class); + var first = client.insert(type, "first").job(); + var second = client.insert(type, "second").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='running', metadata='{\"cancel_attempted_at\":null}' WHERE id IN (?, ?)", + first.id(), + second.id())) { + statement.executeUpdate(); + } + try (var statement = + Sql.prepare( + connection, + Sql.query(database, "cancel_requested"), + Json.encode(List.of(first.id()))); + var rows = statement.executeQuery()) { + assertTrue(rows.next()); + assertEquals(first.id(), rows.getLong(1)); + assertFalse(rows.next()); + } + return null; + }); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void checksOnlyOnHintsAndRetriesFailures(boolean failFirst) throws Exception { + var database = databases.open(directory.resolve("river.db")); + var fail = new AtomicBoolean(); + var queries = new AtomicInteger(); + var checked = new LinkedBlockingQueue(); + var source = + (DataSource) + Proxy.newProxyInstance( + DataSource.class.getClassLoader(), + new Class[] {DataSource.class}, + (proxy, method, args) -> { + if (!method.getName().equals("getConnection")) + throw new UnsupportedOperationException(); + var connection = database.connection(); + return Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + (ignored, operation, arguments) -> { + if (operation.getName().equals("prepareStatement") + && arguments[0].equals(Sql.query(database, "cancel_requested"))) { + checked.add(queries.incrementAndGet()); + if (fail.compareAndSet(true, false)) + throw new SQLException("temporary cancellation read failure"); + } + try { + return operation.invoke(connection, arguments); + } catch (InvocationTargetException error) { + throw error.getCause(); + } + }); + }); + var ready = new CompletableFuture(); + var entered = new CompletableFuture(); + var finish = new CompletableFuture(); + var errors = new LinkedBlockingQueue(); + var type = JobType.of("cancellation_check", String.class); + var client = + Client.builder(new Database(source, database.dialect()).withSchema(database.schema())) + .queue("default", 1) + .leadership(false) + .pollInterval(Duration.ofMillis(10)) + .stopTimeout(Duration.ofSeconds(5)) + .errorHandler(errors::add) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .worker( + type, + context -> { + entered.complete(null); + finish.join(); + context.checkCancelled(); + }) + .build(); + try (var run = client.start()) { + try { + ready.get(5, TimeUnit.SECONDS); + var job = client.insert(type, "working").job(); + entered.get(5, TimeUnit.SECONDS); + // Let the startup hint settle, then leave a running attempt idle for several old ticks. + TimeUnit.MILLISECONDS.sleep(200); + int before = queries.get(); + checked.clear(); + TimeUnit.MILLISECONDS.sleep(200); + assertEquals(before, queries.get()); + fail.set(failFirst); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(job.id(), "default")); + return null; + }); + assertNotNull(checked.poll(5, TimeUnit.SECONDS)); + if (failFirst) assertNotNull(checked.poll(5, TimeUnit.SECONDS)); + TimeUnit.MILLISECONDS.sleep(200); + assertEquals(before + (failFirst ? 2 : 1), queries.get()); + assertEquals(failFirst ? 1 : 0, errors.size()); + } finally { + finish.complete(null); + } + } + } + + @Test + @SuppressWarnings("try") + void identifiedNoticeBypassesBlockedReadAndIgnoresOlderAttempts() throws Exception { + var database = databases.open(directory.resolve("river.db")); + var block = new AtomicBoolean(); + var readStarted = new CompletableFuture(); + var releaseRead = new CompletableFuture(); + var source = + (DataSource) + Proxy.newProxyInstance( + DataSource.class.getClassLoader(), + new Class[] {DataSource.class}, + (proxy, method, args) -> { + if (!method.getName().equals("getConnection")) + throw new UnsupportedOperationException(); + var connection = database.connection(); + return Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + (ignored, operation, arguments) -> { + if (operation.getName().equals("prepareStatement") + && arguments[0].equals(Sql.query(database, "cancel_requested")) + && block.compareAndSet(true, false)) { + readStarted.complete(null); + releaseRead.get(5, TimeUnit.SECONDS); + } + try { + return operation.invoke(connection, arguments); + } catch (InvocationTargetException error) { + throw error.getCause(); + } + }); + }); + var ready = new CompletableFuture(); + var entered = new CompletableFuture>(); + var cancelled = new CompletableFuture(); + var type = JobType.of("cancellation_delivery", String.class); + var client = + Client.builder(new Database(source, database.dialect()).withSchema(database.schema())) + .queue("default", 1) + .leadership(false) + .pollInterval(Duration.ofMillis(10)) + .stopTimeout(Duration.ofSeconds(5)) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .worker( + type, + context -> { + entered.complete(context); + context.awaitCancellation(); + cancelled.complete(null); + context.checkCancelled(); + }) + .build(); + try (var run = client.start()) { + try { + ready.get(5, TimeUnit.SECONDS); + var inserted = client.insert(type, "working").job(); + var context = entered.get(5, TimeUnit.SECONDS); + var job = context.job(); + block.set(true); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(inserted.id(), "default")); + return null; + }); + readStarted.get(5, TimeUnit.SECONDS); + client.transaction( + connection -> { + client.notify( + connection, + Protocol.cancel(job.id(), job.queue(), job.attempt() + 1, job.attemptedAt())); + client.notify( + connection, + Protocol.cancel( + job.id(), job.queue(), job.attempt(), job.attemptedAt().minusSeconds(1))); + return null; + }); + assertThrows( + java.util.concurrent.TimeoutException.class, + () -> cancelled.get(100, TimeUnit.MILLISECONDS)); + client.cancel(job.id()); + cancelled.get(2, TimeUnit.SECONDS); + } finally { + releaseRead.complete(null); + if (entered.isDone()) entered.join().requestCancellation(WorkContext.Cancellation.SHUTDOWN); + } + } + } +} diff --git a/java/river/src/test/java/com/riverqueue/ProtocolTest.java b/java/river/src/test/java/com/riverqueue/ProtocolTest.java index 3496f78a0..994513338 100644 --- a/java/river/src/test/java/com/riverqueue/ProtocolTest.java +++ b/java/river/src/test/java/com/riverqueue/ProtocolTest.java @@ -5,7 +5,6 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Random; @@ -79,27 +78,20 @@ Stream notificationDispatch() throws Exception { DynamicTest.dynamicTest( name + " received by " + clientId, () -> { - var target = context(42); - var unrelated = context(43); + var cancellations = new ArrayList(); var queues = new ArrayList(); var leadership = new ArrayList(); - var pending = new HashMap(); var dispatcher = new Protocol.Dispatcher( clientId, - Map.of(42L, target, 43L, unrelated), - pending, + notice -> cancellations.add(notice.id()), queues::add, leadership::add); dispatcher.dispatch( fixture.required("topic").asString(), Json.encode(fixture.required("payload"))); - assertEquals( - name.equals("cancel") ? WorkContext.Cancellation.REMOTE : null, - target.cancellation()); - assertNull(unrelated.cancellation(), "Notification must not cancel another job"); - assertTrue(pending.isEmpty(), "A registered attempt consumes its cancellation"); + assertEquals(name.equals("cancel") ? List.of(42L) : List.of(), cancellations); switch (name) { case "insert", "metadata_changed", "pause", "resume" -> { assertEquals(List.of(new Protocol.QueueNotice("priority", name)), queues); @@ -130,36 +122,32 @@ Stream notificationDispatch() throws Exception { } @Test - void notificationTopicsAndPendingCancellation() throws Exception { + void notificationTopicsRouteCancellationChecks() throws Exception { var fixture = Conformance.fixture("protocol_values.json").required("notifications"); var cancel = java.util.stream.StreamSupport.stream(fixture.spliterator(), false) .filter(value -> value.required("name").asString().equals("cancel")) .findFirst() .orElseThrow(); - var unrelated = context(43); - var pending = new HashMap(); + var cancellations = new ArrayList(); var queues = new ArrayList(); var leadership = new ArrayList(); var dispatcher = new Protocol.Dispatcher( - "observer", Map.of(43L, unrelated), pending, queues::add, leadership::add); + "observer", notice -> cancellations.add(notice.id()), queues::add, leadership::add); String payload = Json.encode(cancel.required("payload")); dispatcher.dispatch("unrelated", payload); dispatcher.dispatch(Protocol.TOPIC_LEADERSHIP, payload); - assertTrue(pending.isEmpty()); assertTrue(queues.isEmpty()); assertTrue(leadership.isEmpty()); dispatcher.dispatch(Protocol.TOPIC_INSERT, payload); assertEquals(List.of(new Protocol.QueueNotice("priority", "insert")), queues); - assertTrue( - pending.isEmpty(), "A control payload on the insert topic must not cancel an attempt"); queues.clear(); + assertTrue(cancellations.isEmpty()); dispatcher.dispatch(cancel.required("topic").asString(), payload); - assertEquals(java.util.Set.of(42L), pending.keySet()); - assertNull(unrelated.cancellation()); + assertEquals(List.of(42L), cancellations); assertTrue(queues.isEmpty()); assertTrue(leadership.isEmpty()); } @@ -199,7 +187,7 @@ private static WorkContext context(long id, JsonNode metadata) { Stream notificationEncoding() throws Exception { var notifications = Map.of( - "cancel", Protocol.cancel(42, "priority"), + "cancel", Protocol.cancel(42, "priority", 1, Instant.parse("2026-01-02T03:04:05.678Z")), "insert", Protocol.insert("priority"), "metadata_changed", Protocol.queueMetadataChanged("priority", Map.of("owner", "candidate")), diff --git a/java/river/src/test/java/com/riverqueue/RetryTest.java b/java/river/src/test/java/com/riverqueue/RetryTest.java index 683b2a0eb..7aa6696c3 100644 --- a/java/river/src/test/java/com/riverqueue/RetryTest.java +++ b/java/river/src/test/java/com/riverqueue/RetryTest.java @@ -9,6 +9,7 @@ import java.time.Instant; import java.time.ZoneOffset; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; @@ -27,6 +28,55 @@ class RetryTest { @RegisterExtension final TestDatabase databases = new TestDatabase(); @TempDir Path directory; + @Test + void retryClearsCancellation() { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry cancelled").job(); + var cancelled = client.cancel(inserted.id()); + assertTrue(cancelled.metadata().has("cancel_attempted_at")); + var retried = client.retry(inserted.id()); + assertEquals(Job.State.AVAILABLE, retried.state()); + assertFalse(retried.metadata().has("cancel_attempted_at")); + assertFalse(client.get(inserted.id()).metadata().has("cancel_attempted_at")); + } + + @Test + void retryClearsCancellationWithoutReorderingAvailableJob() { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database, Clock.fixed(NOW, ZoneOffset.UTC)); + var inserted = + client + .insert( + TYPE, + "already available", + InsertOptions.builder() + .metadata(Map.of("cancel_attempted_at", NOW.toString(), "preserved", true)) + .scheduledAt(NOW.minusSeconds(3600)) + .build()) + .job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='available', attempt=max_attempts WHERE id=?", + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + return null; + }); + var retried = client.retry(inserted.id()); + assertEquals(inserted.scheduledAt(), retried.scheduledAt()); + assertEquals(inserted.maxAttempts(), retried.maxAttempts()); + var expectedMetadata = (tools.jackson.databind.node.ObjectNode) inserted.metadata().deepCopy(); + expectedMetadata.remove("cancel_attempted_at"); + assertEquals(expectedMetadata, retried.metadata()); + assertEquals(retried, client.get(inserted.id())); + } + @ParameterizedTest @EnumSource(Job.State.class) void retryNotifiesOnlyWhenTheRowChanges(Job.State state) { @@ -67,6 +117,36 @@ void retryNotifiesOnlyWhenTheRowChanges(Job.State state) { } } + @ParameterizedTest + @ValueSource( + strings = {"null", "42", "false", "\"cancel_attempted_at\"", "[\"cancel_attempted_at\"]"}) + void retryPreservesNonObjectMetadata(String metadata) { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry metadata").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='retryable', metadata=" + + (database.dialect() == Database.Dialect.POSTGRES + ? "CAST(? AS jsonb)" + : "jsonb(?)") + + " WHERE id=?", + metadata, + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + return null; + }); + + assertEquals(Job.State.AVAILABLE, client.retry(inserted.id()).state()); + assertEquals(Json.parse(metadata), client.get(inserted.id()).metadata()); + } + @ParameterizedTest @ValueSource(strings = {"local", "peer", "transaction"}) @SuppressWarnings("try") // Resources intentionally scope background activity. @@ -217,4 +297,36 @@ void sqliteNotificationFailureRollsBackTheRetry() { job.metadata().path(Protocol.METADATA_OUTPUT).asString()); assertEquals(0, notifications.get()); } + + @Test + void sqliteRetryQueryPreservesMalformedMetadata() { + var database = TestDatabase.sqlite(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry malformed metadata").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE river_job SET state='retryable', metadata='[not json' WHERE id=?", + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + // The public retry validates the locked row first. Exercise the SQL + // directly so cleanup cannot turn a decode error into a SQL error. + try (var statement = + Sql.prepare( + connection, + Sql.query(database, "retry"), + database.timestamp(NOW), + inserted.id()); + var rows = statement.executeQuery()) { + assertTrue(rows.next()); + assertEquals("[not json", rows.getString("metadata")); + } + return null; + }); + var error = assertThrows(IllegalArgumentException.class, () -> client.get(inserted.id())); + assertTrue(error.getMessage().contains("metadata")); + } } diff --git a/java/river/src/test/java/com/riverqueue/RunHandleTest.java b/java/river/src/test/java/com/riverqueue/RunHandleTest.java index 46efd1c4f..7ae5b42b3 100644 --- a/java/river/src/test/java/com/riverqueue/RunHandleTest.java +++ b/java/river/src/test/java/com/riverqueue/RunHandleTest.java @@ -70,6 +70,125 @@ void builderRequiresServiceIntervalsBelowQueueRetention(String value, boolean va .getMessage()); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void cancellationContinuesAfterGracefulStop(boolean pollOnly) throws Exception { + var client = client(); + var entered = new CompletableFuture>(); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .worker( + TYPE, + context -> { + entered.complete(context); + context.awaitCancellation(); + context.checkCancelled(); + }) + .build() + .start(); + try (var subscription = workers.client().subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel while draining").job(); + entered.get(5, TimeUnit.SECONDS); + // The timeout leaves active work draining after dispatch stops. + // The listener must remain available to deliver peer cancellation. + assertThrows(RiverException.class, workers::stop); + client.cancel(inserted.id()); + assertEquals(Job.State.CANCELLED, outcome.get(5, TimeUnit.SECONDS).job().state()); + workers.stop(); + } finally { + workers.stopAndCancel(); + } + } + + @Test + void cancellationContinuesWhileFetchingIsBlocked() throws Exception { + var entered = new CompletableFuture(); + var fetching = new CompletableFuture(); + var release = new CompletableFuture(); + var cancelled = new CompletableFuture(); + var client = + client() + .withPilot( + new Pilot() { + @Override + public int workerLimit(String queue, int configured) { + if (entered.isDone()) { + fetching.complete(null); + release.join(); + } + return configured; + } + }); + var workers = + workers(client) + .pollInterval(Duration.ofMillis(10)) + .worker( + TYPE, + context -> { + entered.complete(null); + context.awaitCancellation(); + cancelled.complete(context.cancellation()); + context.checkCancelled(); + }) + .build() + .start(); + try { + var inserted = client.insert(TYPE, "cancel while fetching").job(); + fetching.get(5, TimeUnit.SECONDS); + client.cancel(inserted.id()); + assertEquals(WorkContext.Cancellation.REMOTE, cancelled.get(5, TimeUnit.SECONDS)); + } finally { + release.complete(null); + workers.stopAndCancel(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void cancellationDuringArgumentDecodingIsReconciled(boolean pollOnly) throws Exception { + var decoding = new CompletableFuture(); + var release = new CompletableFuture(); + var client = + client() + .withPilot( + new Pilot() { + @Override + public tools.jackson.databind.JsonNode decode( + Job job) { + decoding.complete(null); + release.join(); + return job.args(); + } + }); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .worker( + TYPE, + context -> { + context.awaitCancellation(); + context.checkCancelled(); + }) + .build() + .start(); + try (var subscription = workers.client().subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel before registration").job(); + decoding.get(5, TimeUnit.SECONDS); + client.cancel(inserted.id()); + release.complete(null); + assertEquals(Job.State.CANCELLED, outcome.get(5, TimeUnit.SECONDS).job().state()); + workers.stop(); + } finally { + release.complete(null); + workers.stopAndCancel(); + } + } + @ParameterizedTest @CsvSource({"REMOTE,false", "REMOTE,true", "TIMEOUT,false", "TIMEOUT,true"}) @SuppressWarnings("try") // Resources intentionally scope background activity. @@ -802,6 +921,74 @@ void invalidRetryPoliciesFallBack(String policy) throws Exception { } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void retriedJobIgnoresStaleCancellation(boolean delayedUntilRunning) throws Exception { + var client = client(); + var ready = new CompletableFuture(); + var paused = new LinkedBlockingQueue(); + var entered = new CompletableFuture(); + var finish = new CompletableFuture(); + if (!delayedUntilRunning) finish.complete(null); + var outcome = new CompletableFuture(); + var errors = new LinkedBlockingQueue(); + var workers = + workers(client) + .pollOnly(false) + .errorHandler(errors::add) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .worker( + TYPE, + context -> { + entered.complete(null); + finish.join(); + context.checkCancelled(); + }) + .build() + .start(); + try (var pauseSubscription = + workers.client().subscribe(paused::add, Client.EventKind.QUEUE_PAUSED); + var outcomeSubscription = + workers.client().subscribe(outcome::complete, Client.EventKind.JOB_COMPLETED)) { + ready.get(5, TimeUnit.SECONDS); + var inserted = + client + .insert( + TYPE, + "retry after an idle cancellation", + InsertOptions.builder().scheduledAt(Instant.now().plusSeconds(3600)).build()) + .job(); + client.cancel(inserted.id()); + client.queues().pause("default"); + // The pause follows cancellation on the same listener, so receiving it + // proves the unmatched cancellation was dispatched before retry. + assertNotNull(paused.poll(5, TimeUnit.SECONDS)); + client.retry(inserted.id()); + client.queues().resume("default"); + if (delayedUntilRunning) { + entered.get(5, TimeUnit.SECONDS); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(inserted.id(), "default")); + return null; + }); + client.queues().pause("default"); + assertNotNull(paused.poll(5, TimeUnit.SECONDS)); + finish.complete(null); + } + assertEquals(inserted.id(), outcome.get(5, TimeUnit.SECONDS).job().id()); + workers.stop(); + } finally { + finish.complete(null); + workers.stopAndCancel(); + } + assertTrue(errors.isEmpty(), () -> "Unexpected worker errors: " + errors); + } + @ParameterizedTest @CsvSource({"false,1,5000", "true,1,5000", "false,3600,1", "true,3600,1"}) @SuppressWarnings("try") // Resources intentionally scope background activity. diff --git a/js/driver/pg/src/driver.integration.test.ts b/js/driver/pg/src/driver.integration.test.ts index 99588eebe..9bcc9aa89 100644 --- a/js/driver/pg/src/driver.integration.test.ts +++ b/js/driver/pg/src/driver.integration.test.ts @@ -533,11 +533,23 @@ describe("PgDriver integration", () => { await expect(driver.jobGetCancelRequested([])).resolves.toEqual([]); }); + it("ignores cancellation names stored in non-object metadata", async () => { + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_cancel_array`) + ); + await pool.query( + "UPDATE river_job SET state = 'running', metadata = $1::jsonb WHERE id = $2", + ['["cancel_attempted_at"]', job.id.toString()] + ); + await expect(driver.jobGetCancelRequested([job.id])).resolves.toEqual([]); + }); + it("retries non-running jobs but leaves running jobs unchanged", async () => { const scheduledAt = Temporal.Instant.from("2026-08-30T15:00:00Z"); const now = Temporal.Instant.from("2026-08-30T12:00:00.123456Z"); const retryable = await driver.jobInsert( insertParams(`${filePrefix}_retry`, { + metadata: { cancel_attempted_at: now.toString(), preserved: true }, scheduledAt, state: "retryable", }) @@ -560,10 +572,33 @@ describe("PgDriver integration", () => { }); expect(retried!.state).toBe("available"); + expect(retried!.metadata).toEqual({ preserved: true }); expect(retried!.scheduledAt.toString()).toBe("2026-08-30T12:00:00.123456Z"); expect(unchanged!.state).toBe("running"); }); + it("clears cancellation without reordering an already available retry", async () => { + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_retry_available`, { + metadata: { cancel_attempted_at: now.toString(), preserved: true }, + scheduledAt: now.subtract({ hours: 1 }), + state: "available", + }) + ); + await pool.query( + "UPDATE river_job SET attempt = max_attempts WHERE id = $1", + [job.id.toString()] + ); + const expectedMetadata = { ...job.metadata }; + delete expectedMetadata.cancel_attempted_at; + const retried = await driver.jobRetryWithOptions({ id: job.id, now }); + expect(retried?.metadata).toEqual(expectedMetadata); + expect(retried?.scheduledAt.toString()).toBe(job.scheduledAt.toString()); + expect(retried?.maxAttempts).toBe(job.maxAttempts); + expect((await driver.jobGet(job.id))?.metadata).toEqual(expectedMetadata); + }); + it("returns the committed row to the loser of a cancel or retry race", async () => { const inserted = await driver.jobInsert( insertParams(`${filePrefix}_race`, { @@ -635,6 +670,29 @@ describe("PgDriver integration", () => { ); }); + it.each(["null", "42", '"cancel_attempted_at"', '["cancel_attempted_at"]'])( + "retries jobs without changing non-object metadata: %s", + async (metadata) => { + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_retry_metadata`, { state: "retryable" }) + ); + await pool.query( + "UPDATE river_job SET metadata = $1::jsonb WHERE id = $2::bigint", + [metadata, job.id.toString()] + ); + + expect(await driver.jobRetry(job.id)).toMatchObject({ + id: job.id, + state: "available", + }); + const persisted = await pool.query<{ metadata: unknown }>( + "SELECT metadata FROM river_job WHERE id = $1::bigint", + [job.id.toString()] + ); + expect(persisted.rows[0]?.metadata).toEqual(JSON.parse(metadata)); + } + ); + it("acts on rows River can't fully read by ID and lists them like Go", async () => { const kind = `${filePrefix}_poison`; const uniqueKey = new Uint8Array(32).fill(13); @@ -2437,32 +2495,43 @@ describe("PgDriver integration", () => { await iterator.return(undefined); }); - it("decodes exact bigint IDs from cancellation notifications", async () => { - const exactID = 9_007_199_254_741_001n; - await pool.query( - `INSERT INTO river_job (id, args, kind, max_attempts) - VALUES ($1::bigint, '{}'::jsonb, $2::text, 25)`, - [exactID.toString(10), `${filePrefix}_large_cancel_id`] - ); - const abort = new AbortController(); - const iterator = driver.jobCancellationSubscribe( - `${filePrefix}_cancel_owner`, - abort.signal - ); - const next = iterator.next(); - await waitForListenerPID(pool, "river_control"); - - await driver.jobCancel(exactID); - await expect(next).resolves.toEqual({ - done: false, - value: { - attemptedBy: `${filePrefix}_cancel_owner`, - id: exactID, - }, - }); - abort.abort(); - await iterator.return(undefined); - }); + it.each([false, true])( + "decodes exact bigint IDs from cancellation notifications with identity=%s", + async (identified) => { + const exactID = 9_007_199_254_741_001n; + const attemptedAt = Temporal.Instant.from("2026-01-02T03:04:05.6789Z"); + await pool.query( + `INSERT INTO river_job (id, args, kind, max_attempts, attempt, attempted_at, state) + VALUES ($1::bigint, '{}'::jsonb, $2::text, 25, $3::smallint, $4::timestamptz, $5::river_job_state)`, + [ + exactID.toString(10), + `${filePrefix}_large_cancel_id`, + identified ? 2 : 0, + identified ? attemptedAt.toString() : null, + identified ? "running" : "available", + ] + ); + const abort = new AbortController(); + const iterator = driver.jobCancellationSubscribe( + `${filePrefix}_cancel_owner`, + abort.signal + ); + const next = iterator.next(); + await waitForListenerPID(pool, "river_control"); + + await driver.jobCancel(exactID); + await expect(next).resolves.toEqual({ + done: false, + value: { + ...(identified ? { attempt: 2, attemptedAt } : {}), + attemptedBy: `${filePrefix}_cancel_owner`, + id: exactID, + }, + }); + abort.abort(); + await iterator.return(undefined); + } + ); it("notifies listeners of a committed insert notification", async () => { const abort = new AbortController(); diff --git a/js/driver/pg/src/driver.test.ts b/js/driver/pg/src/driver.test.ts index ec5d186e8..011236b12 100644 --- a/js/driver/pg/src/driver.test.ts +++ b/js/driver/pg/src/driver.test.ts @@ -834,6 +834,52 @@ describe("PgDriver", () => { }); }); + it("fails an overflowing notification stream so its subscriber can recover", async () => { + const leased = mockLeasedPgClient(); + leased.query.mockResolvedValue({ + command: "LISTEN", + fields: [], + oid: 0, + rowCount: 0, + rows: [], + }); + const pool = { + ...mockPgClient(), + connect: vi.fn().mockResolvedValue(leased), + idleCount: 1, + options: { max: 4 }, + totalCount: 1, + } as unknown as Pool; + const driver = testPgDriver(pool, { schema: "river" }); + const abort = new AbortController(); + const ready = Promise.withResolvers(); + const iterator = driver.listen(["river_control"], abort.signal, () => + ready.resolve(undefined) + ); + try { + const first = iterator.next(); + await ready.promise; + leased.emit("notification", { + channel: "river.river_control", + payload: "first", + }); + await first; + for (let id = 0; id < 1_025; id++) { + leased.emit("notification", { + channel: "river.river_control", + payload: JSON.stringify({ action: "cancel", job_id: id }), + }); + } + await expect(iterator.next()).rejects.toThrow( + "notification buffer overflow" + ); + expect(leased.release).toHaveBeenCalledWith(true); + } finally { + abort.abort(); + await iterator.return(undefined); + } + }); + it("pings an idle listener and fails when a ping never answers", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); try { diff --git a/js/driver/pg/src/driver.ts b/js/driver/pg/src/driver.ts index 3287ad65c..01ca76d25 100644 --- a/js/driver/pg/src/driver.ts +++ b/js/driver/pg/src/driver.ts @@ -5,6 +5,7 @@ import type { DriverInsertResult, DriverRecord, InsertDriverOptions, + JobCancellationNotice, JobClaimOptions, JobClaimParams, JobClaimResult, @@ -861,7 +862,7 @@ export class PgRuntime implements RuntimeDriver< attemptedBy: string, signal: AbortSignal, ready?: () => void - ): AsyncGenerator<{ attemptedBy: string; id: bigint }> { + ): AsyncGenerator { return notifySql.jobCancellationSubscribe( this.#db, attemptedBy, diff --git a/js/driver/pg/src/sql/jobs.ts b/js/driver/pg/src/sql/jobs.ts index 37b91ecb6..e675eaa30 100644 --- a/js/driver/pg/src/sql/jobs.ts +++ b/js/driver/pg/src/sql/jobs.ts @@ -72,7 +72,7 @@ export async function jobCancelWithOptions( const sql = { with: ` WITH locked_job AS ( - SELECT id, queue, state, finalized_at + SELECT id, attempt, attempted_at, queue, state, finalized_at FROM ${jobTable} WHERE id = $1::bigint FOR UPDATE @@ -84,6 +84,8 @@ export async function jobCancelWithOptions( concat(coalesce($4::text, current_schema()), '.', $2::text), json_build_object( 'action', 'cancel', + 'attempt', attempt, + 'attempted_at', attempted_at, 'job_id', id, 'queue', queue )::text @@ -294,6 +296,7 @@ export async function jobGetCancelRequested( SELECT id::text AS id FROM ${db.table("river_job")} WHERE id = any($1::bigint[]) + AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' ORDER BY id @@ -822,19 +825,25 @@ export async function jobRetryWithOptions( SET state = 'available', max_attempts = CASE + WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN river_job.max_attempts WHEN river_job.attempt = river_job.max_attempts THEN river_job.max_attempts + 1 ELSE river_job.max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce($2::timestamptz, now()) + finalized_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN river_job.finalized_at ELSE NULL END, + metadata = CASE + WHEN jsonb_typeof(river_job.metadata) = 'object' + THEN river_job.metadata - 'cancel_attempted_at' + ELSE river_job.metadata + END, + scheduled_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN river_job.scheduled_at ELSE coalesce($2::timestamptz, now()) END FROM job_to_update WHERE river_job.id = job_to_update.id AND river_job.state != 'running' - AND NOT ( + AND ((jsonb_typeof(river_job.metadata) = 'object' AND river_job.metadata ? 'cancel_attempted_at') OR NOT ( river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) - ) + )) RETURNING river_job.* ) `, diff --git a/js/driver/pg/src/sql/notify.ts b/js/driver/pg/src/sql/notify.ts index 8f295b946..408ec41c0 100644 --- a/js/driver/pg/src/sql/notify.ts +++ b/js/driver/pg/src/sql/notify.ts @@ -3,7 +3,10 @@ import { LinkedAbortSignal, quoteIdentifier } from "riverqueue/unstable-driver"; import { Buffer } from "node:buffer"; import type { Notification as NodePgNotification } from "pg"; import { ConfigurationError } from "riverqueue"; -import type { RuntimeNotification } from "riverqueue/unstable-driver"; +import type { + JobCancellationNotice, + RuntimeNotification, +} from "riverqueue/unstable-driver"; import type { PgDatabase, PgQueryable } from "../database.js"; import { POSTGRES_IDENTIFIER_MAX_BYTES } from "../database.js"; import { @@ -15,7 +18,7 @@ import { PG_EXACT_TYPES } from "../exact-types.js"; import { abortablePromise, PgClientLease } from "../lease.js"; import type { PgNotification, PgOperationOptions } from "../types.js"; -/** Notifications buffered per connection; beyond it the oldest is dropped. */ +/** Notifications buffered per connection; overflow forces resubscription and recovery. */ const NOTIFICATION_QUEUE_CAPACITY = 1_024; /** Idle time after which a LISTEN connection is pinged, like River's Go notifier. */ @@ -212,7 +215,7 @@ export async function* jobCancellationSubscribe( attemptedBy: string, signal: AbortSignal, ready?: () => void -): AsyncGenerator<{ attemptedBy: string; id: bigint }> { +): AsyncGenerator { for await (const notification of listen( db, ["river_control"], @@ -236,7 +239,23 @@ export async function* jobCancellationSubscribe( } try { const id = exactCancellationID(notification.payload, value.job_id); - if (id !== null) yield { attemptedBy, id }; + if (id !== null) { + if ( + "attempt" in value && + typeof value.attempt === "number" && + Number.isSafeInteger(value.attempt) && + value.attempt > 0 && + "attempted_at" in value && + typeof value.attempted_at === "string" + ) { + yield { + attempt: value.attempt, + attemptedAt: Temporal.Instant.from(value.attempted_at), + attemptedBy, + id, + }; + } else yield { attemptedBy, id }; + } } catch { // Ignore malformed notifications; polling remains the durable path. } @@ -282,7 +301,7 @@ class NotificationQueue { } push(value: PgNotification): void { - if (this.#signal.aborted) return; + if (this.#signal.aborted || this.#failure !== undefined) return; if (this.#waiting !== undefined) { this.#waiting.resolve(value); this.#waiting = undefined; @@ -291,10 +310,12 @@ class NotificationQueue { const key = `${value.topic}\0${value.payload}`; if (this.#keys.has(key)) return; if (this.#values.length - this.#head >= NOTIFICATION_QUEUE_CAPACITY) { - const dropped = this.#values[this.#head++]; - if (dropped !== undefined) { - this.#keys.delete(`${dropped.topic}\0${dropped.payload}`); - } + this.fail( + new Error( + "River notification buffer overflow; resubscribe to recover missed notifications" + ) + ); + return; } this.#keys.add(key); this.#values.push(value); diff --git a/js/driver/sqlite/src/compatibility.test.ts b/js/driver/sqlite/src/compatibility.test.ts index e73c52129..1d413bbf1 100644 --- a/js/driver/sqlite/src/compatibility.test.ts +++ b/js/driver/sqlite/src/compatibility.test.ts @@ -633,7 +633,7 @@ describe("SqliteDriver compatibility with Go's riversqlite", () => { errors: null, errors_null: 1n, finalized_at: null, - metadata: cancelled, + metadata: { m: "x" }, scheduled_at: "2026-01-02 06:06:07.123", state: "available", }); diff --git a/js/driver/sqlite/src/driver.test.ts b/js/driver/sqlite/src/driver.test.ts index a9a41e008..3868059e2 100644 --- a/js/driver/sqlite/src/driver.test.ts +++ b/js/driver/sqlite/src/driver.test.ts @@ -326,15 +326,17 @@ describe("SqliteDriver", () => { const retried = await driver.jobRetryDetailed(inserted.id, { now }); expect(retried.status).toBe("retried"); - if (retried.status === "retried") + if (retried.status === "retried") { expect(retried.job.state).toBe("available"); + expect(retried.job.metadata).not.toHaveProperty("cancel_attempted_at"); + } expect((await driver.jobRetryDetailed(inserted.id, { now })).status).toBe( - "unchanged" + "retried" ); database .prepare( - "UPDATE river_job SET state = 'running', attempt = 1 WHERE id = ?" + "UPDATE river_job SET state = 'running', attempt = 1, attempted_at = '2026-01-02T03:04:05.678Z' WHERE id = ?" ) .run(inserted.id); expect(await driver.jobDelete(inserted.id)).toMatchObject({ @@ -346,17 +348,90 @@ describe("SqliteDriver", () => { // Like River for Go, a retry publishes no insert notification. expect(control.map(({ topic }) => topic)).toEqual(["river_control"]); expect(control[0]?.payload).toBe( - `{"action":"cancel","job_id":${inserted.id},"queue":"default"}` + `{"action":"cancel","attempt":0,"attempted_at":null,"job_id":${inserted.id},"queue":"default"}` ); }); + test("increments exhausted attempts at the exact retry timestamp", async () => { + const { database, driver } = await setup(); + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert({ + args: {}, + kind: "retry_boundary", + maxAttempts: 1, + scheduledAt: now, + state: "available", + }); + database + .prepare("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .run(job.id); + const retried = await driver.jobRetryDetailed(job.id, { now }); + expect(retried).toMatchObject({ + status: "retried", + job: { maxAttempts: 2 }, + }); + }); + + test("clears cancellation without reordering an already available retry", async () => { + const { database, driver } = await setup(); + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert({ + args: {}, + kind: "retry_available", + metadata: { cancel_attempted_at: now.toString(), preserved: true }, + scheduledAt: now.subtract({ hours: 1 }), + state: "available", + }); + database + .prepare("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .run(job.id); + const expectedMetadata = { ...job.metadata }; + delete expectedMetadata.cancel_attempted_at; + const retried = await driver.jobRetryDetailed(job.id, { now }); + expect(retried.status).toBe("retried"); + if (retried.status !== "retried") + throw new Error("expected a metadata update"); + expect(retried.job.metadata).toEqual(expectedMetadata); + expect(retried.job.scheduledAt.toString()).toBe(job.scheduledAt.toString()); + expect(retried.job.maxAttempts).toBe(job.maxAttempts); + expect((await driver.jobGet(job.id))?.metadata).toEqual(expectedMetadata); + }); + + test.each(["[not json", "null", '["cancel_attempted_at"]'])( + "retries jobs without changing malformed metadata: %s", + async (metadata) => { + const { database, driver } = await setup(); + const { job } = await driver.jobInsert({ + args: {}, + finalizedAt: Temporal.Now.instant(), + kind: "retry_metadata", + state: "discarded", + }); + database + .prepare("UPDATE river_job SET metadata = ? WHERE id = ?") + .run(metadata, job.id); + + expect(await driver.jobRetryDetailed(job.id)).toMatchObject({ + job: { id: job.id, state: "available" }, + status: "retried", + }); + expect( + database + .prepare( + "SELECT CASE WHEN typeof(metadata) = 'text' THEN metadata ELSE json(metadata) END AS metadata FROM river_job WHERE id = ?" + ) + .get(job.id)?.metadata + ).toBe(metadata); + } + ); + test("notifies a running job's cancellation in the cancelling transaction, like Go", async () => { const { database, driver } = await setup(); const inserted = (await driver.jobInsert({ args: {}, kind: "control_job" })) .job; database .prepare( - "UPDATE river_job SET state = 'running', attempt = 1 WHERE id = ?" + "UPDATE river_job SET state = 'running', attempt = 1, attempted_at = '2026-01-02T03:04:05.678Z' WHERE id = ?" ) .run(inserted.id); const cancels = () => @@ -385,7 +460,7 @@ describe("SqliteDriver", () => { }); expect(cancels()).toEqual([ { - payload: `{"action":"cancel","job_id":${inserted.id},"queue":"default"}`, + payload: `{"action":"cancel","attempt":1,"attempted_at":"2026-01-02T03:04:05.678Z","job_id":${inserted.id},"queue":"default"}`, topic: "river_control", }, ]); diff --git a/js/driver/sqlite/src/driver.ts b/js/driver/sqlite/src/driver.ts index 98a0891c5..f992a8906 100644 --- a/js/driver/sqlite/src/driver.ts +++ b/js/driver/sqlite/src/driver.ts @@ -623,7 +623,7 @@ export class SqliteRuntime implements RuntimeDriver< insertNotification( database, NOTIFICATION_TOPIC_CONTROL, - `{"action":"cancel","job_id":${updated.id},"queue":${JSON.stringify(updated.queue)}}` + `{"action":"cancel","attempt":${updated.attempt},"attempted_at":${JSON.stringify(updated.attemptedAt?.toString() ?? null)},"job_id":${updated.id},"queue":${JSON.stringify(updated.queue)}}` ); } return { @@ -1438,21 +1438,29 @@ export class SqliteRuntime implements RuntimeDriver< const raw = resultStatement( database, ` + WITH retry_params AS (SELECT ? AS now) UPDATE river_job SET state = 'available', max_attempts = CASE + WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = ? + finalized_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, + metadata = CASE + WHEN ${invalidJsonTextSql("metadata")} THEN metadata + ELSE jsonb_remove(metadata, '$.cancel_attempted_at') + END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END WHERE id = ? AND state != 'running' - AND (state != 'available' OR scheduled_at > ?) + AND ((CASE WHEN ${invalidJsonTextSql("metadata")} THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR NOT (state = 'available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING ${JOB_COLUMNS} ` - ).get(timestamp, id, timestamp); + ).get(timestamp, id); const updated = raw === undefined ? null : decodeJobRowPartial(raw).job; const job = updated ?? getJob(database, id, { partial: true }); if (job === null) return { status: "not_found" }; diff --git a/js/etc/riverqueue.unstable-driver.api.md b/js/etc/riverqueue.unstable-driver.api.md index 9830babf3..7c55c41d6 100644 --- a/js/etc/riverqueue.unstable-driver.api.md +++ b/js/etc/riverqueue.unstable-driver.api.md @@ -541,6 +541,8 @@ export interface JobArgsTransformPlugin extends RiverPlugin { ```ts /** Remote cancellation of an attempt currently owned by this process. */ export interface JobCancellationNotice { + readonly attempt?: number; + readonly attemptedAt?: Temporal.Instant; readonly attemptedBy: string; readonly id: bigint; } diff --git a/js/src/driver.ts b/js/src/driver.ts index 7134cc221..98af925b7 100644 --- a/js/src/driver.ts +++ b/js/src/driver.ts @@ -303,6 +303,8 @@ export function jobCompletionKey( /** Remote cancellation of an attempt currently owned by this process. */ export interface JobCancellationNotice { + readonly attempt?: number; + readonly attemptedAt?: Temporal.Instant; readonly attemptedBy: string; readonly id: bigint; } diff --git a/js/src/pilot-runtime.test.ts b/js/src/pilot-runtime.test.ts index 1d0372194..de64ed451 100644 --- a/js/src/pilot-runtime.test.ts +++ b/js/src/pilot-runtime.test.ts @@ -1067,11 +1067,15 @@ describe("pilot services and teardown", () => { }); const run = await client.start(); - const first = await timer.waitFor(({ kind }) => kind === "delay"); + const first = await timer.waitFor( + ({ kind, ms }) => kind === "delay" && ms !== 2_000 + ); expect(runs).toBe(1); await timer.advance(first.ms); await waitUntil(() => runs === 2); - const second = await timer.waitFor(({ kind }) => kind === "delay"); + const second = await timer.waitFor( + ({ kind, ms }) => kind === "delay" && ms !== 2_000 + ); // Backoff grows while the service keeps failing quickly. expect(second.ms).toBeGreaterThan(first.ms); await timer.advance(second.ms); @@ -1079,7 +1083,9 @@ describe("pilot services and teardown", () => { // A failure after a healthy minute starts backoff over. await timer.advance(60_000); longRun.reject(new Error("service failed after a while")); - const third = await timer.waitFor(({ kind }) => kind === "delay"); + const third = await timer.waitFor( + ({ kind, ms }) => kind === "delay" && ms !== 2_000 + ); expect(third.ms).toBe(first.ms); await timer.advance(third.ms); await waitUntil(() => runs === 4); diff --git a/js/src/runtime.test.ts b/js/src/runtime.test.ts index c44c05346..751914b61 100644 --- a/js/src/runtime.test.ts +++ b/js/src/runtime.test.ts @@ -618,6 +618,7 @@ describe("Client runtime", () => { clientId: "runtime-test", completionBatchSize: 1, leaderElectionDisabled: true, + queueControlPollInterval: { milliseconds: 1234 }, queues: { default: { maxWorkers: 1, pollInterval: { milliseconds: 10_000 } }, }, @@ -633,7 +634,9 @@ describe("Client runtime", () => { expect(driver.completionFailures).toBe(0); expect(driver.completionCalls).toBe(3); expect(driver.completions[0]?.kind).toBe("complete"); - expect(timer.delays.slice(-2)).toEqual([1_000, 2_000]); + expect(timer.delays.filter((delay) => delay !== 1234).slice(-2)).toEqual([ + 1_000, 2_000, + ]); expect(timer.activeTimeouts).toBe(0); expect(run.state).toBe("stopped"); }); @@ -741,8 +744,8 @@ describe("Client runtime", () => { const working = live.size; await run.stop(); - // Only the queue loop's link, on the claim signal and its generation's. - expect(working).toBe(2); + // Only the queue loop's link and the runtime's cancellation link remain. + expect(working).toBe(4); expect(live.size).toBe(0); // Nothing per attempt or completion uses `AbortSignal.any`. expect(any).not.toHaveBeenCalled(); @@ -858,6 +861,7 @@ describe("Client runtime", () => { onMetric: (metric) => void metrics.push(metric), }, leaderElectionDisabled: true, + queueControlPollInterval: { milliseconds: 1234 }, queues: { default: { maxWorkers: 1, pollInterval: { milliseconds: 10_000 } }, }, @@ -875,7 +879,9 @@ describe("Client runtime", () => { name: "job_completion_requeued", }); // Every failed attempt backs off, including the last before a requeue. - expect(timer.delays).toEqual([1_000, 2_000, 4_000, 1_000, 2_000]); + expect(timer.delays.filter((delay) => delay !== 1234)).toEqual([ + 1_000, 2_000, 4_000, 1_000, 2_000, + ]); expect(events).toEqual(["job_started", "job_completed"]); }); @@ -4330,7 +4336,9 @@ describe("Client runtime database faults", () => { await run.stop(); expect(run.state).toBe("stopped"); - expect(timer.delays).toEqual([250, 500, 1_000]); + expect(timer.delays.filter((delay) => delay !== 2_000)).toEqual([ + 250, 500, 1_000, + ]); expect( logs.map(({ attributes, level }) => [ level, @@ -4377,7 +4385,7 @@ describe("Client runtime database faults", () => { await waitUntil(() => events.includes("queue_paused")); expect(run.diagnostics.queues.default?.paused).toBe(true); expect(run.state).toBe("running"); - expect(timer.delays).toEqual([250, 500]); + expect(timer.delays.filter((delay) => delay !== 1)).toEqual([250, 500]); expect(logs.map(({ message }) => message)).toEqual([ "River queue control poll failed; retrying after backoff", "River queue control poll failed; retrying after backoff", @@ -4388,6 +4396,14 @@ describe("Client runtime database faults", () => { it("cancels a job whose cancellation arrived while its claim was in flight", async () => { const driver = new FakeRuntimeDriver(); driver.claim = [fakeJob()]; + driver.listRows = [ + { + ...fakeJob(), + state: "running", + attemptedBy: ["runtime-test"], + metadata: { cancel_attempted_at: null }, + }, + ]; const release = Promise.withResolvers(); driver.claimGate = release.promise; const handled = Promise.withResolvers(); @@ -4516,7 +4532,7 @@ describe("Client runtime database faults", () => { () => driver.completions.length === 1 && subscriptions === 3 ); expect(run.state).toBe("running"); - expect(timer.delays).toEqual([250, 250]); + expect(timer.delays.filter((delay) => delay !== 2_000)).toEqual([250, 250]); expect(logs.map(({ level, message }) => [level, message])).toEqual([ [ "warn", @@ -4551,6 +4567,14 @@ describe("Client runtime database faults", () => { }, }); driver.claim = [fakeJob()]; + driver.listRows = [ + { + ...fakeJob(), + state: "running", + attemptedBy: ["runtime-test"], + metadata: { cancel_attempted_at: null }, + }, + ]; const logs: LogEntry[] = []; const timer = new FakeTimer(); const definition = defineJob({ kind: "test" }); @@ -4582,6 +4606,426 @@ describe("Client runtime database faults", () => { await run.stop(); }); + it("delivers identified cancellations during blocked reads and ignores older attempts", async () => { + const driver = new FakeRuntimeDriver(); + const started = Promise.withResolvers(); + const readStarted = Promise.withResolvers(); + const releaseRead = Promise.withResolvers(); + const releaseWorker = Promise.withResolvers(); + const deliver = Promise.withResolvers(); + const attemptedAt = Temporal.Instant.from("2026-01-02T03:04:05.6789Z"); + driver.claim = [{ ...fakeJob(), attempt: 2, attemptedAt }]; + Object.assign(driver, { + jobGetCancelRequested: () => { + readStarted.resolve(undefined); + return releaseRead.promise; + }, + }); + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + await deliver.promise; + // Legacy hint starts a read, simulating a saturated pool. + yield { + topic: "control", + payload: '{"action":"cancel","job_id":101,"queue":"default"}', + }; + await readStarted.promise; + const workerSignal = await started.promise; + for (const identity of [ + { attempt: 1, attempted_at: attemptedAt.toString() }, + { + attempt: 2, + attempted_at: attemptedAt.subtract({ seconds: 1 }).toString(), + }, + ]) { + yield { + topic: "control", + payload: JSON.stringify({ + action: "cancel", + job_id: 101, + queue: "default", + ...identity, + }), + }; + expect(workerSignal.aborted).toBe(false); + } + yield { + topic: "control", + payload: JSON.stringify({ + action: "cancel", + job_id: 101, + queue: "default", + attempt: 2, + attempted_at: attemptedAt.toString(), + }), + }; + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + queues: { default: { maxWorkers: 1 } }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + started.resolve(signal); + await releaseWorker.promise; + signal.throwIfAborted(); + } + ), + }); + const run = await client.start(); + try { + const signal = await started.promise; + deliver.resolve(undefined); + await waitUntil(() => signal.aborted); + expect(signal.reason).toMatchObject({ name: "JobCancelledError" }); + } finally { + deliver.resolve(undefined); + releaseRead.resolve([]); + releaseWorker.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + + for (const failFirst of [false, true]) { + it(`only checks cancellations on hints and retries failures=${failFirst}`, async () => { + const driver = new FakeRuntimeDriver(); + const deliver = Promise.withResolvers(); + const started = Promise.withResolvers(); + const finish = Promise.withResolvers(); + let checks = 0; + Object.assign(driver, { + jobGetCancelRequested: () => { + checks++; + if (failFirst && checks === 1) + throw new Error("temporary read failure"); + return []; + }, + }); + driver.claim = [fakeJob()]; + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + await deliver.promise; + yield { + topic: "control", + payload: '{"action":"cancel","job_id":101,"queue":"default"}', + }; + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + logger: recordingLogger([]), + queues: { default: { maxWorkers: 1 } }, + workers: new Workers().add(defineJob({ kind: "test" }), async () => { + started.resolve(undefined); + await finish.promise; + }), + }); + const timer = new FakeTimer(); + overrideRuntimeTiming(client, { random: () => 0.5, timer }); + const run = await client.start(); + try { + await started.promise; + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(checks).toBe(0); + deliver.resolve(undefined); + await waitUntil(() => checks === (failFirst ? 2 : 1)); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(checks).toBe(failFirst ? 2 : 1); + } finally { + deliver.resolve(undefined); + finish.resolve(undefined); + await run.stop(); + } + }); + } + + for (const stream of ["runtime", "remote", "local"] as const) { + for (const duringClaim of [false, true]) { + it(`revalidates a delayed ${stream} cancellation ${duringClaim ? "during a retry claim" : "against a running retry"}`, async () => { + const driver = new FakeRuntimeDriver(); + const deliver = Promise.withResolvers(); + const handled = Promise.withResolvers(); + const releaseClaim = Promise.withResolvers(); + const finish = Promise.withResolvers(); + const started = Promise.withResolvers(); + let checks = 0; + Object.assign(driver, { + jobGetCancelRequested: () => { + checks++; + return []; + }, + }); + driver.claim = [{ ...fakeJob(), attempt: 2 }]; + if (duringClaim) driver.claimGate = releaseClaim.promise; + if (stream === "runtime") { + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + await deliver.promise; + yield { + topic: "control", + payload: '{"action":"cancel","job_id":101,"queue":"default"}', + }; + handled.resolve(undefined); + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + } else if (stream === "remote") { + Object.defineProperty(driver, "runtimeNotificationSubscribe", { + value: undefined, + }); + Object.defineProperty(driver, "jobCancellationSubscribe", { + value: async function* (attemptedBy: string, signal: AbortSignal) { + await deliver.promise; + yield { attemptedBy, id: 101n }; + handled.resolve(undefined); + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { + once: true, + }) + ); + }, + }); + } + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + queues: { + default: { maxWorkers: 1, pollInterval: { milliseconds: 10_000 } }, + }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + started.resolve(signal); + await finish.promise; + signal.throwIfAborted(); + } + ), + }); + const run = await client.start(); + try { + await waitUntil(() => driver.claimed.length === 1); + if (!duringClaim) await started.promise; + if (stream === "local") { + // The cancelled row returned by a delayed commit belongs to attempt A. + driver.cancelled = { ...fakeJob(), attempt: 1, state: "running" }; + await client.jobs.cancel(101n); + } else { + deliver.resolve(undefined); + await handled.promise; + } + releaseClaim.resolve(undefined); + const signal = await started.promise; + await waitUntil(() => checks > 0); + expect(signal.aborted).toBe(false); + finish.resolve(undefined); + await waitUntil(() => driver.completions.length === 1); + expect(driver.completions[0]?.kind).toBe("complete"); + } finally { + deliver.resolve(undefined); + releaseClaim.resolve(undefined); + finish.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + } + } + + for (const batched of [false, true]) { + for (const phase of ["running", "claim", "draining"] as const) { + it(`recovers a lost cancellation while ${phase}, using ${batched ? "batches" : "rows"}`, async () => { + const driver = new FakeRuntimeDriver(); + const reconnect = Promise.withResolvers(); + const reconnected = Promise.withResolvers(); + let subscriptions = 0; + // eslint-disable-next-line require-yield -- a lost notice is recovered after resubscription + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + if (subscriptions++ === 0) { + await reconnect.promise; + throw new Error("listener disconnected"); + } + reconnected.resolve(undefined); + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + const claim = Promise.withResolvers(); + const started = Promise.withResolvers(); + let marked = false; + let failures = 0; + const read = (ids: readonly bigint[]) => { + if (!marked) return []; + if (failures++ === 0) + throw new Error("temporary cancellation read failure"); + return ids; + }; + if (batched) { + Object.assign(driver, { jobGetCancelRequested: read }); + } else { + Object.assign(driver, { + jobGet: (id: bigint) => + read([id]).length === 0 + ? null + : { + ...fakeJob(), + id, + attemptedBy: ["runtime-test"], + state: "running", + metadata: { cancel_attempted_at: null }, + }, + }); + } + driver.claim = [fakeJob()]; + if (phase === "claim") driver.claimGate = claim.promise; + let handlerSignal: AbortSignal | undefined; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + logger: recordingLogger([]), + queues: { default: { maxWorkers: 1 } }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + handlerSignal = signal; + started.resolve(undefined); + await new Promise((resolve) => { + if (signal.aborted) resolve(); + else + signal.addEventListener("abort", () => resolve(), { + once: true, + }); + }); + signal.throwIfAborted(); + } + ), + }); + overrideRuntimeTiming(client, { + random: () => 0.5, + timer: new FakeTimer(), + }); + const run = await client.start(); + try { + if (phase === "claim") { + await waitUntil(() => driver.claimRequests.length > 0); + } else { + await started.promise; + } + const stopping = phase === "draining" ? run.stop() : undefined; + marked = true; + reconnect.resolve(undefined); + await reconnected.promise; + claim.resolve(undefined); + await waitUntil(() => driver.completions.length === 1); + if (phase !== "claim") + expect(handlerSignal?.reason).toMatchObject({ + name: "JobCancelledError", + }); + expect(driver.completions[0]?.kind).toBe("cancel"); + expect(failures).toBeGreaterThanOrEqual(2); + await stopping; + await run.stop(); + } finally { + reconnect.resolve(undefined); + claim.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + } + } + + for (const batched of [false, true]) { + it(`does not apply a stale ${batched ? "batch" : "row"} cancellation result to a new attempt`, async () => { + const driver = new FakeRuntimeDriver(); + const firstRead = Promise.withResolvers(); + const readStarted = Promise.withResolvers(); + const finishFirst = Promise.withResolvers(); + const finishSecond = Promise.withResolvers(); + const signals: AbortSignal[] = []; + let reads = 0; + const read = async () => { + reads++; + if (reads !== 1) return false; + readStarted.resolve(undefined); + await firstRead.promise; + return true; + }; + if (batched) { + Object.assign(driver, { + jobGetCancelRequested: async () => ((await read()) ? [7n] : []), + }); + } else { + Object.assign(driver, { + jobGet: async () => + (await read()) + ? { + ...fakeJob(), + attemptedBy: ["runtime-test"], + state: "running", + metadata: { cancel_attempted_at: null }, + } + : null, + }); + } + driver.claim = [fakeJob()]; + const client = new Client(driver, { + clientId: "runtime-test", + pollOnly: true, + completionBatchSize: 1, + leaderElectionDisabled: true, + queues: { + default: { + fetchCooldown: { milliseconds: 1 }, + maxWorkers: 1, + pollInterval: { milliseconds: 1 }, + }, + }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + signals.push(signal); + await (signals.length === 1 + ? finishFirst.promise + : finishSecond.promise); + } + ), + }); + overrideRuntimeTiming(client, { + random: () => 0.5, + timer: new FakeTimer(), + }); + const run = await client.start(); + try { + await readStarted.promise; + finishFirst.resolve(undefined); + await waitUntil(() => driver.completions.length === 1); + driver.claim = [fakeJob()]; + await waitUntil(() => signals.length === 2); + firstRead.resolve(undefined); + await waitUntil(() => reads >= 2); + expect(signals[1]?.aborted).toBe(false); + finishSecond.resolve(undefined); + await run.stop(); + } finally { + firstRead.resolve(undefined); + finishFirst.resolve(undefined); + finishSecond.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + } + for (const [stream, read] of [ ["remote cancellation", "each row"], ["runtime notification", "each row"], diff --git a/js/src/runtime.ts b/js/src/runtime.ts index 1be2c83ad..fed982e24 100644 --- a/js/src/runtime.ts +++ b/js/src/runtime.ts @@ -3,7 +3,7 @@ import type { QueueRow, RuntimeDriver } from "./driver.js"; import { LifecycleError, ValidationError } from "./errors.js"; import type { RiverEvent } from "./events.js"; import type { RiverHooks } from "./extensions.js"; -import { unrefTimeout } from "./internal/abort.js"; +import { LinkedAbortSignal, unrefTimeout } from "./internal/abort.js"; import { SYSTEM_TIMER, type RuntimeTimer } from "./internal/backoff.js"; import { EventLoopDelayMonitor } from "./internal/event-loop-delay-monitor.js"; import type { JobRow } from "./job.js"; @@ -211,11 +211,12 @@ export class RuntimeController { readonly ready: Promise; /** - * Ends polling for cancellations once producers drained, so a job can + * Ends notification delivery and cancellation checks once producers drain, so a job can * still be cancelled while its queue drains. */ - readonly #cancellationPollAbort = new AbortController(); - /** How often a runtime without notifications checks for cancellations. */ + readonly #cancellationAbort = new AbortController(); + readonly #cancellationLink: LinkedAbortSignal; + /** How often the runtime checks for durable cancellation requests. */ readonly #cancellationPollIntervalMs: number; readonly #claimAbort = new AbortController(); readonly #client: Client; @@ -328,7 +329,14 @@ export class RuntimeController { transformJobArgs: (row) => this.#runner.transformJobArgs(row), workerRetryPolicy: (kind) => this.#runner.workerRetryPolicy(kind), }); + this.#cancellationLink = new LinkedAbortSignal([ + this.#cancellationAbort.signal, + this.#runAbort.signal, + ]); + const cancellationSignal = this.#cancellationLink.signal; this.#runner = new AttemptRunner(context, { + cancellationSignal, + cancellationRetryIntervalMs: settings.queueControlPollIntervalMs, completions: this.#completions, errorHandler: settings.errorHandler, hooks: settings.hooks, @@ -401,6 +409,7 @@ export class RuntimeController { ? [] : serviceList(binding.pilot.services(), "services()"); this.#notifications = new NotificationPump(context, { + signal: cancellationSignal, leaderResigned: () => this.#services?.leaderResigned(), producer: this.#producer, resignLeadership: () => this.#services?.resignLeadership(), @@ -618,10 +627,11 @@ export class RuntimeController { new LifecycleError("River runtime is stopping") ); await this.#producer.drainAll(); - this.#cancellationPollAbort.abort( + this.#cancellationAbort.abort( new LifecycleError("River runtime is stopping") ); await Promise.allSettled(this.#tasks); + this.#cancellationLink[Symbol.dispose](); this.#servicesAbort.abort(new LifecycleError("River runtime is stopping")); const [completionClose] = await Promise.allSettled([ this.#completions.close(), @@ -664,15 +674,14 @@ export class RuntimeController { } async #initialize(): Promise { - // Like River for Go's client, learn whether the database delivers - // notifications before anything starts. Without them, running jobs - // learn of cancellations by polling until every producer has drained. + // Only clients without notifications poll. Notification hints and reconnects + // request checks independently, until every producer has drained. const listen = await this.#listens(); if (!listen) { this.#trackTask( this.#runner.pollCancellations( this.#cancellationPollIntervalMs, - this.#cancellationPollAbort.signal + this.#cancellationAbort.signal ) ); } diff --git a/js/src/runtime/attempt-runner.ts b/js/src/runtime/attempt-runner.ts index 7f2a36c6e..c0bfee966 100644 --- a/js/src/runtime/attempt-runner.ts +++ b/js/src/runtime/attempt-runner.ts @@ -75,6 +75,8 @@ import { /** Configuration for an {@link AttemptRunner}. */ export interface AttemptRunnerOptions { + readonly cancellationSignal: AbortSignal; + readonly cancellationRetryIntervalMs: number; readonly completions: CompletionPipeline; readonly errorHandler: RiverErrorHandler | undefined; /** Client-level hooks, plugin hooks first. */ @@ -108,10 +110,18 @@ const CANCEL_POLL_TIMEOUT_MS = 10_000; export class AttemptRunner { readonly #attemptExecutor = new AttemptExecutor(); /** In-flight attempts by `attemptKey`. */ - readonly #attempts = new Map(); + readonly #attempts = new Map< + string, + { readonly controller: AbortController; readonly row: JobRow } + >(); + #cancellationCheck: Promise | undefined; + #cancellationCheckPending = false; + #cancellationGeneration = 0; + readonly #cancellationRetryIntervalMs: number; + readonly #cancellationSignal: AbortSignal; /** * For each claim in flight, the `attemptKey`s of cancellations that - * found no attempt running here. + * need a durable check once the claim registers its attempts. */ readonly #claimCancellations = new Set>(); readonly #completions: CompletionPipeline; @@ -131,6 +141,8 @@ export class AttemptRunner { readonly #workers: Workers; constructor(context: RuntimeContext, options: AttemptRunnerOptions) { + this.#cancellationRetryIntervalMs = options.cancellationRetryIntervalMs; + this.#cancellationSignal = options.cancellationSignal; this.#completions = options.completions; this.#context = context; this.#errorHandler = options.errorHandler; @@ -153,18 +165,30 @@ export class AttemptRunner { /** Abort every in-flight attempt with `reason`. */ abortAttempts(reason: unknown): void { - for (const controller of this.#attempts.values()) controller.abort(reason); + for (const { controller } of this.#attempts.values()) + controller.abort(reason); } - /** Cancel the attempt of job `id` that `attemptedBy` owns, if it runs here. */ - cancelAttempt(id: bigint, attemptedBy: string): void { + /** Cancel a matching execution, or reconcile a legacy hint against durable state. */ + cancelAttempt( + id: bigint, + attemptedBy: string, + identity?: { attempt: number; attemptedAt: Temporal.Instant } + ): void { const key = attemptKey(id, attemptedBy); - const attempt = this.#attempts.get(key); - if (attempt !== undefined) { - attempt.abort(new JobCancelledError(id)); + const active = this.#attempts.get(key); + if (active !== undefined && identity !== undefined) { + if ( + active.row.attempt === identity.attempt && + active.row.attemptedAt?.epochNanoseconds === + identity.attemptedAt.epochNanoseconds + ) { + active.controller.abort(new JobCancelledError(id)); + } return; } for (const pending of this.#claimCancellations) pending.add(key); + if (this.#attempts.has(key)) void this.#checkCancellations(); } /** @@ -178,30 +202,40 @@ export class AttemptRunner { const clientId = this.#context.clientId; const driver = this.#context.driver; const requested = driver.jobGetCancelRequested?.bind(driver); - const ids = this.#ownAttemptIds(); + const attempts = this.#ownAttempts(); if (requested === undefined) { // A backend without the batched read checks each job's row. - for (const id of ids) { + for (const [id, controller] of attempts) { signal.throwIfAborted(); - const job = await raceWithAbort(driver.jobGet(id), signal); - if ( - job?.state === "running" && - job.attemptedBy.at(-1) === clientId && - job.metadata.cancel_attempted_at !== undefined - ) { - this.cancelAttempt(id, clientId); + const timeout = this.#context.timer.timeout( + CANCEL_POLL_TIMEOUT_MS, + () => new Error("River's cancellation check timed out") + ); + const link = new LinkedAbortSignal([signal, timeout.signal]); + try { + const job = await raceWithAbort(driver.jobGet(id), link.signal); + if ( + job?.state === "running" && + job.attemptedBy.at(-1) === clientId && + Object.hasOwn(job.metadata, "cancel_attempted_at") + ) { + controller.abort(new JobCancelledError(id)); + } + } finally { + link[Symbol.dispose](); + timeout.dispose(); } } return; } - await this.#cancelRequested(requested, ids, signal); + await this.#cancelRequested(requested, attempts, signal); } /** * Cancel this client's running attempts whose jobs have a cancellation - * request, checking every `intervalMs` until `signal` aborts, like River - * for Go's producers without a notifier. Each check reads the attempts' - * jobs in batches of 1,000 IDs, each bounded by ten seconds, and reads + * request, checking every `intervalMs` until `signal` aborts. Each check + * reads the attempts' jobs in batches of 1,000 IDs (or one row at a time + * without a batched read), each bounded by ten seconds, and reads * nothing while no attempt runs. A failed check is logged and the next * one tries again. */ @@ -209,9 +243,6 @@ export class AttemptRunner { intervalMs: number, signal: AbortSignal ): Promise { - const driver = this.#context.driver; - const requested = driver.jobGetCancelRequested?.bind(driver); - if (requested === undefined) return; for (;;) { try { await this.#context.timer.delay(intervalMs, signal); @@ -219,7 +250,7 @@ export class AttemptRunner { return; } try { - await this.#cancelRequested(requested, this.#ownAttemptIds(), signal); + await this.#checkCancellations(); } catch (error: unknown) { if (signal.aborted) return; this.#context.logger.error( @@ -230,19 +261,27 @@ export class AttemptRunner { } } + /** Reconcile a reconnect or overflow, including claims not yet registered. */ + requestCancellationCheck(): void { + this.#cancellationGeneration++; + void this.#checkCancellations(); + } + /** Cancel this process's attempts of a job just cancelled through this client. */ cancelLocal(job: JobRow): void { - const keyPrefix = `${job.id.toString(10)}:`; - for (const [key, controller] of this.#attempts) { - if (key.startsWith(keyPrefix)) { - controller.abort(new JobCancelledError(job.id)); - } - } - // The job may be in a claim still in flight here. - const key = attemptKey(job.id, this.#context.clientId); - if (!this.#attempts.has(key)) { - for (const pending of this.#claimCancellations) pending.add(key); + const attempt = this.#attempts.get( + attemptKey(job.id, this.#context.clientId) + ); + if ( + attempt !== undefined && + attempt.row.attempt === job.attempt && + attempt.row.attemptedAt?.epochNanoseconds === + job.attemptedAt?.epochNanoseconds + ) { + attempt.controller.abort(new JobCancelledError(job.id)); + return; } + this.cancelAttempt(job.id, this.#context.clientId); } /** Diagnostics reported by each work executor in use. */ @@ -348,21 +387,21 @@ export class AttemptRunner { * an unknown kind, its attempt fails before hooks or middleware run, so the * error handler sees it and the retry policy retries or discards it. * - * `cancelled` starts the attempt already cancelled, for a job whose - * cancellation arrived while the claim that took it was in flight. + * `checkCancellation` confirms a notice received while the claim was in flight + * before starting work. A stale notice cannot cancel a retried attempt. */ async run( row: JobRow, releaseCapacity: () => void, decodeError?: Error, - cancelled = false + checkCancellation = false ): Promise { const registration = workerRegistration(this.#workers, row.kind); const attemptController = new AbortController(); const timeoutController = new AbortController(); const key = attemptKey(row.id, this.#context.clientId); - this.#attempts.set(key, attemptController); - if (cancelled) attemptController.abort(new JobCancelledError(row.id)); + this.#attempts.set(key, { controller: attemptController, row }); + if (checkCancellation) await this.#checkCancellations(); // Like Go, a remote cancellation or shutdown is tracked apart from the // attempt's own timeout: a remote cancellation that arrives after the // timeout still decides how the attempt ends. @@ -479,7 +518,8 @@ export class AttemptRunner { timers.dispose(); this.#peers?.abandon(metadataState); metadataState.active = false; - this.#attempts.delete(key); + if (this.#attempts.get(key)?.controller === attemptController) + this.#attempts.delete(key); } } @@ -495,18 +535,20 @@ export class AttemptRunner { /** * Keep cancellations of jobs not worked here that arrive until `end()`, - * for one claim, like River for Go's producer does for one fetch: an - * attempt the claim starts for such a job starts cancelled. Other - * cancellations are discarded when the claim ends. + * for one claim. Confirm them against durable state before starting work. + * Other hints are discarded when the claim ends. */ watchCancellations(): { cancelled(id: bigint): boolean; end(): void; } { + const generation = this.#cancellationGeneration; const pending = new Set(); this.#claimCancellations.add(pending); return { - cancelled: (id) => pending.has(attemptKey(id, this.#context.clientId)), + cancelled: (id) => + generation !== this.#cancellationGeneration || + pending.has(attemptKey(id, this.#context.clientId)), end: () => { this.#claimCancellations.delete(pending); }, @@ -872,15 +914,15 @@ export class AttemptRunner { } /** - * Cancel the attempts among `ids` whose jobs have a cancellation request, + * Cancel the snapshotted attempts whose jobs have a cancellation request, * reading them in batches of 1,000 IDs, each bounded by ten seconds. */ async #cancelRequested( requested: NonNullable, - ids: readonly bigint[], + attempts: ReadonlyMap, signal: AbortSignal ): Promise { - const clientId = this.#context.clientId; + const ids = [...attempts.keys()]; for (let start = 0; start < ids.length; start += CANCEL_POLL_BATCH_SIZE) { const timeout = this.#context.timer.timeout( CANCEL_POLL_TIMEOUT_MS, @@ -894,7 +936,11 @@ export class AttemptRunner { }), link.signal ); - for (const id of cancelled) this.cancelAttempt(id, clientId); + // Never look up a fresh attempt by ID after the query: a retry may + // already be running with the same ID and client owner. + for (const id of cancelled) { + attempts.get(id)?.abort(new JobCancelledError(id)); + } } finally { link[Symbol.dispose](); timeout.dispose(); @@ -902,12 +948,53 @@ export class AttemptRunner { } } - /** The job IDs of the attempts this client runs itself. */ - #ownAttemptIds(): bigint[] { + /** Coalesce notification checks while preserving notices received during a read. */ + #checkCancellations(): Promise { + this.#cancellationCheckPending = true; + if (this.#cancellationCheck !== undefined) return this.#cancellationCheck; + const check = (async () => { + try { + while (this.#cancellationCheckPending) { + this.#cancellationCheckPending = false; + try { + await this.recoverCancellations(this.#cancellationSignal); + } catch (error: unknown) { + if (this.#cancellationSignal.aborted) return; + this.#context.logger.warn( + "River could not confirm a cancellation notification", + { error: describeError(error) } + ); + try { + await this.#context.timer.delay( + this.#cancellationRetryIntervalMs, + this.#cancellationSignal + ); + } catch { + return; + } + this.#cancellationCheckPending = true; + } + } + } finally { + this.#cancellationCheck = undefined; + } + })(); + this.#cancellationCheck = check; + this.#context.trackTask(check); + return check; + } + + /** Snapshot the executions themselves, so late results cannot cancel retries. */ + #ownAttempts(): Map { const suffix = `:${this.#context.clientId}`; - return [...this.#attempts.keys()] - .filter((key) => key.endsWith(suffix)) - .map((key) => BigInt(key.slice(0, -suffix.length))); + return new Map( + [...this.#attempts] + .filter(([key]) => key.endsWith(suffix)) + .map(([key, { controller }]) => [ + BigInt(key.slice(0, -suffix.length)), + controller, + ]) + ); } async #reportStuck( diff --git a/js/src/runtime/notification-payloads.ts b/js/src/runtime/notification-payloads.ts index 080b64cd2..72c8f8874 100644 --- a/js/src/runtime/notification-payloads.ts +++ b/js/src/runtime/notification-payloads.ts @@ -19,6 +19,33 @@ export function notificationQueue(payload: string): string | null { return null; } +/** Optional execution identity lets cancellation bypass a blocked database read. */ +export function notificationCancellationAttempt( + payload: string +): { attempt: number; attemptedAt: Temporal.Instant } | undefined { + try { + const value: unknown = JSON.parse(payload); + if ( + value !== null && + typeof value === "object" && + "attempt" in value && + typeof value.attempt === "number" && + Number.isSafeInteger(value.attempt) && + value.attempt > 0 && + "attempted_at" in value && + typeof value.attempted_at === "string" + ) { + return { + attempt: value.attempt, + attemptedAt: Temporal.Instant.from(value.attempted_at), + }; + } + } catch { + // Old or malformed identities require durable reconciliation. + } + return undefined; +} + /** The job ID a control notification asks to cancel, if any. */ export function notificationCancellation(payload: string): bigint | null { if (!/"action"\s*:\s*"cancel"/.test(payload)) return null; diff --git a/js/src/runtime/notification-pump.conformance.test.ts b/js/src/runtime/notification-pump.conformance.test.ts index 05b7c2c88..d657fa49c 100644 --- a/js/src/runtime/notification-pump.conformance.test.ts +++ b/js/src/runtime/notification-pump.conformance.test.ts @@ -106,11 +106,22 @@ describe.each(["postgres", "sqlite"] as const)( }, } as unknown as QueueProducer; const runner = { - cancelAttempt: (id: bigint, owner: string) => { - effects.push(["cancel", id, owner]); + cancelAttempt: ( + id: bigint, + owner: string, + identity: { attempt: number; attemptedAt: Temporal.Instant } + ) => { + effects.push([ + "cancel", + id, + owner, + identity.attempt, + identity.attemptedAt.toString(), + ]); }, } as unknown as AttemptRunner; const pump = new NotificationPump(context, { + signal: context.claimSignal, leaderResigned: () => { effects.push(["leaderResigned"]); }, @@ -128,10 +139,7 @@ describe.each(["postgres", "sqlite"] as const)( // Semantic expectations are independent of the fixture's payload. // Checking all effects also rules out unrelated cancellations/wakeups. const expected: Record = { - cancel: [ - ["cancel", 42n, clientId], - ["refresh", "priority"], - ], + cancel: [["cancel", 42n, clientId, 1, "2026-01-02T03:04:05.678Z"]], insert: [["wake", "priority"]], metadata_changed: [["refresh", "priority"]], pause: [["refresh", "priority"]], diff --git a/js/src/runtime/notification-pump.ts b/js/src/runtime/notification-pump.ts index eae6930b1..958a17205 100644 --- a/js/src/runtime/notification-pump.ts +++ b/js/src/runtime/notification-pump.ts @@ -2,16 +2,16 @@ * The notification pump: supervises the backend's notification streams, * waking queues on inserts, applying queue controls and cancellations, and * forwarding leadership resignation requests. Notifications are hints; - * polling remains the durable path. + * reconnects request durable cancellation checks. */ import type { RuntimeNotification } from "../driver.js"; import { LifecycleError } from "../errors.js"; import type { AttemptRunner } from "./attempt-runner.js"; import type { RuntimeContext } from "./context.js"; import { backOffAfterFailure } from "./context.js"; -import { describeError } from "./failures.js"; import { notificationCancellation, + notificationCancellationAttempt, notificationLeaderResigned, notificationQueue, notificationRequestsLeadershipResignation, @@ -20,6 +20,8 @@ import type { QueueProducer } from "./queue-producer.js"; /** Configuration for a {@link NotificationPump}. */ export interface NotificationPumpOptions { + /** Remains live until every producer has drained. */ + readonly signal: AbortSignal; /** Another client resigned leadership. */ readonly leaderResigned: () => void; readonly producer: QueueProducer; @@ -35,6 +37,7 @@ export class NotificationPump { readonly #producer: QueueProducer; readonly #resignLeadership: () => PromiseLike | undefined; readonly #runner: AttemptRunner; + readonly #signal: AbortSignal; constructor(context: RuntimeContext, options: NotificationPumpOptions) { this.#context = context; @@ -42,6 +45,7 @@ export class NotificationPump { this.#producer = options.producer; this.#resignLeadership = options.resignLeadership; this.#runner = options.runner; + this.#signal = options.signal; } /** @@ -91,7 +95,12 @@ export class NotificationPump { case "control": { const cancellation = notificationCancellation(notification.payload); if (cancellation !== null) { - this.#runner.cancelAttempt(cancellation, this.#context.clientId); + this.#runner.cancelAttempt( + cancellation, + this.#context.clientId, + notificationCancellationAttempt(notification.payload) + ); + return; } const queueName = notificationQueue(notification.payload); if (queueName === "*") { @@ -147,24 +156,7 @@ export class NotificationPump { #recoverMissedNotifications(): void { this.#producer.wakeAll(); this.#producer.wakeControl(); - this.#recoverMissedCancellations(); - } - - /** - * Cancel attempts whose cancellation notice was lost while a stream - * reconnected. A failed check is logged; the attempt then finishes as if - * its notice never arrived, as it would before this recovery. - */ - #recoverMissedCancellations(): void { - const recovery = this.#runner - .recoverCancellations(this.#context.claimSignal) - .catch((error: unknown) => { - this.#context.logger.warn( - "River could not check running jobs for missed cancellations", - { error: describeError(error) } - ); - }); - this.#context.trackTask(recovery); + this.#runner.requestCancellationCheck(); } async #remoteCancellationLoop(): Promise { @@ -177,11 +169,17 @@ export class NotificationPump { "remote cancellation stream", (signal, ready) => subscribe(this.#context.clientId, signal, ready), (notice) => { - this.#runner.cancelAttempt(notice.id, notice.attemptedBy); + this.#runner.cancelAttempt( + notice.id, + notice.attemptedBy, + notice.attempt !== undefined && notice.attemptedAt !== undefined + ? { attempt: notice.attempt, attemptedAt: notice.attemptedAt } + : undefined + ); }, () => { if (subscribed) { - this.#recoverMissedCancellations(); + this.#runner.requestCancellationCheck(); } else { subscribed = true; } @@ -202,7 +200,7 @@ export class NotificationPump { connected: () => void, resubscribes: () => boolean = () => true ): Promise { - const signal = this.#context.claimSignal; + const signal = this.#signal; let failures = 0; while (!signal.aborted) { let failure: unknown; diff --git a/producer.go b/producer.go index fd5a4e9f5..612621afa 100644 --- a/producer.go +++ b/producer.go @@ -45,12 +45,12 @@ const ( // Test-only properties. type producerTestSignals struct { - CancelHandledDuringFetch testsignal.TestSignal[int64] // notifies when a cancellation is handled during a fetch DeletedExpiredQueueRecords testsignal.TestSignal[struct{}] // notifies when the producer deletes expired queue records ExecutorShutdownStarted testsignal.TestSignal[struct{}] // notifies when the producer starts draining active jobs JobFetchTriggered testsignal.TestSignal[struct{}] // notifies when the producer's fetch limiter is triggered via triggerJobFetch MetadataChanged testsignal.TestSignal[struct{}] // notifies when the producer detects a metadata change Paused testsignal.TestSignal[struct{}] // notifies when the producer is paused + PolledJobCancellations testsignal.TestSignal[struct{}] // notifies when a cancellation check completed PolledQueueConfig testsignal.TestSignal[struct{}] // notifies when the producer polls for queue settings QueueControlEventTriggered testsignal.TestSignal[*controlEventPayload] // notifies when a queue control event is triggered via triggerQueueControlEvent ReportedProducerStatus testsignal.TestSignal[struct{}] // notifies when the producer reports its own status @@ -112,8 +112,8 @@ type producerConfig struct { QueueEventCallback func(event *Event) // QueuePollInterval is the amount of time between periodic checks for - // queue setting changes and job cancellation requests. This is only used in - // poll-only mode (when no notifier is provided). + // queue setting changes and job cancellation requests in poll-only mode. + // It also sets the retry interval for failed cancellation checks. QueuePollInterval time.Duration // QueueReportInterval is the amount of time between periodic reports // of the queue status. @@ -198,8 +198,8 @@ type producer struct { baseservice.BaseService startstop.BaseStartStop - // Jobs which are currently being worked. The polling goroutine snapshots - // their IDs; only the main goroutine accesses the executors themselves. + // Jobs currently being worked. Cancellation checks snapshot the executors + // so delayed results cannot affect a newer attempt of the same job. activeJobsMu sync.Mutex activeJobs map[int64]*jobexecutor.JobExecutor @@ -215,9 +215,10 @@ type producer struct { pilot riverpilot.Pilot workers *Workers - // Receives job IDs to cancel. Written by notifier and polling goroutines, - // only read from the main goroutine. - cancelCh chan int64 + // Coalesce cancellation hints. The generation catches hints received while + // a fetch has claimed jobs but hasn't registered their executors yet. + cancelGeneration atomic.Uint64 + cancelPollCh chan struct{} // Set to true when the producer thinks it should trigger another fetch as // soon as slots are available. This is written and read by the main @@ -240,8 +241,10 @@ type producer struct { // Receives control messages from the notifier goroutine. Written by notifier // goroutine, only read from main goroutine. queueControlCh chan *controlEventPayload - retryPolicy ClientRetryPolicy - testSignals producerTestSignals + + retryPolicy ClientRetryPolicy + runState atomic.Pointer[producerRunState] + testSignals producerTestSignals } func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pilot riverpilot.Pilot, config *producerConfig) *producer { @@ -259,7 +262,7 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi producer := baseservice.Init(archetype, &producer{ activeJobs: make(map[int64]*jobexecutor.JobExecutor), - cancelCh: make(chan int64, 1000), + cancelPollCh: make(chan struct{}, 1), completer: config.Completer, config: config.mustValidate(), exec: exec, @@ -307,6 +310,17 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } + // Queue controls delivered during shutdown must not override settings + // loaded when the producer restarts. +drainControls: + for { + select { + case <-p.queueControlCh: + default: + break drainControls + } + } + // Workers can be registered after the client is constructed, so capture // kinds at startup. Keep an empty registry non-nil to avoid fetching all kinds. p.fetchKind = nil @@ -373,8 +387,9 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { p.fetchLimiter = chanutil.NewDebouncedChan(fetchCtx, p.config.FetchCooldown, true) var ( - controlSub *notifier.Subscription - insertSub *notifier.Subscription + controlSub *notifier.Subscription + insertSub *notifier.Subscription + reconnectSub *notifier.RecoverySubscription ) if p.config.Notifier != nil { var err error @@ -400,7 +415,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return err } - controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(workCtx)) + controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(fetchCtx, workCtx)) if err != nil { stopped() if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { @@ -408,11 +423,18 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } return err } + + // Recover cancellations whose NOTIFY was lost while disconnected. + reconnectSub = p.config.Notifier.RegisterRecoveryFunc(p.TriggerCancelCheck) } + subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) + p.runState.Store(&producerRunState{done: subroutineCtx.Done(), fetchDone: fetchCtx.Done()}) + go func() { started() defer stopped() // this defer should come first so it's last out + defer p.runState.Store(nil) p.Logger.DebugContext(fetchCtx, p.Name+": Run loop started", slog.String("queue", p.config.Queue), slog.Bool("paused", p.paused)) defer func() { @@ -427,8 +449,11 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { defer controlSub.Unlisten(fetchCtx) } + if reconnectSub != nil { + defer reconnectSub.Unregister() + } + var subroutineWG sync.WaitGroup - subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) subroutineWG.Add(1) go p.fetchPollLoop(subroutineCtx, &subroutineWG) @@ -447,14 +472,14 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { subroutineWG.Add(1) go p.pollForSettingChanges(subroutineCtx, &subroutineWG, initiallyPaused, initialMetadata) - - subroutineWG.Add(1) - go p.pollForJobCancellations(subroutineCtx, &subroutineWG) } + subroutineWG.Add(1) + go p.pollForJobCancellations(subroutineCtx, &subroutineWG) + p.fetchAndRunLoop(fetchCtx, workCtx) p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) - p.executorShutdownLoop(workCtx) + p.executorShutdownLoop() p.Logger.DebugContext(workCtx, p.Name+": Shutdown loop exited, awaiting subroutines", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) cancelSubroutines(fmt.Errorf("producer stopped: %w", startstop.ErrStop)) @@ -467,6 +492,15 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } +// TriggerCancelCheck requests a cancellation check without blocking notification delivery. +func (p *producer) TriggerCancelCheck() { + p.cancelGeneration.Add(1) + select { + case p.cancelPollCh <- struct{}{}: + default: + } +} + // TriggerJobFetch manually triggers the producer to perform a job fetch // (although it's debounced, so it may not happen immediately if a fetch was // performed very recently). This is used by clients using drivers that don't @@ -484,9 +518,25 @@ func (p *producer) TriggerJobFetch() { // listen/notify. This is used by clients using drivers that don't support // listeners to wake a producer immediately after a queue control event was // known to be performed so the producer doesn't have to wait on polling. -func (p *producer) TriggerQueueControlEvent(controlEvent *controlEventPayload) { - p.queueControlCh <- controlEvent - p.testSignals.QueueControlEventTriggered.Signal(controlEvent) +func (p *producer) TriggerQueueControlEvent(ctx context.Context, controlEvent *controlEventPayload) { + state := p.runState.Load() + if state == nil || controlEvent == nil { + return + } + + // Queue controls are only consumed while fetching. Unlike cancellations, + // they must stop being delivered as soon as the producer starts draining. + select { + case <-state.fetchDone: + return + default: + } + select { + case p.queueControlCh <- controlEvent: + p.testSignals.QueueControlEventTriggered.Signal(controlEvent) + case <-ctx.Done(): + case <-state.fetchDone: + } } type controlAction string @@ -499,17 +549,19 @@ const ( ) type controlEventPayload struct { - Action controlAction `json:"action"` - JobID int64 `json:"job_id,omitempty"` - Metadata json.RawMessage `json:"metadata,omitempty"` - Queue string `json:"queue"` + Action controlAction `json:"action"` + Attempt int `json:"attempt,omitempty"` + AttemptedAt *time.Time `json:"attempted_at,omitempty"` + JobID int64 `json:"job_id,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` + Queue string `json:"queue"` } type insertPayload struct { Queue string `json:"queue"` } -func (p *producer) handleControlNotification(workCtx context.Context) func(notifier.NotificationTopic, string) { +func (p *producer) handleControlNotification(fetchCtx, workCtx context.Context) func(notifier.NotificationTopic, string) { return func(topic notifier.NotificationTopic, payload string) { var decoded controlEventPayload if err := json.Unmarshal([]byte(payload), &decoded); err != nil { @@ -519,12 +571,15 @@ func (p *producer) handleControlNotification(workCtx context.Context) func(notif switch decoded.Action { case controlActionMetadataChanged, controlActionPause, controlActionResume: + if fetchCtx.Err() != nil { + return + } if decoded.Queue != rivercommon.AllQueuesString && decoded.Queue != p.config.Queue { p.Logger.DebugContext(workCtx, p.Name+": Queue control notification for other queue", slog.String("action", string(decoded.Action))) return } select { - case <-workCtx.Done(): + case <-fetchCtx.Done(): case p.queueControlCh <- &decoded: default: p.Logger.WarnContext(workCtx, p.Name+": Queue control notification dropped due to full buffer", slog.String("action", string(decoded.Action))) @@ -538,11 +593,17 @@ func (p *producer) handleControlNotification(workCtx context.Context) func(notif ) return } - select { - case <-workCtx.Done(): - case p.cancelCh <- decoded.JobID: - default: - p.Logger.WarnContext(workCtx, p.Name+": Job cancel notification dropped due to full buffer", slog.Int64("job_id", decoded.JobID)) + // New senders identify the execution, so delivery never needs a pool + // connection. Legacy notices and claims still being registered reconcile. + p.activeJobsMu.Lock() + executor := p.activeJobs[decoded.JobID] + p.activeJobsMu.Unlock() + if executor != nil && decoded.Attempt > 0 && decoded.AttemptedAt != nil { + if executor.JobRow.Attempt == decoded.Attempt && executor.JobRow.AttemptedAt != nil && executor.JobRow.AttemptedAt.Equal(*decoded.AttemptedAt) { + executor.Cancel(workCtx) + } + } else { + p.TriggerCancelCheck() } default: p.Logger.DebugContext(workCtx, p.Name+": Received job control notification with unknown action", @@ -560,6 +621,7 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { p.fetchLimiter.Call() fetchResultCh := make(chan producerFetchResult) + for { select { case <-fetchCtx.Done(): @@ -567,10 +629,7 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { case msg := <-p.queueControlCh: switch msg.Action { case controlActionCancel: - // This path is only expected to take effect in poll-only mode, and - // only works for the case of a single process. Multi-process setups - // will have to wait for the next poll event for a cancel to take effect. - p.maybeCancelJob(workCtx, msg.JobID) + p.TriggerCancelCheck() case controlActionMetadataChanged: p.Logger.DebugContext(workCtx, p.Name+": Queue metadata changed", slog.String("queue", p.config.Queue), slog.String("queue_in_message", msg.Queue)) p.testSignals.MetadataChanged.Signal(struct{}{}) @@ -604,8 +663,6 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { default: p.Logger.DebugContext(workCtx, p.Name+": Unknown queue control action", "action", msg.Action) } - case jobID := <-p.cancelCh: - p.maybeCancelJob(workCtx, jobID) case <-p.fetchLimiter.C(): p.innerFetchLoop(workCtx, fetchResultCh) // Ensure we can't start another fetch when fetchCtx is done, even if @@ -676,9 +733,7 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } - // A fetched job may not have an executor yet. Keep unmatched cancellations - // only for this fetch so unrelated IDs cannot affect later fetches. - pendingCancelJobs := make(map[int64]struct{}) + cancelGeneration := p.cancelGeneration.Load() go p.dispatchWork(workCtx, limit, fetchResultCh) for { @@ -687,7 +742,10 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr if result.err != nil { p.Logger.ErrorContext(workCtx, p.Name+": Error fetching jobs", slog.String("err", result.err.Error()), slog.String("queue", p.config.Queue)) } else if numLocked := len(result.jobs); numLocked > 0 { - p.startNewExecutors(workCtx, result.jobs, result.decodeErrors, pendingCancelJobs) + p.startNewExecutors(workCtx, result.jobs, result.decodeErrors) + if p.cancelGeneration.Load() != cancelGeneration { + p.TriggerCancelCheck() + } if numLocked == limit { // Fetch returned the maximum number of jobs that were requested, @@ -699,20 +757,14 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr return case result := <-p.jobResultCh: p.removeActiveJob(result) - case jobID := <-p.cancelCh: - if !p.maybeCancelJob(workCtx, jobID) { - pendingCancelJobs[jobID] = struct{}{} - } - p.testSignals.CancelHandledDuringFetch.Signal(jobID) } } } -func (p *producer) executorShutdownLoop(ctx context.Context) { +func (p *producer) executorShutdownLoop() { p.testSignals.ExecutorShutdownStarted.Signal(struct{}{}) - // No more jobs will be fetched or executed. However, we must wait for all - // in-progress jobs to complete. + // Cancellation polling stays live until every execution has finished. for { p.activeJobsMu.Lock() numActiveJobs := len(p.activeJobs) @@ -720,13 +772,7 @@ func (p *producer) executorShutdownLoop(ctx context.Context) { if numActiveJobs == 0 { return } - - select { - case jobID := <-p.cancelCh: - p.maybeCancelJob(ctx, jobID) - case result := <-p.jobResultCh: - p.removeActiveJob(result) - } + p.removeActiveJob(<-p.jobResultCh) } } @@ -828,17 +874,6 @@ func (p *producer) handleWorkerUnstuck() { p.config.JobStuckCount.Add(-1) } -func (p *producer) maybeCancelJob(ctx context.Context, id int64) bool { - p.activeJobsMu.Lock() - executor, ok := p.activeJobs[id] - p.activeJobsMu.Unlock() - if !ok { - return false - } - executor.Cancel(ctx) - return true -} - func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit { pluginLookup := p.config.PluginLookupGlobal if pluginLookup == nil { @@ -966,7 +1001,7 @@ func (p *producer) heartbeatLogLoop(ctx context.Context, wg *sync.WaitGroup) { // couldn't be decoded isn't worked, but it still gets an executor that fails // its attempt with the decode error, so that it's retried or discarded // through the normal error handling path instead of being left running. -func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.JobRow, decodeErrors map[int64]error, pendingCancelJobs map[int64]struct{}) { +func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.JobRow, decodeErrors map[int64]error) { defaultClientRetryPolicy := retrypolicy.NewDefault(p.Time) for _, job := range jobs { @@ -1007,10 +1042,6 @@ func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype. WorkUnit: workUnit, }) p.addActiveJob(job.ID, executor) - if _, ok := pendingCancelJobs[job.ID]; ok { - // Cancel before execution so the worker starts with a cancelled context. - executor.Cancel(workCtx) - } go executor.Execute(jobCtx) } @@ -1028,32 +1059,43 @@ func (p *producer) handleWorkerDone(job *rivertype.JobRow) { p.jobResultCh <- job } -// Cancellation polling is independent of fetching: saturated and paused producers -// must still cancel jobs, including while draining during graceful shutdown. +// Cancellation checks run independently of fetching and remain live while draining. +// With notifications, the timer only retries failures; otherwise it polls regularly. func (p *producer) pollForJobCancellations(ctx context.Context, wg *sync.WaitGroup) { defer wg.Done() - ticker := time.NewTicker(p.config.QueuePollInterval) - defer ticker.Stop() + timer := time.NewTimer(p.config.QueuePollInterval) + defer timer.Stop() + if p.config.Notifier != nil { + timer.Stop() + } for { select { case <-ctx.Done(): return - case <-ticker.C: - if err := p.pollForJobCancellationsOnce(ctx); err != nil { - if !errors.Is(context.Cause(ctx), startstop.ErrStop) { - p.Logger.ErrorContext(ctx, p.Name+": Error fetching job cancellation requests", slog.String("err", err.Error())) - } - continue + case <-timer.C: + case <-p.cancelPollCh: + } + timer.Stop() + err := p.pollForJobCancellationsOnce(ctx) + if p.config.Notifier == nil || err != nil { + timer.Reset(p.config.QueuePollInterval) + } + if err != nil { + if !errors.Is(context.Cause(ctx), startstop.ErrStop) { + p.Logger.ErrorContext(ctx, p.Name+": Error fetching job cancellation requests", slog.String("err", err.Error())) } + continue } + p.testSignals.PolledJobCancellations.Signal(struct{}{}) } } func (p *producer) pollForJobCancellationsOnce(ctx context.Context) error { p.activeJobsMu.Lock() - jobIDs := slices.Collect(maps.Keys(p.activeJobs)) + executors := maps.Clone(p.activeJobs) p.activeJobsMu.Unlock() + jobIDs := slices.Collect(maps.Keys(executors)) // Bound query size (including SQLite's parameter count) and skip database // access entirely when there are no active jobs. @@ -1069,10 +1111,10 @@ func (p *producer) pollForJobCancellationsOnce(ctx context.Context) error { } for _, id := range cancelIDs { - select { - case <-ctx.Done(): - return ctx.Err() - case p.cancelCh <- id: + // Signal the snapshotted execution: a retry may already be running + // with the same ID by the time the query returns. + if executor := executors[id]; executor != nil { + executor.Cancel(ctx) } } } @@ -1239,6 +1281,11 @@ type producerFetchResult struct { jobs []*rivertype.JobRow } +type producerRunState struct { + done <-chan struct{} + fetchDone <-chan struct{} +} + type errorHandlerAdapter struct { errorHandler ErrorHandler } diff --git a/producer_test.go b/producer_test.go index 8d8545b7a..b3ff82973 100644 --- a/producer_test.go +++ b/producer_test.go @@ -6,18 +6,22 @@ import ( "errors" "fmt" "slices" + "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/require" "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobexecutor" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" "github.com/riverqueue/river/internal/riverinternaltest" "github.com/riverqueue/river/internal/riverinternaltest/sharedtx" + "github.com/riverqueue/river/internal/util/chanutil" "github.com/riverqueue/river/riverdbtest" "github.com/riverqueue/river/riverdriver" "github.com/riverqueue/river/riverdriver/riverpgxv5" @@ -76,6 +80,21 @@ func (p *beforeJobGetAvailablePilot) JobGetAvailable( return p.Pilot.JobGetAvailable(ctx, exec, state, params) } +type producerTestExecutor struct { + riverdriver.Executor + + jobGetAvailableFunc func(context.Context, *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) + jobGetCancelRequestedFunc func(context.Context, *riverdriver.JobGetCancelRequestedParams) ([]int64, error) +} + +func (e *producerTestExecutor) JobGetAvailable(ctx context.Context, params *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) { + return e.jobGetAvailableFunc(ctx, params) +} + +func (e *producerTestExecutor) JobGetCancelRequested(ctx context.Context, params *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + return e.jobGetCancelRequestedFunc(ctx, params) +} + // undecodableKindPilot reports locked jobs of one kind as undecodable. Postgres' // column types don't allow a job row that can't be decoded, so this simulates // one. @@ -108,6 +127,456 @@ func (p *undecodableKindPilot) JobGetAvailable( return res, nil } +// producerTestListenerMock overrides operations on a real listener. +type producerTestListenerMock struct { + riverdriver.Listener + + connectFunc func(ctx context.Context) error + waitForNotificationFunc func(ctx context.Context) (*riverdriver.Notification, error) +} + +func newProducerTestListenerMock(listener riverdriver.Listener) *producerTestListenerMock { + return &producerTestListenerMock{ + Listener: listener, + connectFunc: listener.Connect, + waitForNotificationFunc: listener.WaitForNotification, + } +} + +func (l *producerTestListenerMock) Connect(ctx context.Context) error { return l.connectFunc(ctx) } + +func (l *producerTestListenerMock) WaitForNotification(ctx context.Context) (*riverdriver.Notification, error) { + return l.waitForNotificationFunc(ctx) +} + +func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + blockedFetch bool + cancelFromLookup bool + queryError bool + remote bool + }{ + {name: "BlockedFetchCancelDelivery", blockedFetch: true}, + {name: "BlockedFetchQueryResult", blockedFetch: true, cancelFromLookup: true}, + {name: "CancelDelivery"}, + {name: "QueryResult", cancelFromLookup: true}, + {name: "RemoteCancelDelivery", remote: true}, + {name: "RetryAfterError", cancelFromLookup: true, queryError: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + var signals struct { + FetchStarted testsignal.TestSignal[struct{}] + LoopStopped testsignal.TestSignal[struct{}] + QueryStarted testsignal.TestSignal[time.Duration] + ReleaseFetch testsignal.TestSignal[struct{}] + } + signals.FetchStarted.Init(t) + signals.LoopStopped.Init(t) + signals.QueryStarted.Init(t) + signals.ReleaseFetch.Init(t) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + release := sync.OnceFunc(func() { signals.ReleaseFetch.Signal(struct{}{}) }) + defer release() + archetype := riversharedtest.BaseServiceArchetype(t) + jobCtx, cancelJob := context.WithCancelCause(ctx) + defer cancelJob(nil) + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelJob, JobRow: &rivertype.JobRow{ID: 1, Attempt: 1, AttemptedAt: new(time.Now().UTC())}}) + var queryAttempts int + exec := &producerTestExecutor{ + jobGetAvailableFunc: func(context.Context, *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) { + signals.FetchStarted.Signal(struct{}{}) + <-signals.ReleaseFetch.WaitC() + return &riverdriver.JobGetAvailableResult{}, nil + }, + jobGetCancelRequestedFunc: func(ctx context.Context, _ *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + deadline, ok := ctx.Deadline() + if ok { + signals.QueryStarted.Signal(time.Until(deadline)) + } else { + signals.QueryStarted.Signal(0) + } + queryAttempts++ + if testCase.queryError && queryAttempts == 1 { + return nil, errors.New("temporary cancellation check failure") + } + if testCase.cancelFromLookup { + return []int64{1}, nil + } + <-ctx.Done() + return nil, ctx.Err() + }, + } + pilot := &riverpilot.StandardPilot{} + _, state, err := pilot.ProducerInit(ctx, exec, nil) + require.NoError(t, err) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{MaxWorkers: 2, Notifier: ¬ifier.Notifier{}, Queue: QueueDefault, QueuePollInterval: 20 * time.Millisecond}, + exec: exec, + fetchLimiter: chanutil.NewDebouncedChan(ctx, time.Millisecond, true), + jobResultCh: make(chan *rivertype.JobRow), + paused: !testCase.blockedFetch, + pilot: pilot, + state: state, + queueControlCh: make(chan *controlEventPayload), + }) + prod.runState.Store(&producerRunState{done: ctx.Done()}) + var pollWG sync.WaitGroup + pollWG.Add(1) + go prod.pollForJobCancellations(ctx, &pollWG) + go func() { + prod.fetchAndRunLoop(ctx, ctx) + prod.executorShutdownLoop() + signals.LoopStopped.Signal(struct{}{}) + }() + defer func() { + release() + cancel() + signals.LoopStopped.WaitOrTimeout() + pollWG.Wait() + }() + + if testCase.blockedFetch { + signals.FetchStarted.WaitOrTimeout() + } + prod.TriggerCancelCheck() + queryTimeout := signals.QueryStarted.WaitOrTimeout() + require.Positive(t, queryTimeout) + require.LessOrEqual(t, queryTimeout, 10*time.Second) + if testCase.queryError { + require.Positive(t, signals.QueryStarted.WaitOrTimeout()) + } + + if testCase.remote { + prod.handleControlNotification(ctx, ctx)(notifier.NotificationTopicControl, fmt.Sprintf(`{"action":"cancel","attempt":1,"attempted_at":%q,"job_id":1,"queue":"default"}`, executor.JobRow.AttemptedAt.Format(time.RFC3339Nano))) + } else if !testCase.cancelFromLookup { + client := &Client[struct{}]{producersByQueueName: map[string]*producer{QueueDefault: prod}} + job := *executor.JobRow + job.Queue = QueueDefault + client.notifyProducerOfCancelWithoutNotifier(ctx, &job) + } + select { + case <-jobCtx.Done(): + require.ErrorIs(t, context.Cause(jobCtx), rivertype.ErrJobCancelledRemotely) + case <-time.After(time.Second): + t.Fatal("blocked cancellation check prevented cancellation delivery") + } + + select { + case prod.jobResultCh <- &rivertype.JobRow{ID: 1}: + case <-time.After(time.Second): + t.Fatal("blocked cancellation check prevented job completion") + } + }) + } +} + +func TestProducer_CancellationChecksWaitForHints(t *testing.T) { + t.Parallel() + + for _, failFirst := range []bool{false, true} { + t.Run(fmt.Sprintf("Retry=%v", failFirst), func(t *testing.T) { + t.Parallel() + + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var queries atomic.Int32 + prod := baseservice.Init(riversharedtest.BaseServiceArchetype(t), &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: {}}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{Notifier: ¬ifier.Notifier{}, QueuePollInterval: time.Second}, + exec: &producerTestExecutor{jobGetCancelRequestedFunc: func(context.Context, *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + query := queries.Add(1) + if failFirst && query == 1 { + return nil, errors.New("temporary failure") + } + return nil, nil + }}, + }) + var wg sync.WaitGroup + wg.Add(1) + go prod.pollForJobCancellations(ctx, &wg) + time.Sleep(10 * time.Second) + require.Zero(t, queries.Load()) + + prod.TriggerCancelCheck() + synctest.Wait() + require.EqualValues(t, 1, queries.Load()) + time.Sleep(10 * time.Second) + want := 1 + if failFirst { + want = 2 + } + require.EqualValues(t, want, queries.Load()) + cancel() + wg.Wait() + }) + }) + } +} + +func TestProducer_CancellationNotificationAttemptIdentity(t *testing.T) { + t.Parallel() + + attemptedAt := time.Date(2026, 1, 2, 3, 4, 5, 678900000, time.UTC) + ctx, cancel := context.WithCancelCause(context.Background()) + defer cancel(nil) + archetype := riversharedtest.BaseServiceArchetype(t) + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancel, JobRow: &rivertype.JobRow{Attempt: 2, AttemptedAt: &attemptedAt}}) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{Queue: QueueDefault}, + }) + deliver := prod.handleControlNotification(context.Background(), context.Background()) + for _, payload := range []string{ + `{"action":"cancel","attempt":1,"attempted_at":"2026-01-02T03:04:05.6789Z","job_id":1,"queue":"default"}`, + `{"action":"cancel","attempt":2,"attempted_at":"2026-01-02T03:04:04.6789Z","job_id":1,"queue":"default"}`, + } { + deliver(notifier.NotificationTopicControl, payload) + require.NoError(t, ctx.Err()) + require.Empty(t, prod.cancelPollCh) + } + deliver(notifier.NotificationTopicControl, `{"action":"cancel","attempt":2,"attempted_at":"2026-01-02T03:04:05.6789Z","job_id":1,"queue":"default"}`) + require.ErrorIs(t, context.Cause(ctx), rivertype.ErrJobCancelledRemotely) + require.Empty(t, prod.cancelPollCh) +} + +func TestProducer_HandleControlNotification(t *testing.T) { + t.Parallel() + + for _, action := range []controlAction{controlActionCancel, controlActionMetadataChanged, controlActionPause, controlActionResume} { + t.Run(string(action), func(t *testing.T) { + t.Parallel() + + prod := baseservice.Init(riversharedtest.BaseServiceArchetype(t), &producer{ + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{Queue: QueueDefault}, + queueControlCh: make(chan *controlEventPayload, 1), + }) + fetchCtx, cancelFetch := context.WithCancel(context.Background()) + handleNotification := prod.handleControlNotification(fetchCtx, context.Background()) + cancelFetch() + payload, err := json.Marshal(&controlEventPayload{Action: action, JobID: 123, Queue: QueueDefault}) + require.NoError(t, err) + + handleNotification(notifier.NotificationTopicControl, string(payload)) + + require.Empty(t, prod.queueControlCh) + if action == controlActionCancel { + require.Len(t, prod.cancelPollCh, 1) + } else { + require.Empty(t, prod.cancelPollCh) + } + }) + } +} + +// A cancel NOTIFY lost during a disconnect must be recovered on reconnect. +func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + fetchPending bool + stop bool + }{ + {name: "DuringFetch", fetchPending: true}, + {name: "DuringFetchAndStop", fetchPending: true, stop: true}, + {name: "DuringGracefulStop", stop: true}, + {name: "Running"}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + var ( + archetype = riversharedtest.BaseServiceArchetype(t) + dbPool = riversharedtest.DBPool(ctx, t) + driver = riverpgxv5.New(dbPool) + exec = driver.GetExecutor() + jobUpdates = make(chan []jobcompleter.CompleterJobUpdated, 10) + schema = riverdbtest.TestSchema(ctx, t, driver, nil) + queueName = fmt.Sprintf("test-reconnect-cancel-%05d", randutil.IntBetween(1, 100_000)) + pilot = &riverpilot.StandardPilot{} + ) + + completer := jobcompleter.NewInlineCompleter(archetype, schema, exec, pilot, jobUpdates) + require.NoError(t, completer.Start(ctx)) + t.Cleanup(completer.Stop) + + realListener := driver.GetListener(&riverdriver.GetListenenerParams{Schema: schema}) + mockListener := newProducerTestListenerMock(realListener) + + var ( + forceDisconnect atomic.Bool // set by the test to inject one simulated connection loss + reconnecting atomic.Bool + reconnectStarted testsignal.TestSignal[struct{}] + releaseReconnect = make(chan struct{}) + releaseReconnectOnce sync.Once + ) + reconnectStarted.Init(t) + closeReleaseReconnect := func() { releaseReconnectOnce.Do(func() { close(releaseReconnect) }) } + mockListener.waitForNotificationFunc = func(ctx context.Context) (*riverdriver.Notification, error) { + if forceDisconnect.CompareAndSwap(true, false) { + reconnecting.Store(true) + return nil, errors.New("simulated connection loss") + } + return realListener.WaitForNotification(ctx) + } + mockListener.connectFunc = func(ctx context.Context) error { + if reconnecting.Load() { + reconnectStarted.Signal(struct{}{}) + select { + case <-ctx.Done(): + return ctx.Err() + case <-releaseReconnect: + } + reconnecting.Store(false) + } + return realListener.Connect(ctx) + } + + notif := notifier.New(archetype, mockListener) + var initiallyReady testsignal.TestSignal[struct{}] + initiallyReady.Init(t) + readySub := notif.RegisterRecoveryFunc(func() { initiallyReady.Signal(struct{}{}) }) + require.NoError(t, notif.Start(ctx)) + initiallyReady.WaitOrTimeout() + readySub.Unregister() + t.Cleanup(notif.Stop) + // Release a blocked Connect before stopping the notifier. + t.Cleanup(closeReleaseReconnect) + + var jobStarted testsignal.TestSignal[int64] + jobStarted.Init(t) + workers := NewWorkers() + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + AddWorker(workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + jobStarted.Signal(job.ID) + <-ctx.Done() + return ctx.Err() + })) + + prod := newProducer(archetype, exec, pilot, &producerConfig{ + ClientID: testClientID, + Completer: completer, + ErrorHandler: newTestErrorHandler(), + FetchCooldown: FetchCooldownDefault, + FetchPollInterval: 50 * time.Millisecond, + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), + JobTimeout: JobTimeoutDefault, + MaxWorkers: 10, + Notifier: notif, + Queue: queueName, + QueuePollInterval: time.Hour, + QueueReportInterval: queueReportIntervalDefault, + RetryPolicy: &DefaultClientRetryPolicy{}, + SchedulerInterval: riverinternaltest.SchedulerShortInterval, + Schema: schema, + StaleProducerRetentionPeriod: time.Minute, + Workers: workers, + }) + + config := newTestConfig(t, schema) + insertParams, err := insertParamsFromConfigArgsAndOptions(&prod.Archetype, config, &JobArgs{}, &InsertOpts{Queue: queueName}) + require.NoError(t, err) + _, err = exec.JobInsertFastMany(ctx, &riverdriver.JobInsertFastManyParams{ + Jobs: []*riverdriver.JobInsertFastParams{(*riverdriver.JobInsertFastParams)(insertParams)}, + Schema: schema, + }) + require.NoError(t, err) + + var fetchReturned testsignal.TestSignal[int64] + var releaseFetch testsignal.TestSignal[struct{}] + fetchReturned.Init(t) + releaseFetch.Init(t) + releaseFetchFunc := sync.OnceFunc(func() { releaseFetch.Signal(struct{}{}) }) + if testCase.fetchPending { + prod.testSignals.PolledJobCancellations.Init(t) + prod.pilot = &afterJobGetAvailablePilot{ + Pilot: pilot, + afterJobGetAvailableFunc: func(res *riverdriver.JobGetAvailableResult, err error) { + if err == nil && res != nil && len(res.Jobs) > 0 { + fetchReturned.Signal(res.Jobs[0].ID) + <-releaseFetch.WaitC() + } + }, + } + } + prod.testSignals.ExecutorShutdownStarted.Init(t) + fetchCtx, fetchCancel := context.WithCancel(ctx) + defer fetchCancel() + workCtx, workCancel := context.WithCancel(ctx) + + require.NoError(t, prod.StartWorkContext(fetchCtx, workCtx)) + // Cancel workers before waiting for producer shutdown. + t.Cleanup(prod.Stop) + t.Cleanup(workCancel) + t.Cleanup(closeReleaseReconnect) + t.Cleanup(releaseFetchFunc) + + var jobID int64 + if testCase.fetchPending { + jobID = fetchReturned.WaitOrTimeout() + } else { + jobID = jobStarted.WaitOrTimeout() + } + require.Positive(t, jobID) + if testCase.stop { + fetchCancel() + if !testCase.fetchPending { + prod.testSignals.ExecutorShutdownStarted.WaitOrTimeout() + } + } + + // Wake the listener so it observes the injected disconnect. + forceDisconnect.Store(true) + require.NoError(t, exec.NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Topic: string(notifier.NotificationTopicInsert), + Payload: []string{`{"queue":"` + queueName + `-unrelated-wakeup"}`}, + Schema: schema, + })) + + reconnectStarted.WaitOrTimeout() + + // Cancel while no listener is connected. + _, err = exec.JobCancel(ctx, &riverdriver.JobCancelParams{ + ID: jobID, + CancelAttemptedAt: time.Now().UTC(), + ControlTopic: string(notifier.NotificationTopicControl), + Schema: schema, + }) + require.NoError(t, err) + + closeReleaseReconnect() + if testCase.fetchPending { + prod.testSignals.PolledJobCancellations.WaitOrTimeout() + releaseFetchFunc() + } + + update := riversharedtest.WaitOrTimeout(t, jobUpdates) + require.Equal(t, rivertype.JobStateCancelled, update[0].Job.State, + "job should have observed cancellation after the notifier reconnected and reconciled active jobs against cancel_attempted_at") + }) + } +} + func TestProducer_MetricEmitHook(t *testing.T) { t.Parallel() @@ -270,6 +739,132 @@ func TestProducer_MetricEmitHook(t *testing.T) { }) } +func TestProducer_PollForJobCancellations(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + jobs int + queries int + }{ + {name: "Batched", jobs: 1001, queries: 2}, + {name: "Idle"}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + seen := make(map[int64]bool) + var queries int + prod := &producer{ + activeJobs: make(map[int64]*jobexecutor.JobExecutor), + config: &producerConfig{Schema: "custom"}, + exec: &producerTestExecutor{jobGetCancelRequestedFunc: func(_ context.Context, params *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + queries++ + require.Equal(t, "custom", params.Schema) + require.NotEmpty(t, params.ID) + require.LessOrEqual(t, len(params.ID), jobCancelPollBatchSize) + for _, id := range params.ID { + require.False(t, seen[id], "job checked more than once") + seen[id] = true + } + return nil, nil + }}, + } + for id := range testCase.jobs { + prod.activeJobs[int64(id)] = &jobexecutor.JobExecutor{} + } + + require.NoError(t, prod.pollForJobCancellationsOnce(context.Background())) + require.Equal(t, testCase.queries, queries) + require.Len(t, seen, testCase.jobs) + }) + } + + for _, replaced := range []bool{false, true} { + t.Run(fmt.Sprintf("ReplacedExecutor_%t", replaced), func(t *testing.T) { + t.Parallel() + + archetype := riversharedtest.BaseServiceArchetype(t) + originalCtx, cancelOriginal := context.WithCancelCause(context.Background()) + defer cancelOriginal(nil) + replacementCtx, cancelReplacement := context.WithCancelCause(context.Background()) + defer cancelReplacement(nil) + original := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelOriginal, JobRow: &rivertype.JobRow{ID: 1}}) + replacement := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelReplacement, JobRow: &rivertype.JobRow{ID: 1}}) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: original}, + config: &producerConfig{}, + }) + prod.exec = &producerTestExecutor{ + jobGetCancelRequestedFunc: func(ctx context.Context, params *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + require.Equal(t, []int64{1}, params.ID) + if replaced { + // The old attempt finishes and is retried before the query returns. + cancelOriginal(nil) + prod.activeJobs[1] = replacement + } + return []int64{1}, nil + }, + } + + require.NoError(t, prod.pollForJobCancellationsOnce(context.Background())) + require.NoError(t, replacementCtx.Err()) + if !replaced { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + } + }) + } +} + +func TestProducer_PollForJobCancellationsWakeDuringQuery(t *testing.T) { + t.Parallel() + + var queryStarted, releaseQuery testsignal.TestSignal[struct{}] + queryStarted.Init(t) + releaseQuery.Init(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + jobCtx, cancelJob := context.WithCancelCause(ctx) + defer cancelJob(nil) + archetype := riversharedtest.BaseServiceArchetype(t) + var queries int + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelJob, JobRow: &rivertype.JobRow{ID: 1}})}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{QueuePollInterval: time.Hour}, + exec: &producerTestExecutor{jobGetCancelRequestedFunc: func(ctx context.Context, _ *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + queries++ + if queries == 1 { + queryStarted.Signal(struct{}{}) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-releaseQuery.WaitC(): + return nil, nil + } + } + return []int64{1}, nil + }}, + }) + var wg sync.WaitGroup + wg.Add(1) + go prod.pollForJobCancellations(ctx, &wg) + defer func() { cancel(); wg.Wait() }() + + prod.TriggerCancelCheck() + queryStarted.WaitOrTimeout() + // A burst during a blocked read coalesces, but still causes a fresh read. + for range 1000 { + prod.TriggerCancelCheck() + } + releaseQuery.Signal(struct{}{}) + riversharedtest.WaitOrTimeout(t, jobCtx.Done()) + require.ErrorIs(t, context.Cause(jobCtx), rivertype.ErrJobCancelledRemotely) + cancel() + wg.Wait() + require.Equal(t, 2, queries) +} + func TestProducer_PollOnly(t *testing.T) { t.Parallel() @@ -421,7 +1016,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin } } - mustInsert := func(ctx context.Context, t *testing.T, producer *producer, bundle *testBundle, args JobArgs) { + mustInsert := func(ctx context.Context, t *testing.T, producer *producer, bundle *testBundle, args JobArgs) int64 { t.Helper() insertParams, err := insertParamsFromConfigArgsAndOptions(bundle.archetype, bundle.config, args, &InsertOpts{ @@ -440,11 +1035,12 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin insertParams.ScheduledAt = &bundle.timeBeforeStart } - _, err = bundle.exec.JobInsertFastMany(ctx, &riverdriver.JobInsertFastManyParams{ + inserted, err := bundle.exec.JobInsertFastMany(ctx, &riverdriver.JobInsertFastManyParams{ Jobs: []*riverdriver.JobInsertFastParams{(*riverdriver.JobInsertFastParams)(insertParams)}, Schema: producer.config.Schema, }) require.NoError(t, err) + return inserted[0].Job.ID } startProducer := func(t *testing.T, fetchCtx, workCtx context.Context, producer *producer) { @@ -501,6 +1097,45 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin producer.testSignals.ReportedQueueStatus.WaitOrTimeout() }) + t.Run("RestartDiscardsStaleControls", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.config.QueuePollInterval = time.Hour + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, _ *Job[noOpArgs]) error { return ctx.Err() })) + startProducer(t, ctx, ctx, producer) + producer.Stop() + + // Model controls buffered just as the previous run stopped fetching. + producer.queueControlCh <- &controlEventPayload{Action: controlActionPause, Queue: bundle.queue} + producer.queueControlCh <- &controlEventPayload{Action: controlActionMetadataChanged, Metadata: []byte(`{"old":true}`), Queue: bundle.queue} + _, err := bundle.exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`{"current":true}`), + MetadataDoUpdate: true, + Name: bundle.queue, + Schema: producer.config.Schema, + }) + require.NoError(t, err) + + id := mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + _, err = bundle.exec.JobCancel(ctx, &riverdriver.JobCancelParams{ + CancelAttemptedAt: time.Now(), + ControlTopic: string(notifier.NotificationTopicControl), + ID: id, + Schema: producer.config.Schema, + }) + require.NoError(t, err) + _, err = bundle.exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: id, Schema: producer.config.Schema}) + require.NoError(t, err) + producer.TriggerCancelCheck() + + startProducer(t, ctx, ctx, producer) + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + producer.testSignals.MetadataChanged.RequireEmpty() + producer.testSignals.Paused.RequireEmpty() + }) + t.Run("UnknownJobKind", func(t *testing.T) { t.Parallel() @@ -542,7 +1177,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin } }) - t.Run("CancellationPolling", func(t *testing.T) { + t.Run("CancellationWhilePausedSaturatedOrStopping", func(t *testing.T) { t.Parallel() for _, state := range []string{"Paused", "Running", "Stopping"} { @@ -550,9 +1185,6 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin t.Parallel() producer, bundle := setup(t) - if producer.config.Notifier != nil { - t.Skip("requires polling without a notifier") - } producer.config.FetchPollInterval = time.Hour producer.config.MaxWorkers = 1 producer.config.QueuePollInterval = 20 * time.Millisecond @@ -588,6 +1220,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin Name: producer.config.Queue, Schema: producer.config.Schema, })) + producer.TriggerQueueControlEvent(ctx, &controlEventPayload{Action: controlActionPause, Queue: producer.config.Queue}) producer.testSignals.Paused.WaitOrTimeout() case "Stopping": fetchCancel() @@ -595,7 +1228,8 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin } // Use the executor directly so no local client shortcut can cancel - // the worker. No more jobs can be fetched in any of these states. + // the worker. Notifications or poll-only checks must work even though + // no more jobs can be fetched in any of these states. _, err := bundle.exec.JobCancel(ctx, &riverdriver.JobCancelParams{ ID: jobID, CancelAttemptedAt: time.Now().UTC(), @@ -683,6 +1317,35 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRetryable, update.Job.State) }) + t.Run("CancellationNoticeWithoutMarkerDoesNotCancelJob", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.testSignals.PolledJobCancellations.Init(t) + var entered testsignal.TestSignal[struct{}] + var release testsignal.TestSignal[struct{}] + entered.Init(t) + release.Init(t) + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, _ *Job[noOpArgs]) error { + entered.Signal(struct{}{}) + <-release.WaitC() + return ctx.Err() + })) + id := mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + startProducer(t, ctx, ctx, producer) + t.Cleanup(func() { release.Signal(struct{}{}) }) + entered.WaitOrTimeout() + + // An old notification can arrive after retry cleared the durable marker. + payload, err := json.Marshal(&controlEventPayload{Action: controlActionCancel, JobID: id, Queue: producer.config.Queue}) + require.NoError(t, err) + producer.handleControlNotification(ctx, ctx)(notifier.NotificationTopicControl, string(payload)) + producer.testSignals.PolledJobCancellations.WaitOrTimeout() + release.Signal(struct{}{}) + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + }) + t.Run("CancellationWhileFetchResultPendingReachesJob", func(t *testing.T) { t.Parallel() @@ -690,7 +1353,8 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin // process its cancellation while it has no active executor. producer, bundle := setup(t) producer.config.FetchPollInterval = time.Hour - producer.testSignals.CancelHandledDuringFetch.Init(t) + producer.config.QueuePollInterval = time.Hour + producer.testSignals.PolledJobCancellations.Init(t) var ( fetchReturned testsignal.TestSignal[int64] @@ -716,6 +1380,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin testutil.JobArgsReflectKind[JobArgs] } AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + <-ctx.Done() workerErr.Signal(ctx.Err()) return ctx.Err() })) @@ -735,8 +1400,8 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRunning, jobAfterCancel.State) // Send directly to avoid depending on notifier delivery timing. - producer.cancelCh <- jobID - require.Equal(t, jobID, producer.testSignals.CancelHandledDuringFetch.WaitOrTimeout()) + producer.TriggerCancelCheck() + producer.testSignals.PolledJobCancellations.WaitOrTimeout() releaseFetch.Signal(struct{}{}) require.ErrorIs(t, workerErr.WaitOrTimeout(), context.Canceled) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 4bd6cbfab..3681ea7b6 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -16,7 +16,7 @@ import ( const jobCancel = `-- name: JobCancel :one WITH locked_job AS ( SELECT - id, queue, state, finalized_at + id, attempt, attempted_at, queue, state, finalized_at FROM /* TEMPLATE: schema */river_job WHERE river_job.id = $1 FOR UPDATE @@ -26,7 +26,7 @@ notification AS ( id, CASE WHEN $2::boolean THEN pg_notify( concat(coalesce($3::text, current_schema()), '.', $4::text), - json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text ) END FROM locked_job @@ -617,6 +617,7 @@ const jobGetCancelRequested = `-- name: JobGetCancelRequested :many SELECT id FROM /* TEMPLATE: schema */river_job WHERE id = any($1::bigint[]) + AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' ORDER BY id @@ -1318,18 +1319,20 @@ updated_job AS ( UPDATE /* TEMPLATE: schema */river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce($2::timestamptz, now()) + max_attempts = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN jsonb_typeof(metadata) = 'object' + THEN metadata - 'cancel_attempted_at' ELSE metadata END, + scheduled_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN scheduled_at ELSE coalesce($2::timestamptz, now()) END FROM job_to_update WHERE river_job.id = job_to_update.id -- Do not touch running jobs: AND river_job.state != 'running' - -- If the job is already available with a prior scheduled_at, leave it alone. - AND NOT ( + -- Already available jobs only need stale cancellation metadata cleared. + AND ((jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at') OR NOT ( river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) - ) + )) RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states ) SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states FROM ( diff --git a/riverdriver/riverdrivertest/job_read.go b/riverdriver/riverdrivertest/job_read.go index 05037a26a..c233843f2 100644 --- a/riverdriver/riverdrivertest/job_read.go +++ b/riverdriver/riverdrivertest/job_read.go @@ -810,6 +810,24 @@ func exerciseJobRead[TTx any](ctx context.Context, t *testing.T, executorWithTx require.Equal(t, []int64{cancelRequestedJob.ID}, ids) require.Equal(t, sqliteMalformedValue, sqliteJobColumnText(ctx, t, exec, invalidJob.ID, "metadata")) }) + + t.Run("NonObjectMetadataIgnored", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + ids := make([]int64, 0, 4) + for _, metadata := range []string{`["cancel_attempted_at"]`, `"cancel_attempted_at"`, `42`, `null`} { + job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{ + Metadata: []byte(metadata), + State: new(rivertype.JobStateRunning), + }) + ids = append(ids, job.ID) + } + + requested, err := exec.JobGetCancelRequested(ctx, &riverdriver.JobGetCancelRequestedParams{ID: ids}) + require.NoError(t, err) + require.Empty(t, requested) + }) }) t.Run("JobGetStuck", func(t *testing.T) { diff --git a/riverdriver/riverdrivertest/job_update.go b/riverdriver/riverdrivertest/job_update.go index 150445d63..a490d2da3 100644 --- a/riverdriver/riverdrivertest/job_update.go +++ b/riverdriver/riverdrivertest/job_update.go @@ -511,13 +511,38 @@ func exerciseJobUpdate[TTx any](ctx context.Context, t *testing.T, executorWithT t.Run("JobRetry", func(t *testing.T) { t.Parallel() + t.Run("ClearsCancellationWithoutReorderingAvailableJob", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + now := time.Now().UTC() + job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{ + Attempt: new(1), + MaxAttempts: new(1), + Metadata: []byte(`{"cancel_attempted_at":"2026-01-01T00:00:00Z","preserved":true}`), + ScheduledAt: new(now.Add(-time.Hour)), + }) + + retried, err := exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: job.ID, Now: &now}) + require.NoError(t, err) + require.JSONEq(t, `{"preserved":true}`, string(retried.Metadata)) + require.Equal(t, job.Attempt, retried.Attempt) + require.Equal(t, job.MaxAttempts, retried.MaxAttempts) + require.Equal(t, job.ScheduledAt, retried.ScheduledAt) + persisted, err := exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) + require.NoError(t, err) + require.Equal(t, retried, persisted) + }) + t.Run("DoesNotUpdateARunningJob", func(t *testing.T) { t.Parallel() exec, _ := setup(ctx, t) + metadata := []byte(`{"cancel_attempted_at":"2026-01-01T00:00:00Z","custom":"preserved"}`) job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{ - State: new(rivertype.JobStateRunning), + Metadata: metadata, + State: new(rivertype.JobStateRunning), }) jobAfter, err := exec.JobRetry(ctx, &riverdriver.JobRetryParams{ @@ -525,11 +550,13 @@ func exerciseJobUpdate[TTx any](ctx context.Context, t *testing.T, executorWithT }) require.NoError(t, err) require.Equal(t, rivertype.JobStateRunning, jobAfter.State) + require.JSONEq(t, string(metadata), string(jobAfter.Metadata)) require.WithinDuration(t, job.ScheduledAt, jobAfter.ScheduledAt, time.Microsecond) jobUpdated, err := exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) require.NoError(t, err) require.Equal(t, rivertype.JobStateRunning, jobUpdated.State) + require.JSONEq(t, string(metadata), string(jobUpdated.Metadata)) }) for _, state := range []rivertype.JobState{ @@ -561,6 +588,7 @@ func exerciseJobUpdate[TTx any](ctx context.Context, t *testing.T, executorWithT job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{ FinalizedAt: finalizedAt, + Metadata: []byte(`{"cancel_attempted_at":"2026-01-01T00:00:00Z","custom":"preserved"}`), ScheduledAt: new(now.Add(1 * time.Hour)), State: &state, }) @@ -571,12 +599,14 @@ func exerciseJobUpdate[TTx any](ctx context.Context, t *testing.T, executorWithT }) require.NoError(t, err) require.Equal(t, rivertype.JobStateAvailable, jobAfter.State) + require.JSONEq(t, `{"custom":"preserved"}`, string(jobAfter.Metadata)) require.WithinDuration(t, now, jobAfter.ScheduledAt, bundle.driver.TimePrecision()) jobUpdated, err := exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) require.NoError(t, err) require.Equal(t, rivertype.JobStateAvailable, jobUpdated.State) require.Nil(t, jobUpdated.FinalizedAt) + require.JSONEq(t, `{"custom":"preserved"}`, string(jobUpdated.Metadata)) }) } @@ -633,6 +663,64 @@ func exerciseJobUpdate[TTx any](ctx context.Context, t *testing.T, executorWithT require.Equal(t, rivertype.JobStateAvailable, jobUpdated.State) }) + t.Run("IncrementsExhaustedAttemptsAtExactSchedule", func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + now := time.Now().UTC().Truncate(time.Second) + job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{ + Attempt: new(1), + MaxAttempts: new(1), + ScheduledAt: &now, + }) + + retried, err := exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: job.ID, Now: &now}) + require.NoError(t, err) + require.Equal(t, 2, retried.MaxAttempts) + require.Equal(t, job.ScheduledAt, retried.ScheduledAt) + }) + + t.Run("PreservesInvalidMetadata", func(t *testing.T) { + t.Parallel() + + exec, bundle := setup(ctx, t) + if bundle.driver.DatabaseName() != riverdriver.DatabaseNameSQLite { + t.Skip("only SQLite's JSON columns can hold invalid JSON") + } + job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{State: new(rivertype.JobStateRetryable)}) + sqliteSetJobColumnMalformed(ctx, t, exec, job.ID, "metadata") + + // Retry still validates its returned row, but the SQL must preserve + // malformed metadata and let the decoder report the invalid field. + _, err := exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: job.ID}) + require.ErrorContains(t, err, "error unmarshaling `metadata`") + require.Equal(t, sqliteMalformedValue, sqliteJobColumnText(ctx, t, exec, job.ID, "metadata")) + }) + + t.Run("PreservesNonObjectMetadata", func(t *testing.T) { + t.Parallel() + + for _, metadata := range []string{`["cancel_attempted_at"]`, `false`, `null`, `42`, `"cancel_attempted_at"`} { + t.Run(metadata, func(t *testing.T) { + t.Parallel() + + exec, _ := setup(ctx, t) + job := testfactory.Job(ctx, t, exec, &testfactory.JobOpts{ + Metadata: []byte(metadata), + State: new(rivertype.JobStateRetryable), + }) + + jobAfter, err := exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: job.ID}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateAvailable, jobAfter.State) + require.JSONEq(t, metadata, string(jobAfter.Metadata)) + persisted, err := exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) + require.NoError(t, err) + require.JSONEq(t, metadata, string(persisted.Metadata)) + }) + } + }) + t.Run("ReturnsErrNotFoundIfJobNotFound", func(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverdrivertest/listener.go b/riverdriver/riverdrivertest/listener.go index 6b8df6de8..e87efe52c 100644 --- a/riverdriver/riverdrivertest/listener.go +++ b/riverdriver/riverdrivertest/listener.go @@ -2,6 +2,7 @@ package riverdrivertest import ( "context" + "encoding/json" "fmt" "testing" "time" @@ -111,8 +112,20 @@ func exerciseListener[TTx any](ctx context.Context, t *testing.T, driverWithPool notification := waitForNotification(ctx, t, listener) require.Equal(t, string(notifier.NotificationTopicControl), notification.Topic) - require.JSONEq(t, fmt.Sprintf(`{"action":"cancel","job_id":%d,"queue":%q}`, - committedJob.ID, committedJob.Queue), notification.Payload) + var payload struct { + Action string `json:"action"` + Attempt int `json:"attempt"` + AttemptedAt *time.Time `json:"attempted_at"` + JobID int64 `json:"job_id"` + Queue string `json:"queue"` + } + require.NoError(t, json.Unmarshal([]byte(notification.Payload), &payload)) + require.Equal(t, "cancel", payload.Action) + require.Equal(t, committedJob.Attempt, payload.Attempt) + require.NotNil(t, payload.AttemptedAt) + require.True(t, committedJob.AttemptedAt.Equal(*payload.AttemptedAt)) + require.Equal(t, committedJob.ID, payload.JobID) + require.Equal(t, committedJob.Queue, payload.Queue) requireNoNotification(ctx, t, listener) }) diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index d87361114..7b4b2ede6 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -40,7 +40,7 @@ CREATE TABLE river_job ( -- name: JobCancel :one WITH locked_job AS ( SELECT - id, queue, state, finalized_at + id, attempt, attempted_at, queue, state, finalized_at FROM /* TEMPLATE: schema */river_job WHERE river_job.id = @id FOR UPDATE @@ -50,7 +50,7 @@ notification AS ( id, CASE WHEN @notify::boolean THEN pg_notify( concat(coalesce(sqlc.narg('schema')::text, current_schema()), '.', @control_topic::text), - json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text ) END FROM locked_job @@ -262,6 +262,7 @@ ORDER BY id; SELECT id FROM /* TEMPLATE: schema */river_job WHERE id = any(@id::bigint[]) + AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' ORDER BY id; @@ -531,18 +532,20 @@ updated_job AS ( UPDATE /* TEMPLATE: schema */river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce(sqlc.narg('now')::timestamptz, now()) + max_attempts = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce(sqlc.narg('now')::timestamptz, now()) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce(sqlc.narg('now')::timestamptz, now()) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN jsonb_typeof(metadata) = 'object' + THEN metadata - 'cancel_attempted_at' ELSE metadata END, + scheduled_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce(sqlc.narg('now')::timestamptz, now()) THEN scheduled_at ELSE coalesce(sqlc.narg('now')::timestamptz, now()) END FROM job_to_update WHERE river_job.id = job_to_update.id -- Do not touch running jobs: AND river_job.state != 'running' - -- If the job is already available with a prior scheduled_at, leave it alone. - AND NOT ( + -- Already available jobs only need stale cancellation metadata cleared. + AND ((jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at') OR NOT ( river_job.state = 'available' AND river_job.scheduled_at < coalesce(sqlc.narg('now')::timestamptz, now()) - ) + )) RETURNING river_job.* ) SELECT * FROM ( diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go index a58f06ff0..af3ef49e5 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -16,7 +16,7 @@ import ( const jobCancel = `-- name: JobCancel :one WITH locked_job AS ( SELECT - id, queue, state, finalized_at + id, attempt, attempted_at, queue, state, finalized_at FROM /* TEMPLATE: schema */river_job WHERE river_job.id = $1 FOR UPDATE @@ -26,7 +26,7 @@ notification AS ( id, CASE WHEN $2::boolean THEN pg_notify( concat(coalesce($3::text, current_schema()), '.', $4::text), - json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text ) END FROM locked_job @@ -599,6 +599,7 @@ const jobGetCancelRequested = `-- name: JobGetCancelRequested :many SELECT id FROM /* TEMPLATE: schema */river_job WHERE id = any($1::bigint[]) + AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' ORDER BY id @@ -1282,18 +1283,20 @@ updated_job AS ( UPDATE /* TEMPLATE: schema */river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce($2::timestamptz, now()) + max_attempts = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN jsonb_typeof(metadata) = 'object' + THEN metadata - 'cancel_attempted_at' ELSE metadata END, + scheduled_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN scheduled_at ELSE coalesce($2::timestamptz, now()) END FROM job_to_update WHERE river_job.id = job_to_update.id -- Do not touch running jobs: AND river_job.state != 'running' - -- If the job is already available with a prior scheduled_at, leave it alone. - AND NOT ( + -- Already available jobs only need stale cancellation metadata cleared. + AND ((jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at') OR NOT ( river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) - ) + )) RETURNING river_job.id, river_job.args, river_job.attempt, river_job.attempted_at, river_job.attempted_by, river_job.created_at, river_job.errors, river_job.finalized_at, river_job.kind, river_job.max_attempts, river_job.metadata, river_job.priority, river_job.queue, river_job.state, river_job.scheduled_at, river_job.tags, river_job.unique_key, river_job.unique_states ) SELECT id, args, attempt, attempted_at, attempted_by, created_at, errors, finalized_at, kind, max_attempts, metadata, priority, queue, state, scheduled_at, tags, unique_key, unique_states FROM ( diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql index 8ef471786..0dc456e7b 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql @@ -586,22 +586,44 @@ WHERE id = @id UPDATE /* TEMPLATE: schema */river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) + max_attempts = CASE WHEN state = 'available' AND scheduled_at < coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = CASE WHEN state = 'available' AND scheduled_at < coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) + THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) THEN scheduled_at ELSE coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) END WHERE id = @id -- Do not touch running jobs: AND state != 'running' - -- If the job is already available with a prior scheduled_at, leave it alone. + -- Preserve the position of already available jobs while clearing stale cancellation metadata. -- -- I had to invert the original 'AND NOT' to 'AND'. Something about -- sqlc's SQLite parser cannot detect a parameter inside an `AND NOT`. An -- unfortunate bug that will hopefully be fixed in the future ... AND ( - state <> 'available' - OR scheduled_at > coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) + (CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR state <> 'available' + OR scheduled_at >= coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) ) -RETURNING *; +RETURNING + id, + cast(CASE WHEN typeof(args) = 'text' AND NOT json_valid(args) THEN args ELSE json(args) END AS blob) AS args, + attempt, + attempted_at, + cast(CASE WHEN typeof(attempted_by) = 'text' AND NOT json_valid(attempted_by) THEN attempted_by ELSE json(attempted_by) END AS blob) AS attempted_by, + created_at, + cast(CASE WHEN typeof(errors) = 'text' AND NOT json_valid(errors) THEN errors ELSE json(errors) END AS blob) AS errors, + finalized_at, + kind, + max_attempts, + cast(CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN metadata ELSE json(metadata) END AS blob) AS metadata, + priority, + queue, + state, + scheduled_at, + cast(CASE WHEN typeof(tags) = 'text' AND NOT json_valid(tags) THEN tags ELSE json(tags) END AS blob) AS tags, + unique_key, + unique_states; -- Selects only the columns needed to schedule jobs so that a job with a JSON -- column that isn't valid JSON doesn't fail scheduling. See JobGetAvailable. diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go index 8fc19993d..180754a1e 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go @@ -1450,22 +1450,44 @@ const jobRetry = `-- name: JobRetry :one UPDATE /* TEMPLATE: schema */river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce(cast(?1 AS text), datetime('now', 'subsec')) + max_attempts = CASE WHEN state = 'available' AND scheduled_at < coalesce(cast(?1 AS text), datetime('now', 'subsec')) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + finalized_at = CASE WHEN state = 'available' AND scheduled_at < coalesce(cast(?1 AS text), datetime('now', 'subsec')) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) + THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < coalesce(cast(?1 AS text), datetime('now', 'subsec')) THEN scheduled_at ELSE coalesce(cast(?1 AS text), datetime('now', 'subsec')) END WHERE id = ?2 -- Do not touch running jobs: AND state != 'running' - -- If the job is already available with a prior scheduled_at, leave it alone. + -- Preserve the position of already available jobs while clearing stale cancellation metadata. -- -- I had to invert the original 'AND NOT' to 'AND'. Something about -- sqlc's SQLite parser cannot detect a parameter inside an ` + "`" + `AND NOT` + "`" + `. An -- unfortunate bug that will hopefully be fixed in the future ... AND ( - state <> 'available' - OR scheduled_at > coalesce(cast(?1 AS text), datetime('now', 'subsec')) + (CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR state <> 'available' + OR scheduled_at >= coalesce(cast(?1 AS text), datetime('now', 'subsec')) ) -RETURNING id, json(args), attempt, attempted_at, json(attempted_by), created_at, json(errors), finalized_at, kind, max_attempts, json(metadata), priority, queue, state, scheduled_at, json(tags), unique_key, unique_states +RETURNING + id, + cast(CASE WHEN typeof(args) = 'text' AND NOT json_valid(args) THEN args ELSE json(args) END AS blob) AS args, + attempt, + attempted_at, + cast(CASE WHEN typeof(attempted_by) = 'text' AND NOT json_valid(attempted_by) THEN attempted_by ELSE json(attempted_by) END AS blob) AS attempted_by, + created_at, + cast(CASE WHEN typeof(errors) = 'text' AND NOT json_valid(errors) THEN errors ELSE json(errors) END AS blob) AS errors, + finalized_at, + kind, + max_attempts, + cast(CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN metadata ELSE json(metadata) END AS blob) AS metadata, + priority, + queue, + state, + scheduled_at, + cast(CASE WHEN typeof(tags) = 'text' AND NOT json_valid(tags) THEN tags ELSE json(tags) END AS blob) AS tags, + unique_key, + unique_states ` type JobRetryParams struct { @@ -1473,6 +1495,27 @@ type JobRetryParams struct { ID int64 } +type JobRetryRow struct { + ID int64 + Column2 []byte + Attempt int64 + AttemptedAt *time.Time + Column5 []byte + CreatedAt time.Time + Column7 []byte + FinalizedAt *time.Time + Kind string + MaxAttempts int64 + Column11 []byte + Priority int64 + Queue string + State string + ScheduledAt time.Time + Column16 []byte + UniqueKey []byte + UniqueStates *int64 +} + // Differs by necessity from other drivers because SQLite doesn't support // `UPDATE` inside CTEs so we can't retry if running but select otherwise. // Instead, the driver uses a transaction to optimisticaly try an update, but @@ -1482,26 +1525,26 @@ type JobRetryParams struct { // sqlc bug. Something about sqlc's SQLite parser cannot detect a parameter // inside an `AND NOT`. I'll try to get this fixed upstream at some point so we // can clean this up and keep it more like the Postgres version. -func (q *Queries) JobRetry(ctx context.Context, db DBTX, arg *JobRetryParams) (*RiverJob, error) { +func (q *Queries) JobRetry(ctx context.Context, db DBTX, arg *JobRetryParams) (*JobRetryRow, error) { row := db.QueryRowContext(ctx, jobRetry, arg.Now, arg.ID) - var i RiverJob + var i JobRetryRow err := row.Scan( &i.ID, - &i.Args, + &i.Column2, &i.Attempt, &i.AttemptedAt, - &i.AttemptedBy, + &i.Column5, &i.CreatedAt, - &i.Errors, + &i.Column7, &i.FinalizedAt, &i.Kind, &i.MaxAttempts, - &i.Metadata, + &i.Column11, &i.Priority, &i.Queue, &i.State, &i.ScheduledAt, - &i.Tags, + &i.Column16, &i.UniqueKey, &i.UniqueStates, ) diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql index 4c78b2a9d..39b8ce561 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql @@ -30,10 +30,11 @@ FROM /* TEMPLATE: schema */river_notification; -- name: NotificationInsertJobCancel :exec INSERT INTO /* TEMPLATE: schema */river_notification (payload, topic) -VALUES ( - json_object('action', 'cancel', 'job_id', cast(@job_id AS integer), 'queue', cast(@queue AS text)), +SELECT + json_object('action', 'cancel', 'attempt', job.attempt, 'attempted_at', strftime('%Y-%m-%dT%H:%M:%fZ', job.attempted_at), 'job_id', job.id, 'queue', cast(@queue AS text)), @topic -); +FROM /* TEMPLATE: schema */river_job AS job +WHERE job.id = @job_id; -- name: NotificationInsertMany :exec INSERT INTO /* TEMPLATE: schema */river_notification ( diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go index a5cf11495..032918c56 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go @@ -91,20 +91,21 @@ func (q *Queries) NotificationGetLastID(ctx context.Context, db DBTX) (int64, er const notificationInsertJobCancel = `-- name: NotificationInsertJobCancel :exec INSERT INTO /* TEMPLATE: schema */river_notification (payload, topic) -VALUES ( - json_object('action', 'cancel', 'job_id', cast(?1 AS integer), 'queue', cast(?2 AS text)), - ?3 -) +SELECT + json_object('action', 'cancel', 'attempt', job.attempt, 'attempted_at', strftime('%Y-%m-%dT%H:%M:%fZ', job.attempted_at), 'job_id', job.id, 'queue', cast(?1 AS text)), + ?2 +FROM /* TEMPLATE: schema */river_job AS job +WHERE job.id = ?3 ` type NotificationInsertJobCancelParams struct { - JobID int64 Queue string Topic string + JobID int64 } func (q *Queries) NotificationInsertJobCancel(ctx context.Context, db DBTX, arg *NotificationInsertJobCancelParams) error { - _, err := db.ExecContext(ctx, notificationInsertJobCancel, arg.JobID, arg.Queue, arg.Topic) + _, err := db.ExecContext(ctx, notificationInsertJobCancel, arg.Queue, arg.Topic, arg.JobID) return err } diff --git a/riverdriver/riversqlite/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index 7d1044bfc..b17c6105f 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -859,7 +859,7 @@ func (e *Executor) JobRetry(ctx context.Context, params *riverdriver.JobRetryPar return nil, interpretError(err) } - return jobRowFromInternal(job) + return jobRowFromInternal(jobFromReturningRow((*dbsqlc.JobGetAvailableRow)(job))) }) } diff --git a/ruby/lib/client_runtime.rb b/ruby/lib/client_runtime.rb index 47cb2acaf..5ab995e8b 100644 --- a/ruby/lib/client_runtime.rb +++ b/ruby/lib/client_runtime.rb @@ -24,6 +24,8 @@ def initialize(started_at, started_monotonic, scheduled_at) attr_reader :periodic_jobs def initialize(client, driver, config) + @cancellation_generation = 0 + @cancellation_pending = false @client = client @condition = ConditionVariable.new @config = config @@ -33,7 +35,7 @@ def initialize(client, driver, config) @maintenance_condition = ConditionVariable.new @maintenance_generation = 0 @mutex = Mutex.new - @notifications_disabled = false + @notifications_disabled = config.poll_only || !driver.respond_to?(:notification_listener) @producer_threads = {} @queue_configs = config.queues.dup @queue_generations = Hash.new(0) @@ -120,7 +122,8 @@ def healthy? @mutex.synchronize do @stop_requested || (@producer_threads.all? { |name, thread| @removed_queues[name] || thread.alive? } && (!@maintenance_thread || @maintenance_thread.alive?) && - (!@notification_thread || @notifications_disabled || @notification_thread.alive?)) + (!@notification_thread || @notifications_disabled || @notification_thread.alive?) && + (!@cancellation_thread || @cancellation_thread.alive?)) end end @@ -197,6 +200,9 @@ def start @started = true @stop_requested = false @stopped = false + thread = Thread.new { cancellation_loop } + @cancellation_thread = thread + @threads << thread end begin @@ -243,6 +249,7 @@ def stop(cancel: false, wait: true) @stopped = true @threads.clear @producer_threads.clear + @cancellation_thread = nil @maintenance_thread = nil @notification_thread = nil @leader = false @@ -271,6 +278,16 @@ def wake end private def begin_work(id) + entry = @mutex.synchronize { @running.fetch(id) } + if entry[:check_before_work] + begin + check_remote_cancellations(entries: {id => entry}) + rescue => error + @config.logger.error("River cancellation check failed: #{error.full_message}") + request_cancellation_check + end + end + interruption = @mutex.synchronize do entry = @running.fetch(id) if @stop_requested @@ -286,12 +303,48 @@ def wake raise interruption if interruption end - private def check_remote_cancellations(queue) - entries = @mutex.synchronize { @running.select { |_id, entry| entry[:queue] == queue && entry[:working] } } + private def cancellation_loop + retrying = false + loop do + @mutex.synchronize do + interval = @queue_configs.values.map { |config| config.resolved_fetch_poll_interval(@config) }.min || @config.fetch_poll_interval + deadline = (@config.poll_only || @notifications_disabled || retrying) ? monotonic_now + interval : nil + loop do + return if @stop_requested && @running.empty? + break if @cancellation_pending + + deadline ||= monotonic_now + interval if @notifications_disabled + remaining = deadline && deadline - monotonic_now + break if remaining && remaining <= 0 + + @condition.wait(@mutex, remaining) + end + @cancellation_pending = false + end + retrying = false + begin + check_remote_cancellations + rescue => error + retrying = true + @config.logger.error("River cancellation check failed: #{error.full_message}") + end + end + end + + private def check_remote_cancellations(entries: nil) + # Include registered attempts whose worker has not started yet: a reconnect + # can recover their cancellation between registration and begin_work. + entries ||= @mutex.synchronize { @running.dup } return if entries.empty? @driver.job_get_cancelled_ids(entries.keys).each do |id| - entries.fetch(id)[:thread].raise(JobCancelError) + @mutex.synchronize do + entry = entries.fetch(id) + next unless @running[id].equal?(entry) + + entry[:cancelled] = true + entry[:thread].raise(JobCancelError) if entry[:working] + end end end @@ -511,7 +564,7 @@ def wake operation.call end - private def launch(row) + private def launch(row, generation = nil) gate = ::Queue.new thread = Thread.new do gate.pop @@ -524,7 +577,7 @@ def wake end @mutex.synchronize do - @running[row.id] = {cancelled: @fetching_queues.fetch(row.queue).include?(row.id), queue: row.queue, thread: thread, working: false} + @running[row.id] = {attempt: row.attempt, attempted_at: row.attempted_at, check_before_work: @fetching_queues.fetch(row.queue).include?(row.id) || (generation && generation != @cancellation_generation), queue: row.queue, thread: thread, working: false} end gate.push(true) end @@ -634,6 +687,7 @@ def wake unless listener listener = @driver.notification_listener @mutex.synchronize { @notifications_disabled = listener.nil? } + request_cancellation_check ready << true if ready ready = nil return unless listener @@ -698,7 +752,6 @@ def wake draining = queue_stopping?(queue) break if draining && running_count(queue).zero? - check_remote_cancellations(queue) if draining wait_for_running_jobs(queue, poll_interval) next @@ -719,11 +772,14 @@ def wake queue_state = @driver.queue_get(queue) publish_queue(queue_state.paused_at ? EVENT_QUEUE_PAUSED : EVENT_QUEUE_RESUMED, queue_state) if queue_state if capacity.positive? && !queue_state&.paused_at - @mutex.synchronize { @fetching_queues[queue] = [] } + generation = @mutex.synchronize do + @fetching_queues[queue] = [] + @cancellation_generation + end begin jobs = @driver.job_get_available(attempted_by: @config.id, max: capacity, queue: queue, **fetch_options) last_fetch = monotonic_now - jobs.each { |job| launch(job) } + jobs.each { |job| launch(job, generation) } next unless jobs.empty? ensure @mutex.synchronize { @fetching_queues.delete(queue) } @@ -757,6 +813,14 @@ def wake @mutex.synchronize { @subscriptions.dup }.each { |subscription| subscription.publish(event) } end + private def request_cancellation_check + @mutex.synchronize do + @cancellation_generation += 1 + @cancellation_pending = true + @condition.broadcast + end + end + private def queue_stopping?(queue) @mutex.synchronize { @stop_requested || @removed_queues[queue] } end @@ -778,9 +842,18 @@ def wake @mutex.synchronize do entry = @running[message["job_id"]] + if entry && message["attempt"].is_a?(Integer) && message["attempt"].positive? && message["attempted_at"].is_a?(String) + if entry[:queue] == queue && entry[:attempt] == message["attempt"] && entry[:attempted_at] == Time.iso8601(message["attempted_at"]) + entry[:cancelled] = true + entry[:thread].raise(JobCancelError) if entry[:working] + end + return + end + @cancellation_generation += 1 + @cancellation_pending = true + @condition.broadcast if entry && entry[:queue] == queue - entry[:cancelled] = true - entry[:thread].raise(JobCancelError) if entry[:working] + entry[:check_before_work] = true unless entry[:working] elsif (pending = @fetching_queues[queue]) # A claim may have committed before its workers are registered. pending << message["job_id"] @@ -803,7 +876,7 @@ def wake wake_maintenance if message["leader_id"].is_a?(String) && message["leader_id"] != @config.id end end - rescue JSON::ParserError, TypeError => error + rescue ArgumentError, JSON::ParserError, TypeError => error @config.logger.warn("River ignored invalid notification: #{error.message}") end diff --git a/ruby/lib/driver/runtime.rb b/ruby/lib/driver/runtime.rb index 39a46c0f0..afc90cf63 100644 --- a/ruby/lib/driver/runtime.rb +++ b/ruby/lib/driver/runtime.rb @@ -60,7 +60,7 @@ def job_cancel(id, now: Time.now.utc) row = job_get_by_id(updated_id || original_id) # Other engines interrupt running workers through this notification; # the durable marker remains Ruby's polling fallback. - runtime_notify("river_control", action: "cancel", job_id: row.id, queue: row.queue) if updated_id + runtime_notify("river_control", action: "cancel", attempt: row.attempt, attempted_at: row.attempted_at&.iso8601(6), job_id: row.id, queue: row.queue) if updated_id row end end @@ -249,27 +249,29 @@ def job_rescue_stuck(horizon:, retry_policy:, now: Time.now.utc, max: 1_000, log end def job_retry(id, now: Time.now.utc) + already_available = "state = 'available' AND scheduled_at < #{runtime_time(now)}" + cancellation_requested = runtime_postgres? ? "(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')" : "(CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL" updated_id, job = transaction do # Check the counter in the write itself so a concurrent claim cannot # advance it between validation and the update. updated_id = runtime_returning_ids(<<~SQL).first UPDATE river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - metadata = #{runtime_postgres? ? "metadata - 'cancel_attempted_at'" : "jsonb_remove(metadata, '$.cancel_attempted_at')"}, - finalized_at = NULL, - scheduled_at = #{runtime_time(now)} + max_attempts = CASE WHEN #{already_available} THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + metadata = #{runtime_postgres? ? "CASE WHEN jsonb_typeof(metadata) = 'object' THEN metadata - 'cancel_attempted_at' ELSE metadata END" : "CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END"}, + finalized_at = CASE WHEN #{already_available} THEN finalized_at ELSE NULL END, + scheduled_at = CASE WHEN #{already_available} THEN scheduled_at ELSE #{runtime_time(now)} END WHERE id = #{Integer(id)} AND state != 'running' - AND (state != 'available' OR scheduled_at > #{runtime_time(now)}) - AND attempt < #{River::MAX_ATTEMPTS_LIMIT} + AND (#{cancellation_requested} OR NOT (#{already_available})) + AND ((#{already_available}) OR attempt < #{River::MAX_ATTEMPTS_LIMIT}) RETURNING id SQL job = job_get_by_id(updated_id || id) runtime_notify("river_insert", queue: job.queue) if updated_id && job [updated_id, job] end - if !updated_id && job && job.state != "running" && (job.state != "available" || job.scheduled_at > now) && job.attempt >= River::MAX_ATTEMPTS_LIMIT + if !updated_id && job && job.state != "running" && (job.state != "available" || job.scheduled_at >= now) && job.attempt >= River::MAX_ATTEMPTS_LIMIT raise ArgumentError, "cannot retry a job with #{River::MAX_ATTEMPTS_LIMIT} or more attempts" end job @@ -484,7 +486,7 @@ def time_precision end private def runtime_cancel_attempted - runtime_postgres? ? "metadata ? 'cancel_attempted_at'" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" + runtime_postgres? ? "(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" end private def runtime_claim_jobs(predicate, max, attempted_by, now) diff --git a/ruby/sig/runtime.rbs b/ruby/sig/runtime.rbs index 2782a3465..ba7212e48 100644 --- a/ruby/sig/runtime.rbs +++ b/ruby/sig/runtime.rbs @@ -360,6 +360,9 @@ module River def initialize: (Time, Float, Time) -> void end + @cancellation_generation: Integer + @cancellation_pending: bool + @cancellation_thread: Thread? @client: Client @condition: ConditionVariable @config: Config @@ -406,7 +409,8 @@ module River private def begin_work: (Integer) -> untyped - def check_remote_cancellations: (String) -> untyped + def cancellation_loop: () -> untyped + def check_remote_cancellations: (?entries: Hash[Integer, untyped]?) -> untyped def error_handler_cancel?: (Exception, Job) -> bool def execute: (JobRow) -> untyped def finish_failed: (JobRow, Job, Exception, JobTiming, ?cancelled: bool, ?worker: untyped) -> untyped @@ -416,7 +420,7 @@ module River def finish_work: (Integer) -> untyped def invoke_plugins: (Symbol, *untyped) -> Array[untyped] def invoke_worker: (untyped, Job) -> untyped - def launch: (JobRow) -> untyped + def launch: (JobRow, ?Integer?) -> untyped def maintenance_loop: () -> untyped def monotonic_now: () -> Float def next_retry: (JobRow, Exception, Time, ?worker: untyped) -> Time @@ -428,6 +432,7 @@ module River def periodic_jobs_changed: () -> untyped def producer_loop: (String, QueueConfig) -> untyped def publish: (Symbol, JobRow, JobTiming) -> untyped + def request_cancellation_check: () -> void def queue_stopping?: (String) -> bool def receive_notification: (untyped, untyped) -> untyped def remove_subscription: (Subscription) -> Subscription? diff --git a/ruby/spec/client_runtime_branch_spec.rb b/ruby/spec/client_runtime_branch_spec.rb index 5510cad55..3cea6e3ad 100644 --- a/ruby/spec/client_runtime_branch_spec.rb +++ b/ruby/spec/client_runtime_branch_spec.rb @@ -243,7 +243,7 @@ def execute(runtime, value) value.periodic_jobs.add(periodic) - expect(value.instance_variable_get(:@threads)).to eq([thread]) + expect(value.instance_variable_get(:@threads)).to eq([value.instance_variable_get(:@cancellation_thread), thread]) expect(entered).to be_empty ensure release << true if release @@ -329,7 +329,10 @@ def execute(runtime, value) value.start - expect(value.instance_variable_get(:@threads)).to be_empty + expect(value.instance_variable_get(:@producer_threads)).to be_empty + ensure + driver.define_singleton_method(:leader_release) { |_| } + value&.stop end it "allows queue changes during startup without launching a producer twice" do @@ -344,9 +347,12 @@ def execute(runtime, value) value.start value.send(:start_producer, "dynamic", River::QueueConfig.new(max_workers: 1)) - expect(value.instance_variable_get(:@threads).length).to eq(2) + expect(value.instance_variable_get(:@producer_threads).length).to eq(2) value.queue_remove("branch") value.queue_remove("dynamic") + ensure + driver.define_singleton_method(:leader_release) { |_| } + value&.stop end it "joins only work belonging to a queue as it is removed" do @@ -439,7 +445,7 @@ def execute(runtime, value) end value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { now } - value.define_singleton_method(:launch) { |_| @stop_requested = true if fetched_at.length == 2 } + value.define_singleton_method(:launch) { |_, _generation| @stop_requested = true if fetched_at.length == 2 } value.define_singleton_method(:wait) do |duration| waits << duration now += (waits.length == 1) ? 0.25 : duration @@ -463,7 +469,7 @@ def execute(runtime, value) end value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { now } - value.define_singleton_method(:launch) { |_| } + value.define_singleton_method(:launch) { |_, _generation| } value.define_singleton_method(:wait) { |_| @stop_requested = true } value.send(:producer_loop, "branch", River::QueueConfig.new(max_workers: 1, fetch_cooldown: 1, fetch_poll_interval: 1)) @@ -471,15 +477,10 @@ def execute(runtime, value) expect(fetched_at).to eq([100.0]) end - it "keeps observing active jobs when shutdown interrupts the fetch cooldown" do - checked = [] + it "drains active jobs when shutdown interrupts the fetch cooldown" do fetched = 0 job = row driver = Object.new - driver.define_singleton_method(:job_get_cancelled_ids) { |ids| - checked << ids - [] - } driver.define_singleton_method(:queue_get) { |_| nil } driver.define_singleton_method(:job_get_available) { |**| fetched += 1 @@ -487,7 +488,7 @@ def execute(runtime, value) } value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { 100.0 } - value.define_singleton_method(:launch) do |job| + value.define_singleton_method(:launch) do |job, _generation| @running[job.id] = {queue: "branch", thread: Thread.current, working: true} end value.define_singleton_method(:wait) { |_| @stop_requested = true } @@ -496,7 +497,22 @@ def execute(runtime, value) value.send(:producer_loop, "branch", River::QueueConfig.new(max_workers: 2, fetch_cooldown: 1, fetch_poll_interval: 1)) expect(fetched).to eq(1) - expect(checked).to eq([[1], [1]]) + expect(value.instance_variable_get(:@running)).to be_empty + end + + it "drains active jobs when a fetch fails during shutdown" do + driver = Object.new + value = runtime(driver: driver) + value.instance_variable_set(:@running, {1 => {queue: "branch", working: true}}) + driver.define_singleton_method(:queue_get) do |_queue| + value.instance_variable_set(:@stop_requested, true) + raise "fetch failed during shutdown" + end + value.define_singleton_method(:wait_for_running_jobs) { |_, _| @running.clear } + + value.send(:producer_loop, "branch", River::QueueConfig.new(max_workers: 2)) + + expect(value.instance_variable_get(:@running)).to be_empty end it "rechecks queue pause state after waiting for its fetch cooldown" do @@ -512,7 +528,7 @@ def execute(runtime, value) end value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { now } - value.define_singleton_method(:launch) { |_| @stop_requested = true if fetched_at.length == 2 } + value.define_singleton_method(:launch) { |_, _generation| @stop_requested = true if fetched_at.length == 2 } value.define_singleton_method(:wait) do |duration| @stop_requested = true if paused_at paused_at = Time.now.utc @@ -746,7 +762,58 @@ def work(_job) expect { execute(value, row) }.not_to raise_error end - it "checks cancellation in one batch, excluding inactive work and other queues" do + [false, true].each do |stopping| + it "cancels work during a blocked fetch#{" and graceful stop" if stopping}" do + started, fetching, release = Queue.new, Queue.new, Queue.new + worker = Object.new + worker.define_singleton_method(:work) do |_job| + started << true + Queue.new.pop + end + job = row + marked = false + claims = 0 + driver = Object.new + driver.define_singleton_method(:queue_upsert) { |_name| } + driver.define_singleton_method(:queue_get) { |_name| nil } + driver.define_singleton_method(:job_get_available) do |**| + claims += 1 + if claims == 1 + [job] + else + fetching << true + release.pop + [] + end + end + driver.define_singleton_method(:job_get_cancelled_ids) { |ids| marked ? ids : [] } + driver.define_singleton_method(:job_set_state_if_running) do |**attributes| + job.state = attributes.fetch(:state) + job + end + value = described_class.new(Object.new, driver, + config(worker: worker, queues: {branch: 2}).with( + fetch_cooldown: 0.001, fetch_poll_interval: 0.01, leader_election_disabled: true + )) + subscription = value.subscribe([:job_cancelled]) + value.start + Timeout.timeout(2) { + started.pop + fetching.pop + } + value.stop(wait: false) if stopping + marked = true + + event = Timeout.timeout(2) { subscription.pop } + expect(event.job.state).to eq(River::JOB_STATE_CANCELLED) + ensure + value&.stop(cancel: true, wait: false) + release << true + value&.stop + end + end + + it "checks cancellation across queues, including registered work that has not started" do raised = [] fake_thread = Object.new fake_thread.define_singleton_method(:raise) { |error| raised << error } @@ -754,7 +821,7 @@ def work(_job) driver = Object.new driver.define_singleton_method(:job_get_cancelled_ids) { |ids| checked << ids - [3] + [2, 3] } value = runtime(driver: driver) value.instance_variable_set( @@ -766,10 +833,11 @@ def work(_job) } ) - value.send(:check_remote_cancellations, "branch") + value.send(:check_remote_cancellations) expect(raised).to eq([River::JobCancelError]) - expect(checked).to eq([[3]]) + expect(checked).to eq([[1, 2, 3]]) + expect { value.send(:begin_work, 2) }.to raise_error(River::JobCancelError) end it "interrupts work that begins after stop was requested" do diff --git a/ruby/spec/conformance_spec.rb b/ruby/spec/conformance_spec.rb index d08025afc..dc68fb3cf 100644 --- a/ruby/spec/conformance_spec.rb +++ b/ruby/spec/conformance_spec.rb @@ -163,7 +163,7 @@ class ConformanceRecord < ActiveRecord::Base row = @client.insert(River::JobArgsHash.new(:notify, {}), queue: "priority").job if name == "cancel" @driver.send(:runtime_execute, "UPDATE river_job SET id = 42 WHERE id = #{row.id}") - @driver.job_claim(id: 42, attempted_by: "conformance") + @driver.job_claim(id: 42, attempted_by: "conformance", allow_scheduled: true, now: Time.iso8601("2026-01-02T03:04:05.678Z")) @client.job_cancel(42) end when "metadata_changed", "pause", "resume" @@ -189,7 +189,13 @@ class ConformanceRecord < ActiveRecord::Base ((name == "insert") ? "insert" : "control") end expect(@driver.send(:runtime_value, notification, :topic)).to eq(protocol.fetch("topics").fetch(topic)) - expect(JSON.parse(@driver.send(:runtime_value, notification, :payload))).to eq(fixture.fetch("payload")) + payload = JSON.parse(@driver.send(:runtime_value, notification, :payload)) + expected = fixture.fetch("payload").dup + if name == "cancel" + payload["attempted_at"] = Time.iso8601(payload.fetch("attempted_at")) + expected["attempted_at"] = Time.iso8601(expected.fetch("attempted_at")) + end + expect(payload).to eq(expected) end end diff --git a/ruby/spec/driver_runtime_feature_spec.rb b/ruby/spec/driver_runtime_feature_spec.rb index 8cfef94f7..2b9aab690 100644 --- a/ruby/spec/driver_runtime_feature_spec.rb +++ b/ruby/spec/driver_runtime_feature_spec.rb @@ -137,7 +137,7 @@ expect(postgres_driver.send(:runtime_state, River::JOB_STATE_AVAILABLE)).to end_with("::river_job_state") expect(postgres_driver.send(:runtime_time, Time.utc(2026, 1, 2))).to end_with("::timestamptz") expect(postgres_driver.send(:runtime_merge_metadata, {"ready" => true})).to include("metadata ||") - expect(postgres_driver.send(:runtime_cancel_attempted)).to eq("metadata ? 'cancel_attempted_at'") + expect(postgres_driver.send(:runtime_cancel_attempted)).to eq("(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')") expect(postgres_driver.send(:runtime_update_value, :attempted_by, ["one"])).to include("ARRAY[", "::text[]") expect(postgres_driver.send(:runtime_update_value, :errors, [7])).to include("ARRAY[", "::jsonb[]") expect(postgres_driver.send(:runtime_queue_columns)).to eq("name, created_at, metadata, paused_at, updated_at") diff --git a/ruby/spec/driver_runtime_shared_examples.rb b/ruby/spec/driver_runtime_shared_examples.rb index 785b8a906..b6f80b857 100644 --- a/ruby/spec/driver_runtime_shared_examples.rb +++ b/ruby/spec/driver_runtime_shared_examples.rb @@ -422,6 +422,13 @@ def work(_job) expect(driver.job_get_by_id(row.id).state).to eq("running") end + it "ignores cancellation names stored in non-object metadata" do + job = insert_job(1) + metadata = driver.send(:runtime_json, ["cancel_attempted_at"]) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = #{metadata} WHERE id = #{job.id}") + expect(driver.job_get_cancelled_ids([job.id])).to eq([]) + end + it "looks up cancellation markers for a batch, excluding missing and uncancelled jobs" do cancelled = insert_job(1) unmarked = insert_job(2) @@ -612,6 +619,51 @@ def insert_job(value, **options) expect(driver.job_retry(-1)).to be_nil end + it "increments exhausted attempts at the exact retry timestamp" do + job = insert_job(1, scheduled_at: now, max_attempts: 1) + driver.send(:runtime_execute, "UPDATE river_job SET attempt = max_attempts WHERE id = #{job.id}") + expect(driver.job_retry(job.id, now: now)).to have_attributes(max_attempts: 2, scheduled_at: job.scheduled_at) + end + + it "clears cancellation metadata without reordering an already available retry" do + job = insert_job(1, scheduled_at: now - 60, max_attempts: 1, + metadata: {"cancel_attempted_at" => now.iso8601, "preserved" => true}) + driver.send(:runtime_execute, "UPDATE river_job SET attempt = max_attempts WHERE id = #{job.id}") + + retried = driver.job_retry(job.id, now: now) + expect(retried).to have_attributes(state: "available", attempt: 1, max_attempts: 1, + scheduled_at: job.scheduled_at, metadata: job.metadata.except("cancel_attempted_at")) + expect(driver.job_get_by_id(job.id).metadata).to eq(retried.metadata) + end + + it "preserves malformed SQLite metadata and reports its decoding error", database: :sqlite do + skip "requires SQLite" if driver.send(:runtime_postgres?) + + job = insert_job(1, scheduled_at: now + 60) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = '[not json' WHERE id = #{job.id}") + + expect { driver.job_retry(job.id, now: now) }.to raise_error(River::JobRowDecodeError, /metadata/) + row = driver.send(:runtime_query_rows, "SELECT metadata FROM river_job WHERE id = #{job.id}").first + expect(driver.send(:runtime_value, row, :metadata)).to eq("[not json") + end + + [nil, 42, false, "cancel_attempted_at", ["cancel_attempted_at"]].each do |metadata| + it "preserves non-object metadata #{metadata.inspect} when retrying" do + job = insert_job(1, scheduled_at: now + 60) + encoded = driver.send(:runtime_json, JSON.generate(metadata)) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = #{encoded} WHERE id = #{job.id}") + + if metadata.nil? + expect(driver.job_retry(job.id, now: now)).to have_attributes(state: "available") + else + expect { driver.job_retry(job.id, now: now) }.to raise_error(River::JobRowDecodeError) + end + column = driver.send(:runtime_postgres?) ? "metadata::text" : "json(metadata)" + row = driver.send(:runtime_query_rows, "SELECT #{column} AS metadata FROM river_job WHERE id = #{job.id}").first + expect(JSON.parse(driver.send(:runtime_value, row, :metadata))).to eq(metadata) + end + end + it "rejects manual retries after the last representable attempt without changing the job" do job = insert_job(1, max_attempts: 32_767) original = client.job_update(job.id, attempt: 32_767, finalized_at: now, state: :discarded) diff --git a/ruby/spec/insert_notification_shared_examples.rb b/ruby/spec/insert_notification_shared_examples.rb index 2ac9fce40..3805f0c45 100644 --- a/ruby/spec/insert_notification_shared_examples.rb +++ b/ruby/spec/insert_notification_shared_examples.rb @@ -73,12 +73,12 @@ def insert_notifications it "publishes Go-compatible cancellation only when the transaction commits" do client = River::Client.new(@driver) row = client.insert(River::JobArgsHash.new(:notify, {})).job - @driver.job_claim(id: row.id, attempted_by: "test") + claimed = @driver.job_claim(id: row.id, attempted_by: "test") @driver.transaction do client.job_cancel(row.id) expect(insert_notifications).to be_empty end - expect(insert_notifications).to eq([{"action" => "cancel", "job_id" => row.id, "queue" => "default"}]) + expect(insert_notifications).to eq([{"action" => "cancel", "attempt" => claimed.attempt, "attempted_at" => claimed.attempted_at.iso8601(6), "job_id" => row.id, "queue" => "default"}]) end it "does not notify rolled-back, missing, or already-finalized cancellations" do @@ -107,14 +107,14 @@ def cancellation_notifications it "makes cancellation and its outbox notification visible together on commit" do client = River::Client.new(@driver) row = client.insert(River::JobArgsHash.new(:notify, {})).job - @driver.job_claim(id: row.id, attempted_by: "test") + claimed = @driver.job_claim(id: row.id, attempted_by: "test") @driver.transaction do client.job_cancel(row.id) expect(cancellation_notifications.length).to eq(1) observer = Thread.new { cancellation_notifications } expect(Timeout.timeout(5) { observer.value }).to be_empty end - expect(cancellation_notifications).to eq([{"action" => "cancel", "job_id" => row.id, "queue" => "default"}]) + expect(cancellation_notifications).to eq([{"action" => "cancel", "attempt" => claimed.attempt, "attempted_at" => claimed.attempted_at.iso8601(6), "job_id" => row.id, "queue" => "default"}]) end it "rolls back both cancellation and the outbox row" do diff --git a/ruby/spec/notification_runtime_spec.rb b/ruby/spec/notification_runtime_spec.rb index d3e6e6d57..58043d8f3 100644 --- a/ruby/spec/notification_runtime_spec.rb +++ b/ruby/spec/notification_runtime_spec.rb @@ -12,6 +12,69 @@ )) end + [false, true].each do |fail_first| + it "waits for cancellation hints and retries failures=#{fail_first}" do + driver.define_singleton_method(:notification_listener) { Object.new } + value = River::ClientRuntime.new(Object.new, driver, River::Config.new( + fetch_cooldown: 0.001, fetch_poll_interval: 0.01, leader_election_disabled: true, logger: Logger.new(output) + )) + queries = Queue.new + failures = 0 + driver.define_singleton_method(:job_get_cancelled_ids) do |_ids| + queries << true + if fail_first && failures.zero? + failures += 1 + raise "temporary read failure" + end + [] + end + value.instance_variable_set(:@running, 42 => {queue: "one", thread: Thread.current, working: true}) + checker = Thread.new { value.send(:cancellation_loop) } + sleep 0.05 + expect(queries).to be_empty + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) + Timeout.timeout(2) { queries.pop } + Timeout.timeout(2) { queries.pop } if fail_first + sleep 0.05 + expect(queries).to be_empty + ensure + value&.instance_variable_get(:@mutex)&.synchronize do + value.instance_variable_set(:@stop_requested, true) + value.instance_variable_get(:@running).clear + value.instance_variable_get(:@condition).broadcast + end + checker&.join + end + end + + it "delivers an identified cancellation while a read is blocked and ignores older attempts" do + value = runtime + queried = Queue.new + release = Queue.new + raised = Queue.new + driver.define_singleton_method(:job_get_cancelled_ids) do |_ids| + queried << true + release.pop + [] + end + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + attempted_at = Time.iso8601("2026-01-02T03:04:05.6789Z") + entry = {attempt: 2, attempted_at: attempted_at, queue: "one", thread: thread, working: true} + value.instance_variable_set(:@running, 42 => entry) + checker = Thread.new { value.send(:check_remote_cancellations) } + Timeout.timeout(2) { queried.pop } + [{attempt: 1, attempted_at: attempted_at}, {attempt: 2, attempted_at: attempted_at - 1}].each do |identity| + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", job_id: 42, queue: "one", attempt: identity[:attempt], attempted_at: identity[:attempted_at].iso8601(6))) + expect(raised).to be_empty + end + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", job_id: 42, queue: "one", attempt: 2, attempted_at: attempted_at.iso8601(6))) + expect(Timeout.timeout(2) { raised.pop }).to eq(River::JobCancelError) + ensure + release&.push(true) + checker&.join + end + it "backs off repeated connection failures without indefinitely delaying shutdown" do listener = Object.new listener.define_singleton_method(:poll) { |_| raise "connection lost" } @@ -30,6 +93,7 @@ end it "cancels an attempt before its worker starts but ignores another queue's job" do + driver.define_singleton_method(:job_get_cancelled_ids) { |ids| ids } entry = {queue: "one", thread: Thread.current, working: false} runtime.instance_variable_set(:@running, 42 => entry) runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "two", job_id: 42)) @@ -40,8 +104,79 @@ expect(entry[:working]).to be false end + [false, true].each do |working| + it "ignores a delayed cancellation notice after retry#{" while working" if working}" do + raised = [] + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + entry = {queue: "one", thread: thread, working: working} + runtime.instance_variable_set(:@running, 42 => entry) + driver.define_singleton_method(:job_get_cancelled_ids) { |_ids| [] } + + runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) + runtime.send(:check_remote_cancellations) + runtime.send(:begin_work, 42) unless working + + expect(raised).to be_empty + expect(entry[:cancelled]).to be_nil + expect(entry[:working]).to be true + end + end + + it "does not interrupt a replacement registered during a cancellation check" do + raised = [] + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + original = {queue: "one", thread: thread, working: true} + replacement = original.dup + value = runtime + value.instance_variable_set(:@running, 42 => original) + driver.define_singleton_method(:job_get_cancelled_ids) do |ids| + value.instance_variable_set(:@running, 42 => replacement) + ids + end + + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) + value.send(:check_remote_cancellations) + + expect(raised).to be_empty + expect(replacement[:cancelled]).to be_nil + end + + it "retries a failed cancellation check made before the worker starts" do + raised = Queue.new + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + entry = {check_before_work: true, queue: "one", thread: thread, working: false} + runtime.instance_variable_set(:@running, 42 => entry) + queries = 0 + driver.define_singleton_method(:job_get_cancelled_ids) do |ids| + queries += 1 + raise "temporary read failure" if queries == 1 + + ids + end + + runtime.send(:begin_work, 42) + checker = Thread.new { runtime.send(:cancellation_loop) } + + expect(Timeout.timeout(2) { raised.pop }).to eq(River::JobCancelError) + expect(entry[:cancelled]).to be true + expect(output.string).to include("temporary read failure") + ensure + runtime.instance_variable_get(:@mutex).synchronize do + runtime.instance_variable_set(:@stop_requested, true) + runtime.instance_variable_get(:@running).clear + runtime.instance_variable_get(:@condition).broadcast + end + checker&.join + end + it "ignores invalid payloads without disrupting the receiver" do [nil, "{", "null", "[]", "{}"].each { |payload| runtime.send(:receive_notification, "river_insert", payload) } + runtime.instance_variable_set(:@running, 42 => {attempt: 1, queue: "one"}) + runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", attempt: 1, attempted_at: "invalid", job_id: 42, queue: "one")) + runtime.instance_variable_set(:@running, {}) runtime.send(:receive_notification, "unknown", "{}") runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: "42")) runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) diff --git a/ruby/spec/runtime_draining_shared_examples.rb b/ruby/spec/runtime_draining_shared_examples.rb index 879b6b742..587f58cf1 100644 --- a/ruby/spec/runtime_draining_shared_examples.rb +++ b/ruby/spec/runtime_draining_shared_examples.rb @@ -77,13 +77,21 @@ runtime = client.instance_variable_get(:@runtime) signalled = false failures = 0 + poll_failed = Queue.new @driver.define_singleton_method(:job_get_cancelled_ids) do |ids| if poll_failure && signalled && failures.zero? failures += 1 + poll_failed << true raise "temporary cancellation query failure" end super(ids) end + if poll_failure + # Legacy senders require a durable check, including its retry after a failed read. + @driver.define_singleton_method(:runtime_notify) do |topic, payload| + super(topic, payload.except(:attempt, :attempted_at)) + end + end runtime.define_singleton_method(:queue_stopping?) do |queue| stopping = super(queue) if stopping && !signalled @@ -106,9 +114,9 @@ nil end Timeout.timeout(5) { draining.pop } - - # Persist through another client so only database polling can deliver it. + # A peer notification must still be delivered while this client drains. River::Client.new(@driver).job_cancel(job.id) + Timeout.timeout(5) { poll_failed.pop } if poll_failure event = Timeout.timeout(3) { subscription.pop } expect(event).to have_attributes(kind: :job_cancelled, job: have_attributes(id: job.id, state: "cancelled")) diff --git a/rust/riverqueue/src/client/attempts.rs b/rust/riverqueue/src/client/attempts.rs index 7483935d5..b211a3a8a 100644 --- a/rust/riverqueue/src/client/attempts.rs +++ b/rust/riverqueue/src/client/attempts.rs @@ -2,14 +2,21 @@ use std::collections::HashMap; use std::sync::Mutex; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::Ordering; use tokio_util::sync::CancellationToken; -use crate::client::{ClientInner, PENDING_CANCELLATION_LIMIT, PENDING_CANCELLATION_RETENTION}; +use crate::client::ClientInner; + +#[derive(Clone, Debug, PartialEq)] +pub(super) struct RunningAttempt { + pub(super) attempt: i32, + pub(super) attempted_at: Option>, + pub(super) cancellation: CancellationToken, +} pub(super) fn remove_running_attempt( - running: &Mutex>, + running: &Mutex>, job_id: i64, cancellation: &CancellationToken, ) { @@ -18,106 +25,33 @@ pub(super) fn remove_running_attempt( .unwrap_or_else(std::sync::PoisonError::into_inner); if running .get(&job_id) - .is_some_and(|active| active == cancellation) + .is_some_and(|active| active.cancellation == *cancellation) { running.remove(&job_id); } } pub(super) fn register_running_attempt( - running: &Mutex>, - pending_cancellations: &Mutex>, + inner: &ClientInner, job_id: i64, cancellation: &CancellationToken, + generation: u64, + attempt: i32, + attempted_at: Option>, ) { - // Keep the locks in this order here and in `signal_running_attempt` so a - // cancellation cannot fall between checking the active map and recording - // a just-fetched attempt. - let mut running = running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - running.insert(job_id, cancellation.clone()); - let should_cancel = pending_cancellations + inner + .running .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .remove(&job_id) - .is_some(); - drop(running); - if should_cancel { - cancellation.cancel(); - } -} - -pub(super) fn signal_running_attempt( - running: &Mutex>, - pending_cancellations: &Mutex>, - fetch_registration_windows: &AtomicU64, - job_id: i64, -) { - let running = running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(cancellation) = running.get(&job_id).cloned() { - drop(running); - cancellation.cancel(); - return; - } - if fetch_registration_windows.load(Ordering::SeqCst) == 0 { - return; - } - - let now = std::time::Instant::now(); - let mut pending = pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - pending.retain(|_, received_at| { - now.saturating_duration_since(*received_at) <= PENDING_CANCELLATION_RETENTION - }); - if pending.len() >= PENDING_CANCELLATION_LIMIT - && let Some(oldest_job_id) = pending - .iter() - .min_by_key(|(_, received_at)| **received_at) - .map(|(job_id, _)| *job_id) - { - pending.remove(&oldest_job_id); - } - pending.insert(job_id, now); -} - -pub(super) struct FetchRegistrationGuard<'a> { - pub(super) inner: &'a ClientInner, -} - -impl<'a> FetchRegistrationGuard<'a> { - pub(super) fn new(inner: &'a ClientInner) -> Self { - inner - .fetch_registration_windows - .fetch_add(1, Ordering::SeqCst); - Self { inner } - } -} - -impl Drop for FetchRegistrationGuard<'_> { - fn drop(&mut self) { - // Synchronize the last-window transition with - // `signal_running_attempt`, which holds this lock while deciding - // whether to retain an unmatched cancellation. - let _running = self - .inner - .running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if self - .inner - .fetch_registration_windows - .fetch_sub(1, Ordering::SeqCst) - == 1 - { - self.inner - .pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clear(); - } + .insert( + job_id, + RunningAttempt { + attempt, + attempted_at, + cancellation: cancellation.clone(), + }, + ); + if inner.cancellation_generation.load(Ordering::Acquire) != generation { + inner.cancellation_wakeup.notify_one(); } } diff --git a/rust/riverqueue/src/client/builder.rs b/rust/riverqueue/src/client/builder.rs index 0c378c36d..22c6f3eb2 100644 --- a/rust/riverqueue/src/client/builder.rs +++ b/rust/riverqueue/src/client/builder.rs @@ -1,12 +1,13 @@ //! Client configuration and construction. use std::collections::HashMap; -use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; use std::sync::{Arc, Mutex, RwLock}; use std::time::Duration; use serde_json::{Map, Value}; -use tokio::sync::{broadcast, watch}; +use tokio::sync::{Notify, broadcast, watch}; use crate::__private::Pilot; #[cfg(feature = "postgres")] @@ -802,6 +803,8 @@ impl ClientBuilder { inner: Arc::new(ClientInner { allow_legacy_job_kinds: self.allow_legacy_job_kinds, allow_unregistered_job_kinds: self.allow_unregistered_job_kinds, + cancellation_generation: AtomicU64::new(0), + cancellation_wakeup: Notify::new(), completion_sender: Mutex::new(None), database: self.database, default_max_attempts: self.default_max_attempts, @@ -809,7 +812,6 @@ impl ClientBuilder { events, fetch_cooldown: self.fetch_cooldown, fetch_kinds, - fetch_registration_windows: AtomicU64::new(0), hooks: self.hooks, id: self.id, insert_middleware: self.insert_middleware, @@ -823,7 +825,6 @@ impl ClientBuilder { #[cfg(test)] notifier_start_panics: AtomicU64::new(0), peer_owners: Mutex::new(HashMap::new()), - pending_cancellations: Mutex::new(HashMap::new()), periodic_jobs, pilot: self.pilot, poll_only: self.poll_only, diff --git a/rust/riverqueue/src/client/jobs.rs b/rust/riverqueue/src/client/jobs.rs index ac8482870..adf913f19 100644 --- a/rust/riverqueue/src/client/jobs.rs +++ b/rust/riverqueue/src/client/jobs.rs @@ -9,7 +9,6 @@ use std::{ use serde_json::Map; use crate::__private::DatabaseConnection as PilotDatabaseConnection; -use crate::client::attempts::signal_running_attempt; use crate::client::request::{Target, request_type}; use crate::database::DatabaseTransactionExecutor; use crate::storage::Access; @@ -257,17 +256,10 @@ impl JobCancelRequest<'_> { let row = session.storage(inner).job_cancel(self.id).await?; session.commit().await?; // Without a listener (no backend listener, or a poll-only client), - // wake this client's running attempt directly, like Go's - // `notifyProducerWithoutListenerQueueControlEvent`. Other clients - // observe the committed notification through their own listener or - // outbox poll. + // wake the durable check. The post-commit signal can be delayed until + // a retry is running, so it must not cancel by job ID alone. if own_transaction && !inner.listens_for_notifications() { - signal_running_attempt( - &inner.running, - &inner.pending_cancellations, - &inner.fetch_registration_windows, - self.id, - ); + inner.request_cancellation_check(); } Ok(row) } diff --git a/rust/riverqueue/src/client/mod.rs b/rust/riverqueue/src/client/mod.rs index 68b517853..6b4c813ff 100644 --- a/rust/riverqueue/src/client/mod.rs +++ b/rust/riverqueue/src/client/mod.rs @@ -57,8 +57,7 @@ use std::{ }; use chrono::{DateTime, Utc}; -use tokio::sync::{broadcast, mpsc, watch}; -use tokio_util::sync::CancellationToken; +use tokio::sync::{Notify, broadcast, mpsc, watch}; use self::completer::CompletionUpdate; pub(crate) use self::notify_limiter::InsertNotifyLimiter; @@ -79,8 +78,6 @@ use crate::{ pub(crate) const ATTEMPTED_BY_MAX: i32 = 100; const EVENT_BUFFER_CAPACITY: usize = 10_000; -const PENDING_CANCELLATION_LIMIT: usize = 10_000; -const PENDING_CANCELLATION_RETENTION: Duration = Duration::from_mins(1); // Large queues otherwise become limited by a single Postgres claim round trip. // Concurrent `SKIP LOCKED` claims safely divide the available worker slots. const PARALLEL_FETCH_MINIMUM: usize = 1_000; @@ -103,6 +100,8 @@ const PRODUCER_STALE_RETENTION: Duration = Duration::from_mins(5); pub(crate) struct ClientInner { allow_legacy_job_kinds: bool, allow_unregistered_job_kinds: bool, + cancellation_generation: AtomicU64, + cancellation_wakeup: Notify, completion_sender: Mutex>>, pub(crate) database: Database, default_max_attempts: i32, @@ -115,7 +114,6 @@ pub(crate) struct ClientInner { /// Kinds claims are restricted to, including aliases, or `None` to claim /// every kind. pub(crate) fetch_kinds: Option>, - fetch_registration_windows: AtomicU64, pub(crate) hooks: Vec>, pub(crate) id: String, insert_middleware: Vec>, @@ -139,7 +137,6 @@ pub(crate) struct ClientInner { notifier_start_panics: AtomicU64, /// Peer jobs owned by running attempts, mapped to their ledger. peer_owners: Mutex>, - pending_cancellations: Mutex>, pub(crate) periodic_jobs: PeriodicJobs, pub(crate) pilot: Arc, poll_only: bool, @@ -148,7 +145,7 @@ pub(crate) struct ClientInner { queue_notifications: broadcast::Sender, queues: RwLock>, pub(crate) retry_policy: Arc, - running: Mutex>, + running: Mutex>, #[cfg(feature = "postgres")] pub(crate) schema: SchemaName, soft_stop_timeout: Option, @@ -204,6 +201,11 @@ impl ClientInner { !self.poll_only && self.database.delivers_notifications() } + fn request_cancellation_check(&self) { + self.cancellation_generation.fetch_add(1, Ordering::AcqRel); + self.cancellation_wakeup.notify_one(); + } + /// Whether this client receives notifications from other clients, which /// needs a backend listener and a client that isn't poll-only. When it /// doesn't, it wakes its own runtime directly after committing a change, diff --git a/rust/riverqueue/src/client/notifier.rs b/rust/riverqueue/src/client/notifier.rs index be9aa976e..d8cacad27 100644 --- a/rust/riverqueue/src/client/notifier.rs +++ b/rust/riverqueue/src/client/notifier.rs @@ -22,7 +22,6 @@ use tracing::{debug, error, warn}; use crate::Error; use crate::client::ClientInner; -use crate::client::attempts::signal_running_attempt; use crate::client::backoff::exponential_backoff; #[cfg(feature = "sqlite")] use crate::client::sqlite_backend_error; @@ -40,6 +39,8 @@ const LISTENER_PING_INTERVAL: Duration = Duration::from_secs(5); #[derive(Deserialize)] pub(super) struct ControlNotification { pub(super) action: String, + pub(super) attempt: Option, + pub(super) attempted_at: Option>, pub(super) job_id: Option, pub(super) queue: Option, } @@ -125,13 +126,25 @@ pub(super) fn dispatch_notification( }; match payload.action.as_str() { "cancel" => { - if let Some(job_id) = payload.job_id { - signal_running_attempt( - &inner.running, - &inner.pending_cancellations, - &inner.fetch_registration_windows, - job_id, - ); + if let Some(id) = payload.job_id { + if let (Some(attempt), Some(attempted_at)) = ( + payload.attempt.filter(|attempt| *attempt > 0), + payload.attempted_at, + ) { + let running = inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(active) = running.get(&id) { + if active.attempt == attempt + && active.attempted_at == Some(attempted_at) + { + active.cancellation.cancel(); + } + return; + } + } + inner.request_cancellation_check(); } } // Like Go's producer, a metadata change is handled at once @@ -227,6 +240,7 @@ async fn idles_without_listen_notify( } info!("River's database does not support LISTEN/NOTIFY; polling instead"); report_ready(ready); + inner.request_cancellation_check(); cancel.cancelled().await; Ok(true) } @@ -303,6 +317,7 @@ async fn listen_until_error( "timed out connecting and subscribing".to_owned(), ) })??; + inner.request_cancellation_check(); debug!("River notification listener healthy"); *attempt = 0; report_ready(ready); @@ -331,6 +346,7 @@ async fn listen_until_error( Ok(Ok(None)) => { warn!("River notification listener reconnected; waking producers"); wake_all_producers(queue_notifications); + inner.request_cancellation_check(); } Ok(Err(listener_error)) => return Err(listener_error.into()), Err(_) => { @@ -390,6 +406,7 @@ pub(super) async fn run_sqlite_notifications( .connect(&mut connection) .await .map_err(sqlite_backend_error)?; + inner.request_cancellation_check(); } for topic in [ crate::NOTIFICATION_TOPIC_CONTROL, diff --git a/rust/riverqueue/src/client/producer.rs b/rust/riverqueue/src/client/producer.rs index fa68924f9..9d3f3be68 100644 --- a/rust/riverqueue/src/client/producer.rs +++ b/rust/riverqueue/src/client/producer.rs @@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::Duration; use chrono::Utc; @@ -16,9 +17,7 @@ use tokio_util::sync::CancellationToken; use tracing::{debug, error, warn}; use crate::__private::DatabaseConnection as PilotDatabaseConnection; -use crate::client::attempts::{ - FetchRegistrationGuard, register_running_attempt, signal_running_attempt, -}; +use crate::client::attempts::register_running_attempt; use crate::client::backoff::{RestartBackoff, exponential_backoff}; use crate::client::completer::CompletionUpdate; use crate::client::executor::{AbortOnDrop, execute_job}; @@ -56,10 +55,7 @@ pub(super) async fn run_dynamic_queues( mut changes: watch::Receiver, queues_ready: oneshot::Sender<()>, ) -> Result<(), Error> { - // Without a notifier, running jobs learn of cancellation requests by - // polling until every producer has drained. Whether the client has one is - // known only once a Postgres server's capabilities are detected, so the - // poll checks each time. + // Cancellation checks remain independent of fetching until every producer drains. let _cancellation_poll = AbortOnDrop(tokio::spawn(poll_job_cancellations( Arc::clone(&inner), fetch_cancel.clone(), @@ -142,59 +138,49 @@ fn jittered_fetch_poll_interval(interval: Duration) -> Duration { crate::maintenance::random_duration(interval, interval.saturating_add(jitter)) } -/// Cancels this client's running attempts whose jobs have a cancellation -/// request, checking every [`QUEUE_CONFIG_POLL_INTERVAL`] while the client -/// has no notifier, like River Go's producers without one. That's checked -/// each time, since a Postgres server's lack of `LISTEN`/`NOTIFY` is -/// detected only once the client runs. -/// It keeps checking after `stopping` fires, so jobs can still be cancelled -/// while producers drain; the caller aborts it once they have. -async fn poll_job_cancellations(inner: Arc, stopping: CancellationToken) { - let mut poll = tokio::time::interval(QUEUE_CONFIG_POLL_INTERVAL); - poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // The first tick completes at once; running jobs start later. - poll.tick().await; +/// Checks cancellations on notification/reconnect hints, retrying failed reads. +/// Without a notifier, polls at the queue-control interval instead. +pub(super) async fn poll_job_cancellations(inner: Arc, stopping: CancellationToken) { + let mut retry = false; loop { - poll.tick().await; - if inner.has_notifier() { - continue; + tokio::select! { + () = tokio::time::sleep(QUEUE_CONFIG_POLL_INTERVAL), if retry || !inner.has_notifier() => {}, + () = inner.cancellation_wakeup.notified() => {}, } - let mut ids = inner - .running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .keys() - .copied() - .collect::>(); - ids.sort_unstable(); - for batch in ids.chunks(JOB_CANCEL_POLL_BATCH_SIZE) { - let requested = tokio::time::timeout( - JOB_CANCEL_POLL_TIMEOUT, - crate::storage::job_cancel_requested(&inner, batch), - ) - .await - .map_err(|_| Error::runtime_context("job cancellation poll", "timed out")) - .and_then(|requested| requested); - match requested { - Ok(requested) => { - for id in requested { - signal_running_attempt( - &inner.running, - &inner.pending_cancellations, - &inner.fetch_registration_windows, - id, - ); - } - } - Err(poll_error) => { - if !stopping.is_cancelled() { - error!(error = %crate::error::Chain(&poll_error), "River failed to check for job cancellation requests"); - } - break; - } + let result = check_job_cancellations(&inner).await; + retry = result.is_err(); + if let Err(check_error) = result + && !stopping.is_cancelled() + { + error!(error = %crate::error::Chain(&check_error), "River failed to check for job cancellation requests"); + } + } +} + +pub(super) async fn check_job_cancellations(inner: &ClientInner) -> Result<(), Error> { + let attempts = inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let mut ids = attempts.keys().copied().collect::>(); + ids.sort_unstable(); + for batch in ids.chunks(JOB_CANCEL_POLL_BATCH_SIZE) { + let requested = tokio::time::timeout( + JOB_CANCEL_POLL_TIMEOUT, + crate::storage::job_cancel_requested(inner, batch), + ) + .await + .map_err(|_| Error::runtime_context("job cancellation poll", "timed out"))??; + for id in requested { + // A query can outlive an attempt and its retry. Only signal the + // execution that was present in its snapshot. + if let Some(attempt) = attempts.get(&id) { + attempt.cancellation.cancel(); } } } + Ok(()) } type ProducerOutcome = (String, u64, CancellationToken, Result<(), Error>); @@ -858,7 +844,7 @@ pub(super) async fn run_queue( if available == 0 { continue; } - let registration_guard = FetchRegistrationGuard::new(&inner); + let cancellation_generation = inner.cancellation_generation.load(Ordering::Acquire); let fetched = match (&session, claims_through_session) { (Some(session), true) => { match claim_through_session( @@ -908,10 +894,12 @@ pub(super) async fn run_queue( let hard_cancel = work_cancel.child_token(); let cancellation = hard_cancel.child_token(); register_running_attempt( - &inner.running, - &inner.pending_cancellations, + &inner, row.id, &cancellation, + cancellation_generation, + row.attempt, + row.attempted_at, ); let task_inner = Arc::clone(&inner); let completion_sender = completion_sender.clone(); @@ -928,7 +916,6 @@ pub(super) async fn run_queue( .await; }); } - drop(registration_guard); while let Some(joined) = attempts.tasks.try_join_next_with_id() { attempts.exited(joined.map(|(task_id, ())| task_id), false); } diff --git a/rust/riverqueue/src/client/run.rs b/rust/riverqueue/src/client/run.rs index f8c42bac0..fb91a1e92 100644 --- a/rust/riverqueue/src/client/run.rs +++ b/rust/riverqueue/src/client/run.rs @@ -175,8 +175,8 @@ impl Service { /// When each part of a started client stops, matching River Go's order. /// /// A stop request cancels the client's fetch token. Claims, leadership with -/// its maintenance services, the notifier, and extension runtime services -/// all stop then, like Go's services started on the fetch context. Running +/// its maintenance services, and extension runtime services stop then. +/// The notifier remains available until every queue has drained. Running /// jobs keep going until they finish or `work_cancel` fires, and each /// producer keeps reporting to its extension session until its last attempt /// has left, from a token of its own that the stop request doesn't cancel. @@ -185,7 +185,9 @@ struct StopTokens { claims: CancellationToken, /// Leader election and leader-owned maintenance. leadership: CancellationToken, - /// The notifier and extension runtime services. + /// Notification delivery remains live until queues finish draining. + notifier: CancellationToken, + /// Extension runtime services. services: CancellationToken, } @@ -194,6 +196,7 @@ impl StopTokens { Self { claims: fetch_cancel.child_token(), leadership: fetch_cancel.child_token(), + notifier: CancellationToken::new(), services: fetch_cancel.child_token(), } } @@ -247,7 +250,14 @@ impl Supervisor { let Some((service, started_at)) = self.services.remove(&task_id) else { continue; }; - let stopping = self.fetch_cancel.is_cancelled(); + if service == Service::Queues { + self.stop.notifier.cancel(); + } + let stopping = if service == Service::Notifier { + self.stop.notifier.is_cancelled() + } else { + self.fetch_cancel.is_cancelled() + }; match outcome { // The completer ends once every producer has dropped its sender. Ok(()) if stopping || service == Service::Completer => {} @@ -374,7 +384,7 @@ impl Supervisor { )) } Service::Notifier => { - cancel = self.stop.services.child_token(); + cancel = self.stop.notifier.child_token(); let ready = self.notifier_ready.clone(); #[cfg(test)] if inner @@ -526,9 +536,9 @@ impl Stopper { /// Requests a soft stop. /// /// The client stops fetching new jobs and lets running jobs finish before - /// it stops. Leader election, maintenance, and the notification listener - /// stop at once, while each queue's producer keeps reporting its - /// running jobs until they finish. When the builder's `soft_stop_timeout` + /// it stops. Leader election and maintenance stop at once. The notification + /// listener keeps delivering cancellations, and each queue's producer keeps + /// reporting its running jobs until they finish. When the builder's `soft_stop_timeout` /// is set, jobs still running after that timeout are cancelled as if by /// [`Stopper::stop_and_cancel`]. pub fn stop(&self) { diff --git a/rust/riverqueue/src/client/tests.rs b/rust/riverqueue/src/client/tests.rs index d52f83f76..e09504450 100644 --- a/rust/riverqueue/src/client/tests.rs +++ b/rust/riverqueue/src/client/tests.rs @@ -1,7 +1,8 @@ use serde::Deserialize; use serde_json::Map; +use tokio_util::sync::CancellationToken; -use super::attempts::{register_running_attempt, remove_running_attempt, signal_running_attempt}; +use super::attempts::{RunningAttempt, register_running_attempt, remove_running_attempt}; #[cfg(feature = "postgres")] use super::completer::CompletionBatcher; #[cfg(feature = "sqlite")] @@ -94,12 +95,26 @@ fn completion_cleanup_preserves_newer_attempt() { let job_id = 42; let first = CancellationToken::new(); let second = CancellationToken::new(); - let running = Mutex::new(HashMap::from([(job_id, first.clone())])); + let running = Mutex::new(HashMap::from([( + job_id, + RunningAttempt { + attempt: 1, + attempted_at: None, + cancellation: first.clone(), + }, + )])); running .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(job_id, second.clone()); + .insert( + job_id, + RunningAttempt { + attempt: 2, + attempted_at: None, + cancellation: second.clone(), + }, + ); remove_running_attempt(&running, job_id, &first); assert_eq!( @@ -107,7 +122,11 @@ fn completion_cleanup_preserves_newer_attempt() { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .get(&job_id), - Some(&second) + Some(&RunningAttempt { + attempt: 2, + attempted_at: None, + cancellation: second.clone() + }) ); remove_running_attempt(&running, job_id, &second); @@ -171,11 +190,13 @@ async fn go_notification_fixtures_dispatch_correctly() { let cancellation = CancellationToken::new(); let unrelated = CancellationToken::new(); for (job_id, token) in [(42, &cancellation), (43, &unrelated)] { - register_running_attempt( - &inner.running, - &inner.pending_cancellations, + inner.running.lock().unwrap().insert( job_id, - token, + RunningAttempt { + attempt: 1, + attempted_at: Some("2026-01-02T03:04:05.678Z".parse().unwrap()), + cancellation: token.clone(), + }, ); } @@ -211,21 +232,18 @@ async fn go_notification_fixtures_dispatch_correctly() { } assert_eq!(cancellation.is_cancelled(), name == "cancel", "{name}"); assert!( - !unrelated.is_cancelled(), - "{name}: cancelled an unrelated job" + tokio::time::timeout(Duration::ZERO, inner.cancellation_wakeup.notified()) + .await + .is_err(), + "{name}" ); + assert!(!unrelated.is_cancelled(), "{name}"); assert!( - matches!( - producer.try_recv(), - Err(broadcast::error::TryRecvError::Empty) - ), + producer.try_recv().is_err(), "{name}: unexpected producer notification" ); assert!( - matches!( - leadership.try_recv(), - Err(broadcast::error::TryRecvError::Empty) - ), + leadership.try_recv().is_err(), "{name}: unexpected leadership notification" ); } @@ -233,50 +251,185 @@ async fn go_notification_fixtures_dispatch_correctly() { pool.close().await; } -#[test] -fn pending_cancellation_reaches_fetched_attempt() { - let job_id = 42; +#[tokio::test] +async fn identified_cancellation_needs_no_database_and_ignores_older_attempts() { + #[cfg(feature = "postgres")] + let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); + #[cfg(not(feature = "postgres"))] + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); let cancellation = CancellationToken::new(); - let fetch_registration_windows = AtomicU64::new(1); - let pending_cancellations = Mutex::new(HashMap::new()); - let running = Mutex::new(HashMap::new()); - - signal_running_attempt( - &running, - &pending_cancellations, - &fetch_registration_windows, - job_id, + let attempted_at: chrono::DateTime = "2026-01-02T03:04:05.6789Z".parse().unwrap(); + register_running_attempt(&client.inner, 42, &cancellation, 0, 2, Some(attempted_at)); + for (attempt, at) in [ + (1, attempted_at), + (2, attempted_at - chrono::Duration::seconds(1)), + (2, attempted_at), + ] { + let payload = serde_json::json!({"action":"cancel", "attempt":attempt, "attempted_at":at, "job_id":42,"queue":"default"}); + dispatch_notification( + &client.inner, + &client.inner.queue_notifications, + crate::NOTIFICATION_TOPIC_CONTROL, + &payload.to_string(), + ); + assert_eq!( + cancellation.is_cancelled(), + attempt == 2 && at == attempted_at + ); + } + assert!( + tokio::time::timeout(Duration::ZERO, client.inner.cancellation_wakeup.notified()) + .await + .is_err() ); - register_running_attempt(&running, &pending_cancellations, job_id, &cancellation); +} - assert!(cancellation.is_cancelled()); +#[tokio::test] +async fn registration_without_a_hint_does_not_query_cancellations() { + #[cfg(feature = "postgres")] + let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); + #[cfg(not(feature = "postgres"))] + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); + let cancellation = CancellationToken::new(); + register_running_attempt(&client.inner, 42, &cancellation, 0, 1, None); assert!( - pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() + tokio::time::timeout(Duration::ZERO, client.inner.cancellation_wakeup.notified(),) + .await + .is_err() ); + // A hint consumed before registration still requires a fresh snapshot afterward. + client.inner.request_cancellation_check(); + client.inner.cancellation_wakeup.notified().await; + register_running_attempt(&client.inner, 43, &cancellation, 0, 1, None); + tokio::time::timeout( + Duration::from_secs(1), + client.inner.cancellation_wakeup.notified(), + ) + .await + .unwrap(); + assert!(!cancellation.is_cancelled()); } -#[test] -fn unmatched_cancellation_is_not_retained_without_fetch() { - let fetch_registration_windows = AtomicU64::new(0); - let pending_cancellations = Mutex::new(HashMap::new()); - let running = Mutex::new(HashMap::new()); - - signal_running_attempt( - &running, - &pending_cancellations, - &fetch_registration_windows, - 42, - ); +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn cancellation_checks_wait_for_hints_and_retry_failed_reads() { + let queries = Arc::new(AtomicU64::new(0)); + let observed = Arc::clone(&queries); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .before_acquire(move |_, _| { + let observed = Arc::clone(&observed); + Box::pin(async move { + observed.fetch_add(1, Ordering::AcqRel); + Ok(true) + }) + }) + .connect("sqlite::memory:") + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + let cancellation = CancellationToken::new(); + register_running_attempt(&client.inner, 42, &cancellation, 0, 1, None); + let task = tokio::spawn(super::producer::poll_job_cancellations( + Arc::clone(&client.inner), + CancellationToken::new(), + )); + let before = queries.load(Ordering::Acquire); + tokio::time::sleep(Duration::from_millis(2100)).await; + assert_eq!(before, queries.load(Ordering::Acquire)); + client.inner.request_cancellation_check(); + tokio::time::timeout(Duration::from_secs(5), async { + while queries.load(Ordering::Acquire) == before { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + // The first check fails because the table is absent. Acquiring its sole connection + // waits for that check to finish; repairing the database must be enough to recover. + let mut connection = pool.acquire().await.unwrap(); + sqlx::query("CREATE TABLE river_job (id INTEGER, state TEXT, metadata TEXT)") + .execute(&mut *connection) + .await + .unwrap(); + sqlx::query(r#"INSERT INTO river_job VALUES (42, 'running', '{"cancel_attempted_at":"2026-01-01T00:00:00Z"}')"#).execute(&mut *connection).await.unwrap(); + drop(connection); + tokio::time::timeout(Duration::from_secs(5), cancellation.cancelled()) + .await + .unwrap(); + let after = queries.load(Ordering::Acquire); + tokio::time::sleep(Duration::from_millis(2100)).await; + assert_eq!(after, queries.load(Ordering::Acquire)); + task.abort(); + let _ = task.await; + pool.close().await; +} - assert!( - pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn delayed_cancellation_signals_recheck_retry_marker() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + riverqueue_migrate::SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()) + .without_notifications() + .build() + .unwrap(); + let id: i64 = sqlx::query_scalar("INSERT INTO river_job (kind, args, state, attempt, max_attempts, metadata) VALUES ('test', jsonb('{}'), 'running', 1, 2, jsonb('{}')) RETURNING id") + .fetch_one(&pool).await.unwrap(); + let original = CancellationToken::new(); + register_running_attempt( + &client.inner, + id, + &original, + client.inner.cancellation_generation.load(Ordering::Acquire), + 1, + None, + ); + client.jobs().cancel(id).await.unwrap(); + sqlx::query("UPDATE river_job SET state = 'cancelled', finalized_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = ?").bind(id).execute(&pool).await.unwrap(); + client.jobs().retry(id).await.unwrap(); + sqlx::query("UPDATE river_job SET state = 'running', attempt = 2 WHERE id = ?") + .bind(id) + .execute(&pool) + .await + .unwrap(); + let replacement = CancellationToken::new(); + register_running_attempt( + &client.inner, + id, + &replacement, + client.inner.cancellation_generation.load(Ordering::Acquire), + 1, + None, + ); + dispatch_notification( + &client.inner, + &client.inner.queue_notifications, + crate::NOTIFICATION_TOPIC_CONTROL, + &format!(r#"{{"action":"cancel","job_id":{id},"queue":"default"}}"#), ); + super::producer::check_job_cancellations(&client.inner) + .await + .unwrap(); + assert!(!replacement.is_cancelled()); + // The local post-commit signal is also only a hint, and remains useful + // when the current attempt actually has a durable cancellation request. + client.jobs().cancel(id).await.unwrap(); + assert!(!replacement.is_cancelled()); + super::producer::check_job_cancellations(&client.inner) + .await + .unwrap(); + assert!(replacement.is_cancelled()); + pool.close().await; } fn retry_row(error_count: usize) -> JobRow { @@ -638,11 +791,14 @@ async fn completer_abandons_its_backlog_when_a_batch_fails_during_shutdown() { let now = std::time::Instant::now(); for job_id in 0..i64::try_from(updates).unwrap() { let cancellation = CancellationToken::new(); - inner - .running - .lock() - .unwrap() - .insert(job_id, cancellation.clone()); + inner.running.lock().unwrap().insert( + job_id, + RunningAttempt { + attempt: 0, + attempted_at: None, + cancellation: cancellation.clone(), + }, + ); sender .try_send(CompletionUpdate { attempt: None, diff --git a/rust/riverqueue/src/database/sqlite.rs b/rust/riverqueue/src/database/sqlite.rs index 2152e1341..9526ccf2e 100644 --- a/rust/riverqueue/src/database/sqlite.rs +++ b/rust/riverqueue/src/database/sqlite.rs @@ -949,25 +949,30 @@ pub(crate) async fn retry( ) -> Result, BackendError> { let sql = format!( r#" + WITH retry_params AS (SELECT ? AS now) UPDATE river_job SET state = 'available', max_attempts = CASE + WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = ? + finalized_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) + THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END WHERE id = ? AND state != 'running' - AND (state != 'available' OR scheduled_at > ?) + AND ((CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR NOT (state = 'available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {JOB_COLUMNS} "# ); sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) .bind(sqlite_time(now)) .bind(id) - .bind(sqlite_time(now)) .fetch_optional(&mut *connection) .await? .map(JobRecord::into_job) @@ -2537,6 +2542,46 @@ mod tests { assert_eq!(row.attempted_by, ["three", "four", "five"]); } + #[tokio::test] + async fn retry_at_exact_schedule_increments_max_attempts() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Available, now).await; + sqlx::query("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .bind(job.id) + .execute(&mut *connection) + .await + .unwrap(); + let retried = retry(&mut connection, job.id, now).await.unwrap().unwrap(); + assert_eq!(retried.max_attempts, job.max_attempts + 1); + assert_eq!(retried.scheduled_at, job.scheduled_at); + } + + #[tokio::test] + async fn retry_preserves_malformed_metadata_for_row_validation() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Retryable, now).await; + sqlx::query("UPDATE river_job SET metadata = '[not json' WHERE id = ?") + .bind(job.id) + .execute(&mut *connection) + .await + .unwrap(); + + let error = retry(&mut connection, job.id, now).await.unwrap_err(); + assert!( + matches!(error, BackendError::InvalidRow(ref message) if message.contains("metadata")) + ); + let metadata: String = sqlx::query_scalar("SELECT metadata FROM river_job WHERE id = ?") + .bind(job.id) + .fetch_one(&mut *connection) + .await + .unwrap(); + assert_eq!(metadata, "[not json"); + } + /// Inserts a job for the undecodable row tests. async fn insert_test_job( connection: &mut SqliteConnection, diff --git a/rust/riverqueue/src/storage/postgres.rs b/rust/riverqueue/src/storage/postgres.rs index 316eb3587..4673aac8a 100644 --- a/rust/riverqueue/src/storage/postgres.rs +++ b/rust/riverqueue/src/storage/postgres.rs @@ -45,9 +45,9 @@ impl Backend for PostgresBackend<'_> { let table = self.schema.qualify("river_job"); let sql = format!( "WITH locked AS (\ - SELECT id, queue, state, finalized_at FROM {table} WHERE id = $1 FOR UPDATE\ + SELECT id, attempt, attempted_at, queue, state, finalized_at FROM {table} WHERE id = $1 FOR UPDATE\ ), notified AS (\ - SELECT id, CASE WHEN $5::boolean THEN pg_notify(concat(coalesce($2::text, current_schema()), '.', $3::text), json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text) END \ + SELECT id, CASE WHEN $5::boolean THEN pg_notify(concat(coalesce($2::text, current_schema()), '.', $3::text), json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text) END \ FROM locked WHERE state NOT IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NULL\ ), updated AS (\ UPDATE {table} AS job SET \ @@ -80,7 +80,7 @@ impl Backend for PostgresBackend<'_> { let table = self.schema.qualify("river_job"); let sql = format!( "SELECT id FROM {table} \ - WHERE id = any($1) AND metadata ? 'cancel_attempted_at' AND state = 'running' \ + WHERE id = any($1) AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' \ ORDER BY id" ); Ok(sqlx::query_scalar(AssertSqlSafe(sql)) @@ -238,10 +238,11 @@ impl Backend for PostgresBackend<'_> { let sql = format!( "WITH locked AS (SELECT id FROM {table} WHERE id = $1 FOR UPDATE), \ updated AS (UPDATE {table} AS job SET state = 'available', \ - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, \ - finalized_at = NULL, scheduled_at = now() \ + max_attempts = CASE WHEN job.state = 'available' AND job.scheduled_at < now() THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, \ + finalized_at = CASE WHEN job.state = 'available' AND job.scheduled_at < now() THEN finalized_at ELSE NULL END, metadata = CASE WHEN jsonb_typeof(metadata) = 'object' \ + THEN metadata - 'cancel_attempted_at' ELSE metadata END, scheduled_at = CASE WHEN job.state = 'available' AND job.scheduled_at < now() THEN scheduled_at ELSE now() END \ FROM locked WHERE job.id = locked.id AND job.state != 'running' \ - AND NOT (job.state = 'available' AND job.scheduled_at < now()) RETURNING job.*) \ + AND ((jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at') OR NOT (job.state = 'available' AND job.scheduled_at < now())) RETURNING job.*) \ SELECT {}, false AS unique_skipped_as_duplicate FROM updated AS job \ UNION ALL SELECT {}, false AS unique_skipped_as_duplicate \ FROM (SELECT * FROM {table} WHERE id = $1 FOR UPDATE) AS job \ diff --git a/rust/riverqueue/src/storage/sqlite.rs b/rust/riverqueue/src/storage/sqlite.rs index 57694909e..1f01ba3c4 100644 --- a/rust/riverqueue/src/storage/sqlite.rs +++ b/rust/riverqueue/src/storage/sqlite.rs @@ -70,6 +70,8 @@ impl Backend for SqliteBackend<'_> { }; let payload = serde_json::json!({ "action": "cancel", + "attempt": row.attempt, + "attempted_at": row.attempted_at, "job_id": id, "queue": row.queue, }) diff --git a/rust/riverqueue/tests/client_handles.rs b/rust/riverqueue/tests/client_handles.rs index 8f00521a1..a3b8ccc51 100644 --- a/rust/riverqueue/tests/client_handles.rs +++ b/rust/riverqueue/tests/client_handles.rs @@ -546,6 +546,7 @@ macro_rules! scenarios { .unwrap(); let retried = jobs.retry(cancelled).tx(&mut tx).await.unwrap(); assert_eq!(retried.state, JobState::Available); + assert!(!retried.metadata.contains_key("cancel_attempted_at")); tx.commit().await.unwrap(); assert_eq!( @@ -557,6 +558,35 @@ macro_rules! scenarios { fixture.cleanup().await; } + #[tokio::test(flavor = "multi_thread")] + async fn retry_clears_cancellation_without_reordering_an_available_job() { + let fixture = Fixture::new().await; + let metadata: riverqueue::JobMetadata = + r#"{"cancel_attempted_at":"2026-01-01T00:00:00Z","preserved":true}"# + .parse() + .unwrap(); + let id = fixture + .client + .insert(args("already available")) + .opts(InsertOpts::default().with_metadata(metadata)) + .await + .unwrap() + .id(); + let jobs = fixture.client.jobs(); + let before = jobs.get(id).await.unwrap(); + assert_eq!(before.state, JobState::Available); + + let retried = jobs.retry(id).await.unwrap(); + assert!(!retried.metadata.contains_key("cancel_attempted_at")); + assert_eq!(retried.metadata.get::("preserved").unwrap(), Some(true)); + assert_eq!(retried.scheduled_at, before.scheduled_at); + assert_eq!(retried.max_attempts, before.max_attempts); + assert!( + !jobs.get(id).await.unwrap().metadata.contains_key("cancel_attempted_at") + ); + fixture.cleanup().await; + } + #[tokio::test(flavor = "multi_thread")] async fn job_requests_run_on_the_pool_without_a_transaction() { let fixture = Fixture::new().await; diff --git a/rust/riverqueue/tests/poll_only_cancellation.rs b/rust/riverqueue/tests/poll_only_cancellation.rs index 76b6928d1..457ff1fed 100644 --- a/rust/riverqueue/tests/poll_only_cancellation.rs +++ b/rust/riverqueue/tests/poll_only_cancellation.rs @@ -1,5 +1,5 @@ -//! Cancellation of running jobs in clients without notifications, on every -//! backend. +//! Durable cancellation of running jobs, with notifications enabled or disabled, +//! including during graceful shutdown, on every backend. //! //! Postgres tests run in a unique schema and fail rather than skip when //! `RIVER_RUST_DATABASE_URL` is unset; SQLite tests use a temporary file. @@ -56,11 +56,11 @@ impl Backend { } } -/// A client without notifications cancels its running job once it polls the -/// cancellation another client requested, including while it's stopping and -/// waiting for that job. +/// Both notification delivery and poll-only checks remain live during graceful shutdown. async fn polls_for_remote_cancellation(backend: Backend) { - for while_stopping in [false, true] { + for (notifications, while_stopping) in + [(false, false), (false, true), (true, false), (true, true)] + { let started = Arc::new(Semaphore::new(0)); let mut workers = Workers::new(); let worker_started = Arc::clone(&started); @@ -75,13 +75,17 @@ async fn polls_for_remote_cancellation(backend: Backend) { } }) .unwrap(); - let client = Client::builder(backend.database()) + let builder = Client::builder(backend.database()) .queue(riverqueue::QUEUE_DEFAULT, QueueConfig::new(1)) .without_leader_election() - .without_notifications() - .workers(workers) - .build() - .unwrap(); + .workers(workers); + let client = if notifications { + builder + } else { + builder.without_notifications() + } + .build() + .unwrap(); let other = Client::builder(backend.database()).build().unwrap(); let mut events = client.subscribe(&[EventKind::JobCancelled]).unwrap(); let mut run = client.start().unwrap();