diff --git a/CHANGELOG.md b/CHANGELOG.md index b0724cece..55be01b14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- Recover `JobCancel()` notifications lost during notifier disconnects by reconciling running jobs after reconnect. [PR #1361](https://github.com/riverqueue/river/pull/1361). +- Deliver `JobCancel()` locally in `PollOnly` mode without blocking unstarted or stopping clients. [PR #1361](https://github.com/riverqueue/river/pull/1361). +- Clear previous cancellation metadata when retrying a job. [PR #1361](https://github.com/riverqueue/river/pull/1361). - Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440). - Fixed the notifier retrying without backoff and flooding logs when the context passed to `Client.Start` reaches its deadline. [PR #1495](https://github.com/riverqueue/river/pull/1495). - Fixed `JobSnooze` accepting negative durations despite documenting that they panic. [PR #1529](https://github.com/riverqueue/river/pull/1529). diff --git a/client.go b/client.go index e06af4b0c..3a117a053 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() @@ -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 } @@ -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,20 +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 } - if producer, ok := c.producersByQueueName[queue]; ok { - producer.TriggerQueueControlEvent(controlEvent) + 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() + + // Keep the matching execution attached even if delivery waits on the buffer. + select { + case producer.cancelCh <- producerCancelRequest{executor: executor, jobID: job.ID}: + producer.testSignals.QueueControlEventTriggered.Signal(&controlEventPayload{ + Action: controlActionCancel, + JobID: job.ID, + Queue: job.Queue, + }) + case <-ctx.Done(): + case <-state.done: } } diff --git a/client_test.go b/client_test.go index 472bac298..f0df01026 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.cancelCh) + 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,274 @@ 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 + stopped func() + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + targetProducer := &producer{cancelCh: make(chan producerCancelRequest, 1)} + ctx, cancel := context.WithCancel(context.Background()) + ctx, _, started, stopped := targetProducer.StartInit(ctx) + targetProducer.runState.Store(&producerRunState{done: ctx.Done()}) + started() + targetProducer.cancelCh <- producerCancelRequest{jobID: 123} + stop := sync.OnceFunc(func() { + cancel() + stopped() + }) + t.Cleanup(stop) + + return &testBundle{ + client: &Client[pgx.Tx]{ + producersByQueueName: map[string]*producer{QueueDefault: targetProducer}, + }, + producer: targetProducer, + stopped: stop, + } + } + + for _, testCase := range []struct { + name string + replaceBeforeDelivery bool + reuseAttempt bool + }{ + {name: "BufferedCancellationKeepsOriginalExecutor"}, + {name: "NewAttemptBeforeDeliveryIsReconciled", replaceBeforeDelivery: true}, + {name: "NewAttemptedAtBeforeDeliveryIsReconciled", replaceBeforeDelivery: true, reuseAttempt: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + <-bundle.producer.cancelCh + 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.replaceBeforeDelivery { + bundle.producer.activeJobs[job.ID] = replacement + } + + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), job) + bundle.producer.activeJobs[job.ID] = replacement + request := <-bundle.producer.cancelCh + require.Equal(t, !testCase.replaceBeforeDelivery, request.cancel(context.Background())) + require.NoError(t, replacementCtx.Err()) + if !testCase.replaceBeforeDelivery { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + } + }) + } + + 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.notifyProducerOfCancelWithoutNotifier(ctx, &rivertype.JobRow{Queue: QueueDefault}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.cancelCh, 1) + }) + + t.Run("FullBufferProducerStopped", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{Queue: QueueDefault}) + returned.Signal(struct{}{}) + }() + bundle.stopped() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.cancelCh, 1) + }) + + t.Run("RunningProducer", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + <-bundle.producer.cancelCh + + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{ID: 456, Queue: QueueDefault}) + + require.Equal(t, int64(456), (<-bundle.producer.cancelCh).jobID) + }) + + t.Run("StopInProgress", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + shouldStop, _, finalizeStop := bundle.producer.StopInit() + require.True(t, shouldStop) + defer func() { + bundle.stopped() + finalizeStop(true) + }() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerOfCancelWithoutNotifier(ctx, &rivertype.JobRow{Queue: QueueDefault}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.cancelCh, 1) + }) +} + +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{ + cancelCh: make(chan producerCancelRequest, 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.Equal(t, int64(123), (<-bundle.producer.cancelCh).jobID) + 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 +9244,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 +11176,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/internal/notifier/notifier.go b/internal/notifier/notifier.go index e860cb17b..93d3d34bc 100644 --- a/internal/notifier/notifier.go +++ b/internal/notifier/notifier.go @@ -22,6 +22,21 @@ import ( "github.com/riverqueue/river/rivershared/util/timeoututil" ) +// ListenerReadyFunc runs after initial connection and each successful reconnect. +type ListenerReadyFunc func() + +// ListenerReadySubscription unregisters a listener-ready callback. +type ListenerReadySubscription struct { + notifier *Notifier +} + +// Unregister stops future callbacks; an invocation already in progress may finish. +func (s *ListenerReadySubscription) Unregister() { + s.notifier.mu.Lock() + defer s.notifier.mu.Unlock() + delete(s.notifier.listenerReadyFuncs, s) +} + type NotificationTopic string const ( @@ -88,12 +103,13 @@ type Notifier struct { testSignals notifierTestSignals waitInterruptChan chan func() - mu sync.RWMutex - isConnected bool - isStarted bool - isWaiting bool - subscriptions map[NotificationTopic][]*Subscription - waitCancel context.CancelFunc + mu sync.RWMutex + isConnected bool + isStarted bool + isWaiting bool + listenerReadyFuncs map[*ListenerReadySubscription]ListenerReadyFunc + subscriptions map[NotificationTopic][]*Subscription + waitCancel context.CancelFunc } func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notifier { @@ -102,12 +118,29 @@ func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notif notificationBuf: make(chan *riverdriver.Notification, 1000), waitInterruptChan: make(chan func(), 10), - subscriptions: make(map[NotificationTopic][]*Subscription), + listenerReadyFuncs: make(map[*ListenerReadySubscription]ListenerReadyFunc), + subscriptions: make(map[NotificationTopic][]*Subscription), }) return notifier } +// RegisterListenerReadyFunc registers a callback for initial connection and reconnects. +func (n *Notifier) RegisterListenerReadyFunc(listenerReadyFunc ListenerReadyFunc) *ListenerReadySubscription { + n.mu.Lock() + defer n.mu.Unlock() + + sub := &ListenerReadySubscription{notifier: n} + n.listenerReadyFuncs[sub] = listenerReadyFunc + 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 +148,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 +252,16 @@ func (n *Notifier) listenAndWait(ctx context.Context) error { n.testSignals.ListeningBegin.Signal(struct{}{}) defer n.testSignals.ListeningEnd.Signal(struct{}{}) + listenerReadyFuncs := func() []ListenerReadyFunc { + n.mu.RLock() + defer n.mu.RUnlock() + + return maputil.Values(n.listenerReadyFuncs) + }() + for _, listenerReadyFunc := range listenerReadyFuncs { + listenerReadyFunc() + } + drainInterrupts := func() { for { select { diff --git a/internal/notifier/notifier_test.go b/internal/notifier/notifier_test.go index 76ad0b820..0d138ca61 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("RegisterListenerReadyFunc", 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.RegisterListenerReadyFunc(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() diff --git a/java/river/src/main/java/com/riverqueue/Client.java b/java/river/src/main/java/com/riverqueue/Client.java index b3a63512f..a867856d2 100644 --- a/java/river/src/main/java/com/riverqueue/Client.java +++ b/java/river/src/main/java/com/riverqueue/Client.java @@ -781,11 +781,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..f13bdaae5 100644 --- a/java/river/src/main/java/com/riverqueue/Protocol.java +++ b/java/river/src/main/java/com/riverqueue/Protocol.java @@ -2,6 +2,7 @@ import java.util.Map; import java.util.function.Consumer; +import java.util.function.LongConsumer; import java.util.regex.Pattern; /** Stored and transmitted values shared with the other River implementations. */ @@ -57,8 +58,7 @@ 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, + LongConsumer cancellation, Consumer queues, Consumer leadership) { void dispatch(String topic, String payload) { @@ -70,14 +70,9 @@ 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); - } + // Notices identify a job, not an attempt. Confirm its durable marker + // before cancelling so a delayed notice cannot undo a retry. + cancellation.accept(value.path("job_id").asLong()); } else if ((action.equals("metadata_changed") || action.equals("pause") || action.equals("resume")) diff --git a/java/river/src/main/java/com/riverqueue/Workers.java b/java/river/src/main/java/com/riverqueue/Workers.java index 567d515e1..7dc8e0dbc 100644 --- a/java/river/src/main/java/com/riverqueue/Workers.java +++ b/java/river/src/main/java/com/riverqueue/Workers.java @@ -31,7 +31,6 @@ public final class Workers 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() @@ -40,6 +39,7 @@ public final class Workers implements AutoCloseable { .plusDays(1) .atStartOfDay() .toInstant(java.time.ZoneOffset.UTC); + private final Semaphore cancellationWake = new Semaphore(0); private final Builder config; private final RetryPolicy defaultRetryPolicy = RetryPolicy.defaults(); private final java.util.concurrent.BlockingQueue completions = @@ -81,7 +81,10 @@ private Workers(Builder config) { id = config.id == null ? "java-" + UUID.randomUUID() : config.id; notifications = new Protocol.Dispatcher( - id, active, pendingCancellation, this::queueNotification, this::leadershipNotification); + id, + ignored -> cancellationWake.release(), + this::queueNotification, + this::leadershipNotification); queues.putAll(config.queues); periodic = new PeriodicJobBundle( @@ -99,6 +102,7 @@ private Workers(Builder config) { 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); if (config.leadership) executor.submit(this::elect); @@ -244,7 +248,6 @@ public Subscription subscribe(Consumer subscriber, EventKind... kinds) { private void dispatch() { long nextPoll = 0; long nextMaintenance = 0; - long nextCancelPoll = 0; while (running.get()) { try { long now = System.nanoTime(); @@ -272,25 +275,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(); @@ -443,8 +427,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( () -> { @@ -1394,6 +1376,41 @@ 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. + while (running.get() || dispatcherStopped.getCount() > 0 || !active.isEmpty()) { + try { + cancellationWake.drainPermits(); + var attempts = Map.copyOf(active); + if (!attempts.isEmpty()) { + river.transaction( + c -> { + try (var statement = + Sql.prepare(c, Sql.query(river.database(), "cancel_requested")); + 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 (Exception e) { + report(e); + } + try { + cancellationWake.tryAcquire(50, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + } + } + private void publish(Event event) { for (var subscriber : subscribers) try { diff --git a/java/river/src/main/resources/com/riverqueue/postgres.sql b/java/river/src/main/resources/com/riverqueue/postgres.sql index d5a92e3ca..10e7df6f5 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 diff --git a/java/river/src/main/resources/com/riverqueue/sqlite.sql b/java/river/src/main/resources/com/riverqueue/sqlite.sql index f38cfd267..2ccbff7de 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 diff --git a/java/river/src/test/java/com/riverqueue/ProtocolTest.java b/java/river/src/test/java/com/riverqueue/ProtocolTest.java index eb0424058..116f0b250 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,17 @@ 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, - queues::add, - leadership::add); + clientId, cancellations::add, 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,44 +119,35 @@ 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); + new Protocol.Dispatcher("observer", cancellations::add, 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()); } - private static WorkContext context(long id) { - return context(id, Json.object()); - } - private static WorkContext context(long id, JsonNode metadata) { var now = Instant.parse("2026-01-02T03:04:05Z"); var job = diff --git a/java/river/src/test/java/com/riverqueue/RetryTest.java b/java/river/src/test/java/com/riverqueue/RetryTest.java index 52385d523..77577ec99 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"}) void retryWakesWorkersAfterCommit(String mode) throws Exception { @@ -216,4 +296,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/WorkersTest.java b/java/river/src/test/java/com/riverqueue/WorkersTest.java index fc3a7c56d..cc1f36e32 100644 --- a/java/river/src/test/java/com/riverqueue/WorkersTest.java +++ b/java/river/src/test/java/com/riverqueue/WorkersTest.java @@ -71,6 +71,120 @@ void builderRequiresServiceIntervalsBelowQueueRetention(String value, boolean va .getMessage()); } + @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)) + .add( + TYPE, + context -> { + entered.complete(null); + context.awaitCancellation(); + cancelled.complete(context.cancellation()); + context.checkCancelled(); + }) + .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}) + void cancellationContinuesAfterGracefulStop(boolean pollOnly) throws Exception { + var client = client(); + var entered = new CompletableFuture>(); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .add( + TYPE, + context -> { + entered.complete(context); + context.awaitCancellation(); + context.checkCancelled(); + }) + .start(); + try (var subscription = workers.subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel while draining").job(); + entered.get(5, TimeUnit.SECONDS); + // The timeout leaves active work draining after the dispatcher and + // listener have stopped. A peer can still cancel that work. + 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(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + 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) + .add( + TYPE, + context -> { + context.awaitCancellation(); + context.checkCancelled(); + }) + .start(); + try (var subscription = workers.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"}) void cancellationRespectsTheHandlersOutcome(WorkContext.Cancellation cause, boolean interrupted) @@ -694,6 +808,71 @@ void invalidRetryPoliciesFallBack(String policy) throws Exception { } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + 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); + }) + .add( + TYPE, + context -> { + entered.complete(null); + finish.join(); + context.checkCancelled(); + }) + .start(); + try (var pauseSubscription = workers.subscribe(paused::add, Workers.EventKind.QUEUE_PAUSED); + var outcomeSubscription = + workers.subscribe(outcome::complete, Workers.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"}) void shortDelaysWakeFetchingBeforeThePollInterval( diff --git a/js/driver/pg/src/driver.integration.test.ts b/js/driver/pg/src/driver.integration.test.ts index 25ecd0f93..4ef3e1c25 100644 --- a/js/driver/pg/src/driver.integration.test.ts +++ b/js/driver/pg/src/driver.integration.test.ts @@ -531,11 +531,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", }) @@ -558,10 +570,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`, { @@ -633,6 +668,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); diff --git a/js/driver/pg/src/sql/jobs.ts b/js/driver/pg/src/sql/jobs.ts index 37b91ecb6..f95c75659 100644 --- a/js/driver/pg/src/sql/jobs.ts +++ b/js/driver/pg/src/sql/jobs.ts @@ -294,6 +294,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 +823,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/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 3506895e0..70d0f92b8 100644 --- a/js/driver/sqlite/src/driver.test.ts +++ b/js/driver/sqlite/src/driver.test.ts @@ -326,10 +326,12 @@ 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 @@ -350,6 +352,79 @@ describe("SqliteDriver", () => { ); }); + 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" })) diff --git a/js/driver/sqlite/src/driver.ts b/js/driver/sqlite/src/driver.ts index 98a0891c5..ead745c98 100644 --- a/js/driver/sqlite/src/driver.ts +++ b/js/driver/sqlite/src/driver.ts @@ -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/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..781973d07 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"); }); @@ -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,264 @@ describe("Client runtime database faults", () => { 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 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); + marked = true; + claim.resolve(undefined); + } + await started.promise; + const stopping = phase === "draining" ? run.stop() : undefined; + marked = true; + await waitUntil(() => driver.completions.length === 1); + expect(handlerSignal?.reason).toMatchObject({ + name: "JobCancelledError", + }); + expect(driver.completions[0]?.kind).toBe("cancel"); + expect(failures).toBeGreaterThanOrEqual(2); + await stopping; + await run.stop(); + } finally { + 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", + 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..a8ce8d425 100644 --- a/js/src/runtime.ts +++ b/js/src/runtime.ts @@ -215,7 +215,7 @@ export class RuntimeController { * still be cancelled while its queue drains. */ readonly #cancellationPollAbort = new AbortController(); - /** How often a runtime without notifications checks for cancellations. */ + /** How often the runtime checks for durable cancellation requests. */ readonly #cancellationPollIntervalMs: number; readonly #claimAbort = new AbortController(); readonly #client: Client; @@ -664,18 +664,15 @@ 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. + // Notifications deliver cancellations promptly. Keep a durable fallback + // until every producer drains, including during reconnects and shutdown. const listen = await this.#listens(); - if (!listen) { - this.#trackTask( - this.#runner.pollCancellations( - this.#cancellationPollIntervalMs, - this.#cancellationPollAbort.signal - ) - ); - } + this.#trackTask( + this.#runner.pollCancellations( + this.#cancellationPollIntervalMs, + this.#cancellationPollAbort.signal + ) + ); // Like River for Go's notifier, notifications are listening before any // queue claims, and a failure to listen fails the start. await this.#notifications.start(listen); diff --git a/js/src/runtime/attempt-runner.ts b/js/src/runtime/attempt-runner.ts index cca2ea5a5..29423aee0 100644 --- a/js/src/runtime/attempt-runner.ts +++ b/js/src/runtime/attempt-runner.ts @@ -108,10 +108,15 @@ 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; /** * 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; @@ -153,18 +158,15 @@ 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. */ + /** Check a cancellation hint against durable state before aborting work. */ cancelAttempt(id: bigint, attemptedBy: string): void { const key = attemptKey(id, attemptedBy); - const attempt = this.#attempts.get(key); - if (attempt !== undefined) { - attempt.abort(new JobCancelledError(id)); - return; - } for (const pending of this.#claimCancellations) pending.add(key); + if (this.#attempts.has(key)) void this.#checkCancellations(); } /** @@ -178,30 +180,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 +221,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 +228,7 @@ export class AttemptRunner { return; } try { - await this.#cancelRequested(requested, this.#ownAttemptIds(), signal); + await this.recoverCancellations(signal); } catch (error: unknown) { if (signal.aborted) return; this.#context.logger.error( @@ -232,17 +241,19 @@ export class AttemptRunner { /** 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 +359,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 +490,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,9 +507,8 @@ 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; @@ -871,15 +882,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, @@ -893,7 +904,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(); @@ -901,12 +916,44 @@ 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.#runSignal); + } catch (error: unknown) { + if (this.#runSignal.aborted) return; + this.#context.logger.warn( + "River could not confirm a cancellation notification", + { error: describeError(error) } + ); + } + } + } 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/producer.go b/producer.go index fd5a4e9f5..a5f1dcfd0 100644 --- a/producer.go +++ b/producer.go @@ -45,18 +45,20 @@ 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 - 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 - ReportedQueueStatus testsignal.TestSignal[struct{}] // notifies when the producer reports queue status - Resumed testsignal.TestSignal[struct{}] // notifies when the producer is resumed - StartedExecutors testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass + 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 + 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 + Reconciled testsignal.TestSignal[struct{}] // notifies when durable cancellation state was applied + ReconnectHandledDuringFetch testsignal.TestSignal[struct{}] // notifies when a reconnect is handled during a fetch + ReportedProducerStatus testsignal.TestSignal[struct{}] // notifies when the producer reports its own status + ReportedQueueStatus testsignal.TestSignal[struct{}] // notifies when the producer reports queue status + Resumed testsignal.TestSignal[struct{}] // notifies when the producer is resumed + StartedExecutors testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass } func (ts *producerTestSignals) Init(tb testutil.TestingTB) { @@ -67,6 +69,7 @@ func (ts *producerTestSignals) Init(tb testutil.TestingTB) { ts.Paused.Init(tb) ts.PolledQueueConfig.Init(tb) ts.QueueControlEventTriggered.Init(tb) + ts.ReconnectHandledDuringFetch.Init(tb) ts.ReportedQueueStatus.Init(tb) ts.ReportedProducerStatus.Init(tb) ts.Resumed.Init(tb) @@ -215,9 +218,9 @@ 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 + // Receives job IDs from notifier, polling, and local cancellation calls. + // Only the main goroutine reads it. + cancelCh chan producerCancelRequest // 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 +243,13 @@ 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 + + // Coalesces reconnect notifications for reconciliation on the main goroutine. + reconnectCh chan struct{} + + retryPolicy ClientRetryPolicy + runState atomic.Pointer[producerRunState] + testSignals producerTestSignals } func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pilot riverpilot.Pilot, config *producerConfig) *producer { @@ -259,7 +267,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), + cancelCh: make(chan producerCancelRequest, 1000), completer: config.Completer, config: config.mustValidate(), exec: exec, @@ -268,6 +276,7 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi jobTimeout: config.JobTimeout, pilot: pilot, queueControlCh: make(chan *controlEventPayload, 100), + reconnectCh: make(chan struct{}, 1), retryPolicy: config.RetryPolicy, workers: config.Workers, }) @@ -307,6 +316,18 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } + // Controls left over from the previous run must not override current queue + // settings or cancel a retried attempt. Delivery can race with shutdown. +drainControls: + for { + select { + case <-p.cancelCh: + 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 +394,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.ListenerReadySubscription ) if p.config.Notifier != nil { var err error @@ -400,7 +422,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 +430,23 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } return err } + + // Recover cancellations whose NOTIFY was lost while disconnected. + reconnectSub = p.config.Notifier.RegisterListenerReadyFunc(func() { + select { + case p.reconnectCh <- struct{}{}: + default: + } + }) } + 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 +461,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) @@ -453,8 +490,6 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } 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.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)) @@ -484,9 +519,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 @@ -509,7 +560,7 @@ 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 +570,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))) @@ -540,7 +594,7 @@ func (p *producer) handleControlNotification(workCtx context.Context) func(notif } select { case <-workCtx.Done(): - case p.cancelCh <- decoded.JobID: + case p.cancelCh <- producerCancelRequest{jobID: decoded.JobID}: default: p.Logger.WarnContext(workCtx, p.Name+": Job cancel notification dropped due to full buffer", slog.Int64("job_id", decoded.JobID)) } @@ -560,17 +614,75 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { p.fetchLimiter.Call() fetchResultCh := make(chan producerFetchResult) + // Reconciliation must remain live while workers drain after fetching stops. + reconcileCtx, cancelReconcile := context.WithCancel(context.WithoutCancel(fetchCtx)) + var reconcileWG sync.WaitGroup + defer func() { + cancelReconcile() + reconcileWG.Wait() + }() + + reconcileResultCh := make(chan producerReconcileResult) + var ( + reconcileAttempt int + reconcilePending bool + reconciling bool + retryTimer *time.Timer + ) + defer func() { + if retryTimer != nil { + retryTimer.Stop() + } + }() + startReconcile := func() { + if reconciling { + reconcilePending = true + return + } + if retryTimer != nil { + retryTimer.Stop() + retryTimer = nil + } + p.activeJobsMu.Lock() + numActiveJobs := len(p.activeJobs) + p.activeJobsMu.Unlock() + if numActiveJobs == 0 { + reconcileAttempt = 0 + return + } + reconciling = true + reconcileWG.Add(1) + p.handleNotifierReconnect(reconcileCtx, reconcileResultCh, &reconcileWG) + } + handleReconcileResult := func(result producerReconcileResult) { + reconciling = false + if result.err != nil { + p.Logger.ErrorContext(workCtx, p.Name+": Error reconciling cancellations; will retry", slog.String("err", result.err.Error())) + if retryTimer != nil { + retryTimer.Stop() + } + retryTimer = p.scheduleReconnectRetry(reconcileCtx, reconcileAttempt) + reconcileAttempt++ + } else { + reconcileAttempt = 0 + p.applyReconcileResult(workCtx, result) + p.testSignals.Reconciled.Signal(struct{}{}) + } + if reconcilePending { + reconcilePending = false + startReconcile() + } + } + +fetchLoop: for { select { case <-fetchCtx.Done(): - return + break fetchLoop 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) + startReconcile() 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,15 +716,21 @@ 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 request := <-p.cancelCh: + if !request.cancel(workCtx) { + startReconcile() + } + case <-p.reconnectCh: + startReconcile() + case result := <-reconcileResultCh: + handleReconcileResult(result) case <-p.fetchLimiter.C(): - p.innerFetchLoop(workCtx, fetchResultCh) + p.innerFetchLoop(workCtx, fetchResultCh, reconcileResultCh, handleReconcileResult, startReconcile) // Ensure we can't start another fetch when fetchCtx is done, even if // the fetchLimiter is also ready to fire: select { case <-fetchCtx.Done(): - return + break fetchLoop default: } case result := <-p.jobResultCh: @@ -627,6 +745,9 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { } } } + + p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) + p.executorShutdownLoop(workCtx, reconcileResultCh, handleReconcileResult, startReconcile) } // Loops every FetchPollInterval to check for jobs. This is meant as a back up @@ -661,7 +782,7 @@ func (p *producer) jitteredFetchPollInterval() time.Duration { return randutil.DurationBetween(p.config.FetchPollInterval, p.config.FetchPollInterval+jitterRange) } -func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan producerFetchResult) { +func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan producerFetchResult, reconcileResultCh <-chan producerReconcileResult, handleReconcileResult func(producerReconcileResult), startReconcile func()) { var limit int if p.paused { limit = 0 @@ -676,9 +797,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{}) + var reconcileAfterFetch bool go p.dispatchWork(workCtx, limit, fetchResultCh) for { @@ -687,7 +806,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 reconcileAfterFetch { + startReconcile() + } if numLocked == limit { // Fetch returned the maximum number of jobs that were requested, @@ -697,18 +819,28 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } return + case <-p.reconnectCh: + // Claimed jobs may still be absent from activeJobs until this fetch + // returns. Recheck them after starting their executors too. + reconcileAfterFetch = true + startReconcile() + p.testSignals.ReconnectHandledDuringFetch.Signal(struct{}{}) + case result := <-reconcileResultCh: + handleReconcileResult(result) case result := <-p.jobResultCh: p.removeActiveJob(result) - case jobID := <-p.cancelCh: - if !p.maybeCancelJob(workCtx, jobID) { - pendingCancelJobs[jobID] = struct{}{} + case request := <-p.cancelCh: + if !request.cancel(workCtx) { + // The claim may not be registered yet. Check durable state once it is. + reconcileAfterFetch = true + startReconcile() } - p.testSignals.CancelHandledDuringFetch.Signal(jobID) + p.testSignals.CancelHandledDuringFetch.Signal(request.jobID) } } } -func (p *producer) executorShutdownLoop(ctx context.Context) { +func (p *producer) executorShutdownLoop(ctx context.Context, reconcileResultCh <-chan producerReconcileResult, handleReconcileResult func(producerReconcileResult), startReconcile func()) { p.testSignals.ExecutorShutdownStarted.Signal(struct{}{}) // No more jobs will be fetched or executed. However, we must wait for all @@ -722,10 +854,16 @@ func (p *producer) executorShutdownLoop(ctx context.Context) { } select { - case jobID := <-p.cancelCh: - p.maybeCancelJob(ctx, jobID) + case request := <-p.cancelCh: + if !request.cancel(ctx) { + startReconcile() + } case result := <-p.jobResultCh: p.removeActiveJob(result) + case <-p.reconnectCh: + startReconcile() + case result := <-reconcileResultCh: + handleReconcileResult(result) } } } @@ -787,6 +925,18 @@ func (p *producer) addActiveJob(id int64, executor *jobexecutor.JobExecutor) { p.activeJobsMu.Unlock() } +// applyReconcileResult runs on the producer goroutine and ignores replaced executions. +func (p *producer) applyReconcileResult(ctx context.Context, result producerReconcileResult) { + for _, id := range result.cancelRequestedIDs { + p.activeJobsMu.Lock() + executor := p.activeJobs[id] + p.activeJobsMu.Unlock() + if executor != nil && executor == result.executors[id] { + executor.Cancel(ctx) + } + } +} + func (p *producer) removeActiveJob(job *rivertype.JobRow) { p.activeJobsMu.Lock() executor := p.activeJobs[job.ID] @@ -828,15 +978,46 @@ func (p *producer) handleWorkerUnstuck() { p.config.JobStuckCount.Add(-1) } -func (p *producer) maybeCancelJob(ctx context.Context, id int64) bool { +// handleNotifierReconnect snapshots active executions before querying asynchronously. +func (p *producer) handleNotifierReconnect(ctx context.Context, resultCh chan<- producerReconcileResult, wg *sync.WaitGroup) { p.activeJobsMu.Lock() - executor, ok := p.activeJobs[id] - p.activeJobsMu.Unlock() - if !ok { - return false + executors := make(map[int64]*jobexecutor.JobExecutor, len(p.activeJobs)) + ids := make([]int64, 0, len(p.activeJobs)) + for id, executor := range p.activeJobs { + executors[id] = executor + ids = append(ids, id) } - executor.Cancel(ctx) - return true + + p.activeJobsMu.Unlock() + + go func() { + defer wg.Done() + + cancelRequestedIDs, err := timeoututil.WithTimeoutV(ctx, 10*time.Second, p.Name+".handleNotifierReconnect", func(ctx context.Context) ([]int64, error) { + return p.exec.JobGetCancelRequested(ctx, &riverdriver.JobGetCancelRequestedParams{ + ID: ids, + Schema: p.config.Schema, + }) + }) + select { + case <-ctx.Done(): + case resultCh <- producerReconcileResult{cancelRequestedIDs: cancelRequestedIDs, err: err, executors: executors}: + } + }() +} + +// scheduleReconnectRetry retries a failed lookup without blocking the producer. +func (p *producer) scheduleReconnectRetry(ctx context.Context, attempt int) *time.Timer { + sleepDuration := serviceutil.ExponentialBackoff(attempt, serviceutil.MaxAttemptsBeforeResetDefault) + return time.AfterFunc(sleepDuration, func() { + if ctx.Err() != nil { + return + } + select { + case p.reconnectCh <- struct{}{}: + default: + } + }) } func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit { @@ -966,7 +1147,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 +1188,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) } @@ -1052,8 +1229,9 @@ func (p *producer) pollForJobCancellations(ctx context.Context, wg *sync.WaitGro 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 +1247,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) } } } @@ -1233,12 +1411,38 @@ func (p *producer) reportQueueStatusOnce(ctx context.Context) { p.testSignals.ReportedQueueStatus.Signal(struct{}{}) } +// A local request retains the executor that matched the cancelled database row. +// Requests without that identity must be confirmed against durable state. +type producerCancelRequest struct { + executor *jobexecutor.JobExecutor + jobID int64 +} + +func (r producerCancelRequest) cancel(ctx context.Context) bool { + if r.executor == nil { + return false + } + r.executor.Cancel(ctx) + return true +} + type producerFetchResult struct { decodeErrors map[int64]error err error jobs []*rivertype.JobRow } +type producerReconcileResult struct { + cancelRequestedIDs []int64 + err error + executors map[int64]*jobexecutor.JobExecutor +} + +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..4c51e50dc 100644 --- a/producer_test.go +++ b/producer_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "slices" + "sync" "sync/atomic" "testing" "time" @@ -13,11 +14,13 @@ import ( "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 +79,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 +126,298 @@ 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_ApplyReconcileResult(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + cancelRequested bool + replaced bool + wantCancel bool + }{ + {name: "Cancelled", cancelRequested: true, wantCancel: true}, + {name: "NoCancellation"}, + {name: "ReplacedExecutor", cancelRequested: true, replaced: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + archetype := riversharedtest.BaseServiceArchetype(t) + ctx, cancel := context.WithCancelCause(context.Background()) + t.Cleanup(func() { cancel(nil) }) + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancel, JobRow: &rivertype.JobRow{ID: 1, Attempt: 1}}) + snapshotExecutor := executor + if testCase.replaced { + snapshotExecutor = &jobexecutor.JobExecutor{} + } + prod := baseservice.Init(archetype, &producer{activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}}) + + var cancelRequestedIDs []int64 + if testCase.cancelRequested { + cancelRequestedIDs = []int64{1} + } + prod.applyReconcileResult(context.Background(), producerReconcileResult{ + cancelRequestedIDs: cancelRequestedIDs, + executors: map[int64]*jobexecutor.JobExecutor{1: snapshotExecutor}, + }) + if testCase.wantCancel { + require.ErrorIs(t, context.Cause(ctx), rivertype.ErrJobCancelledRemotely) + } else { + require.NoError(t, ctx.Err()) + } + }) + } +} + +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{ + cancelCh: make(chan producerCancelRequest, 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.cancelCh, 1) + require.Equal(t, int64(123), (<-prod.cancelCh).jobID) + } else { + require.Empty(t, prod.cancelCh) + } + }) + } +} + +// 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.RegisterListenerReadyFunc(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: queuePollIntervalDefault, + 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.ReconnectHandledDuringFetch.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.ReconnectHandledDuringFetch.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 +580,52 @@ func TestProducer_MetricEmitHook(t *testing.T) { }) } +func TestProducer_PollForJobCancellations(t *testing.T) { + t.Parallel() + + 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}, + cancelCh: make(chan producerCancelRequest, 1), + 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())) + // Process any queued cancellations as the producer loop would. + select { + case request := <-prod.cancelCh: + request.cancel(context.Background()) + default: + } + require.NoError(t, replacementCtx.Err()) + if !replaced { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + } + }) + } +} + func TestProducer_PollOnly(t *testing.T) { t.Parallel() @@ -322,6 +678,124 @@ func TestProducer_PollOnly(t *testing.T) { }) } +func TestProducer_ReconcileDoesNotBlock(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + blockedFetch bool + cancelFromLookup bool + queryError bool + }{ + {name: "BlockedFetchCancelDelivery", blockedFetch: true}, + {name: "BlockedFetchReconcileResult", blockedFetch: true, cancelFromLookup: true}, + {name: "CancelDelivery"}, + {name: "ReconcileResult", cancelFromLookup: 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] + } + signals.FetchStarted.Init(t) + signals.LoopStopped.Init(t) + signals.QueryStarted.Init(t) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + releaseFetch := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseFetch) }) } + 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}}) + var queryAttempts int + exec := &producerTestExecutor{ + jobGetAvailableFunc: func(context.Context, *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) { + signals.FetchStarted.Signal(struct{}{}) + <-releaseFetch + 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 reconciliation 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}, + cancelCh: make(chan producerCancelRequest, 1), + config: &producerConfig{MaxWorkers: 2}, + 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), + reconnectCh: make(chan struct{}, 1), + }) + go func() { + prod.fetchAndRunLoop(ctx, ctx) + signals.LoopStopped.Signal(struct{}{}) + }() + defer func() { + release() + cancel() + signals.LoopStopped.WaitOrTimeout() + }() + + if testCase.blockedFetch { + signals.FetchStarted.WaitOrTimeout() + } + prod.reconnectCh <- struct{}{} + 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.cancelFromLookup { + prod.cancelCh <- producerCancelRequest{executor: executor, jobID: 1} + } + select { + case <-jobCtx.Done(): + require.ErrorIs(t, context.Cause(jobCtx), rivertype.ErrJobCancelledRemotely) + case <-time.After(time.Second): + t.Fatal("blocked reconciliation prevented cancellation delivery") + } + + select { + case prod.jobResultCh <- &rivertype.JobRow{ID: 1}: + case <-time.After(time.Second): + t.Fatal("blocked reconciliation prevented job completion") + } + }) + } +} + func TestProducer_WithNotifier(t *testing.T) { t.Parallel() @@ -421,7 +895,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 +914,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 +976,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.cancelCh <- producerCancelRequest{jobID: id} + + 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() @@ -683,6 +1197,33 @@ 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.Reconciled.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. + producer.cancelCh <- producerCancelRequest{jobID: id} + producer.testSignals.Reconciled.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() @@ -716,6 +1257,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,7 +1277,7 @@ 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 + producer.cancelCh <- producerCancelRequest{jobID: jobID} require.Equal(t, jobID, producer.testSignals.CancelHandledDuringFetch.WaitOrTimeout()) releaseFetch.Signal(struct{}{}) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 4bd6cbfab..05b4fd854 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -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 a4de886f8..0b05ee242 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/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index d87361114..ac038c1c4 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -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..069001869 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -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 6e651cbfe..13271c6a8 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 0ff9958d2..11c26f388 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/river_sqlite_driver.go b/riverdriver/riversqlite/river_sqlite_driver.go index 1fefd422a..ff7eb13ac 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 ac064ca8e..30d702ec8 100644 --- a/ruby/lib/client_runtime.rb +++ b/ruby/lib/client_runtime.rb @@ -120,7 +120,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 +198,9 @@ def start @started = true @stop_requested = false @stopped = false + thread = Thread.new { cancellation_loop } + @cancellation_thread = thread + @threads << thread end begin @@ -243,6 +247,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 +276,15 @@ 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}") + end + end + interruption = @mutex.synchronize do entry = @running.fetch(id) if @stop_requested @@ -286,12 +300,45 @@ 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 + loop do + begin + check_remote_cancellations + rescue => error + @config.logger.error("River cancellation check failed: #{error.full_message}") + end + + finished = @mutex.synchronize do + if @stop_requested && @running.empty? + true + else + interval = @queue_configs.values.map { |config| config.resolved_fetch_poll_interval(@config) }.min || @config.fetch_poll_interval + # Continue polling while producers drain, independently of fetches + # and their cooldowns. Completions and queue changes wake us early. + @condition.wait(@mutex, interval) unless @running.any? { |_id, entry| entry[:check_cancellation] } + false + end + end + break if finished + end + end + + private def check_remote_cancellations(entries: nil) + entries ||= @mutex.synchronize do + @running.select { |_id, entry| entry[:working] || entry[:check_cancellation] }.tap do |snapshot| + snapshot.each_value { |entry| entry[:check_cancellation] = false } + end + end 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 @@ -524,7 +571,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] = {check_before_work: @fetching_queues.fetch(row.queue).include?(row.id), queue: row.queue, thread: thread, working: false} end gate.push(true) end @@ -698,7 +745,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 @@ -779,8 +825,9 @@ def wake @mutex.synchronize do entry = @running[message["job_id"]] if entry && entry[:queue] == queue - entry[:cancelled] = true - entry[:thread].raise(JobCancelError) if entry[:working] + entry[:check_before_work] = true unless entry[:working] + entry[:check_cancellation] = true + @condition.broadcast elsif (pending = @fetching_queues[queue]) # A claim may have committed before its workers are registered. pending << message["job_id"] diff --git a/ruby/lib/driver/runtime.rb b/ruby/lib/driver/runtime.rb index f353230ca..a4ab6d395 100644 --- a/ruby/lib/driver/runtime.rb +++ b/ruby/lib/driver/runtime.rb @@ -242,27 +242,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 @@ -489,7 +491,7 @@ def request_resign 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 ab94b3a7e..6b7edbcc5 100644 --- a/ruby/sig/runtime.rbs +++ b/ruby/sig/runtime.rbs @@ -358,6 +358,7 @@ module River def initialize: (Time, Float, Time) -> void end + @cancellation_thread: Thread? @client: Client @condition: ConditionVariable @config: Config @@ -404,7 +405,8 @@ module River private def begin_work: (Integer) -> untyped - def check_remote_cancellations: (String) -> untyped + def cancellation_loop: () -> untyped + def check_remote_cancellations: () -> untyped def error_handler_cancel?: (Exception, Job) -> bool def execute: (JobRow) -> untyped def finish_failed: (JobRow, Job, Exception, JobTiming, ?cancelled: bool, ?worker: untyped) -> untyped diff --git a/ruby/spec/client_runtime_branch_spec.rb b/ruby/spec/client_runtime_branch_spec.rb index 999be512c..03394a6af 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 @@ -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 @@ -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 @@ -728,7 +744,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 in one batch, excluding inactive work" do raised = [] fake_thread = Object.new fake_thread.define_singleton_method(:raise) { |error| raised << error } @@ -748,10 +815,10 @@ 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, 3]]) end it "interrupts work that begins after stop was requested" do diff --git a/ruby/spec/driver_runtime_feature_spec.rb b/ruby/spec/driver_runtime_feature_spec.rb index 3c793561f..4771894d3 100644 --- a/ruby/spec/driver_runtime_feature_spec.rb +++ b/ruby/spec/driver_runtime_feature_spec.rb @@ -121,7 +121,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 058f54ce4..4218f86df 100644 --- a/ruby/spec/driver_runtime_shared_examples.rb +++ b/ruby/spec/driver_runtime_shared_examples.rb @@ -293,6 +293,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) @@ -483,6 +490,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/notification_runtime_spec.rb b/ruby/spec/notification_runtime_spec.rb index d3e6e6d57..2e9d18b45 100644 --- a/ruby/spec/notification_runtime_spec.rb +++ b/ruby/spec/notification_runtime_spec.rb @@ -30,6 +30,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,6 +41,45 @@ 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 "ignores invalid payloads without disrupting the receiver" do [nil, "{", "null", "[]", "{}"].each { |payload| runtime.send(:receive_notification, "river_insert", payload) } runtime.send(:receive_notification, "unknown", "{}") diff --git a/ruby/spec/runtime_draining_shared_examples.rb b/ruby/spec/runtime_draining_shared_examples.rb index 879b6b742..66e0cea34 100644 --- a/ruby/spec/runtime_draining_shared_examples.rb +++ b/ruby/spec/runtime_draining_shared_examples.rb @@ -77,9 +77,11 @@ 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) @@ -106,6 +108,7 @@ nil end Timeout.timeout(5) { draining.pop } + Timeout.timeout(5) { poll_failed.pop } if poll_failure # Persist through another client so only database polling can deliver it. River::Client.new(@driver).job_cancel(job.id) diff --git a/rust/riverqueue/src/client/attempts.rs b/rust/riverqueue/src/client/attempts.rs index 7483935d5..f51ffe204 100644 --- a/rust/riverqueue/src/client/attempts.rs +++ b/rust/riverqueue/src/client/attempts.rs @@ -2,11 +2,10 @@ use std::collections::HashMap; use std::sync::Mutex; -use std::sync::atomic::{AtomicU64, Ordering}; use tokio_util::sync::CancellationToken; -use crate::client::{ClientInner, PENDING_CANCELLATION_LIMIT, PENDING_CANCELLATION_RETENTION}; +use crate::client::ClientInner; pub(super) fn remove_running_attempt( running: &Mutex>, @@ -25,99 +24,15 @@ pub(super) fn remove_running_attempt( } pub(super) fn register_running_attempt( - running: &Mutex>, - pending_cancellations: &Mutex>, + inner: &ClientInner, job_id: i64, cancellation: &CancellationToken, ) { - // 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, cancellation.clone()); + // Recheck after registration even if a notice arrived during the claim. + inner.cancellation_wakeup.notify_one(); } diff --git a/rust/riverqueue/src/client/builder.rs b/rust/riverqueue/src/client/builder.rs index 0c378c36d..058214cf1 100644 --- a/rust/riverqueue/src/client/builder.rs +++ b/rust/riverqueue/src/client/builder.rs @@ -1,12 +1,14 @@ //! Client configuration and construction. use std::collections::HashMap; -use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::sync::atomic::AtomicBool; +#[cfg(test)] +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 +804,7 @@ 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_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..94ed79cce 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.cancellation_wakeup.notify_one(); } Ok(row) } diff --git a/rust/riverqueue/src/client/mod.rs b/rust/riverqueue/src/client/mod.rs index 68b517853..66a71958e 100644 --- a/rust/riverqueue/src/client/mod.rs +++ b/rust/riverqueue/src/client/mod.rs @@ -57,7 +57,7 @@ use std::{ }; use chrono::{DateTime, Utc}; -use tokio::sync::{broadcast, mpsc, watch}; +use tokio::sync::{Notify, broadcast, mpsc, watch}; use tokio_util::sync::CancellationToken; use self::completer::CompletionUpdate; @@ -79,8 +79,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 +101,7 @@ const PRODUCER_STALE_RETENTION: Duration = Duration::from_mins(5); pub(crate) struct ClientInner { allow_legacy_job_kinds: bool, allow_unregistered_job_kinds: bool, + 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, diff --git a/rust/riverqueue/src/client/notifier.rs b/rust/riverqueue/src/client/notifier.rs index be9aa976e..07ab86e67 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; @@ -125,13 +124,8 @@ 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 payload.job_id.is_some() { + inner.cancellation_wakeup.notify_one(); } } // Like Go's producer, a metadata change is handled at once diff --git a/rust/riverqueue/src/client/producer.rs b/rust/riverqueue/src/client/producer.rs index fa68924f9..e6045da78 100644 --- a/rust/riverqueue/src/client/producer.rs +++ b/rust/riverqueue/src/client/producer.rs @@ -16,9 +16,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 +54,8 @@ 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. + // Recover lost cancellation notices, including after the notifier stops + // during graceful shutdown, until every producer has drained. let _cancellation_poll = AbortOnDrop(tokio::spawn(poll_job_cancellations( Arc::clone(&inner), fetch_cancel.clone(), @@ -143,10 +139,9 @@ fn jittered_fetch_poll_interval(interval: Duration) -> Duration { } /// 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. +/// request, checking every [`QUEUE_CONFIG_POLL_INTERVAL`]. Notifications +/// deliver cancellations promptly; polling also recovers notices lost during +/// reconnects or while a fetched attempt was still being registered. /// 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) { @@ -155,46 +150,42 @@ async fn poll_job_cancellations(inner: Arc, stopping: CancellationT // The first tick completes at once; running jobs start later. poll.tick().await; loop { - poll.tick().await; - if inner.has_notifier() { - continue; + tokio::select! { + _ = poll.tick() => {}, + () = 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; - } + if let Err(poll_error) = check_job_cancellations(&inner).await + && !stopping.is_cancelled() + { + error!(error = %crate::error::Chain(&poll_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(cancellation) = attempts.get(&id) { + cancellation.cancel(); } } } + Ok(()) } type ProducerOutcome = (String, u64, CancellationToken, Result<(), Error>); @@ -858,7 +849,6 @@ pub(super) async fn run_queue( if available == 0 { continue; } - let registration_guard = FetchRegistrationGuard::new(&inner); let fetched = match (&session, claims_through_session) { (Some(session), true) => { match claim_through_session( @@ -907,12 +897,7 @@ pub(super) async fn run_queue( for (row, decode_error) in claimed { let hard_cancel = work_cancel.child_token(); let cancellation = hard_cancel.child_token(); - register_running_attempt( - &inner.running, - &inner.pending_cancellations, - row.id, - &cancellation, - ); + register_running_attempt(&inner, row.id, &cancellation); let task_inner = Arc::clone(&inner); let completion_sender = completion_sender.clone(); let session_row = attempts.session.is_some().then(|| row.clone()); @@ -928,7 +913,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/tests.rs b/rust/riverqueue/src/client/tests.rs index d52f83f76..45e3ae72a 100644 --- a/rust/riverqueue/src/client/tests.rs +++ b/rust/riverqueue/src/client/tests.rs @@ -1,7 +1,7 @@ use serde::Deserialize; use serde_json::Map; -use super::attempts::{register_running_attempt, remove_running_attempt, signal_running_attempt}; +use super::attempts::{register_running_attempt, remove_running_attempt}; #[cfg(feature = "postgres")] use super::completer::CompletionBatcher; #[cfg(feature = "sqlite")] @@ -171,12 +171,7 @@ 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, - job_id, - token, - ); + inner.running.lock().unwrap().insert(job_id, token.clone()); } dispatch_notification( @@ -209,11 +204,18 @@ async fn go_notification_fixtures_dispatch_correctly() { ), _ => unreachable!(), } - assert_eq!(cancellation.is_cancelled(), name == "cancel", "{name}"); assert!( - !unrelated.is_cancelled(), - "{name}: cancelled an unrelated job" + !cancellation.is_cancelled(), + "{name}: notice cancelled without a durable check" ); + assert_eq!( + tokio::time::timeout(Duration::ZERO, inner.cancellation_wakeup.notified()) + .await + .is_ok(), + name == "cancel", + "{name}" + ); + assert!(!unrelated.is_cancelled(), "{name}"); assert!( matches!( producer.try_recv(), @@ -233,50 +235,73 @@ 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 registration_wakes_cancellation_check() { + #[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, - ); - register_running_attempt(&running, &pending_cancellations, job_id, &cancellation); - - assert!(cancellation.is_cancelled()); - assert!( - pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() - ); + register_running_attempt(&client.inner, 42, &cancellation); + 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, - ); - - 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.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); + 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 { diff --git a/rust/riverqueue/src/database/sqlite.rs b/rust/riverqueue/src/database/sqlite.rs index 4e8b26e21..40dfe3444 100644 --- a/rust/riverqueue/src/database/sqlite.rs +++ b/rust/riverqueue/src/database/sqlite.rs @@ -935,25 +935,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) @@ -2523,6 +2528,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..08b0b7fe0 100644 --- a/rust/riverqueue/src/storage/postgres.rs +++ b/rust/riverqueue/src/storage/postgres.rs @@ -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/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..a01234f70 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, including when notifications are +//! disabled or a cancellation notice is lost, 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. @@ -11,8 +11,8 @@ mod support; use std::{convert::Infallible, sync::Arc, time::Duration}; use riverqueue::{ - __private::Database, Client, EventKind, Job, JobArgs, JobState, QueueConfig, WorkContext, - WorkOutcome, Workers, + __private::Database, Client, EventKind, Job, JobArgs, JobState, JobUpdateParams, QueueConfig, + WorkContext, WorkOutcome, Workers, }; use serde::{Deserialize, Serialize}; use tokio::sync::Semaphore; @@ -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. +/// Durable cancellation survives lost notifications and 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(); @@ -96,7 +100,27 @@ async fn polls_for_remote_cancellation(backend: Backend) { if while_stopping { run.stopper().stop(); } - other.jobs().cancel(id).await.unwrap(); + if notifications { + // Persist a cancellation without its notification, as though that + // notice was lost while the listener reconnected. + other + .jobs() + .update( + id, + JobUpdateParams::default().metadata( + [( + "cancel_attempted_at".to_owned(), + chrono::Utc::now().to_rfc3339().into(), + )] + .into_iter() + .collect(), + ), + ) + .await + .unwrap(); + } else { + other.jobs().cancel(id).await.unwrap(); + } let event = tokio::time::timeout(TIMEOUT, events.recv()) .await