From 25d74d2a3db4c2b84f626b85578879be09747b2f Mon Sep 17 00:00:00 2001 From: Chris George Date: Sun, 23 Aug 2026 14:32:18 -0700 Subject: [PATCH 1/6] Reconcile active jobs against cancel_attempted_at after notifier reconnects --- CHANGELOG.md | 1 + internal/notifier/notifier.go | 59 ++++++++-- internal/notifier/notifier_test.go | 48 ++++++++ producer.go | 119 ++++++++++++++++++- producer_test.go | 177 +++++++++++++++++++++++++++++ 5 files changed, 393 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a6b73c75..7582f2680 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Fixed `Client.InsertManyFast` panicking after its transaction committed on a client configured with queues when the driver doesn't support a listener, such as a `riverdatabasesql` driver created with `New` instead of `NewWithPgxListener`. It now wakes producers for the inserted jobs' queues right away, like `Client.InsertMany`. [PR #1558](https://github.com/riverqueue/river/pull/1558). - Fixed `Client.QueueUpdate` without metadata crashing the program when the client works that queue and the driver doesn't support a listener. [PR #1558](https://github.com/riverqueue/river/pull/1558). - Fixed the notifier's run loop blocking forever when a topic `LISTEN` or `UNLISTEN` handed to it finished after its caller had stopped waiting (after 5 seconds, with `timed out waiting for interrupt operation`), as can happen on a slow or unresponsive connection. `Client.Stop` then never completed, and neither did a `Client.Start` that failed this way. [PR #1559](https://github.com/riverqueue/river/pull/1559). +- A `JobCancel()` sent while a client's notifier is disconnected/reconnecting is no longer lost. Previously, the cancellation's `NOTIFY` could commit while nothing was `LISTEN`ing (Postgres `NOTIFY` is fire-and-forget), leaving the running job unaware of the cancellation until `JobRescuer`'s much longer stuck-job sweep. Each client now reconciles its currently running jobs against `cancel_attempted_at` by primary key every time its notifier (re)establishes healthy listening, closing the gap with no added steady-state query load. [PR #1361](https://github.com/riverqueue/river/pull/1361), as discussed in [#1358](https://github.com/riverqueue/river/issues/1358). ## [0.50.0] - 2026-10-26 diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index d453eddc3..8fba56743 100644 --- a/internal/notifier/notifier.go +++ b/internal/notifier/notifier.go @@ -88,12 +88,46 @@ 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 +} + +// ListenerReadyFunc is invoked each time the notifier establishes healthy +// listening on all of its currently subscribed topics. This includes the +// very first connect as well as every subsequent reconnect that follows a +// connection loss, so callers whose work is only meaningful after a genuine +// reconnect (as opposed to the initial connect) should make that check +// themselves. +type ListenerReadyFunc func() + +// ListenerReadySubscription is a handle to a func registered with +// RegisterListenerReadyFunc. Call Unregister to stop receiving invocations. +type ListenerReadySubscription struct { + notifier *Notifier +} + +func (s *ListenerReadySubscription) Unregister() { + s.notifier.mu.Lock() + defer s.notifier.mu.Unlock() + delete(s.notifier.listenerReadyFuncs, s) +} + +// RegisterListenerReadyFunc registers a function to be invoked every time the +// notifier (re)establishes healthy listening on all of its subscribed topics. +// It's used to let callers reconcile any state that could have missed a +// notification while the notifier was disconnected/reconnecting. +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 New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notifier { @@ -102,7 +136,8 @@ 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 } @@ -214,6 +249,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 8d5c21cc3..c36016f59 100644 --- a/internal/notifier/notifier_test.go +++ b/internal/notifier/notifier_test.go @@ -623,6 +623,54 @@ func TestNotifier(t *testing.T) { // Subscription should still work. 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 + + readyChan := make(chan struct{}, 10) + sub := notifier.RegisterListenerReadyFunc(func() { readyChan <- struct{}{} }) + t.Cleanup(sub.Unregister) + + start(t, notifier) + + // Fires once on the initial connect. + riversharedtest.WaitOrTimeout(t, readyChan) + + // And again once the notifier reconnects after the injected failure. + require.EqualError(t, notifier.testSignals.BackoffError.WaitOrTimeout(), "error during wait") + riversharedtest.WaitOrTimeout(t, readyChan) + + // Unregister stops further invocations. There's no direct signal for + // "would have fired again", so drain what's pending, unregister, and + // confirm a subsequent restart doesn't add anything new. + for len(readyChan) > 0 { + <-readyChan + } + sub.Unregister() + + notifier.Stop() + require.NoError(t, notifier.Start(ctx)) + notifier.testSignals.ListeningBegin.WaitOrTimeout() + require.Empty(t, readyChan) + }) } func TestNotifier_Start(t *testing.T) { diff --git a/producer.go b/producer.go index fd5a4e9f5..814cd2697 100644 --- a/producer.go +++ b/producer.go @@ -240,8 +240,22 @@ 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 + + // Receives a signal every time the notifier (re)establishes healthy + // listening. Written by the notifier goroutine (via a registered + // ListenerReadyFunc), only read from main goroutine. Used to reconcile + // activeJobs against cancel_attempted_at in case a cancel NOTIFY was lost + // while the notifier was disconnected/reconnecting. Buffered to 1 and + // written to non-blockingly because reconnects are rare and coalescing + // multiple signals into one reconciliation pass is fine. + reconnectCh chan struct{} + + // Only used by main goroutine. Tracks the current backoff attempt for + // retrying the reconnect reconciliation query after a failure. + reconcileAttempt int + + retryPolicy ClientRetryPolicy + testSignals producerTestSignals } func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pilot riverpilot.Pilot, config *producerConfig) *producer { @@ -268,6 +282,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, }) @@ -373,8 +388,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 @@ -408,6 +424,21 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } return err } + + // Reconcile activeJobs against cancel_attempted_at every time the + // notifier (re)establishes healthy listening. This closes the window + // described in https://github.com/riverqueue/river/issues/1358: a + // cancel NOTIFY sent while the notifier was disconnected/reconnecting + // is otherwise lost for good (Postgres NOTIFY is fire-and-forget). + // Firing on the very first connect too is harmless: activeJobs is + // always empty at that point in startup, so the fetchAndRunLoop + // handler below is a no-op query-free check. + reconnectSub = p.config.Notifier.RegisterListenerReadyFunc(func() { + select { + case p.reconnectCh <- struct{}{}: + default: + } + }) } go func() { @@ -427,6 +458,10 @@ 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)) @@ -606,6 +641,8 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { } case jobID := <-p.cancelCh: p.maybeCancelJob(workCtx, jobID) + case <-p.reconnectCh: + p.handleNotifierReconnect(workCtx) case <-p.fetchLimiter.C(): p.innerFetchLoop(workCtx, fetchResultCh) // Ensure we can't start another fetch when fetchCtx is done, even if @@ -839,6 +876,80 @@ func (p *producer) maybeCancelJob(ctx context.Context, id int64) bool { return true } +// jobMetadataWithCancelAttemptedAt mirrors the identical unexported struct in +// internal/maintenance/job_rescuer.go, which reads the same field for the +// same reason. +type jobMetadataWithCancelAttemptedAt struct { + CancelAttemptedAt time.Time `json:"cancel_attempted_at"` +} + +// handleNotifierReconnect reconciles activeJobs against cancel_attempted_at +// after the notifier has (re)established healthy listening. It exists to +// close the window described in +// https://github.com/riverqueue/river/issues/1358: a cancel NOTIFY sent while +// the notifier was disconnected/reconnecting is lost for good (Postgres +// NOTIFY is fire-and-forget), so a job cancelled during that window would +// otherwise run to completion unaware, undetected until JobRescuer's much +// longer stuck-job sweep. +// +// Only ever called from the main goroutine, so activeJobs can be read +// without a lock. +func (p *producer) handleNotifierReconnect(ctx context.Context) { + if len(p.activeJobs) == 0 { + p.reconcileAttempt = 0 + return + } + + ids := make([]int64, 0, len(p.activeJobs)) + for id := range p.activeJobs { + ids = append(ids, id) + } + + jobs, err := p.exec.JobGetByIDMany(ctx, &riverdriver.JobGetByIDManyParams{ + ID: ids, + Schema: p.config.Schema, + }) + if err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Error reconciling active jobs against cancel_attempted_at after notifier reconnect; will retry", + slog.String("err", err.Error())) + p.scheduleReconnectRetry(ctx) + return + } + p.reconcileAttempt = 0 + + for _, job := range jobs { + var metadata jobMetadataWithCancelAttemptedAt + if err := json.Unmarshal(job.Metadata, &metadata); err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Error unmarshaling job metadata while reconciling cancellations after notifier reconnect", + slog.Int64("job_id", job.ID), slog.String("err", err.Error())) + continue + } + if !metadata.CancelAttemptedAt.IsZero() { + p.maybeCancelJob(ctx, job.ID) + } + } +} + +// scheduleReconnectRetry re-arms p.reconnectCh after a backoff so a +// reconciliation query that failed (e.g. a transient DB error) doesn't +// silently drop the only signal that a reconnect happened, which would +// recreate the lost-cancel window this mechanism exists to close. +func (p *producer) scheduleReconnectRetry(ctx context.Context) { + if ctx.Err() != nil { + return + } + + sleepDuration := serviceutil.ExponentialBackoff(p.reconcileAttempt, serviceutil.MaxAttemptsBeforeResetDefault) + p.reconcileAttempt++ + + time.AfterFunc(sleepDuration, func() { + select { + case p.reconnectCh <- struct{}{}: + default: + } + }) +} + func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit { pluginLookup := p.config.PluginLookupGlobal if pluginLookup == nil { diff --git a/producer_test.go b/producer_test.go index 8d8545b7a..d7a9eca24 100644 --- a/producer_test.go +++ b/producer_test.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "slices" + "sync" "sync/atomic" "testing" "time" @@ -108,6 +109,182 @@ func (p *undecodableKindPilot) JobGetAvailable( return res, nil } +// producerTestListenerMock wraps a real riverdriver.Listener so tests can +// override individual operations. Unlike notifier.ListenerMock (unexported to +// the internal/notifier package), this lives in package river so it can be +// used to drive a real *notifier.Notifier from producer-level tests. +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) +} + +// TestProducer_JobCancelSurvivesNotifierReconnect is a deterministic repro of +// https://github.com/riverqueue/river/issues/1358: a `NOTIFY` sent while the +// notifier is disconnected/reconnecting is lost (Postgres `NOTIFY` is +// fire-and-forget), so a job cancelled during that window previously ran to +// completion unaware. It forces exactly one simulated connection loss between +// the job starting and `JobCancel` being applied, holds the notifier at the +// start of its reconnect attempt until the cancellation has landed, then lets +// it finish reconnecting and asserts the job observes the cancellation. +func TestProducer_JobCancelSurvivesNotifierReconnect(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 // set once that loss has been injected; cleared once Connect is released + releaseReconnect = make(chan struct{}) + releaseReconnectOnce sync.Once + ) + 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() { + <-releaseReconnect + reconnecting.Store(false) + } + return realListener.Connect(ctx) + } + + notif := notifier.New(archetype, mockListener) + require.NoError(t, notif.Start(ctx)) + t.Cleanup(notif.Stop) + // Cleanups run LIFO, so registering this after notif.Stop makes it run + // first: an assertion failure anywhere in the test body must not leave + // notif.Stop() deadlocked waiting on a mocked Connect call that's + // holding the reconnect gate open forever. + t.Cleanup(closeReleaseReconnect) + + jobStarted := make(chan int64, 1) + workers := NewWorkers() + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + AddWorker(workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + jobStarted <- 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) + + workCtx, workCancel := context.WithCancel(ctx) + + require.NoError(t, prod.StartWorkContext(ctx, workCtx)) + // Cleanups run LIFO: cancel workCtx (unblocking the worker) before + // prod.Stop tries to wait for it to exit, so cleanup can't deadlock + // regardless of whether the fix under test actually cancelled the job + // itself. + t.Cleanup(prod.Stop) + t.Cleanup(workCancel) + + jobID := riversharedtest.WaitOrTimeout(t, jobStarted) + + // Arm the injected disconnect, then wake the notifier's currently-blocked + // WaitForNotification call with a harmless notification on an unrelated + // queue so it loops back around and picks up the injected error + // immediately (rather than waiting on the notifier's own ping interval). + forceDisconnect.Store(true) + require.NoError(t, exec.NotifyMany(ctx, &riverdriver.NotifyManyParams{ + Topic: string(notifier.NotificationTopicInsert), + Payload: []string{`{"queue":"` + queueName + `-unrelated-wakeup"}`}, + Schema: schema, + })) + + // Wait until the notifier has actually observed the injected error and is + // blocked trying to reconnect. + require.Eventually(t, reconnecting.Load, 5*time.Second, 5*time.Millisecond, + "notifier never reached its reconnect attempt after the injected connection loss") + + // While the notifier is down, cancel the running job. The NOTIFY this + // sends is lost because nothing is LISTENing right now. + _, err = exec.JobCancel(ctx, &riverdriver.JobCancelParams{ + ID: jobID, + CancelAttemptedAt: time.Now().UTC(), + ControlTopic: string(notifier.NotificationTopicControl), + Schema: schema, + }) + require.NoError(t, err) + + // Now let the notifier finish reconnecting. + closeReleaseReconnect() + + 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() From 376c5a379aa2065edd02a474a6a240d0c5d8054f Mon Sep 17 00:00:00 2001 From: Chris George Date: Sun, 23 Aug 2026 14:53:09 -0700 Subject: [PATCH 2/6] Deliver cancellation locally when the client has no notifier --- client.go | 73 +++++++++++++++++++++++++++--- client_test.go | 118 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 186 insertions(+), 5 deletions(-) diff --git a/client.go b/client.go index e0abd8956..3bf0a32d4 100644 --- a/client.go +++ b/client.go @@ -1533,11 +1533,22 @@ 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, - }) + // Only a running job has an in-process executor to interrupt (an + // available/scheduled/retryable job was cancelled directly above, and a + // finalized job needed no cancellation at all), so there's nothing for a + // non-running job to dispatch. This also matters for correctness, not + // just efficiency: it keeps a burst of cancellations against jobs that + // were never running (e.g. cancelling many freshly-inserted jobs before + // the client has even been started) from filling the producer's + // control-event channel, which TriggerQueueControlEvent below sends to + // unconditionally and would otherwise block this call indefinitely. + if job.State == rivertype.JobStateRunning { + c.notifyProducerOfCancelWithoutNotifier(job.Queue, &controlEventPayload{ + Action: controlActionCancel, + JobID: job.ID, + Queue: job.Queue, + }) + } return job, nil } @@ -3006,6 +3017,58 @@ func (c *Client[TTx]) notifyProducerWithoutListenerQueueControlEvent(queue strin } } +// Notifies an internal producer of a job cancellation so a job running in +// this same process observes it immediately, for clients that have no +// notifier to deliver it via listen/notify. +// +// This deliberately does NOT share notifyProducerWithoutListenerQueueControlEvent +// above despite the near-identical dispatch: the two helpers guard on +// different conditions for a reason, and unifying them would either reopen +// this gap or introduce new bugs in the other one. +// +// - This helper guards on c.notifier == nil (this client has no notifier, +// for any reason, including an explicit Config.PollOnly on an otherwise +// listener-capable driver). Queue control events guard on +// driver.SupportsListener() instead, because pause/resume/metadata +// changes already have an independent, documented convergence path in +// poll-only mode (producer.pollForSettingChanges, on QueuePollInterval) — +// using this same client.notifier == nil guard there would double-deliver +// metadata changes not deduplicated the way pause/resume are, and would +// panic on QueueUpdate's nil control event when metadata is empty. +// - Job cancellation has no such fallback: nothing else ever re-checks a +// running job's cancel_attempted_at in poll-only mode, so a lost signal +// here is only caught by JobRescuer's much longer stuck-job sweep. +// +// The dispatch below (via TriggerQueueControlEvent) sends unconditionally and +// blocks if the target producer's control-event channel is full — the same +// characteristic notifyProducerWithoutListenerQueueControlEvent above already +// has for genuinely listener-incapable drivers. The caller (JobCancel) is +// responsible for only invoking this for jobs it knows are actually running, +// both because there's nothing to dispatch for a job with no in-process +// executor and to avoid a burst of cancellations against non-running jobs +// (e.g. many freshly-inserted jobs cancelled before the client is started) +// filling that channel and blocking indefinitely. +// +// 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]) notifyProducerOfCancelWithoutNotifier(queue string, controlEvent *controlEventPayload) { + if c.notifier != nil { + return + } + + c.producersMu.RLock() + defer c.producersMu.RUnlock() + + if len(c.producersByQueueName) < 1 { + return + } + + if producer, ok := c.producersByQueueName[queue]; ok { + producer.TriggerQueueControlEvent(controlEvent) + } +} + func (c *Client[TTx]) queueUpdate(ctx context.Context, executorTx riverdriver.ExecutorTx, name string, params *QueueUpdateParams) (*rivertype.Queue, *controlEventPayload, error) { updateMetadata := len(params.Metadata) > 0 diff --git a/client_test.go b/client_test.go index 8ee93b6b2..dbadae3e7 100644 --- a/client_test.go +++ b/client_test.go @@ -1465,6 +1465,112 @@ func Test_Client_Common(t *testing.T) { require.WithinDuration(t, time.Now(), *event.Job.FinalizedAt, 2*time.Second) }) + // Unlike CancelRunningJobPollOnly above (which uses a fake driver that + // reports SupportsListener() == false to simulate a driver that can't + // listen/notify at all), this uses the real pgx driver with an explicit + // Config.PollOnly, which is the actually-common "I don't want a notifier" + // case. driver.SupportsListener() is true here, so a same-process + // cancellation delivered via JobCancel() has nothing to do with the + // listen/notify path at all: it must go through the client's local, + // no-notifier dispatch instead. + t.Run("CancelRunningJobExplicitPollOnly", func(t *testing.T) { + t.Parallel() + + config, bundle := setupConfig(t) + config.PollOnly = true + + client, err := NewClient(bundle.driver, config) + require.NoError(t, err) + + 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) + + 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) + }) + + // A non-running job (available, scheduled, or retryable) is cancelled + // directly in the database and has no in-process executor to notify, so + // JobCancel must not dispatch a local control event for it. This matters + // beyond efficiency: notifyProducerOfCancelWithoutNotifier's dispatch + // blocks if the target producer's control-event channel (buffered to + // 100) fills up, which an unstarted (or otherwise not-yet-draining) + // producer never does. Cancelling more than that many never-started jobs + // against an explicit-PollOnly client (which has no notifier and so + // would otherwise take the local dispatch path) must not block. + 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 + } + + // Deliberately cancelling before the client is started: nothing is + // running yet, and nothing is consuming the (not-yet-existent) + // producer's control-event channel either. + 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") + } + }) + t.Run("CancelScheduledJob", func(t *testing.T) { t.Parallel() @@ -1530,15 +1636,27 @@ 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) + // Only a running job dispatches a local control event (see + // notifyProducerOfCancelWithoutNotifier), so make sure the job has + // actually started before cancelling it. + startedJobID := riversharedtest.WaitOrTimeout(t, jobStartedChan) + require.Equal(t, insertRes.Job.ID, startedJobID) + _, err = client.JobCancel(ctx, insertRes.Job.ID) require.NoError(t, err) From 42fa20b0f9537e442eccee8c8fe2c16b2d4d6e84 Mon Sep 17 00:00:00 2001 From: Chris George Date: Tue, 20 Oct 2026 07:35:49 -0700 Subject: [PATCH 3/6] Address cancellation reconciliation review feedback Run reconnect reconciliation asynchronously with a bounded context and reject stale execution results. Keep local cancellation lifecycle-aware without blocking on shutdown locks, and clear old cancellation metadata on retry across drivers. Add deterministic regression coverage and shorten comments and changelog entries. --- client.go | 72 ++--- client_test.go | 249 ++++++++++++++++-- internal/notifier/notifier.go | 58 ++-- internal/notifier/notifier_test.go | 100 +++---- producer.go | 223 ++++++++++------ producer_test.go | 233 +++++++++++++--- .../internal/dbsqlc/river_job.sql.go | 1 + riverdriver/riverdrivertest/job_update.go | 9 +- .../riverpgxv5/internal/dbsqlc/river_job.sql | 1 + .../internal/dbsqlc/river_job.sql.go | 1 + .../riversqlite/internal/dbsqlc/river_job.sql | 1 + .../internal/dbsqlc/river_job.sql.go | 1 + 12 files changed, 673 insertions(+), 276 deletions(-) diff --git a/client.go b/client.go index 3bf0a32d4..9a642b535 100644 --- a/client.go +++ b/client.go @@ -1533,17 +1533,8 @@ func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.Jo return nil, err } - // Only a running job has an in-process executor to interrupt (an - // available/scheduled/retryable job was cancelled directly above, and a - // finalized job needed no cancellation at all), so there's nothing for a - // non-running job to dispatch. This also matters for correctness, not - // just efficiency: it keeps a burst of cancellations against jobs that - // were never running (e.g. cancelling many freshly-inserted jobs before - // the client has even been started) from filling the producer's - // control-event channel, which TriggerQueueControlEvent below sends to - // unconditionally and would otherwise block this call indefinitely. if job.State == rivertype.JobStateRunning { - c.notifyProducerOfCancelWithoutNotifier(job.Queue, &controlEventPayload{ + c.notifyProducerOfCancelWithoutNotifier(ctx, job.Queue, &controlEventPayload{ Action: controlActionCancel, JobID: job.ID, Queue: job.Queue, @@ -3017,55 +3008,36 @@ func (c *Client[TTx]) notifyProducerWithoutListenerQueueControlEvent(queue strin } } -// Notifies an internal producer of a job cancellation so a job running in -// this same process observes it immediately, for clients that have no -// notifier to deliver it via listen/notify. -// -// This deliberately does NOT share notifyProducerWithoutListenerQueueControlEvent -// above despite the near-identical dispatch: the two helpers guard on -// different conditions for a reason, and unifying them would either reopen -// this gap or introduce new bugs in the other one. -// -// - This helper guards on c.notifier == nil (this client has no notifier, -// for any reason, including an explicit Config.PollOnly on an otherwise -// listener-capable driver). Queue control events guard on -// driver.SupportsListener() instead, because pause/resume/metadata -// changes already have an independent, documented convergence path in -// poll-only mode (producer.pollForSettingChanges, on QueuePollInterval) — -// using this same client.notifier == nil guard there would double-deliver -// metadata changes not deduplicated the way pause/resume are, and would -// panic on QueueUpdate's nil control event when metadata is empty. -// - Job cancellation has no such fallback: nothing else ever re-checks a -// running job's cancel_attempted_at in poll-only mode, so a lost signal -// here is only caught by JobRescuer's much longer stuck-job sweep. -// -// The dispatch below (via TriggerQueueControlEvent) sends unconditionally and -// blocks if the target producer's control-event channel is full — the same -// characteristic notifyProducerWithoutListenerQueueControlEvent above already -// has for genuinely listener-incapable drivers. The caller (JobCancel) is -// responsible for only invoking this for jobs it knows are actually running, -// both because there's nothing to dispatch for a job with no in-process -// executor and to avoid a burst of cancellations against non-running jobs -// (e.g. many freshly-inserted jobs cancelled before the client is started) -// filling that channel and blocking indefinitely. -// -// 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]) notifyProducerOfCancelWithoutNotifier(queue string, controlEvent *controlEventPayload) { +// 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, queue string, controlEvent *controlEventPayload) { if c.notifier != nil { return } c.producersMu.RLock() - defer c.producersMu.RUnlock() + producer := c.producersByQueueName[queue] + c.producersMu.RUnlock() + if producer == nil { + return + } - if len(c.producersByQueueName) < 1 { + state := producer.cancelState.Load() + if state == nil { return } + select { + case <-state.done: + return + default: + } - if producer, ok := c.producersByQueueName[queue]; ok { - producer.TriggerQueueControlEvent(controlEvent) + // cancelCh is also drained while a database fetch is in progress. + select { + case producer.cancelCh <- controlEvent.JobID: + producer.testSignals.QueueControlEventTriggered.Signal(controlEvent) + case <-ctx.Done(): + case <-state.done: } } diff --git a/client_test.go b/client_test.go index dbadae3e7..1372c60a3 100644 --- a/client_test.go +++ b/client_test.go @@ -1465,22 +1465,17 @@ func Test_Client_Common(t *testing.T) { require.WithinDuration(t, time.Now(), *event.Job.FinalizedAt, 2*time.Second) }) - // Unlike CancelRunningJobPollOnly above (which uses a fake driver that - // reports SupportsListener() == false to simulate a driver that can't - // listen/notify at all), this uses the real pgx driver with an explicit - // Config.PollOnly, which is the actually-common "I don't want a notifier" - // case. driver.SupportsListener() is true here, so a same-process - // cancellation delivered via JobCancel() has nothing to do with the - // listen/notify path at all: it must go through the client's local, - // no-notifier dispatch instead. + // 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) @@ -1507,6 +1502,9 @@ func Test_Client_Common(t *testing.T) { 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) @@ -1514,15 +1512,66 @@ func Test_Client_Common(t *testing.T) { require.WithinDuration(t, time.Now(), *event.Job.FinalizedAt, 2*time.Second) }) - // A non-running job (available, scheduled, or retryable) is cancelled - // directly in the database and has no in-process executor to notify, so - // JobCancel must not dispatch a local control event for it. This matters - // beyond efficiency: notifyProducerOfCancelWithoutNotifier's dispatch - // blocks if the target producer's control-event channel (buffered to - // 100) fills up, which an unstarted (or otherwise not-yet-draining) - // producer never does. Cancelling more than that many never-started jobs - // against an explicit-PollOnly client (which has no notifier and so - // would otherwise take the local dispatch path) must not block. + 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("CancelManyNonRunningJobsExplicitPollOnlyDoesNotBlock", func(t *testing.T) { t.Parallel() @@ -1549,9 +1598,7 @@ func Test_Client_Common(t *testing.T) { jobIDs[i] = insertRes.Job.ID } - // Deliberately cancelling before the client is started: nothing is - // running yet, and nothing is consuming the (not-yet-existent) - // producer's control-event channel either. + // Producers exist before Start, but have no consumer yet. cancelDone := make(chan error, 1) go func() { for _, jobID := range jobIDs { @@ -1571,6 +1618,56 @@ func Test_Client_Common(t *testing.T) { } }) + 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() @@ -1651,9 +1748,7 @@ func Test_Client_Common(t *testing.T) { insertRes, err := client.Insert(ctx, &JobArgs{}, nil) require.NoError(t, err) - // Only a running job dispatches a local control event (see - // notifyProducerOfCancelWithoutNotifier), so make sure the job has - // actually started before cancelling it. + // Wait for a running attempt before testing local cancellation delivery. startedJobID := riversharedtest.WaitOrTimeout(t, jobStartedChan) require.Equal(t, insertRes.Job.ID, startedJobID) @@ -2764,6 +2859,114 @@ 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 int64, 1)} + ctx, cancel := context.WithCancel(context.Background()) + ctx, _, started, stopped := targetProducer.StartInit(ctx) + targetProducer.cancelState.Store(&producerCancelState{done: ctx.Done()}) + started() + targetProducer.cancelCh <- 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, + } + } + + 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, QueueDefault, &controlEventPayload{}) + 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(), QueueDefault, &controlEventPayload{}) + 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(), QueueDefault, &controlEventPayload{ + Action: controlActionCancel, + JobID: 456, + Queue: QueueDefault, + }) + + require.Equal(t, int64(456), <-bundle.producer.cancelCh) + }) + + 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, QueueDefault, &controlEventPayload{}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.cancelCh, 1) + }) +} + func Test_Client_Stop_Common(t *testing.T) { t.Parallel() diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index 8fba56743..6bc5a7731 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 ( @@ -97,39 +112,6 @@ type Notifier struct { waitCancel context.CancelFunc } -// ListenerReadyFunc is invoked each time the notifier establishes healthy -// listening on all of its currently subscribed topics. This includes the -// very first connect as well as every subsequent reconnect that follows a -// connection loss, so callers whose work is only meaningful after a genuine -// reconnect (as opposed to the initial connect) should make that check -// themselves. -type ListenerReadyFunc func() - -// ListenerReadySubscription is a handle to a func registered with -// RegisterListenerReadyFunc. Call Unregister to stop receiving invocations. -type ListenerReadySubscription struct { - notifier *Notifier -} - -func (s *ListenerReadySubscription) Unregister() { - s.notifier.mu.Lock() - defer s.notifier.mu.Unlock() - delete(s.notifier.listenerReadyFuncs, s) -} - -// RegisterListenerReadyFunc registers a function to be invoked every time the -// notifier (re)establishes healthy listening on all of its subscribed topics. -// It's used to let callers reconcile any state that could have missed a -// notification while the notifier was disconnected/reconnecting. -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 New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notifier { notifier := baseservice.Init(archetype, &Notifier{ listener: listener, @@ -142,6 +124,16 @@ func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notif 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 { ctx, shouldStart, started, stopped := n.StartInit(ctx) if !shouldStart { diff --git a/internal/notifier/notifier_test.go b/internal/notifier/notifier_test.go index c36016f59..4af6927c3 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() @@ -623,54 +675,6 @@ func TestNotifier(t *testing.T) { // Subscription should still work. 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 - - readyChan := make(chan struct{}, 10) - sub := notifier.RegisterListenerReadyFunc(func() { readyChan <- struct{}{} }) - t.Cleanup(sub.Unregister) - - start(t, notifier) - - // Fires once on the initial connect. - riversharedtest.WaitOrTimeout(t, readyChan) - - // And again once the notifier reconnects after the injected failure. - require.EqualError(t, notifier.testSignals.BackoffError.WaitOrTimeout(), "error during wait") - riversharedtest.WaitOrTimeout(t, readyChan) - - // Unregister stops further invocations. There's no direct signal for - // "would have fired again", so drain what's pending, unregister, and - // confirm a subsequent restart doesn't add anything new. - for len(readyChan) > 0 { - <-readyChan - } - sub.Unregister() - - notifier.Stop() - require.NoError(t, notifier.Start(ctx)) - notifier.testSignals.ListeningBegin.WaitOrTimeout() - require.Empty(t, readyChan) - }) } func TestNotifier_Start(t *testing.T) { diff --git a/producer.go b/producer.go index 814cd2697..4b00d16a6 100644 --- a/producer.go +++ b/producer.go @@ -215,9 +215,10 @@ type producer struct { pilot riverpilot.Pilot workers *Workers - // Receives job IDs to cancel. Written by notifier and polling goroutines, - // only read from the main goroutine. - cancelCh chan int64 + // Receives job IDs from notifier, polling, and local cancellation calls. + // Only the main goroutine reads it. + cancelCh chan int64 + cancelState atomic.Pointer[producerCancelState] // 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 @@ -241,19 +242,9 @@ type producer struct { // goroutine, only read from main goroutine. queueControlCh chan *controlEventPayload - // Receives a signal every time the notifier (re)establishes healthy - // listening. Written by the notifier goroutine (via a registered - // ListenerReadyFunc), only read from main goroutine. Used to reconcile - // activeJobs against cancel_attempted_at in case a cancel NOTIFY was lost - // while the notifier was disconnected/reconnecting. Buffered to 1 and - // written to non-blockingly because reconnects are rare and coalescing - // multiple signals into one reconciliation pass is fine. + // Coalesces reconnect notifications for reconciliation on the main goroutine. reconnectCh chan struct{} - // Only used by main goroutine. Tracks the current backoff attempt for - // retrying the reconnect reconciliation query after a failure. - reconcileAttempt int - retryPolicy ClientRetryPolicy testSignals producerTestSignals } @@ -425,14 +416,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return err } - // Reconcile activeJobs against cancel_attempted_at every time the - // notifier (re)establishes healthy listening. This closes the window - // described in https://github.com/riverqueue/river/issues/1358: a - // cancel NOTIFY sent while the notifier was disconnected/reconnecting - // is otherwise lost for good (Postgres NOTIFY is fire-and-forget). - // Firing on the very first connect too is harmless: activeJobs is - // always empty at that point in startup, so the fetchAndRunLoop - // handler below is a no-op query-free check. + // Recover cancellations whose NOTIFY was lost while disconnected. reconnectSub = p.config.Notifier.RegisterListenerReadyFunc(func() { select { case p.reconnectCh <- struct{}{}: @@ -441,9 +425,13 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { }) } + subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) + p.cancelState.Store(&producerCancelState{done: subroutineCtx.Done()}) + go func() { started() defer stopped() // this defer should come first so it's last out + defer p.cancelState.Store(nil) p.Logger.DebugContext(fetchCtx, p.Name+": Run loop started", slog.String("queue", p.config.Queue), slog.Bool("paused", p.paused)) defer func() { @@ -463,7 +451,6 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } var subroutineWG sync.WaitGroup - subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) subroutineWG.Add(1) go p.fetchPollLoop(subroutineCtx, &subroutineWG) @@ -595,6 +582,64 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { p.fetchLimiter.Call() fetchResultCh := make(chan producerFetchResult) + reconcileCtx, cancelReconcile := context.WithCancel(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) + } + if reconcilePending { + reconcilePending = false + startReconcile() + } + } + for { select { case <-fetchCtx.Done(): @@ -642,9 +687,11 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { case jobID := <-p.cancelCh: p.maybeCancelJob(workCtx, jobID) case <-p.reconnectCh: - p.handleNotifierReconnect(workCtx) + 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 { @@ -698,7 +745,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 @@ -734,6 +781,10 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } return + case <-p.reconnectCh: + startReconcile() + case result := <-reconcileResultCh: + handleReconcileResult(result) case result := <-p.jobResultCh: p.removeActiveJob(result) case jobID := <-p.cancelCh: @@ -824,6 +875,30 @@ 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 _, job := range result.jobs { + p.activeJobsMu.Lock() + executor := p.activeJobs[job.ID] + p.activeJobsMu.Unlock() + if executor == nil || executor != result.executors[job.ID] || executor.JobRow.Attempt != job.Attempt { + continue + } + + var metadata struct { + CancelAttemptedAt time.Time `json:"cancel_attempted_at"` + } + if err := json.Unmarshal(job.Metadata, &metadata); err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Error unmarshaling cancellation metadata", + slog.Int64("job_id", job.ID), slog.String("err", err.Error())) + continue + } + if !metadata.CancelAttemptedAt.IsZero() { + p.maybeCancelJob(ctx, job.ID) + } + } +} + func (p *producer) removeActiveJob(job *rivertype.JobRow) { p.activeJobsMu.Lock() executor := p.activeJobs[job.ID] @@ -876,73 +951,41 @@ func (p *producer) maybeCancelJob(ctx context.Context, id int64) bool { return true } -// jobMetadataWithCancelAttemptedAt mirrors the identical unexported struct in -// internal/maintenance/job_rescuer.go, which reads the same field for the -// same reason. -type jobMetadataWithCancelAttemptedAt struct { - CancelAttemptedAt time.Time `json:"cancel_attempted_at"` -} - -// handleNotifierReconnect reconciles activeJobs against cancel_attempted_at -// after the notifier has (re)established healthy listening. It exists to -// close the window described in -// https://github.com/riverqueue/river/issues/1358: a cancel NOTIFY sent while -// the notifier was disconnected/reconnecting is lost for good (Postgres -// NOTIFY is fire-and-forget), so a job cancelled during that window would -// otherwise run to completion unaware, undetected until JobRescuer's much -// longer stuck-job sweep. -// -// Only ever called from the main goroutine, so activeJobs can be read -// without a lock. -func (p *producer) handleNotifierReconnect(ctx context.Context) { - if len(p.activeJobs) == 0 { - p.reconcileAttempt = 0 - return - } - +// handleNotifierReconnect snapshots active executions before querying asynchronously. +func (p *producer) handleNotifierReconnect(ctx context.Context, resultCh chan<- producerReconcileResult, wg *sync.WaitGroup) { + p.activeJobsMu.Lock() + executors := make(map[int64]*jobexecutor.JobExecutor, len(p.activeJobs)) ids := make([]int64, 0, len(p.activeJobs)) - for id := range p.activeJobs { + for id, executor := range p.activeJobs { + executors[id] = executor ids = append(ids, id) } - jobs, err := p.exec.JobGetByIDMany(ctx, &riverdriver.JobGetByIDManyParams{ - ID: ids, - Schema: p.config.Schema, - }) - if err != nil { - p.Logger.ErrorContext(ctx, p.Name+": Error reconciling active jobs against cancel_attempted_at after notifier reconnect; will retry", - slog.String("err", err.Error())) - p.scheduleReconnectRetry(ctx) - return - } - p.reconcileAttempt = 0 + p.activeJobsMu.Unlock() - for _, job := range jobs { - var metadata jobMetadataWithCancelAttemptedAt - if err := json.Unmarshal(job.Metadata, &metadata); err != nil { - p.Logger.ErrorContext(ctx, p.Name+": Error unmarshaling job metadata while reconciling cancellations after notifier reconnect", - slog.Int64("job_id", job.ID), slog.String("err", err.Error())) - continue - } - if !metadata.CancelAttemptedAt.IsZero() { - p.maybeCancelJob(ctx, job.ID) + go func() { + defer wg.Done() + + jobs, err := timeoututil.WithTimeoutV(ctx, 10*time.Second, p.Name+".handleNotifierReconnect", func(ctx context.Context) ([]*rivertype.JobRow, error) { + return p.exec.JobGetByIDMany(ctx, &riverdriver.JobGetByIDManyParams{ + ID: ids, + Schema: p.config.Schema, + }) + }) + select { + case <-ctx.Done(): + case resultCh <- producerReconcileResult{err: err, executors: executors, jobs: jobs}: } - } + }() } -// scheduleReconnectRetry re-arms p.reconnectCh after a backoff so a -// reconciliation query that failed (e.g. a transient DB error) doesn't -// silently drop the only signal that a reconnect happened, which would -// recreate the lost-cancel window this mechanism exists to close. -func (p *producer) scheduleReconnectRetry(ctx context.Context) { - if ctx.Err() != nil { - return - } - - sleepDuration := serviceutil.ExponentialBackoff(p.reconcileAttempt, serviceutil.MaxAttemptsBeforeResetDefault) - p.reconcileAttempt++ - - time.AfterFunc(sleepDuration, func() { +// 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: @@ -1344,12 +1387,22 @@ func (p *producer) reportQueueStatusOnce(ctx context.Context) { p.testSignals.ReportedQueueStatus.Signal(struct{}{}) } +type producerCancelState struct { + done <-chan struct{} +} + type producerFetchResult struct { decodeErrors map[int64]error err error jobs []*rivertype.JobRow } +type producerReconcileResult struct { + err error + executors map[int64]*jobexecutor.JobExecutor + jobs []*rivertype.JobRow +} + type errorHandlerAdapter struct { errorHandler ErrorHandler } diff --git a/producer_test.go b/producer_test.go index d7a9eca24..c63e04325 100644 --- a/producer_test.go +++ b/producer_test.go @@ -14,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" @@ -77,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) + jobGetByIDManyFunc func(context.Context, *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) +} + +func (e *producerTestExecutor) JobGetAvailable(ctx context.Context, params *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) { + return e.jobGetAvailableFunc(ctx, params) +} + +func (e *producerTestExecutor) JobGetByIDMany(ctx context.Context, params *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) { + return e.jobGetByIDManyFunc(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. @@ -109,10 +126,7 @@ func (p *undecodableKindPilot) JobGetAvailable( return res, nil } -// producerTestListenerMock wraps a real riverdriver.Listener so tests can -// override individual operations. Unlike notifier.ListenerMock (unexported to -// the internal/notifier package), this lives in package river so it can be -// used to drive a real *notifier.Notifier from producer-level tests. +// producerTestListenerMock overrides operations on a real listener. type producerTestListenerMock struct { riverdriver.Listener @@ -134,14 +148,49 @@ func (l *producerTestListenerMock) WaitForNotification(ctx context.Context) (*ri return l.waitForNotificationFunc(ctx) } -// TestProducer_JobCancelSurvivesNotifierReconnect is a deterministic repro of -// https://github.com/riverqueue/river/issues/1358: a `NOTIFY` sent while the -// notifier is disconnected/reconnecting is lost (Postgres `NOTIFY` is -// fire-and-forget), so a job cancelled during that window previously ran to -// completion unaware. It forces exactly one simulated connection loss between -// the job starting and `JobCancel` being applied, holds the notifier at the -// start of its reconnect attempt until the cancellation has landed, then lets -// it finish reconnecting and asserts the job observes the cancellation. +func TestProducer_ApplyReconcileResult(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + attempt int + metadata string + replaced bool + wantCancel bool + }{ + {name: "Cancelled", attempt: 1, metadata: `{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`, wantCancel: true}, + {name: "DifferentAttempt", attempt: 2, metadata: `{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`}, + {name: "MalformedMetadata", attempt: 1, metadata: `{"cancel_attempted_at":true}`}, + {name: "NoCancellation", attempt: 1, metadata: `{}`}, + {name: "ReplacedExecutor", attempt: 1, metadata: `{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`, 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}}) + + prod.applyReconcileResult(context.Background(), producerReconcileResult{ + executors: map[int64]*jobexecutor.JobExecutor{1: snapshotExecutor}, + jobs: []*rivertype.JobRow{{ID: 1, Attempt: testCase.attempt, Metadata: []byte(testCase.metadata)}}, + }) + if testCase.wantCancel { + require.ErrorIs(t, context.Cause(ctx), rivertype.ErrJobCancelledRemotely) + } else { + require.NoError(t, ctx.Err()) + } + }) + } +} + +// A cancel NOTIFY lost during a disconnect must be recovered on reconnect. func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { t.Parallel() @@ -167,10 +216,12 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { var ( forceDisconnect atomic.Bool // set by the test to inject one simulated connection loss - reconnecting atomic.Bool // set once that loss has been injected; cleared once Connect is released + 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) { @@ -181,7 +232,12 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { } mockListener.connectFunc = func(ctx context.Context) error { if reconnecting.Load() { - <-releaseReconnect + reconnectStarted.Signal(struct{}{}) + select { + case <-ctx.Done(): + return ctx.Err() + case <-releaseReconnect: + } reconnecting.Store(false) } return realListener.Connect(ctx) @@ -190,13 +246,11 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { notif := notifier.New(archetype, mockListener) require.NoError(t, notif.Start(ctx)) t.Cleanup(notif.Stop) - // Cleanups run LIFO, so registering this after notif.Stop makes it run - // first: an assertion failure anywhere in the test body must not leave - // notif.Stop() deadlocked waiting on a mocked Connect call that's - // holding the reconnect gate open forever. + // Release a blocked Connect before stopping the notifier. t.Cleanup(closeReleaseReconnect) - jobStarted := make(chan int64, 1) + var jobStarted testsignal.TestSignal[int64] + jobStarted.Init(t) workers := NewWorkers() type JobArgs struct { @@ -204,7 +258,7 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { } AddWorker(workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { - jobStarted <- job.ID + jobStarted.Signal(job.ID) <-ctx.Done() return ctx.Err() })) @@ -242,19 +296,13 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { workCtx, workCancel := context.WithCancel(ctx) require.NoError(t, prod.StartWorkContext(ctx, workCtx)) - // Cleanups run LIFO: cancel workCtx (unblocking the worker) before - // prod.Stop tries to wait for it to exit, so cleanup can't deadlock - // regardless of whether the fix under test actually cancelled the job - // itself. + // Cancel workers before waiting for producer shutdown. t.Cleanup(prod.Stop) t.Cleanup(workCancel) - jobID := riversharedtest.WaitOrTimeout(t, jobStarted) + jobID := jobStarted.WaitOrTimeout() - // Arm the injected disconnect, then wake the notifier's currently-blocked - // WaitForNotification call with a harmless notification on an unrelated - // queue so it loops back around and picks up the injected error - // immediately (rather than waiting on the notifier's own ping interval). + // Wake the listener so it observes the injected disconnect. forceDisconnect.Store(true) require.NoError(t, exec.NotifyMany(ctx, &riverdriver.NotifyManyParams{ Topic: string(notifier.NotificationTopicInsert), @@ -262,13 +310,9 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { Schema: schema, })) - // Wait until the notifier has actually observed the injected error and is - // blocked trying to reconnect. - require.Eventually(t, reconnecting.Load, 5*time.Second, 5*time.Millisecond, - "notifier never reached its reconnect attempt after the injected connection loss") + reconnectStarted.WaitOrTimeout() - // While the notifier is down, cancel the running job. The NOTIFY this - // sends is lost because nothing is LISTENing right now. + // Cancel while no listener is connected. _, err = exec.JobCancel(ctx, &riverdriver.JobCancelParams{ ID: jobID, CancelAttemptedAt: time.Now().UTC(), @@ -277,7 +321,6 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { }) require.NoError(t, err) - // Now let the notifier finish reconnecting. closeReleaseReconnect() update := riversharedtest.WaitOrTimeout(t, jobUpdates) @@ -499,6 +542,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 + }, + jobGetByIDManyFunc: func(ctx context.Context, _ *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, 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 []*rivertype.JobRow{{ID: 1, Attempt: 1, Metadata: []byte(`{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`)}}, 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 int64, 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 <- 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() diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 4bd6cbfab..53dae3d6f 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -1320,6 +1320,7 @@ updated_job AS ( state = 'available', max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, finalized_at = NULL, + metadata = metadata - 'cancel_attempted_at', scheduled_at = coalesce($2::timestamptz, now()) FROM job_to_update WHERE river_job.id = job_to_update.id diff --git a/riverdriver/riverdrivertest/job_update.go b/riverdriver/riverdrivertest/job_update.go index 150445d63..1346c5008 100644 --- a/riverdriver/riverdrivertest/job_update.go +++ b/riverdriver/riverdrivertest/job_update.go @@ -516,8 +516,10 @@ func exerciseJobUpdate[TTx any](ctx context.Context, t *testing.T, executorWithT 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 +527,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 +565,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 +576,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)) }) } diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index d87361114..6b11a3bc5 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -533,6 +533,7 @@ updated_job AS ( state = 'available', max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, finalized_at = NULL, + metadata = metadata - 'cancel_attempted_at', scheduled_at = coalesce(sqlc.narg('now')::timestamptz, now()) FROM job_to_update WHERE river_job.id = job_to_update.id diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go index a58f06ff0..57bf72f8b 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -1284,6 +1284,7 @@ updated_job AS ( state = 'available', max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, finalized_at = NULL, + metadata = metadata - 'cancel_attempted_at', scheduled_at = coalesce($2::timestamptz, now()) FROM job_to_update WHERE river_job.id = job_to_update.id diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql index 8ef471786..3380dc0dd 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql @@ -588,6 +588,7 @@ SET state = 'available', max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, finalized_at = NULL, + metadata = jsonb_remove(metadata, '$.cancel_attempted_at'), scheduled_at = coalesce(cast(sqlc.narg('now') AS text), datetime('now', 'subsec')) WHERE id = @id -- Do not touch running jobs: diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go index 8fc19993d..2a7472f1d 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go @@ -1452,6 +1452,7 @@ SET state = 'available', max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, finalized_at = NULL, + metadata = jsonb_remove(metadata, '$.cancel_attempted_at'), scheduled_at = coalesce(cast(?1 AS text), datetime('now', 'subsec')) WHERE id = ?2 -- Do not touch running jobs: From 63847c5febe82f9beafe2e80838eb8a8a2661948 Mon Sep 17 00:00:00 2001 From: Brandur Date: Sat, 24 Oct 2026 12:20:20 -0500 Subject: [PATCH 4/6] =?UTF-8?q?=F0=9F=A4=96=20Keep=20cancellation=20recove?= =?UTF-8?q?ry=20active=20while=20workers=20drain?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep Go's notifier and reconciliation running through graceful shutdown and queue removal, while preserving cancellable startup. Reconcile jobs claimed across reconnects and prevent delayed polling results from cancelling newer attempts. Keep cancellation polling active in Java, JavaScript, Ruby, and Rust through fetch stalls, missed notifications, and shutdown. Clear stale cancellation markers when retrying jobs. Add regression coverage for draining, reconnects, in-flight fetches, and stale cancellation results. --- client.go | 84 +- client_test.go | 397 ++++++++- internal/notifier/notifier.go | 13 +- .../src/main/java/com/riverqueue/Client.java | 6 +- .../main/java/com/riverqueue/Protocol.java | 15 +- .../main/java/com/riverqueue/RunHandle.java | 65 +- .../resources/com/riverqueue/postgres.sql | 8 +- .../main/resources/com/riverqueue/sqlite.sql | 8 +- .../java/com/riverqueue/ProtocolTest.java | 32 +- .../test/java/com/riverqueue/RetryTest.java | 112 +++ .../java/com/riverqueue/RunHandleTest.java | 187 ++++ js/driver/pg/src/driver.integration.test.ts | 58 ++ js/driver/pg/src/sql/jobs.ts | 15 +- js/driver/sqlite/src/compatibility.test.ts | 2 +- js/driver/sqlite/src/driver.test.ts | 79 +- js/driver/sqlite/src/driver.ts | 16 +- js/src/pilot-runtime.test.ts | 12 +- js/src/runtime.test.ts | 292 ++++++- js/src/runtime.ts | 21 +- js/src/runtime/attempt-runner.ts | 155 ++-- producer.go | 318 ++----- producer_test.go | 799 ++++++++++++------ .../internal/dbsqlc/river_job.sql.go | 16 +- riverdriver/riverdrivertest/job_read.go | 18 + riverdriver/riverdrivertest/job_update.go | 81 ++ .../riverpgxv5/internal/dbsqlc/river_job.sql | 16 +- .../internal/dbsqlc/river_job.sql.go | 16 +- .../riversqlite/internal/dbsqlc/river_job.sql | 37 +- .../internal/dbsqlc/river_job.sql.go | 72 +- .../riversqlite/river_sqlite_driver.go | 2 +- ruby/lib/client_runtime.rb | 63 +- ruby/lib/driver/runtime.rb | 18 +- ruby/sig/runtime.rbs | 4 +- ruby/spec/client_runtime_branch_spec.rb | 93 +- ruby/spec/driver_runtime_feature_spec.rb | 2 +- ruby/spec/driver_runtime_shared_examples.rb | 52 ++ ruby/spec/notification_runtime_spec.rb | 40 + ruby/spec/runtime_draining_shared_examples.rb | 3 + rust/riverqueue/src/client/attempts.rs | 99 +-- rust/riverqueue/src/client/builder.rs | 9 +- rust/riverqueue/src/client/jobs.rs | 14 +- rust/riverqueue/src/client/mod.rs | 7 +- rust/riverqueue/src/client/notifier.rs | 10 +- rust/riverqueue/src/client/producer.rs | 94 +-- rust/riverqueue/src/client/tests.rs | 125 +-- rust/riverqueue/src/database/sqlite.rs | 53 +- rust/riverqueue/src/storage/postgres.rs | 9 +- rust/riverqueue/tests/client_handles.rs | 30 + .../tests/poll_only_cancellation.rs | 52 +- 49 files changed, 2674 insertions(+), 1055 deletions(-) diff --git a/client.go b/client.go index 9a642b535..4cbb64edd 100644 --- a/client.go +++ b/client.go @@ -1231,6 +1231,16 @@ func (c *Client[TTx]) Start(ctx context.Context) error { fetchCtx := withClient(fetchCtx, c) workCtx = withClient(workCtx, c) + // Keep cancellation notifications and reconnects live until producers + // finish draining. The notifier is stopped with the other services below. + if c.notifier != nil { + if err := c.notifier.StartWithConnectContext(context.WithoutCancel(fetchCtx), fetchCtx); err != nil { + workCancel(err) + stopServicesOnError() + return err + } + } + if err := startstop.StartAll(fetchCtx, c.services...); err != nil { workCancel(err) stopServicesOnError() @@ -1503,8 +1513,8 @@ func (c *Client[TTx]) Driver() riverdriver.Driver[TTx] { // // If the job is currently running, it is not immediately cancelled, but is // instead marked for cancellation. The client running the job will also be -// notified to cancel the running job's context. When running without a notifier, -// clients poll for cancellation requests every two seconds. Although the job's +// notified to cancel the running job's context. Clients also poll for cancellation +// requests every two seconds to recover missed notifications. Although the job's // context will be cancelled, since Go does not provide a mechanism to // interrupt a running goroutine the job will continue running until it returns. // As always, it is important for workers to respect context cancellation and @@ -1534,11 +1544,7 @@ func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.Jo } if job.State == rivertype.JobStateRunning { - c.notifyProducerOfCancelWithoutNotifier(ctx, job.Queue, &controlEventPayload{ - Action: controlActionCancel, - JobID: job.ID, - Queue: job.Queue, - }) + c.notifyProducerOfCancelWithoutNotifier(ctx, job) } return job, nil @@ -1562,8 +1568,8 @@ func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.Jo // // If the job is currently running, it is not immediately cancelled, but is // instead marked for cancellation. The client running the job will also be -// notified to cancel the running job's context. When running without a notifier, -// clients poll for cancellation requests every two seconds. Although the job's +// notified to cancel the running job's context. Clients also poll for cancellation +// requests every two seconds to recover missed notifications. Although the job's // context will be cancelled, since Go does not provide a mechanism to // interrupt a running goroutine the job will continue running until it returns. // As always, it is important for workers to respect context cancellation and @@ -2428,14 +2434,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 { @@ -2447,7 +2454,9 @@ func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) erro } } + c.producersMu.Lock() delete(c.producersByQueueName, queueName) + c.producersMu.Unlock() return nil } @@ -2835,7 +2844,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 } @@ -2905,7 +2914,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 } @@ -2968,7 +2977,7 @@ func (c *Client[TTx]) QueueUpdate(ctx context.Context, name string, params *Queu // No control event is returned when metadata isn't being changed, in which // case there's nothing for the producer to act on. if controlEvent != nil { - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) } return queue, nil @@ -2991,38 +3000,35 @@ 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() - - if len(c.producersByQueueName) < 1 { - return - } + producer := c.producersByQueueName[queue] + c.producersMu.RUnlock() - if producer, ok := c.producersByQueueName[queue]; ok { - producer.TriggerQueueControlEvent(controlEvent) + 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, queue string, controlEvent *controlEventPayload) { +func (c *Client[TTx]) notifyProducerOfCancelWithoutNotifier(ctx context.Context, job *rivertype.JobRow) { if c.notifier != nil { return } c.producersMu.RLock() - producer := c.producersByQueueName[queue] + producer := c.producersByQueueName[job.Queue] c.producersMu.RUnlock() if producer == nil { return } - state := producer.cancelState.Load() + state := producer.runState.Load() if state == nil { return } @@ -3032,13 +3038,27 @@ func (c *Client[TTx]) notifyProducerOfCancelWithoutNotifier(ctx context.Context, default: } - // cancelCh is also drained while a database fetch is in progress. - select { - case producer.cancelCh <- controlEvent.JobID: - producer.testSignals.QueueControlEventTriggered.Signal(controlEvent) - case <-ctx.Done(): - case <-state.done: + producer.activeJobsMu.Lock() + executor := producer.activeJobs[job.ID] + if executor != nil { + sameAttemptedAt := executor.JobRow.AttemptedAt == job.AttemptedAt || + (executor.JobRow.AttemptedAt != nil && job.AttemptedAt != nil && executor.JobRow.AttemptedAt.Equal(*job.AttemptedAt)) + if executor.JobRow.Attempt != job.Attempt || !sameAttemptedAt { + executor = nil + } } + producer.activeJobsMu.Unlock() + + if executor != nil { + executor.Cancel(ctx) + } else { + producer.TriggerCancelCheck() + } + producer.testSignals.QueueControlEventTriggered.Signal(&controlEventPayload{ + Action: controlActionCancel, + JobID: job.ID, + Queue: job.Queue, + }) } func (c *Client[TTx]) queueUpdate(ctx context.Context, executorTx riverdriver.ExecutorTx, name string, params *QueueUpdateParams) (*rivertype.Queue, *controlEventPayload, error) { diff --git a/client_test.go b/client_test.go index 1372c60a3..829b0fa1d 100644 --- a/client_test.go +++ b/client_test.go @@ -1572,6 +1572,67 @@ func Test_Client_Common(t *testing.T) { 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() @@ -1663,7 +1724,7 @@ func Test_Client_Common(t *testing.T) { }() require.NoError(t, cancelled.WaitOrTimeout()) - require.Empty(t, producer.cancelCh) + require.Empty(t, producer.cancelPollCh) require.Empty(t, producer.queueControlCh) }) } @@ -2765,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() @@ -2865,30 +3015,137 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { type testBundle struct { client *Client[pgx.Tx] producer *producer - stopped func() + stopFunc context.CancelFunc + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + prod := &producer{cancelPollCh: make(chan struct{}, 1)} + prod.runState.Store(&producerRunState{done: ctx.Done()}) + return &testBundle{ + client: &Client[pgx.Tx]{producersByQueueName: map[string]*producer{QueueDefault: prod}}, + producer: prod, + stopFunc: cancel, + } + } + + t.Run("FullWakeBuffer", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.producer.TriggerCancelCheck() + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{Queue: QueueDefault}) + returned.Signal(struct{}{}) + }() + returned.WaitOrTimeout() + require.Len(t, bundle.producer.cancelPollCh, 1) + }) + + for _, stopped := range []bool{false, true} { + t.Run(fmt.Sprintf("InactiveProducer_Stopped_%t", stopped), func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + if stopped { + bundle.stopFunc() + } else { + bundle.producer.runState.Store(nil) + } + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{Queue: QueueDefault}) + require.Empty(t, bundle.producer.cancelPollCh) + }) + } + + for _, testCase := range []struct { + name string + replaced bool + reuseAttempt bool + }{ + {name: "MatchingAttempt"}, + {name: "NewAttempt", replaced: true}, + {name: "NewAttemptedAt", replaced: true, reuseAttempt: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + archetype := riversharedtest.BaseServiceArchetype(t) + originalCtx, cancelOriginal := context.WithCancelCause(context.Background()) + defer cancelOriginal(nil) + replacementCtx, cancelReplacement := context.WithCancelCause(context.Background()) + defer cancelReplacement(nil) + attemptedAt := time.Now().UTC() + job := &rivertype.JobRow{ID: 456, Attempt: 1, AttemptedAt: &attemptedAt, Queue: QueueDefault} + original := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelOriginal, JobRow: job}) + replacementJob := *job + replacementJob.AttemptedAt = new(attemptedAt.Add(time.Second)) + if !testCase.reuseAttempt { + replacementJob.Attempt++ + } + replacement := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelReplacement, JobRow: &replacementJob}) + bundle.producer.activeJobs = map[int64]*jobexecutor.JobExecutor{job.ID: original} + if testCase.replaced { + bundle.producer.activeJobs[job.ID] = replacement + } + + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), job) + + require.NoError(t, replacementCtx.Err()) + if testCase.replaced { + require.NoError(t, originalCtx.Err()) + require.Len(t, bundle.producer.cancelPollCh, 1) + } else { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + require.Empty(t, bundle.producer.cancelPollCh) + } + }) + } +} + +func Test_Client_NotifyProducerWithoutListenerQueueControlEvent(t *testing.T) { + t.Parallel() + + type testBundle struct { + client *Client[pgx.Tx] + fetchStopped <-chan struct{} + producer *producer + stopFunc func() } setup := func(t *testing.T) *testBundle { t.Helper() - targetProducer := &producer{cancelCh: make(chan int64, 1)} ctx, cancel := context.WithCancel(context.Background()) + runCtx, cancelRun := context.WithCancel(context.Background()) + targetProducer := &producer{ + cancelPollCh: make(chan struct{}, 1), + queueControlCh: make(chan *controlEventPayload, 1), + } ctx, _, started, stopped := targetProducer.StartInit(ctx) - targetProducer.cancelState.Store(&producerCancelState{done: ctx.Done()}) + targetProducer.runState.Store(&producerRunState{done: runCtx.Done(), fetchDone: ctx.Done()}) started() - targetProducer.cancelCh <- 123 - stop := sync.OnceFunc(func() { + targetProducer.queueControlCh <- &controlEventPayload{Action: controlActionPause} + stopFunc := sync.OnceFunc(func() { cancel() + cancelRun() stopped() }) - t.Cleanup(stop) + t.Cleanup(stopFunc) return &testBundle{ client: &Client[pgx.Tx]{ + driver: &DriverPollOnly{}, producersByQueueName: map[string]*producer{QueueDefault: targetProducer}, }, - producer: targetProducer, - stopped: stop, + fetchStopped: ctx.Done(), + producer: targetProducer, + stopFunc: stopFunc, } } @@ -2902,68 +3159,72 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { var returned testsignal.TestSignal[struct{}] returned.Init(t) go func() { - bundle.client.notifyProducerOfCancelWithoutNotifier(ctx, QueueDefault, &controlEventPayload{}) + bundle.client.notifyProducerWithoutListenerQueueControlEvent(ctx, QueueDefault, &controlEventPayload{}) returned.Signal(struct{}{}) }() returned.WaitOrTimeout() - require.Len(t, bundle.producer.cancelCh, 1) + require.Len(t, bundle.producer.queueControlCh, 1) }) - t.Run("FullBufferProducerStopped", func(t *testing.T) { + 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.notifyProducerOfCancelWithoutNotifier(context.Background(), QueueDefault, &controlEventPayload{}) + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, &controlEventPayload{}) returned.Signal(struct{}{}) }() - bundle.stopped() returned.WaitOrTimeout() - require.Len(t, bundle.producer.cancelCh, 1) + require.Len(t, bundle.producer.queueControlCh, 1) }) - t.Run("RunningProducer", func(t *testing.T) { + t.Run("QueueRemoval", func(t *testing.T) { t.Parallel() bundle := setup(t) - <-bundle.producer.cancelCh + 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) - bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), QueueDefault, &controlEventPayload{ - Action: controlActionCancel, - JobID: 456, - Queue: QueueDefault, - }) + // 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) - require.Equal(t, int64(456), <-bundle.producer.cancelCh) + // Cancellation delivery remains live until the running attempts drain. + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{ID: 123, Queue: QueueDefault}) + require.Len(t, bundle.producer.cancelPollCh, 1) + bundle.stopFunc() + require.NoError(t, queueRemoved.WaitOrTimeout()) + require.Empty(t, bundle.client.producersByQueueName) }) - t.Run("StopInProgress", func(t *testing.T) { + t.Run("RunningProducer", 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() + <-bundle.producer.queueControlCh + controlEvent := &controlEventPayload{Action: controlActionResume, Queue: QueueDefault} + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, controlEvent) - var returned testsignal.TestSignal[struct{}] - returned.Init(t) - go func() { - bundle.client.notifyProducerOfCancelWithoutNotifier(ctx, QueueDefault, &controlEventPayload{}) - returned.Signal(struct{}{}) - }() - - returned.WaitOrTimeout() - require.Len(t, bundle.producer.cancelCh, 1) + require.Same(t, controlEvent, <-bundle.producer.queueControlCh) }) } @@ -9023,6 +9284,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() @@ -10914,6 +11216,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 6bc5a7731..799b9ae1e 100644 --- a/internal/notifier/notifier.go +++ b/internal/notifier/notifier.go @@ -135,6 +135,12 @@ func (n *Notifier) RegisterListenerReadyFunc(listenerReadyFunc ListenerReadyFunc } 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 @@ -142,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 diff --git a/java/river/src/main/java/com/riverqueue/Client.java b/java/river/src/main/java/com/riverqueue/Client.java index 752f5a9d8..ceafd0d1b 100644 --- a/java/river/src/main/java/com/riverqueue/Client.java +++ b/java/river/src/main/java/com/riverqueue/Client.java @@ -1182,11 +1182,7 @@ public Job retry(Connection connection, long id) { var before = locked(c, id); try (var statement = Sql.prepare( - c, - Sql.query(database, "retry"), - database.timestamp(clock.instant()), - id, - database.timestamp(clock.instant())); + c, Sql.query(database, "retry"), database.timestamp(clock.instant()), id); var rows = statement.executeQuery()) { boolean updated = rows.next(); var job = updated ? read(rows) : before; diff --git a/java/river/src/main/java/com/riverqueue/Protocol.java b/java/river/src/main/java/com/riverqueue/Protocol.java index 1e0a1da43..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/RunHandle.java b/java/river/src/main/java/com/riverqueue/RunHandle.java index 8346562da..4ddfcc3be 100644 --- a/java/river/src/main/java/com/riverqueue/RunHandle.java +++ b/java/river/src/main/java/com/riverqueue/RunHandle.java @@ -33,7 +33,6 @@ public final class RunHandle implements AutoCloseable { private final Map> active = new ConcurrentHashMap<>(); private final Map activeThreads = new ConcurrentHashMap<>(); - private final Map pendingCancellation = new ConcurrentHashMap<>(); private final AtomicBoolean maintenanceBusy = new AtomicBoolean(); private volatile Instant nextReindex = Instant.now() @@ -42,6 +41,7 @@ public final class RunHandle implements AutoCloseable { .plusDays(1) .atStartOfDay() .toInstant(java.time.ZoneOffset.UTC); + private final Semaphore cancellationWake = new Semaphore(0); private final Client.Builder config; private final RetryPolicy defaultRetryPolicy = RetryPolicy.defaults(); private final java.util.concurrent.BlockingQueue completions = @@ -83,7 +83,10 @@ public final class RunHandle implements AutoCloseable { id = river.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( @@ -164,6 +167,7 @@ void launch() { unsubscribeInsert = river.onInsertCommit(wake::release); if (!config.pollOnly) executor.submit(this::listen); executor.submit(this::completeBatches); + executor.submit(this::pollCancellations); executor.submit(this::supervise); executor.submit(this::dispatch); dispatcherStarted = true; @@ -267,7 +271,6 @@ public void stopAndCancel() { private void dispatch() { long nextPoll = 0; long nextMaintenance = 0; - long nextCancelPoll = 0; while (running.get()) { try { long now = System.nanoTime(); @@ -295,25 +298,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(); @@ -466,8 +450,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( () -> { @@ -1417,6 +1399,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 : river.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 75e2d53b0..c859aa7a2 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 3496f78a0..ca198346e 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,36 +119,31 @@ 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()); } diff --git a/java/river/src/test/java/com/riverqueue/RetryTest.java b/java/river/src/test/java/com/riverqueue/RetryTest.java index 683b2a0eb..7aa6696c3 100644 --- a/java/river/src/test/java/com/riverqueue/RetryTest.java +++ b/java/river/src/test/java/com/riverqueue/RetryTest.java @@ -9,6 +9,7 @@ import java.time.Instant; import java.time.ZoneOffset; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; @@ -27,6 +28,55 @@ class RetryTest { @RegisterExtension final TestDatabase databases = new TestDatabase(); @TempDir Path directory; + @Test + void retryClearsCancellation() { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry cancelled").job(); + var cancelled = client.cancel(inserted.id()); + assertTrue(cancelled.metadata().has("cancel_attempted_at")); + var retried = client.retry(inserted.id()); + assertEquals(Job.State.AVAILABLE, retried.state()); + assertFalse(retried.metadata().has("cancel_attempted_at")); + assertFalse(client.get(inserted.id()).metadata().has("cancel_attempted_at")); + } + + @Test + void retryClearsCancellationWithoutReorderingAvailableJob() { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database, Clock.fixed(NOW, ZoneOffset.UTC)); + var inserted = + client + .insert( + TYPE, + "already available", + InsertOptions.builder() + .metadata(Map.of("cancel_attempted_at", NOW.toString(), "preserved", true)) + .scheduledAt(NOW.minusSeconds(3600)) + .build()) + .job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='available', attempt=max_attempts WHERE id=?", + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + return null; + }); + var retried = client.retry(inserted.id()); + assertEquals(inserted.scheduledAt(), retried.scheduledAt()); + assertEquals(inserted.maxAttempts(), retried.maxAttempts()); + var expectedMetadata = (tools.jackson.databind.node.ObjectNode) inserted.metadata().deepCopy(); + expectedMetadata.remove("cancel_attempted_at"); + assertEquals(expectedMetadata, retried.metadata()); + assertEquals(retried, client.get(inserted.id())); + } + @ParameterizedTest @EnumSource(Job.State.class) void retryNotifiesOnlyWhenTheRowChanges(Job.State state) { @@ -67,6 +117,36 @@ void retryNotifiesOnlyWhenTheRowChanges(Job.State state) { } } + @ParameterizedTest + @ValueSource( + strings = {"null", "42", "false", "\"cancel_attempted_at\"", "[\"cancel_attempted_at\"]"}) + void retryPreservesNonObjectMetadata(String metadata) { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry metadata").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='retryable', metadata=" + + (database.dialect() == Database.Dialect.POSTGRES + ? "CAST(? AS jsonb)" + : "jsonb(?)") + + " WHERE id=?", + metadata, + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + return null; + }); + + assertEquals(Job.State.AVAILABLE, client.retry(inserted.id()).state()); + assertEquals(Json.parse(metadata), client.get(inserted.id()).metadata()); + } + @ParameterizedTest @ValueSource(strings = {"local", "peer", "transaction"}) @SuppressWarnings("try") // Resources intentionally scope background activity. @@ -217,4 +297,36 @@ void sqliteNotificationFailureRollsBackTheRetry() { job.metadata().path(Protocol.METADATA_OUTPUT).asString()); assertEquals(0, notifications.get()); } + + @Test + void sqliteRetryQueryPreservesMalformedMetadata() { + var database = TestDatabase.sqlite(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry malformed metadata").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE river_job SET state='retryable', metadata='[not json' WHERE id=?", + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + // The public retry validates the locked row first. Exercise the SQL + // directly so cleanup cannot turn a decode error into a SQL error. + try (var statement = + Sql.prepare( + connection, + Sql.query(database, "retry"), + database.timestamp(NOW), + inserted.id()); + var rows = statement.executeQuery()) { + assertTrue(rows.next()); + assertEquals("[not json", rows.getString("metadata")); + } + return null; + }); + var error = assertThrows(IllegalArgumentException.class, () -> client.get(inserted.id())); + assertTrue(error.getMessage().contains("metadata")); + } } diff --git a/java/river/src/test/java/com/riverqueue/RunHandleTest.java b/java/river/src/test/java/com/riverqueue/RunHandleTest.java index 46efd1c4f..f5b650d39 100644 --- a/java/river/src/test/java/com/riverqueue/RunHandleTest.java +++ b/java/river/src/test/java/com/riverqueue/RunHandleTest.java @@ -70,6 +70,125 @@ void builderRequiresServiceIntervalsBelowQueueRetention(String value, boolean va .getMessage()); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void cancellationContinuesAfterGracefulStop(boolean pollOnly) throws Exception { + var client = client(); + var entered = new CompletableFuture>(); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .worker( + TYPE, + context -> { + entered.complete(context); + context.awaitCancellation(); + context.checkCancelled(); + }) + .build() + .start(); + try (var subscription = workers.client().subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel while draining").job(); + entered.get(5, TimeUnit.SECONDS); + // The timeout leaves active work draining after 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(); + } + } + + @Test + void cancellationContinuesWhileFetchingIsBlocked() throws Exception { + var entered = new CompletableFuture(); + var fetching = new CompletableFuture(); + var release = new CompletableFuture(); + var cancelled = new CompletableFuture(); + var client = + client() + .withPilot( + new Pilot() { + @Override + public int workerLimit(String queue, int configured) { + if (entered.isDone()) { + fetching.complete(null); + release.join(); + } + return configured; + } + }); + var workers = + workers(client) + .pollInterval(Duration.ofMillis(10)) + .worker( + TYPE, + context -> { + entered.complete(null); + context.awaitCancellation(); + cancelled.complete(context.cancellation()); + context.checkCancelled(); + }) + .build() + .start(); + try { + var inserted = client.insert(TYPE, "cancel while fetching").job(); + fetching.get(5, TimeUnit.SECONDS); + client.cancel(inserted.id()); + assertEquals(WorkContext.Cancellation.REMOTE, cancelled.get(5, TimeUnit.SECONDS)); + } finally { + release.complete(null); + workers.stopAndCancel(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void cancellationDuringArgumentDecodingIsReconciled(boolean pollOnly) throws Exception { + var decoding = new CompletableFuture(); + var release = new CompletableFuture(); + var client = + client() + .withPilot( + new Pilot() { + @Override + public tools.jackson.databind.JsonNode decode( + Job job) { + decoding.complete(null); + release.join(); + return job.args(); + } + }); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .worker( + TYPE, + context -> { + context.awaitCancellation(); + context.checkCancelled(); + }) + .build() + .start(); + try (var subscription = workers.client().subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel before registration").job(); + decoding.get(5, TimeUnit.SECONDS); + client.cancel(inserted.id()); + release.complete(null); + assertEquals(Job.State.CANCELLED, outcome.get(5, TimeUnit.SECONDS).job().state()); + workers.stop(); + } finally { + release.complete(null); + workers.stopAndCancel(); + } + } + @ParameterizedTest @CsvSource({"REMOTE,false", "REMOTE,true", "TIMEOUT,false", "TIMEOUT,true"}) @SuppressWarnings("try") // Resources intentionally scope background activity. @@ -802,6 +921,74 @@ void invalidRetryPoliciesFallBack(String policy) throws Exception { } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void retriedJobIgnoresStaleCancellation(boolean delayedUntilRunning) throws Exception { + var client = client(); + var ready = new CompletableFuture(); + var paused = new LinkedBlockingQueue(); + var entered = new CompletableFuture(); + var finish = new CompletableFuture(); + if (!delayedUntilRunning) finish.complete(null); + var outcome = new CompletableFuture(); + var errors = new LinkedBlockingQueue(); + var workers = + workers(client) + .pollOnly(false) + .errorHandler(errors::add) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .worker( + TYPE, + context -> { + entered.complete(null); + finish.join(); + context.checkCancelled(); + }) + .build() + .start(); + try (var pauseSubscription = + workers.client().subscribe(paused::add, Client.EventKind.QUEUE_PAUSED); + var outcomeSubscription = + workers.client().subscribe(outcome::complete, Client.EventKind.JOB_COMPLETED)) { + ready.get(5, TimeUnit.SECONDS); + var inserted = + client + .insert( + TYPE, + "retry after an idle cancellation", + InsertOptions.builder().scheduledAt(Instant.now().plusSeconds(3600)).build()) + .job(); + client.cancel(inserted.id()); + client.queues().pause("default"); + // The pause follows cancellation on the same listener, so receiving it + // proves the unmatched cancellation was dispatched before retry. + assertNotNull(paused.poll(5, TimeUnit.SECONDS)); + client.retry(inserted.id()); + client.queues().resume("default"); + if (delayedUntilRunning) { + entered.get(5, TimeUnit.SECONDS); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(inserted.id(), "default")); + return null; + }); + client.queues().pause("default"); + assertNotNull(paused.poll(5, TimeUnit.SECONDS)); + finish.complete(null); + } + assertEquals(inserted.id(), outcome.get(5, TimeUnit.SECONDS).job().id()); + workers.stop(); + } finally { + finish.complete(null); + workers.stopAndCancel(); + } + assertTrue(errors.isEmpty(), () -> "Unexpected worker errors: " + errors); + } + @ParameterizedTest @CsvSource({"false,1,5000", "true,1,5000", "false,3600,1", "true,3600,1"}) @SuppressWarnings("try") // Resources intentionally scope background activity. diff --git a/js/driver/pg/src/driver.integration.test.ts b/js/driver/pg/src/driver.integration.test.ts index 99588eebe..084b0ae7f 100644 --- a/js/driver/pg/src/driver.integration.test.ts +++ b/js/driver/pg/src/driver.integration.test.ts @@ -533,11 +533,23 @@ describe("PgDriver integration", () => { await expect(driver.jobGetCancelRequested([])).resolves.toEqual([]); }); + it("ignores cancellation names stored in non-object metadata", async () => { + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_cancel_array`) + ); + await pool.query( + "UPDATE river_job SET state = 'running', metadata = $1::jsonb WHERE id = $2", + ['["cancel_attempted_at"]', job.id.toString()] + ); + await expect(driver.jobGetCancelRequested([job.id])).resolves.toEqual([]); + }); + it("retries non-running jobs but leaves running jobs unchanged", async () => { const scheduledAt = Temporal.Instant.from("2026-08-30T15:00:00Z"); const now = Temporal.Instant.from("2026-08-30T12:00:00.123456Z"); const retryable = await driver.jobInsert( insertParams(`${filePrefix}_retry`, { + metadata: { cancel_attempted_at: now.toString(), preserved: true }, scheduledAt, state: "retryable", }) @@ -560,10 +572,33 @@ describe("PgDriver integration", () => { }); expect(retried!.state).toBe("available"); + expect(retried!.metadata).toEqual({ preserved: true }); expect(retried!.scheduledAt.toString()).toBe("2026-08-30T12:00:00.123456Z"); expect(unchanged!.state).toBe("running"); }); + it("clears cancellation without reordering an already available retry", async () => { + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_retry_available`, { + metadata: { cancel_attempted_at: now.toString(), preserved: true }, + scheduledAt: now.subtract({ hours: 1 }), + state: "available", + }) + ); + await pool.query( + "UPDATE river_job SET attempt = max_attempts WHERE id = $1", + [job.id.toString()] + ); + const expectedMetadata = { ...job.metadata }; + delete expectedMetadata.cancel_attempted_at; + const retried = await driver.jobRetryWithOptions({ id: job.id, now }); + expect(retried?.metadata).toEqual(expectedMetadata); + expect(retried?.scheduledAt.toString()).toBe(job.scheduledAt.toString()); + expect(retried?.maxAttempts).toBe(job.maxAttempts); + expect((await driver.jobGet(job.id))?.metadata).toEqual(expectedMetadata); + }); + it("returns the committed row to the loser of a cancel or retry race", async () => { const inserted = await driver.jobInsert( insertParams(`${filePrefix}_race`, { @@ -635,6 +670,29 @@ describe("PgDriver integration", () => { ); }); + it.each(["null", "42", '"cancel_attempted_at"', '["cancel_attempted_at"]'])( + "retries jobs without changing non-object metadata: %s", + async (metadata) => { + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_retry_metadata`, { state: "retryable" }) + ); + await pool.query( + "UPDATE river_job SET metadata = $1::jsonb WHERE id = $2::bigint", + [metadata, job.id.toString()] + ); + + expect(await driver.jobRetry(job.id)).toMatchObject({ + id: job.id, + state: "available", + }); + const persisted = await pool.query<{ metadata: unknown }>( + "SELECT metadata FROM river_job WHERE id = $1::bigint", + [job.id.toString()] + ); + expect(persisted.rows[0]?.metadata).toEqual(JSON.parse(metadata)); + } + ); + it("acts on rows River can't fully read by ID and lists them like Go", async () => { const kind = `${filePrefix}_poison`; const uniqueKey = new Uint8Array(32).fill(13); 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 a9a41e008..ed9c4b77a 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 7f2a36c6e..99190a97a 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; @@ -872,15 +883,15 @@ export class AttemptRunner { } /** - * Cancel the attempts among `ids` whose jobs have a cancellation request, + * Cancel the snapshotted attempts whose jobs have a cancellation request, * reading them in batches of 1,000 IDs, each bounded by ten seconds. */ async #cancelRequested( requested: NonNullable, - ids: readonly bigint[], + attempts: ReadonlyMap, signal: AbortSignal ): Promise { - const clientId = this.#context.clientId; + const ids = [...attempts.keys()]; for (let start = 0; start < ids.length; start += CANCEL_POLL_BATCH_SIZE) { const timeout = this.#context.timer.timeout( CANCEL_POLL_TIMEOUT_MS, @@ -894,7 +905,11 @@ export class AttemptRunner { }), link.signal ); - for (const id of cancelled) this.cancelAttempt(id, clientId); + // Never look up a fresh attempt by ID after the query: a retry may + // already be running with the same ID and client owner. + for (const id of cancelled) { + attempts.get(id)?.abort(new JobCancelledError(id)); + } } finally { link[Symbol.dispose](); timeout.dispose(); @@ -902,12 +917,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 4b00d16a6..070aa5f17 100644 --- a/producer.go +++ b/producer.go @@ -45,12 +45,12 @@ const ( // Test-only properties. type producerTestSignals struct { - CancelHandledDuringFetch testsignal.TestSignal[int64] // notifies when a cancellation is handled during a fetch DeletedExpiredQueueRecords testsignal.TestSignal[struct{}] // notifies when the producer deletes expired queue records ExecutorShutdownStarted testsignal.TestSignal[struct{}] // notifies when the producer starts draining active jobs JobFetchTriggered testsignal.TestSignal[struct{}] // notifies when the producer's fetch limiter is triggered via triggerJobFetch MetadataChanged testsignal.TestSignal[struct{}] // notifies when the producer detects a metadata change Paused testsignal.TestSignal[struct{}] // notifies when the producer is paused + PolledJobCancellations testsignal.TestSignal[struct{}] // notifies when a cancellation check completed PolledQueueConfig testsignal.TestSignal[struct{}] // notifies when the producer polls for queue settings QueueControlEventTriggered testsignal.TestSignal[*controlEventPayload] // notifies when a queue control event is triggered via triggerQueueControlEvent ReportedProducerStatus testsignal.TestSignal[struct{}] // notifies when the producer reports its own status @@ -198,8 +198,8 @@ type producer struct { baseservice.BaseService startstop.BaseStartStop - // Jobs which are currently being worked. The polling goroutine snapshots - // their IDs; only the main goroutine accesses the executors themselves. + // Jobs currently being worked. Cancellation checks snapshot the executors + // so delayed results cannot affect a newer attempt of the same job. activeJobsMu sync.Mutex activeJobs map[int64]*jobexecutor.JobExecutor @@ -215,10 +215,10 @@ type producer struct { pilot riverpilot.Pilot workers *Workers - // Receives job IDs from notifier, polling, and local cancellation calls. - // Only the main goroutine reads it. - cancelCh chan int64 - cancelState atomic.Pointer[producerCancelState] + // Coalesce cancellation hints. The generation catches hints received while + // a fetch has claimed jobs but hasn't registered their executors yet. + cancelGeneration atomic.Uint64 + cancelPollCh chan struct{} // Set to true when the producer thinks it should trigger another fetch as // soon as slots are available. This is written and read by the main @@ -242,10 +242,8 @@ type producer struct { // goroutine, only read from main goroutine. queueControlCh chan *controlEventPayload - // Coalesces reconnect notifications for reconciliation on the main goroutine. - reconnectCh chan struct{} - retryPolicy ClientRetryPolicy + runState atomic.Pointer[producerRunState] testSignals producerTestSignals } @@ -264,7 +262,7 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi producer := baseservice.Init(archetype, &producer{ activeJobs: make(map[int64]*jobexecutor.JobExecutor), - cancelCh: make(chan int64, 1000), + cancelPollCh: make(chan struct{}, 1), completer: config.Completer, config: config.mustValidate(), exec: exec, @@ -273,7 +271,6 @@ 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, }) @@ -313,6 +310,17 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } + // Queue controls delivered during shutdown must not override settings + // loaded when the producer restarts. +drainControls: + for { + select { + case <-p.queueControlCh: + default: + break drainControls + } + } + // Workers can be registered after the client is constructed, so capture // kinds at startup. Keep an empty registry non-nil to avoid fetching all kinds. p.fetchKind = nil @@ -407,7 +415,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return err } - controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(workCtx)) + controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(fetchCtx, workCtx)) if err != nil { stopped() if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { @@ -417,21 +425,16 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } // Recover cancellations whose NOTIFY was lost while disconnected. - reconnectSub = p.config.Notifier.RegisterListenerReadyFunc(func() { - select { - case p.reconnectCh <- struct{}{}: - default: - } - }) + reconnectSub = p.config.Notifier.RegisterListenerReadyFunc(p.TriggerCancelCheck) } subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) - p.cancelState.Store(&producerCancelState{done: subroutineCtx.Done()}) + 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.cancelState.Store(nil) + 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() { @@ -469,14 +472,14 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { subroutineWG.Add(1) go p.pollForSettingChanges(subroutineCtx, &subroutineWG, initiallyPaused, initialMetadata) - - subroutineWG.Add(1) - go p.pollForJobCancellations(subroutineCtx, &subroutineWG) } + subroutineWG.Add(1) + go p.pollForJobCancellations(subroutineCtx, &subroutineWG) + p.fetchAndRunLoop(fetchCtx, workCtx) p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) - p.executorShutdownLoop(workCtx) + p.executorShutdownLoop() p.Logger.DebugContext(workCtx, p.Name+": Shutdown loop exited, awaiting subroutines", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) cancelSubroutines(fmt.Errorf("producer stopped: %w", startstop.ErrStop)) @@ -489,6 +492,15 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } +// TriggerCancelCheck wakes the cancellation poller without blocking notification delivery. +func (p *producer) TriggerCancelCheck() { + p.cancelGeneration.Add(1) + select { + case p.cancelPollCh <- struct{}{}: + default: + } +} + // TriggerJobFetch manually triggers the producer to perform a job fetch // (although it's debounced, so it may not happen immediately if a fetch was // performed very recently). This is used by clients using drivers that don't @@ -506,9 +518,25 @@ func (p *producer) TriggerJobFetch() { // listen/notify. This is used by clients using drivers that don't support // listeners to wake a producer immediately after a queue control event was // known to be performed so the producer doesn't have to wait on polling. -func (p *producer) TriggerQueueControlEvent(controlEvent *controlEventPayload) { - p.queueControlCh <- controlEvent - p.testSignals.QueueControlEventTriggered.Signal(controlEvent) +func (p *producer) TriggerQueueControlEvent(ctx context.Context, controlEvent *controlEventPayload) { + state := p.runState.Load() + if state == nil || controlEvent == nil { + return + } + + // Queue controls are only consumed while fetching. Unlike cancellations, + // they must stop being delivered as soon as the producer starts draining. + select { + case <-state.fetchDone: + return + default: + } + select { + case p.queueControlCh <- controlEvent: + p.testSignals.QueueControlEventTriggered.Signal(controlEvent) + case <-ctx.Done(): + case <-state.fetchDone: + } } type controlAction string @@ -531,7 +559,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 { @@ -541,12 +569,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))) @@ -560,12 +591,7 @@ func (p *producer) handleControlNotification(workCtx context.Context) func(notif ) return } - select { - case <-workCtx.Done(): - case p.cancelCh <- decoded.JobID: - default: - p.Logger.WarnContext(workCtx, p.Name+": Job cancel notification dropped due to full buffer", slog.Int64("job_id", decoded.JobID)) - } + p.TriggerCancelCheck() default: p.Logger.DebugContext(workCtx, p.Name+": Received job control notification with unknown action", slog.String("action", string(decoded.Action)), @@ -582,63 +608,6 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { p.fetchLimiter.Call() fetchResultCh := make(chan producerFetchResult) - reconcileCtx, cancelReconcile := context.WithCancel(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) - } - if reconcilePending { - reconcilePending = false - startReconcile() - } - } for { select { @@ -647,10 +616,7 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { case msg := <-p.queueControlCh: switch msg.Action { case controlActionCancel: - // This path is only expected to take effect in poll-only mode, and - // only works for the case of a single process. Multi-process setups - // will have to wait for the next poll event for a cancel to take effect. - p.maybeCancelJob(workCtx, msg.JobID) + p.TriggerCancelCheck() case controlActionMetadataChanged: p.Logger.DebugContext(workCtx, p.Name+": Queue metadata changed", slog.String("queue", p.config.Queue), slog.String("queue_in_message", msg.Queue)) p.testSignals.MetadataChanged.Signal(struct{}{}) @@ -684,14 +650,8 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { default: p.Logger.DebugContext(workCtx, p.Name+": Unknown queue control action", "action", msg.Action) } - case jobID := <-p.cancelCh: - p.maybeCancelJob(workCtx, jobID) - case <-p.reconnectCh: - startReconcile() - case result := <-reconcileResultCh: - handleReconcileResult(result) case <-p.fetchLimiter.C(): - p.innerFetchLoop(workCtx, fetchResultCh, reconcileResultCh, handleReconcileResult, startReconcile) + p.innerFetchLoop(workCtx, fetchResultCh) // Ensure we can't start another fetch when fetchCtx is done, even if // the fetchLimiter is also ready to fire: select { @@ -745,7 +705,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, reconcileResultCh <-chan producerReconcileResult, handleReconcileResult func(producerReconcileResult), startReconcile func()) { +func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan producerFetchResult) { var limit int if p.paused { limit = 0 @@ -760,9 +720,7 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } - // A fetched job may not have an executor yet. Keep unmatched cancellations - // only for this fetch so unrelated IDs cannot affect later fetches. - pendingCancelJobs := make(map[int64]struct{}) + cancelGeneration := p.cancelGeneration.Load() go p.dispatchWork(workCtx, limit, fetchResultCh) for { @@ -771,7 +729,10 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr if result.err != nil { p.Logger.ErrorContext(workCtx, p.Name+": Error fetching jobs", slog.String("err", result.err.Error()), slog.String("queue", p.config.Queue)) } else if numLocked := len(result.jobs); numLocked > 0 { - p.startNewExecutors(workCtx, result.jobs, result.decodeErrors, pendingCancelJobs) + p.startNewExecutors(workCtx, result.jobs, result.decodeErrors) + if p.cancelGeneration.Load() != cancelGeneration { + p.TriggerCancelCheck() + } if numLocked == limit { // Fetch returned the maximum number of jobs that were requested, @@ -781,26 +742,16 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } return - case <-p.reconnectCh: - startReconcile() - case result := <-reconcileResultCh: - handleReconcileResult(result) case result := <-p.jobResultCh: p.removeActiveJob(result) - case jobID := <-p.cancelCh: - if !p.maybeCancelJob(workCtx, jobID) { - pendingCancelJobs[jobID] = struct{}{} - } - p.testSignals.CancelHandledDuringFetch.Signal(jobID) } } } -func (p *producer) executorShutdownLoop(ctx context.Context) { +func (p *producer) executorShutdownLoop() { p.testSignals.ExecutorShutdownStarted.Signal(struct{}{}) - // No more jobs will be fetched or executed. However, we must wait for all - // in-progress jobs to complete. + // Cancellation polling stays live until every execution has finished. for { p.activeJobsMu.Lock() numActiveJobs := len(p.activeJobs) @@ -808,13 +759,7 @@ func (p *producer) executorShutdownLoop(ctx context.Context) { if numActiveJobs == 0 { return } - - select { - case jobID := <-p.cancelCh: - p.maybeCancelJob(ctx, jobID) - case result := <-p.jobResultCh: - p.removeActiveJob(result) - } + p.removeActiveJob(<-p.jobResultCh) } } @@ -875,30 +820,6 @@ 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 _, job := range result.jobs { - p.activeJobsMu.Lock() - executor := p.activeJobs[job.ID] - p.activeJobsMu.Unlock() - if executor == nil || executor != result.executors[job.ID] || executor.JobRow.Attempt != job.Attempt { - continue - } - - var metadata struct { - CancelAttemptedAt time.Time `json:"cancel_attempted_at"` - } - if err := json.Unmarshal(job.Metadata, &metadata); err != nil { - p.Logger.ErrorContext(ctx, p.Name+": Error unmarshaling cancellation metadata", - slog.Int64("job_id", job.ID), slog.String("err", err.Error())) - continue - } - if !metadata.CancelAttemptedAt.IsZero() { - p.maybeCancelJob(ctx, job.ID) - } - } -} - func (p *producer) removeActiveJob(job *rivertype.JobRow) { p.activeJobsMu.Lock() executor := p.activeJobs[job.ID] @@ -940,59 +861,6 @@ func (p *producer) handleWorkerUnstuck() { p.config.JobStuckCount.Add(-1) } -func (p *producer) maybeCancelJob(ctx context.Context, id int64) bool { - p.activeJobsMu.Lock() - executor, ok := p.activeJobs[id] - p.activeJobsMu.Unlock() - if !ok { - return false - } - executor.Cancel(ctx) - return true -} - -// handleNotifierReconnect snapshots active executions before querying asynchronously. -func (p *producer) handleNotifierReconnect(ctx context.Context, resultCh chan<- producerReconcileResult, wg *sync.WaitGroup) { - p.activeJobsMu.Lock() - 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) - } - - p.activeJobsMu.Unlock() - - go func() { - defer wg.Done() - - jobs, err := timeoututil.WithTimeoutV(ctx, 10*time.Second, p.Name+".handleNotifierReconnect", func(ctx context.Context) ([]*rivertype.JobRow, error) { - return p.exec.JobGetByIDMany(ctx, &riverdriver.JobGetByIDManyParams{ - ID: ids, - Schema: p.config.Schema, - }) - }) - select { - case <-ctx.Done(): - case resultCh <- producerReconcileResult{err: err, executors: executors, jobs: jobs}: - } - }() -} - -// 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 { pluginLookup := p.config.PluginLookupGlobal if pluginLookup == nil { @@ -1120,7 +988,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 { @@ -1161,10 +1029,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) } @@ -1194,20 +1058,23 @@ func (p *producer) pollForJobCancellations(ctx context.Context, wg *sync.WaitGro case <-ctx.Done(): return case <-ticker.C: - if err := p.pollForJobCancellationsOnce(ctx); err != nil { - if !errors.Is(context.Cause(ctx), startstop.ErrStop) { - p.Logger.ErrorContext(ctx, p.Name+": Error fetching job cancellation requests", slog.String("err", err.Error())) - } - continue + case <-p.cancelPollCh: + } + if err := p.pollForJobCancellationsOnce(ctx); err != nil { + if !errors.Is(context.Cause(ctx), startstop.ErrStop) { + p.Logger.ErrorContext(ctx, p.Name+": Error fetching job cancellation requests", slog.String("err", err.Error())) } + continue } + p.testSignals.PolledJobCancellations.Signal(struct{}{}) } } func (p *producer) pollForJobCancellationsOnce(ctx context.Context) error { p.activeJobsMu.Lock() - jobIDs := slices.Collect(maps.Keys(p.activeJobs)) + executors := maps.Clone(p.activeJobs) p.activeJobsMu.Unlock() + jobIDs := slices.Collect(maps.Keys(executors)) // Bound query size (including SQLite's parameter count) and skip database // access entirely when there are no active jobs. @@ -1223,10 +1090,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) } } } @@ -1387,20 +1254,15 @@ func (p *producer) reportQueueStatusOnce(ctx context.Context) { p.testSignals.ReportedQueueStatus.Signal(struct{}{}) } -type producerCancelState struct { - done <-chan struct{} -} - type producerFetchResult struct { decodeErrors map[int64]error err error jobs []*rivertype.JobRow } -type producerReconcileResult struct { - err error - executors map[int64]*jobexecutor.JobExecutor - jobs []*rivertype.JobRow +type producerRunState struct { + done <-chan struct{} + fetchDone <-chan struct{} } type errorHandlerAdapter struct { diff --git a/producer_test.go b/producer_test.go index c63e04325..46e920fe8 100644 --- a/producer_test.go +++ b/producer_test.go @@ -82,16 +82,16 @@ func (p *beforeJobGetAvailablePilot) JobGetAvailable( type producerTestExecutor struct { riverdriver.Executor - jobGetAvailableFunc func(context.Context, *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) - jobGetByIDManyFunc func(context.Context, *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) + 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) JobGetByIDMany(ctx context.Context, params *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) { - return e.jobGetByIDManyFunc(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' @@ -148,184 +148,356 @@ func (l *producerTestListenerMock) WaitForNotification(ctx context.Context) (*ri return l.waitForNotificationFunc(ctx) } -func TestProducer_ApplyReconcileResult(t *testing.T) { +func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { t.Parallel() for _, testCase := range []struct { - name string - attempt int - metadata string - replaced bool - wantCancel bool + name string + blockedFetch bool + cancelFromLookup bool + queryError bool }{ - {name: "Cancelled", attempt: 1, metadata: `{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`, wantCancel: true}, - {name: "DifferentAttempt", attempt: 2, metadata: `{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`}, - {name: "MalformedMetadata", attempt: 1, metadata: `{"cancel_attempted_at":true}`}, - {name: "NoCancellation", attempt: 1, metadata: `{}`}, - {name: "ReplacedExecutor", attempt: 1, metadata: `{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`, replaced: true}, + {name: "BlockedFetchCancelDelivery", blockedFetch: true}, + {name: "BlockedFetchQueryResult", blockedFetch: true, cancelFromLookup: true}, + {name: "CancelDelivery"}, + {name: "QueryResult", cancelFromLookup: true}, + {name: "RetryAfterError", cancelFromLookup: true, queryError: 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{} + var signals struct { + FetchStarted testsignal.TestSignal[struct{}] + LoopStopped testsignal.TestSignal[struct{}] + QueryStarted testsignal.TestSignal[time.Duration] + ReleaseFetch testsignal.TestSignal[struct{}] } - prod := baseservice.Init(archetype, &producer{activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}}) + signals.FetchStarted.Init(t) + signals.LoopStopped.Init(t) + signals.QueryStarted.Init(t) + signals.ReleaseFetch.Init(t) - prod.applyReconcileResult(context.Background(), producerReconcileResult{ - executors: map[int64]*jobexecutor.JobExecutor{1: snapshotExecutor}, - jobs: []*rivertype.JobRow{{ID: 1, Attempt: testCase.attempt, Metadata: []byte(testCase.metadata)}}, + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + release := sync.OnceFunc(func() { signals.ReleaseFetch.Signal(struct{}{}) }) + defer release() + archetype := riversharedtest.BaseServiceArchetype(t) + jobCtx, cancelJob := context.WithCancelCause(ctx) + defer cancelJob(nil) + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelJob, JobRow: &rivertype.JobRow{ID: 1, Attempt: 1}}) + var queryAttempts int + exec := &producerTestExecutor{ + jobGetAvailableFunc: func(context.Context, *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) { + signals.FetchStarted.Signal(struct{}{}) + <-signals.ReleaseFetch.WaitC() + return &riverdriver.JobGetAvailableResult{}, nil + }, + jobGetCancelRequestedFunc: func(ctx context.Context, _ *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + deadline, ok := ctx.Deadline() + if ok { + signals.QueryStarted.Signal(time.Until(deadline)) + } else { + signals.QueryStarted.Signal(0) + } + queryAttempts++ + if testCase.queryError && queryAttempts == 1 { + return nil, errors.New("temporary cancellation check failure") + } + if testCase.cancelFromLookup { + return []int64{1}, nil + } + <-ctx.Done() + return nil, ctx.Err() + }, + } + pilot := &riverpilot.StandardPilot{} + _, state, err := pilot.ProducerInit(ctx, exec, nil) + require.NoError(t, err) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{MaxWorkers: 2, QueuePollInterval: 20 * time.Millisecond}, + exec: exec, + fetchLimiter: chanutil.NewDebouncedChan(ctx, time.Millisecond, true), + jobResultCh: make(chan *rivertype.JobRow), + paused: !testCase.blockedFetch, + pilot: pilot, + state: state, + queueControlCh: make(chan *controlEventPayload), }) - if testCase.wantCancel { - require.ErrorIs(t, context.Cause(ctx), rivertype.ErrJobCancelledRemotely) - } else { - require.NoError(t, ctx.Err()) + prod.runState.Store(&producerRunState{done: ctx.Done()}) + var pollWG sync.WaitGroup + pollWG.Add(1) + go prod.pollForJobCancellations(ctx, &pollWG) + go func() { + prod.fetchAndRunLoop(ctx, ctx) + prod.executorShutdownLoop() + signals.LoopStopped.Signal(struct{}{}) + }() + defer func() { + release() + cancel() + signals.LoopStopped.WaitOrTimeout() + pollWG.Wait() + }() + + if testCase.blockedFetch { + signals.FetchStarted.WaitOrTimeout() + } + prod.TriggerCancelCheck() + queryTimeout := signals.QueryStarted.WaitOrTimeout() + require.Positive(t, queryTimeout) + require.LessOrEqual(t, queryTimeout, 10*time.Second) + if testCase.queryError { + require.Positive(t, signals.QueryStarted.WaitOrTimeout()) + } + + if !testCase.cancelFromLookup { + client := &Client[struct{}]{producersByQueueName: map[string]*producer{QueueDefault: prod}} + job := *executor.JobRow + job.Queue = QueueDefault + client.notifyProducerOfCancelWithoutNotifier(ctx, &job) + } + select { + case <-jobCtx.Done(): + require.ErrorIs(t, context.Cause(jobCtx), rivertype.ErrJobCancelledRemotely) + case <-time.After(time.Second): + t.Fatal("blocked cancellation check prevented cancellation delivery") + } + + select { + case prod.jobResultCh <- &rivertype.JobRow{ID: 1}: + case <-time.After(time.Second): + t.Fatal("blocked cancellation check prevented job completion") } }) } } -// A cancel NOTIFY lost during a disconnect must be recovered on reconnect. -func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { +func TestProducer_HandleControlNotification(t *testing.T) { t.Parallel() - ctx := context.Background() + for _, action := range []controlAction{controlActionCancel, controlActionMetadataChanged, controlActionPause, controlActionResume} { + t.Run(string(action), func(t *testing.T) { + t.Parallel() - 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: + prod := baseservice.Init(riversharedtest.BaseServiceArchetype(t), &producer{ + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{Queue: QueueDefault}, + queueControlCh: make(chan *controlEventPayload, 1), + }) + fetchCtx, cancelFetch := context.WithCancel(context.Background()) + handleNotification := prod.handleControlNotification(fetchCtx, context.Background()) + cancelFetch() + payload, err := json.Marshal(&controlEventPayload{Action: action, JobID: 123, Queue: QueueDefault}) + require.NoError(t, err) + + handleNotification(notifier.NotificationTopicControl, string(payload)) + + require.Empty(t, prod.queueControlCh) + if action == controlActionCancel { + require.Len(t, prod.cancelPollCh, 1) + } else { + require.Empty(t, prod.cancelPollCh) } - reconnecting.Store(false) - } - return realListener.Connect(ctx) + }) } +} - notif := notifier.New(archetype, mockListener) - require.NoError(t, notif.Start(ctx)) - t.Cleanup(notif.Stop) - // Release a blocked Connect before stopping the notifier. - t.Cleanup(closeReleaseReconnect) +// A cancel NOTIFY lost during a disconnect must be recovered on reconnect. +func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { + t.Parallel() - var jobStarted testsignal.TestSignal[int64] - jobStarted.Init(t) - workers := NewWorkers() + 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() - type JobArgs struct { - testutil.JobArgsReflectKind[JobArgs] - } + ctx := context.Background() - 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, - }) + 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{} + ) - 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) + completer := jobcompleter.NewInlineCompleter(archetype, schema, exec, pilot, jobUpdates) + require.NoError(t, completer.Start(ctx)) + t.Cleanup(completer.Stop) - workCtx, workCancel := context.WithCancel(ctx) + 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) + } - require.NoError(t, prod.StartWorkContext(ctx, workCtx)) - // Cancel workers before waiting for producer shutdown. - t.Cleanup(prod.Stop) - t.Cleanup(workCancel) + 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] + } - jobID := jobStarted.WaitOrTimeout() + AddWorker(workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + jobStarted.Signal(job.ID) + <-ctx.Done() + return ctx.Err() + })) + + prod := newProducer(archetype, exec, pilot, &producerConfig{ + ClientID: testClientID, + Completer: completer, + ErrorHandler: newTestErrorHandler(), + FetchCooldown: FetchCooldownDefault, + FetchPollInterval: 50 * time.Millisecond, + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), + JobTimeout: JobTimeoutDefault, + MaxWorkers: 10, + Notifier: notif, + Queue: queueName, + QueuePollInterval: time.Hour, + QueueReportInterval: queueReportIntervalDefault, + RetryPolicy: &DefaultClientRetryPolicy{}, + SchedulerInterval: riverinternaltest.SchedulerShortInterval, + Schema: schema, + StaleProducerRetentionPeriod: time.Minute, + Workers: workers, + }) - // 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, - })) + 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) - reconnectStarted.WaitOrTimeout() + var fetchReturned testsignal.TestSignal[int64] + var releaseFetch testsignal.TestSignal[struct{}] + fetchReturned.Init(t) + releaseFetch.Init(t) + releaseFetchFunc := sync.OnceFunc(func() { releaseFetch.Signal(struct{}{}) }) + if testCase.fetchPending { + prod.testSignals.PolledJobCancellations.Init(t) + prod.pilot = &afterJobGetAvailablePilot{ + Pilot: pilot, + afterJobGetAvailableFunc: func(res *riverdriver.JobGetAvailableResult, err error) { + if err == nil && res != nil && len(res.Jobs) > 0 { + fetchReturned.Signal(res.Jobs[0].ID) + <-releaseFetch.WaitC() + } + }, + } + } + prod.testSignals.ExecutorShutdownStarted.Init(t) + fetchCtx, fetchCancel := context.WithCancel(ctx) + defer fetchCancel() + workCtx, workCancel := context.WithCancel(ctx) + + require.NoError(t, prod.StartWorkContext(fetchCtx, workCtx)) + // Cancel workers before waiting for producer shutdown. + t.Cleanup(prod.Stop) + t.Cleanup(workCancel) + t.Cleanup(closeReleaseReconnect) + t.Cleanup(releaseFetchFunc) + + var jobID int64 + if testCase.fetchPending { + jobID = fetchReturned.WaitOrTimeout() + } else { + jobID = jobStarted.WaitOrTimeout() + } + require.Positive(t, jobID) + if testCase.stop { + fetchCancel() + if !testCase.fetchPending { + prod.testSignals.ExecutorShutdownStarted.WaitOrTimeout() + } + } - // 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) + // 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() + closeReleaseReconnect() + if testCase.fetchPending { + prod.testSignals.PolledJobCancellations.WaitOrTimeout() + releaseFetchFunc() + } - update := riversharedtest.WaitOrTimeout(t, jobUpdates) - require.Equal(t, rivertype.JobStateCancelled, update[0].Job.State, - "job should have observed cancellation after the notifier reconnected and reconciled active jobs against cancel_attempted_at") + 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) { @@ -490,6 +662,132 @@ func TestProducer_MetricEmitHook(t *testing.T) { }) } +func TestProducer_PollForJobCancellations(t *testing.T) { + t.Parallel() + + for _, testCase := range []struct { + name string + jobs int + queries int + }{ + {name: "Batched", jobs: 1001, queries: 2}, + {name: "Idle"}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + seen := make(map[int64]bool) + var queries int + prod := &producer{ + activeJobs: make(map[int64]*jobexecutor.JobExecutor), + config: &producerConfig{Schema: "custom"}, + exec: &producerTestExecutor{jobGetCancelRequestedFunc: func(_ context.Context, params *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + queries++ + require.Equal(t, "custom", params.Schema) + require.NotEmpty(t, params.ID) + require.LessOrEqual(t, len(params.ID), jobCancelPollBatchSize) + for _, id := range params.ID { + require.False(t, seen[id], "job checked more than once") + seen[id] = true + } + return nil, nil + }}, + } + for id := range testCase.jobs { + prod.activeJobs[int64(id)] = &jobexecutor.JobExecutor{} + } + + require.NoError(t, prod.pollForJobCancellationsOnce(context.Background())) + require.Equal(t, testCase.queries, queries) + require.Len(t, seen, testCase.jobs) + }) + } + + for _, replaced := range []bool{false, true} { + t.Run(fmt.Sprintf("ReplacedExecutor_%t", replaced), func(t *testing.T) { + t.Parallel() + + archetype := riversharedtest.BaseServiceArchetype(t) + originalCtx, cancelOriginal := context.WithCancelCause(context.Background()) + defer cancelOriginal(nil) + replacementCtx, cancelReplacement := context.WithCancelCause(context.Background()) + defer cancelReplacement(nil) + original := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelOriginal, JobRow: &rivertype.JobRow{ID: 1}}) + replacement := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelReplacement, JobRow: &rivertype.JobRow{ID: 1}}) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: original}, + config: &producerConfig{}, + }) + prod.exec = &producerTestExecutor{ + jobGetCancelRequestedFunc: func(ctx context.Context, params *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + require.Equal(t, []int64{1}, params.ID) + if replaced { + // The old attempt finishes and is retried before the query returns. + cancelOriginal(nil) + prod.activeJobs[1] = replacement + } + return []int64{1}, nil + }, + } + + require.NoError(t, prod.pollForJobCancellationsOnce(context.Background())) + require.NoError(t, replacementCtx.Err()) + if !replaced { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + } + }) + } +} + +func TestProducer_PollForJobCancellationsWakeDuringQuery(t *testing.T) { + t.Parallel() + + var queryStarted, releaseQuery testsignal.TestSignal[struct{}] + queryStarted.Init(t) + releaseQuery.Init(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + jobCtx, cancelJob := context.WithCancelCause(ctx) + defer cancelJob(nil) + archetype := riversharedtest.BaseServiceArchetype(t) + var queries int + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelJob, JobRow: &rivertype.JobRow{ID: 1}})}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{QueuePollInterval: time.Hour}, + exec: &producerTestExecutor{jobGetCancelRequestedFunc: func(ctx context.Context, _ *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + queries++ + if queries == 1 { + queryStarted.Signal(struct{}{}) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-releaseQuery.WaitC(): + return nil, nil + } + } + return []int64{1}, nil + }}, + }) + var wg sync.WaitGroup + wg.Add(1) + go prod.pollForJobCancellations(ctx, &wg) + defer func() { cancel(); wg.Wait() }() + + prod.TriggerCancelCheck() + queryStarted.WaitOrTimeout() + // A burst during a blocked read coalesces, but still causes a fresh read. + for range 1000 { + prod.TriggerCancelCheck() + } + releaseQuery.Signal(struct{}{}) + riversharedtest.WaitOrTimeout(t, jobCtx.Done()) + require.ErrorIs(t, context.Cause(jobCtx), rivertype.ErrJobCancelledRemotely) + cancel() + wg.Wait() + require.Equal(t, 2, queries) +} + func TestProducer_PollOnly(t *testing.T) { t.Parallel() @@ -542,124 +840,6 @@ 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 - }, - jobGetByIDManyFunc: func(ctx context.Context, _ *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, 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 []*rivertype.JobRow{{ID: 1, Attempt: 1, Metadata: []byte(`{"cancel_attempted_at":"2026-10-20T00:00:00Z"}`)}}, 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 int64, 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 <- 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() @@ -759,7 +939,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{ @@ -778,11 +958,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) { @@ -839,6 +1020,45 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin producer.testSignals.ReportedQueueStatus.WaitOrTimeout() }) + t.Run("RestartDiscardsStaleControls", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.config.QueuePollInterval = time.Hour + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, _ *Job[noOpArgs]) error { return ctx.Err() })) + startProducer(t, ctx, ctx, producer) + producer.Stop() + + // Model controls buffered just as the previous run stopped fetching. + producer.queueControlCh <- &controlEventPayload{Action: controlActionPause, Queue: bundle.queue} + producer.queueControlCh <- &controlEventPayload{Action: controlActionMetadataChanged, Metadata: []byte(`{"old":true}`), Queue: bundle.queue} + _, err := bundle.exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`{"current":true}`), + MetadataDoUpdate: true, + Name: bundle.queue, + Schema: producer.config.Schema, + }) + require.NoError(t, err) + + id := mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + _, err = bundle.exec.JobCancel(ctx, &riverdriver.JobCancelParams{ + CancelAttemptedAt: time.Now(), + ControlTopic: string(notifier.NotificationTopicControl), + ID: id, + Schema: producer.config.Schema, + }) + require.NoError(t, err) + _, err = bundle.exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: id, Schema: producer.config.Schema}) + require.NoError(t, err) + producer.TriggerCancelCheck() + + startProducer(t, ctx, ctx, producer) + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + producer.testSignals.MetadataChanged.RequireEmpty() + producer.testSignals.Paused.RequireEmpty() + }) + t.Run("UnknownJobKind", func(t *testing.T) { t.Parallel() @@ -888,9 +1108,6 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin t.Parallel() producer, bundle := setup(t) - if producer.config.Notifier != nil { - t.Skip("requires polling without a notifier") - } producer.config.FetchPollInterval = time.Hour producer.config.MaxWorkers = 1 producer.config.QueuePollInterval = 20 * time.Millisecond @@ -926,6 +1143,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin Name: producer.config.Queue, Schema: producer.config.Schema, })) + producer.TriggerQueueControlEvent(ctx, &controlEventPayload{Action: controlActionPause, Queue: producer.config.Queue}) producer.testSignals.Paused.WaitOrTimeout() case "Stopping": fetchCancel() @@ -937,7 +1155,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin _, err := bundle.exec.JobCancel(ctx, &riverdriver.JobCancelParams{ ID: jobID, CancelAttemptedAt: time.Now().UTC(), - ControlTopic: string(notifier.NotificationTopicControl), + ControlTopic: "unsubscribed_control", Schema: producer.config.Schema, }) require.NoError(t, err) @@ -1021,6 +1239,35 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRetryable, update.Job.State) }) + t.Run("CancellationNoticeWithoutMarkerDoesNotCancelJob", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.testSignals.PolledJobCancellations.Init(t) + var entered testsignal.TestSignal[struct{}] + var release testsignal.TestSignal[struct{}] + entered.Init(t) + release.Init(t) + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, _ *Job[noOpArgs]) error { + entered.Signal(struct{}{}) + <-release.WaitC() + return ctx.Err() + })) + id := mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + startProducer(t, ctx, ctx, producer) + t.Cleanup(func() { release.Signal(struct{}{}) }) + entered.WaitOrTimeout() + + // An old notification can arrive after retry cleared the durable marker. + payload, err := json.Marshal(&controlEventPayload{Action: controlActionCancel, JobID: id, Queue: producer.config.Queue}) + require.NoError(t, err) + producer.handleControlNotification(ctx, ctx)(notifier.NotificationTopicControl, string(payload)) + producer.testSignals.PolledJobCancellations.WaitOrTimeout() + release.Signal(struct{}{}) + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + }) + t.Run("CancellationWhileFetchResultPendingReachesJob", func(t *testing.T) { t.Parallel() @@ -1028,7 +1275,8 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin // process its cancellation while it has no active executor. producer, bundle := setup(t) producer.config.FetchPollInterval = time.Hour - producer.testSignals.CancelHandledDuringFetch.Init(t) + producer.config.QueuePollInterval = time.Hour + producer.testSignals.PolledJobCancellations.Init(t) var ( fetchReturned testsignal.TestSignal[int64] @@ -1054,6 +1302,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() })) @@ -1073,8 +1322,8 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRunning, jobAfterCancel.State) // Send directly to avoid depending on notifier delivery timing. - producer.cancelCh <- jobID - require.Equal(t, jobID, producer.testSignals.CancelHandledDuringFetch.WaitOrTimeout()) + producer.TriggerCancelCheck() + producer.testSignals.PolledJobCancellations.WaitOrTimeout() releaseFetch.Signal(struct{}{}) require.ErrorIs(t, workerErr.WaitOrTimeout(), context.Canceled) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 53dae3d6f..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,19 +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, - metadata = metadata - 'cancel_attempted_at', - 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 1346c5008..a490d2da3 100644 --- a/riverdriver/riverdrivertest/job_update.go +++ b/riverdriver/riverdrivertest/job_update.go @@ -511,6 +511,29 @@ 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() @@ -640,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 6b11a3bc5..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,19 +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, - metadata = metadata - 'cancel_attempted_at', - 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 57bf72f8b..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,19 +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, - metadata = metadata - 'cancel_attempted_at', - 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 3380dc0dd..0dc456e7b 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql @@ -586,23 +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, - metadata = jsonb_remove(metadata, '$.cancel_attempted_at'), - 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 2a7472f1d..180754a1e 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_job.sql.go @@ -1450,23 +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, - metadata = jsonb_remove(metadata, '$.cancel_attempted_at'), - 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 { @@ -1474,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 @@ -1483,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 7d1044bfc..b17c6105f 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -859,7 +859,7 @@ func (e *Executor) JobRetry(ctx context.Context, params *riverdriver.JobRetryPar return nil, interpretError(err) } - return jobRowFromInternal(job) + return jobRowFromInternal(jobFromReturningRow((*dbsqlc.JobGetAvailableRow)(job))) }) } diff --git a/ruby/lib/client_runtime.rb b/ruby/lib/client_runtime.rb index 47cb2acaf..9ada74962 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 39a46c0f0..ed64584f1 100644 --- a/ruby/lib/driver/runtime.rb +++ b/ruby/lib/driver/runtime.rb @@ -249,27 +249,29 @@ def job_rescue_stuck(horizon:, retry_policy:, now: Time.now.utc, max: 1_000, log end def job_retry(id, now: Time.now.utc) + already_available = "state = 'available' AND scheduled_at < #{runtime_time(now)}" + cancellation_requested = runtime_postgres? ? "(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')" : "(CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL" updated_id, job = transaction do # Check the counter in the write itself so a concurrent claim cannot # advance it between validation and the update. updated_id = runtime_returning_ids(<<~SQL).first UPDATE river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - metadata = #{runtime_postgres? ? "metadata - 'cancel_attempted_at'" : "jsonb_remove(metadata, '$.cancel_attempted_at')"}, - finalized_at = NULL, - scheduled_at = #{runtime_time(now)} + max_attempts = CASE WHEN #{already_available} THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + metadata = #{runtime_postgres? ? "CASE WHEN jsonb_typeof(metadata) = 'object' THEN metadata - 'cancel_attempted_at' ELSE metadata END" : "CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END"}, + finalized_at = CASE WHEN #{already_available} THEN finalized_at ELSE NULL END, + scheduled_at = CASE WHEN #{already_available} THEN scheduled_at ELSE #{runtime_time(now)} END WHERE id = #{Integer(id)} AND state != 'running' - AND (state != 'available' OR scheduled_at > #{runtime_time(now)}) - AND attempt < #{River::MAX_ATTEMPTS_LIMIT} + AND (#{cancellation_requested} OR NOT (#{already_available})) + AND ((#{already_available}) OR attempt < #{River::MAX_ATTEMPTS_LIMIT}) RETURNING id SQL job = job_get_by_id(updated_id || id) runtime_notify("river_insert", queue: job.queue) if updated_id && job [updated_id, job] end - if !updated_id && job && job.state != "running" && (job.state != "available" || job.scheduled_at > now) && job.attempt >= River::MAX_ATTEMPTS_LIMIT + if !updated_id && job && job.state != "running" && (job.state != "available" || job.scheduled_at >= now) && job.attempt >= River::MAX_ATTEMPTS_LIMIT raise ArgumentError, "cannot retry a job with #{River::MAX_ATTEMPTS_LIMIT} or more attempts" end job @@ -484,7 +486,7 @@ def time_precision end private def runtime_cancel_attempted - runtime_postgres? ? "metadata ? 'cancel_attempted_at'" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" + runtime_postgres? ? "(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" end private def runtime_claim_jobs(predicate, max, attempted_by, now) diff --git a/ruby/sig/runtime.rbs b/ruby/sig/runtime.rbs index 2782a3465..1f2b8d73c 100644 --- a/ruby/sig/runtime.rbs +++ b/ruby/sig/runtime.rbs @@ -360,6 +360,7 @@ module River def initialize: (Time, Float, Time) -> void end + @cancellation_thread: Thread? @client: Client @condition: ConditionVariable @config: Config @@ -406,7 +407,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 5510cad55..39eaefcfe 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 @@ -746,7 +762,58 @@ def work(_job) expect { execute(value, row) }.not_to raise_error end - it "checks cancellation in one batch, excluding inactive work and other queues" do + [false, true].each do |stopping| + it "cancels work during a blocked fetch#{" and graceful stop" if stopping}" do + started, fetching, release = Queue.new, Queue.new, Queue.new + worker = Object.new + worker.define_singleton_method(:work) do |_job| + started << true + Queue.new.pop + end + job = row + marked = false + claims = 0 + driver = Object.new + driver.define_singleton_method(:queue_upsert) { |_name| } + driver.define_singleton_method(:queue_get) { |_name| nil } + driver.define_singleton_method(:job_get_available) do |**| + claims += 1 + if claims == 1 + [job] + else + fetching << true + release.pop + [] + end + end + driver.define_singleton_method(:job_get_cancelled_ids) { |ids| marked ? ids : [] } + driver.define_singleton_method(:job_set_state_if_running) do |**attributes| + job.state = attributes.fetch(:state) + job + end + value = described_class.new(Object.new, driver, + config(worker: worker, queues: {branch: 2}).with( + fetch_cooldown: 0.001, fetch_poll_interval: 0.01, leader_election_disabled: true + )) + subscription = value.subscribe([:job_cancelled]) + value.start + Timeout.timeout(2) { + started.pop + fetching.pop + } + value.stop(wait: false) if stopping + marked = true + + event = Timeout.timeout(2) { subscription.pop } + expect(event.job.state).to eq(River::JOB_STATE_CANCELLED) + ensure + value&.stop(cancel: true, wait: false) + release << true + value&.stop + end + end + + it "checks cancellation across queues in one batch, excluding inactive work" do raised = [] fake_thread = Object.new fake_thread.define_singleton_method(:raise) { |error| raised << error } @@ -766,10 +833,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 8cfef94f7..2b9aab690 100644 --- a/ruby/spec/driver_runtime_feature_spec.rb +++ b/ruby/spec/driver_runtime_feature_spec.rb @@ -137,7 +137,7 @@ expect(postgres_driver.send(:runtime_state, River::JOB_STATE_AVAILABLE)).to end_with("::river_job_state") expect(postgres_driver.send(:runtime_time, Time.utc(2026, 1, 2))).to end_with("::timestamptz") expect(postgres_driver.send(:runtime_merge_metadata, {"ready" => true})).to include("metadata ||") - expect(postgres_driver.send(:runtime_cancel_attempted)).to eq("metadata ? 'cancel_attempted_at'") + expect(postgres_driver.send(:runtime_cancel_attempted)).to eq("(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')") expect(postgres_driver.send(:runtime_update_value, :attempted_by, ["one"])).to include("ARRAY[", "::text[]") expect(postgres_driver.send(:runtime_update_value, :errors, [7])).to include("ARRAY[", "::jsonb[]") expect(postgres_driver.send(:runtime_queue_columns)).to eq("name, created_at, metadata, paused_at, updated_at") diff --git a/ruby/spec/driver_runtime_shared_examples.rb b/ruby/spec/driver_runtime_shared_examples.rb index 785b8a906..b6f80b857 100644 --- a/ruby/spec/driver_runtime_shared_examples.rb +++ b/ruby/spec/driver_runtime_shared_examples.rb @@ -422,6 +422,13 @@ def work(_job) expect(driver.job_get_by_id(row.id).state).to eq("running") end + it "ignores cancellation names stored in non-object metadata" do + job = insert_job(1) + metadata = driver.send(:runtime_json, ["cancel_attempted_at"]) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = #{metadata} WHERE id = #{job.id}") + expect(driver.job_get_cancelled_ids([job.id])).to eq([]) + end + it "looks up cancellation markers for a batch, excluding missing and uncancelled jobs" do cancelled = insert_job(1) unmarked = insert_job(2) @@ -612,6 +619,51 @@ def insert_job(value, **options) expect(driver.job_retry(-1)).to be_nil end + it "increments exhausted attempts at the exact retry timestamp" do + job = insert_job(1, scheduled_at: now, max_attempts: 1) + driver.send(:runtime_execute, "UPDATE river_job SET attempt = max_attempts WHERE id = #{job.id}") + expect(driver.job_retry(job.id, now: now)).to have_attributes(max_attempts: 2, scheduled_at: job.scheduled_at) + end + + it "clears cancellation metadata without reordering an already available retry" do + job = insert_job(1, scheduled_at: now - 60, max_attempts: 1, + metadata: {"cancel_attempted_at" => now.iso8601, "preserved" => true}) + driver.send(:runtime_execute, "UPDATE river_job SET attempt = max_attempts WHERE id = #{job.id}") + + retried = driver.job_retry(job.id, now: now) + expect(retried).to have_attributes(state: "available", attempt: 1, max_attempts: 1, + scheduled_at: job.scheduled_at, metadata: job.metadata.except("cancel_attempted_at")) + expect(driver.job_get_by_id(job.id).metadata).to eq(retried.metadata) + end + + it "preserves malformed SQLite metadata and reports its decoding error", database: :sqlite do + skip "requires SQLite" if driver.send(:runtime_postgres?) + + job = insert_job(1, scheduled_at: now + 60) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = '[not json' WHERE id = #{job.id}") + + expect { driver.job_retry(job.id, now: now) }.to raise_error(River::JobRowDecodeError, /metadata/) + row = driver.send(:runtime_query_rows, "SELECT metadata FROM river_job WHERE id = #{job.id}").first + expect(driver.send(:runtime_value, row, :metadata)).to eq("[not json") + end + + [nil, 42, false, "cancel_attempted_at", ["cancel_attempted_at"]].each do |metadata| + it "preserves non-object metadata #{metadata.inspect} when retrying" do + job = insert_job(1, scheduled_at: now + 60) + encoded = driver.send(:runtime_json, JSON.generate(metadata)) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = #{encoded} WHERE id = #{job.id}") + + if metadata.nil? + expect(driver.job_retry(job.id, now: now)).to have_attributes(state: "available") + else + expect { driver.job_retry(job.id, now: now) }.to raise_error(River::JobRowDecodeError) + end + column = driver.send(:runtime_postgres?) ? "metadata::text" : "json(metadata)" + row = driver.send(:runtime_query_rows, "SELECT #{column} AS metadata FROM river_job WHERE id = #{job.id}").first + expect(JSON.parse(driver.send(:runtime_value, row, :metadata))).to eq(metadata) + end + end + it "rejects manual retries after the last representable attempt without changing the job" do job = insert_job(1, max_attempts: 32_767) original = client.job_update(job.id, attempt: 32_767, finalized_at: now, state: :discarded) diff --git a/ruby/spec/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 2152e1341..9526ccf2e 100644 --- a/rust/riverqueue/src/database/sqlite.rs +++ b/rust/riverqueue/src/database/sqlite.rs @@ -949,25 +949,30 @@ pub(crate) async fn retry( ) -> Result, BackendError> { let sql = format!( r#" + WITH retry_params AS (SELECT ? AS now) UPDATE river_job SET state = 'available', max_attempts = CASE + WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = ? + finalized_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) + THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END WHERE id = ? AND state != 'running' - AND (state != 'available' OR scheduled_at > ?) + AND ((CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR NOT (state = 'available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {JOB_COLUMNS} "# ); sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) .bind(sqlite_time(now)) .bind(id) - .bind(sqlite_time(now)) .fetch_optional(&mut *connection) .await? .map(JobRecord::into_job) @@ -2537,6 +2542,46 @@ mod tests { assert_eq!(row.attempted_by, ["three", "four", "five"]); } + #[tokio::test] + async fn retry_at_exact_schedule_increments_max_attempts() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Available, now).await; + sqlx::query("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .bind(job.id) + .execute(&mut *connection) + .await + .unwrap(); + let retried = retry(&mut connection, job.id, now).await.unwrap().unwrap(); + assert_eq!(retried.max_attempts, job.max_attempts + 1); + assert_eq!(retried.scheduled_at, job.scheduled_at); + } + + #[tokio::test] + async fn retry_preserves_malformed_metadata_for_row_validation() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Retryable, now).await; + sqlx::query("UPDATE river_job SET metadata = '[not json' WHERE id = ?") + .bind(job.id) + .execute(&mut *connection) + .await + .unwrap(); + + let error = retry(&mut connection, job.id, now).await.unwrap_err(); + assert!( + matches!(error, BackendError::InvalidRow(ref message) if message.contains("metadata")) + ); + let metadata: String = sqlx::query_scalar("SELECT metadata FROM river_job WHERE id = ?") + .bind(job.id) + .fetch_one(&mut *connection) + .await + .unwrap(); + assert_eq!(metadata, "[not json"); + } + /// Inserts a job for the undecodable row tests. async fn insert_test_job( connection: &mut SqliteConnection, diff --git a/rust/riverqueue/src/storage/postgres.rs b/rust/riverqueue/src/storage/postgres.rs index 316eb3587..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 From 70d514565bf1688e54bb2026686fc4fe3d6f49e7 Mon Sep 17 00:00:00 2001 From: Brandur Date: Tue, 27 Oct 2026 17:49:40 -0500 Subject: [PATCH 5/6] =?UTF-8?q?=F0=9F=A4=96=20Reconcile=20cancellations=20?= =?UTF-8?q?on=20demand=20when=20notifications=20are=20available?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wake cancellation checks on notices, reconnects, and notification buffer overflow across all five runtimes. Retry failed reads until they succeed, and retain periodic checks only for clients without notifications. Preserve cancellation hints across claims and registration, and keep notification delivery and cancellation checks active until running jobs drain. Cover idle behavior, transient failures, overflow, and shutdown recovery. --- client.go | 8 +- internal/notifier/notifier.go | 80 +++++++------ internal/notifier/notifier_test.go | 34 +++++- .../main/java/com/riverqueue/RunHandle.java | 53 +++++++-- .../com/riverqueue/CancellationCheckTest.java | 111 ++++++++++++++++++ .../java/com/riverqueue/RunHandleTest.java | 4 +- js/driver/pg/src/driver.test.ts | 46 ++++++++ js/driver/pg/src/sql/notify.ts | 14 ++- js/src/runtime.test.ts | 93 +++++++++++++-- js/src/runtime.ts | 36 ++++-- js/src/runtime/attempt-runner.ts | 33 +++++- .../notification-pump.conformance.test.ts | 1 + js/src/runtime/notification-pump.ts | 30 ++--- producer.go | 30 +++-- producer_test.go | 57 ++++++++- ruby/lib/client_runtime.rb | 69 +++++++---- ruby/sig/runtime.rbs | 7 +- ruby/spec/client_runtime_branch_spec.rb | 15 +-- ruby/spec/notification_runtime_spec.rb | 64 ++++++++++ ruby/spec/runtime_draining_shared_examples.rb | 5 +- rust/riverqueue/src/client/attempts.rs | 7 +- rust/riverqueue/src/client/builder.rs | 2 +- rust/riverqueue/src/client/jobs.rs | 2 +- rust/riverqueue/src/client/mod.rs | 6 + rust/riverqueue/src/client/notifier.rs | 6 +- rust/riverqueue/src/client/producer.rs | 30 ++--- rust/riverqueue/src/client/run.rs | 26 ++-- rust/riverqueue/src/client/tests.rs | 82 ++++++++++++- .../tests/poll_only_cancellation.rs | 32 +---- 29 files changed, 761 insertions(+), 222 deletions(-) create mode 100644 java/river/src/test/java/com/riverqueue/CancellationCheckTest.java diff --git a/client.go b/client.go index 4cbb64edd..d3fde61f7 100644 --- a/client.go +++ b/client.go @@ -1513,8 +1513,8 @@ func (c *Client[TTx]) Driver() riverdriver.Driver[TTx] { // // If the job is currently running, it is not immediately cancelled, but is // instead marked for cancellation. The client running the job will also be -// notified to cancel the running job's context. Clients also poll for cancellation -// requests every two seconds to recover missed notifications. Although the job's +// notified to cancel the running job's context. Clients reconcile missed requests +// after reconnecting, or poll every two seconds without a notifier. Although the job's // context will be cancelled, since Go does not provide a mechanism to // interrupt a running goroutine the job will continue running until it returns. // As always, it is important for workers to respect context cancellation and @@ -1568,8 +1568,8 @@ func (c *Client[TTx]) JobCancel(ctx context.Context, jobID int64) (*rivertype.Jo // // If the job is currently running, it is not immediately cancelled, but is // instead marked for cancellation. The client running the job will also be -// notified to cancel the running job's context. Clients also poll for cancellation -// requests every two seconds to recover missed notifications. Although the job's +// notified to cancel the running job's context. Clients reconcile missed requests +// after reconnecting, or poll every two seconds without a notifier. Although the job's // context will be cancelled, since Go does not provide a mechanism to // interrupt a running goroutine the job will continue running until it returns. // As always, it is important for workers to respect context cancellation and diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index 799b9ae1e..d6b7ab175 100644 --- a/internal/notifier/notifier.go +++ b/internal/notifier/notifier.go @@ -22,21 +22,6 @@ 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 ( @@ -63,6 +48,22 @@ var NotificationTopicLongest = func() NotificationTopic { //nolint:gochecknoglob type NotifyFunc func(topic NotificationTopic, payload string) +// RecoveryFunc requests reconciliation after listening starts or notifications are +// dropped. It must not block; calls may overlap with notification delivery. +type RecoveryFunc func() + +// RecoverySubscription unregisters a recovery callback. +type RecoverySubscription struct { + notifier *Notifier +} + +// Unregister stops future callbacks; an invocation already in progress may finish. +func (s *RecoverySubscription) Unregister() { + s.notifier.mu.Lock() + defer s.notifier.mu.Unlock() + delete(s.notifier.recoveryFuncs, s) +} + type Subscription struct { notifyFunc NotifyFunc notifier *Notifier @@ -103,13 +104,13 @@ type Notifier struct { testSignals notifierTestSignals waitInterruptChan chan func() - mu sync.RWMutex - isConnected bool - isStarted bool - isWaiting bool - listenerReadyFuncs map[*ListenerReadySubscription]ListenerReadyFunc - subscriptions map[NotificationTopic][]*Subscription - waitCancel context.CancelFunc + mu sync.RWMutex + isConnected bool + isStarted bool + isWaiting bool + recoveryFuncs map[*RecoverySubscription]RecoveryFunc + subscriptions map[NotificationTopic][]*Subscription + waitCancel context.CancelFunc } func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notifier { @@ -118,19 +119,19 @@ func New(archetype *baseservice.Archetype, listener riverdriver.Listener) *Notif notificationBuf: make(chan *riverdriver.Notification, 1000), waitInterruptChan: make(chan func(), 10), - listenerReadyFuncs: make(map[*ListenerReadySubscription]ListenerReadyFunc), - subscriptions: make(map[NotificationTopic][]*Subscription), + recoveryFuncs: make(map[*RecoverySubscription]RecoveryFunc), + subscriptions: make(map[NotificationTopic][]*Subscription), }) return notifier } -// RegisterListenerReadyFunc registers a callback for initial connection and reconnects. -func (n *Notifier) RegisterListenerReadyFunc(listenerReadyFunc ListenerReadyFunc) *ListenerReadySubscription { +// RegisterRecoveryFunc registers a callback for connection recovery and buffer overflow. +func (n *Notifier) RegisterRecoveryFunc(recoveryFunc RecoveryFunc) *RecoverySubscription { n.mu.Lock() defer n.mu.Unlock() - sub := &ListenerReadySubscription{notifier: n} - n.listenerReadyFuncs[sub] = listenerReadyFunc + sub := &RecoverySubscription{notifier: n} + n.recoveryFuncs[sub] = recoveryFunc return sub } @@ -252,15 +253,7 @@ 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() - } + n.requestRecovery() drainInterrupts := func() { for { @@ -322,6 +315,18 @@ func (n *Notifier) listenerClose(ctx context.Context, skipLock bool) { n.isConnected = false } +func (n *Notifier) requestRecovery() { + recoveryFuncs := func() []RecoveryFunc { + n.mu.RLock() + defer n.mu.RUnlock() + + return maputil.Values(n.recoveryFuncs) + }() + for _, recoveryFunc := range recoveryFuncs { + recoveryFunc() + } +} + // shouldIgnoreListenerError returns true if the error is a certain type of // common error that we see when trying to close a listener. // @@ -445,6 +450,7 @@ func (n *Notifier) waitOnce(ctx context.Context) error { case n.notificationBuf <- notification: default: n.Logger.WarnContext(ctx, n.Name+": Dropping notification due to full buffer", "payload", notification.Payload) + n.requestRecovery() } } }() diff --git a/internal/notifier/notifier_test.go b/internal/notifier/notifier_test.go index 4af6927c3..20e5ed123 100644 --- a/internal/notifier/notifier_test.go +++ b/internal/notifier/notifier_test.go @@ -282,7 +282,7 @@ func TestNotifier(t *testing.T) { require.Equal(t, TopicAndPayload{testTopic1, "msg1"}, riversharedtest.WaitOrTimeout(t, notifyChan)) }) - t.Run("RegisterListenerReadyFunc", func(t *testing.T) { + t.Run("RegisterRecoveryFunc", func(t *testing.T) { t.Parallel() notifier, _ := setup(t, nil) @@ -305,7 +305,7 @@ func TestNotifier(t *testing.T) { var ready testsignal.TestSignal[struct{}] ready.Init(t) - sub := notifier.RegisterListenerReadyFunc(func() { ready.Signal(struct{}{}) }) + sub := notifier.RegisterRecoveryFunc(func() { ready.Signal(struct{}{}) }) t.Cleanup(sub.Unregister) start(t, notifier) @@ -677,6 +677,36 @@ func TestNotifier(t *testing.T) { }) } +func TestNotifier_OverflowRequestsRecovery(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var recovered testsignal.TestSignal[struct{}] + recovered.Init(t) + delivered := false + notif := New(riversharedtest.BaseServiceArchetype(t), &ListenerMock{ + waitForNotificationFunc: func(ctx context.Context) (*riverdriver.Notification, error) { + if !delivered { + delivered = true + return &riverdriver.Notification{Topic: string(NotificationTopicControl), Payload: `{"action":"cancel","job_id":1}`}, nil + } + <-ctx.Done() + return nil, ctx.Err() + }, + }) + notif.notificationBuf = make(chan *riverdriver.Notification, 1) + notif.notificationBuf <- &riverdriver.Notification{} + sub := notif.RegisterRecoveryFunc(func() { recovered.Signal(struct{}{}) }) + defer sub.Unregister() + var done testsignal.TestSignal[error] + done.Init(t) + go func() { done.Signal(notif.waitOnce(ctx)) }() + recovered.WaitOrTimeout() + cancel() + require.ErrorIs(t, done.WaitOrTimeout(), context.Canceled) +} + func TestNotifier_Start(t *testing.T) { t.Parallel() diff --git a/java/river/src/main/java/com/riverqueue/RunHandle.java b/java/river/src/main/java/com/riverqueue/RunHandle.java index 4ddfcc3be..c956ddc68 100644 --- a/java/river/src/main/java/com/riverqueue/RunHandle.java +++ b/java/river/src/main/java/com/riverqueue/RunHandle.java @@ -19,6 +19,7 @@ import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; import org.postgresql.PGConnection; @@ -41,6 +42,8 @@ public final class RunHandle implements AutoCloseable { .plusDays(1) .atStartOfDay() .toInstant(java.time.ZoneOffset.UTC); + private final AtomicLong cancellationGeneration = new AtomicLong(); + private volatile boolean cancellationPollOnly; private final Semaphore cancellationWake = new Semaphore(0); private final Client.Builder config; private final RetryPolicy defaultRetryPolicy = RetryPolicy.defaults(); @@ -79,12 +82,13 @@ public final class RunHandle implements AutoCloseable { RunHandle(Client river, Client.Builder config) { this.config = config; + cancellationPollOnly = config.pollOnly; this.river = river; id = river.id(); notifications = new Protocol.Dispatcher( id, - ignored -> cancellationWake.release(), + ignored -> requestCancellationCheck(), this::queueNotification, this::leadershipNotification); queues.putAll(config.queues); @@ -315,6 +319,7 @@ private void dispatch() { } } dispatcherStopped.countDown(); + cancellationWake.release(); } private record FetchWakeup(Instant scheduledAt, String notifyQueue) @@ -352,6 +357,7 @@ private boolean wakeDueFetches() { } private void claim(String queue, int count) { + long generation = cancellationGeneration.get(); List jobs = river.transaction( c -> @@ -397,6 +403,7 @@ private void claim(String queue, int count) { return result; })); for (var decoded : jobs) launch(decoded, config.handlers.get(decoded.job().kind())); + if (!jobs.isEmpty() && cancellationGeneration.get() != generation) cancellationWake.release(); } private void launch(Client.Decoded decoded, Registration registration) { @@ -494,6 +501,7 @@ private void launch(Job row, WorkContext context, WorkContext.Step queueLock.lock(); try { active.remove(row.id(), context); + if (!running.get() && active.isEmpty()) cancellationWake.release(); queueChanged.signalAll(); } finally { queueLock.unlock(); @@ -515,6 +523,7 @@ private void workMiddleware(WorkContext context, int index, WorkContext.Step List> claimPeers( WorkContext parent, JobType type, Client.Transaction> operation) { parent.checkRuntimeCancellation(); + long generation = cancellationGeneration.get(); var rows = river.transaction( connection -> { @@ -533,6 +542,7 @@ List> claimPeers( // Publish before checking the parent so a concurrent cancellation cannot miss both. context.attach(this, parent); } + if (!rows.isEmpty() && cancellationGeneration.get() != generation) cancellationWake.release(); var contexts = new ArrayList>(); for (var decoded : rows) { if (decoded.failure() != null) throw decoded.failure(); @@ -562,6 +572,7 @@ void completePeer(WorkContext parent, WorkContext context, Exception failu queueLock.lock(); try { active.remove(context.job().id(), context); + if (!running.get() && active.isEmpty()) cancellationWake.release(); queueChanged.signalAll(); } finally { queueLock.unlock(); @@ -729,6 +740,10 @@ private void completeBatches() { } } + private boolean hasActiveWork() { + return running.get() || dispatcherStopped.getCount() > 0 || !active.isEmpty(); + } + private void heartbeat(String queue) { river.transaction( c -> { @@ -762,9 +777,13 @@ private void listen() { listenSqlite(); return; } - while (running.get()) { + while (hasActiveWork()) { try (var connection = river.database().connection()) { - if (!river.database().supportsNotifications(connection)) return; + if (!river.database().supportsNotifications(connection)) { + cancellationPollOnly = true; + cancellationWake.release(); + return; + } listener = connection; String schema = river.database().schema(); if (schema.isEmpty()) @@ -779,10 +798,11 @@ private void listen() { statement.execute("LISTEN \"" + schema.replace("\"", "\"\"") + "." + topic + "\""); } wake.release(); + requestCancellationCheck(); observe("listen_ready"); var postgres = connection.unwrap(PGConnection.class); try { - while (running.get()) { + while (hasActiveWork()) { var notifications = postgres.getNotifications(25); if (notifications != null) for (var notification : notifications) @@ -797,7 +817,7 @@ private void listen() { } } } catch (Exception e) { - if (running.get()) { + if (hasActiveWork()) { report(e); try { TimeUnit.MILLISECONDS.sleep(100); @@ -824,7 +844,7 @@ private void deliver(String topic, String payload) { private void listenSqlite() { long after = 0; boolean initialized = false; - while (running.get()) { + while (hasActiveWork()) { try { if (!initialized) { after = @@ -841,6 +861,7 @@ private void listenSqlite() { // The initial cursor skips earlier notifications, including inserts since the first // fetch. Poll again once listening starts, as the Postgres listener does. wake.release(); + requestCancellationCheck(); observe("listen_ready"); } final long cursor = after; @@ -1402,9 +1423,14 @@ 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()) { + boolean retry = false; + while (hasActiveWork()) { try { + if (cancellationPollOnly || retry) cancellationWake.tryAcquire(50, TimeUnit.MILLISECONDS); + else cancellationWake.acquire(); cancellationWake.drainPermits(); + if (!hasActiveWork()) return; + retry = false; var attempts = Map.copyOf(active); if (!attempts.isEmpty()) { river.transaction( @@ -1422,14 +1448,12 @@ private void pollCancellations() { return null; }); } - } catch (Exception e) { - report(e); - } - try { - cancellationWake.tryAcquire(50, TimeUnit.MILLISECONDS); } catch (InterruptedException e) { Thread.currentThread().interrupt(); return; + } catch (Exception e) { + retry = true; + report(e); } } } @@ -1454,6 +1478,11 @@ private void report(Throwable error) { } } + private void requestCancellationCheck() { + cancellationGeneration.incrementAndGet(); + cancellationWake.release(); + } + private Duration retryDelay(Job job, Instant now) { try { var delay = Objects.requireNonNull(config.retryPolicy.delay(job), "Retry delay"); diff --git a/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java b/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java new file mode 100644 index 000000000..e42814197 --- /dev/null +++ b/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java @@ -0,0 +1,111 @@ +package com.riverqueue; + +import static org.junit.jupiter.api.Assertions.*; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.SQLException; +import java.time.Duration; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import javax.sql.DataSource; +import org.junit.jupiter.api.extension.RegisterExtension; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +class CancellationCheckTest { + @RegisterExtension final TestDatabase databases = new TestDatabase(); + @TempDir Path directory; + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + @SuppressWarnings("try") // Resources intentionally scope background activity. + void checksOnlyOnHintsAndRetriesFailures(boolean failFirst) throws Exception { + var database = databases.open(directory.resolve("river.db")); + var fail = new AtomicBoolean(); + var queries = new AtomicInteger(); + var checked = new LinkedBlockingQueue(); + var source = + (DataSource) + Proxy.newProxyInstance( + DataSource.class.getClassLoader(), + new Class[] {DataSource.class}, + (proxy, method, args) -> { + if (!method.getName().equals("getConnection")) + throw new UnsupportedOperationException(); + var connection = database.connection(); + return Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + (ignored, operation, arguments) -> { + if (operation.getName().equals("prepareStatement") + && arguments[0].equals(Sql.query(database, "cancel_requested"))) { + checked.add(queries.incrementAndGet()); + if (fail.compareAndSet(true, false)) + throw new SQLException("temporary cancellation read failure"); + } + try { + return operation.invoke(connection, arguments); + } catch (InvocationTargetException error) { + throw error.getCause(); + } + }); + }); + var ready = new CompletableFuture(); + var entered = new CompletableFuture(); + var finish = new CompletableFuture(); + var errors = new LinkedBlockingQueue(); + var type = JobType.of("cancellation_check", String.class); + var client = + Client.builder(new Database(source, database.dialect()).withSchema(database.schema())) + .queue("default", 1) + .leadership(false) + .pollInterval(Duration.ofMillis(10)) + .stopTimeout(Duration.ofSeconds(5)) + .errorHandler(errors::add) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .worker( + type, + context -> { + entered.complete(null); + finish.join(); + context.checkCancelled(); + }) + .build(); + try (var run = client.start()) { + try { + ready.get(5, TimeUnit.SECONDS); + var job = client.insert(type, "working").job(); + entered.get(5, TimeUnit.SECONDS); + // Let the startup hint settle, then leave a running attempt idle for several old ticks. + TimeUnit.MILLISECONDS.sleep(200); + int before = queries.get(); + checked.clear(); + TimeUnit.MILLISECONDS.sleep(200); + assertEquals(before, queries.get()); + fail.set(failFirst); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(job.id(), "default")); + return null; + }); + assertNotNull(checked.poll(5, TimeUnit.SECONDS)); + if (failFirst) assertNotNull(checked.poll(5, TimeUnit.SECONDS)); + TimeUnit.MILLISECONDS.sleep(200); + assertEquals(before + (failFirst ? 2 : 1), queries.get()); + assertEquals(failFirst ? 1 : 0, errors.size()); + } finally { + finish.complete(null); + } + } + } +} diff --git a/java/river/src/test/java/com/riverqueue/RunHandleTest.java b/java/river/src/test/java/com/riverqueue/RunHandleTest.java index f5b650d39..7ae5b42b3 100644 --- a/java/river/src/test/java/com/riverqueue/RunHandleTest.java +++ b/java/river/src/test/java/com/riverqueue/RunHandleTest.java @@ -92,8 +92,8 @@ void cancellationContinuesAfterGracefulStop(boolean pollOnly) throws Exception { try (var subscription = workers.client().subscribe(outcome::complete)) { var inserted = client.insert(TYPE, "cancel while draining").job(); entered.get(5, TimeUnit.SECONDS); - // The timeout leaves active work draining after the dispatcher and - // listener have stopped. A peer can still cancel that work. + // The timeout leaves active work draining after dispatch stops. + // The listener must remain available to deliver peer cancellation. assertThrows(RiverException.class, workers::stop); client.cancel(inserted.id()); assertEquals(Job.State.CANCELLED, outcome.get(5, TimeUnit.SECONDS).job().state()); diff --git a/js/driver/pg/src/driver.test.ts b/js/driver/pg/src/driver.test.ts index ec5d186e8..011236b12 100644 --- a/js/driver/pg/src/driver.test.ts +++ b/js/driver/pg/src/driver.test.ts @@ -834,6 +834,52 @@ describe("PgDriver", () => { }); }); + it("fails an overflowing notification stream so its subscriber can recover", async () => { + const leased = mockLeasedPgClient(); + leased.query.mockResolvedValue({ + command: "LISTEN", + fields: [], + oid: 0, + rowCount: 0, + rows: [], + }); + const pool = { + ...mockPgClient(), + connect: vi.fn().mockResolvedValue(leased), + idleCount: 1, + options: { max: 4 }, + totalCount: 1, + } as unknown as Pool; + const driver = testPgDriver(pool, { schema: "river" }); + const abort = new AbortController(); + const ready = Promise.withResolvers(); + const iterator = driver.listen(["river_control"], abort.signal, () => + ready.resolve(undefined) + ); + try { + const first = iterator.next(); + await ready.promise; + leased.emit("notification", { + channel: "river.river_control", + payload: "first", + }); + await first; + for (let id = 0; id < 1_025; id++) { + leased.emit("notification", { + channel: "river.river_control", + payload: JSON.stringify({ action: "cancel", job_id: id }), + }); + } + await expect(iterator.next()).rejects.toThrow( + "notification buffer overflow" + ); + expect(leased.release).toHaveBeenCalledWith(true); + } finally { + abort.abort(); + await iterator.return(undefined); + } + }); + it("pings an idle listener and fails when a ping never answers", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); try { diff --git a/js/driver/pg/src/sql/notify.ts b/js/driver/pg/src/sql/notify.ts index 8f295b946..ac966c5d5 100644 --- a/js/driver/pg/src/sql/notify.ts +++ b/js/driver/pg/src/sql/notify.ts @@ -15,7 +15,7 @@ import { PG_EXACT_TYPES } from "../exact-types.js"; import { abortablePromise, PgClientLease } from "../lease.js"; import type { PgNotification, PgOperationOptions } from "../types.js"; -/** Notifications buffered per connection; beyond it the oldest is dropped. */ +/** Notifications buffered per connection; overflow forces resubscription and recovery. */ const NOTIFICATION_QUEUE_CAPACITY = 1_024; /** Idle time after which a LISTEN connection is pinged, like River's Go notifier. */ @@ -282,7 +282,7 @@ class NotificationQueue { } push(value: PgNotification): void { - if (this.#signal.aborted) return; + if (this.#signal.aborted || this.#failure !== undefined) return; if (this.#waiting !== undefined) { this.#waiting.resolve(value); this.#waiting = undefined; @@ -291,10 +291,12 @@ class NotificationQueue { const key = `${value.topic}\0${value.payload}`; if (this.#keys.has(key)) return; if (this.#values.length - this.#head >= NOTIFICATION_QUEUE_CAPACITY) { - const dropped = this.#values[this.#head++]; - if (dropped !== undefined) { - this.#keys.delete(`${dropped.topic}\0${dropped.payload}`); - } + this.fail( + new Error( + "River notification buffer overflow; resubscribe to recover missed notifications" + ) + ); + return; } this.#keys.add(key); this.#values.push(value); diff --git a/js/src/runtime.test.ts b/js/src/runtime.test.ts index 781973d07..7ddfbb19b 100644 --- a/js/src/runtime.test.ts +++ b/js/src/runtime.test.ts @@ -744,8 +744,8 @@ describe("Client runtime", () => { const working = live.size; await run.stop(); - // Only the queue loop's link, on the claim signal and its generation's. - expect(working).toBe(2); + // Only the queue loop's link and the runtime's cancellation link remain. + expect(working).toBe(4); expect(live.size).toBe(0); // Nothing per attempt or completion uses `AbortSignal.any`. expect(any).not.toHaveBeenCalled(); @@ -4606,6 +4606,63 @@ describe("Client runtime database faults", () => { await run.stop(); }); + for (const failFirst of [false, true]) { + it(`only checks cancellations on hints and retries failures=${failFirst}`, async () => { + const driver = new FakeRuntimeDriver(); + const deliver = Promise.withResolvers(); + const started = Promise.withResolvers(); + const finish = Promise.withResolvers(); + let checks = 0; + Object.assign(driver, { + jobGetCancelRequested: () => { + checks++; + if (failFirst && checks === 1) + throw new Error("temporary read failure"); + return []; + }, + }); + driver.claim = [fakeJob()]; + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + await deliver.promise; + yield { + topic: "control", + payload: '{"action":"cancel","job_id":101,"queue":"default"}', + }; + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + logger: recordingLogger([]), + queues: { default: { maxWorkers: 1 } }, + workers: new Workers().add(defineJob({ kind: "test" }), async () => { + started.resolve(undefined); + await finish.promise; + }), + }); + const timer = new FakeTimer(); + overrideRuntimeTiming(client, { random: () => 0.5, timer }); + const run = await client.start(); + try { + await started.promise; + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(checks).toBe(0); + deliver.resolve(undefined); + await waitUntil(() => checks === (failFirst ? 2 : 1)); + await new Promise((resolve) => setTimeout(resolve, 30)); + expect(checks).toBe(failFirst ? 2 : 1); + } finally { + deliver.resolve(undefined); + finish.resolve(undefined); + await run.stop(); + } + }); + } + for (const stream of ["runtime", "remote", "local"] as const) { for (const duringClaim of [false, true]) { it(`revalidates a delayed ${stream} cancellation ${duringClaim ? "during a retry claim" : "against a running retry"}`, async () => { @@ -4703,6 +4760,21 @@ describe("Client runtime database faults", () => { for (const phase of ["running", "claim", "draining"] as const) { it(`recovers a lost cancellation while ${phase}, using ${batched ? "batches" : "rows"}`, async () => { const driver = new FakeRuntimeDriver(); + const reconnect = Promise.withResolvers(); + const reconnected = Promise.withResolvers(); + let subscriptions = 0; + // eslint-disable-next-line require-yield -- a lost notice is recovered after resubscription + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + if (subscriptions++ === 0) { + await reconnect.promise; + throw new Error("listener disconnected"); + } + reconnected.resolve(undefined); + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; const claim = Promise.withResolvers(); const started = Promise.withResolvers(); let marked = false; @@ -4762,21 +4834,25 @@ describe("Client runtime database faults", () => { try { if (phase === "claim") { await waitUntil(() => driver.claimRequests.length > 0); - marked = true; - claim.resolve(undefined); + } else { + await started.promise; } - await started.promise; const stopping = phase === "draining" ? run.stop() : undefined; marked = true; + reconnect.resolve(undefined); + await reconnected.promise; + claim.resolve(undefined); await waitUntil(() => driver.completions.length === 1); - expect(handlerSignal?.reason).toMatchObject({ - name: "JobCancelledError", - }); + if (phase !== "claim") + expect(handlerSignal?.reason).toMatchObject({ + name: "JobCancelledError", + }); expect(driver.completions[0]?.kind).toBe("cancel"); expect(failures).toBeGreaterThanOrEqual(2); await stopping; await run.stop(); } finally { + reconnect.resolve(undefined); claim.resolve(undefined); await run.stop({ mode: "cancel" }); } @@ -4820,6 +4896,7 @@ describe("Client runtime database faults", () => { driver.claim = [fakeJob()]; const client = new Client(driver, { clientId: "runtime-test", + pollOnly: true, completionBatchSize: 1, leaderElectionDisabled: true, queues: { diff --git a/js/src/runtime.ts b/js/src/runtime.ts index a8ce8d425..fed982e24 100644 --- a/js/src/runtime.ts +++ b/js/src/runtime.ts @@ -3,7 +3,7 @@ import type { QueueRow, RuntimeDriver } from "./driver.js"; import { LifecycleError, ValidationError } from "./errors.js"; import type { RiverEvent } from "./events.js"; import type { RiverHooks } from "./extensions.js"; -import { unrefTimeout } from "./internal/abort.js"; +import { LinkedAbortSignal, unrefTimeout } from "./internal/abort.js"; import { SYSTEM_TIMER, type RuntimeTimer } from "./internal/backoff.js"; import { EventLoopDelayMonitor } from "./internal/event-loop-delay-monitor.js"; import type { JobRow } from "./job.js"; @@ -211,10 +211,11 @@ export class RuntimeController { readonly ready: Promise; /** - * Ends polling for cancellations once producers drained, so a job can + * Ends notification delivery and cancellation checks once producers drain, so a job can * still be cancelled while its queue drains. */ - readonly #cancellationPollAbort = new AbortController(); + readonly #cancellationAbort = new AbortController(); + readonly #cancellationLink: LinkedAbortSignal; /** How often the runtime checks for durable cancellation requests. */ readonly #cancellationPollIntervalMs: number; readonly #claimAbort = new AbortController(); @@ -328,7 +329,14 @@ export class RuntimeController { transformJobArgs: (row) => this.#runner.transformJobArgs(row), workerRetryPolicy: (kind) => this.#runner.workerRetryPolicy(kind), }); + this.#cancellationLink = new LinkedAbortSignal([ + this.#cancellationAbort.signal, + this.#runAbort.signal, + ]); + const cancellationSignal = this.#cancellationLink.signal; this.#runner = new AttemptRunner(context, { + cancellationSignal, + cancellationRetryIntervalMs: settings.queueControlPollIntervalMs, completions: this.#completions, errorHandler: settings.errorHandler, hooks: settings.hooks, @@ -401,6 +409,7 @@ export class RuntimeController { ? [] : serviceList(binding.pilot.services(), "services()"); this.#notifications = new NotificationPump(context, { + signal: cancellationSignal, leaderResigned: () => this.#services?.leaderResigned(), producer: this.#producer, resignLeadership: () => this.#services?.resignLeadership(), @@ -618,10 +627,11 @@ export class RuntimeController { new LifecycleError("River runtime is stopping") ); await this.#producer.drainAll(); - this.#cancellationPollAbort.abort( + this.#cancellationAbort.abort( new LifecycleError("River runtime is stopping") ); await Promise.allSettled(this.#tasks); + this.#cancellationLink[Symbol.dispose](); this.#servicesAbort.abort(new LifecycleError("River runtime is stopping")); const [completionClose] = await Promise.allSettled([ this.#completions.close(), @@ -664,15 +674,17 @@ export class RuntimeController { } async #initialize(): Promise { - // Notifications deliver cancellations promptly. Keep a durable fallback - // until every producer drains, including during reconnects and shutdown. + // Only clients without notifications poll. Notification hints and reconnects + // request checks independently, until every producer has drained. const listen = await this.#listens(); - this.#trackTask( - this.#runner.pollCancellations( - this.#cancellationPollIntervalMs, - this.#cancellationPollAbort.signal - ) - ); + if (!listen) { + this.#trackTask( + this.#runner.pollCancellations( + this.#cancellationPollIntervalMs, + this.#cancellationAbort.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 99190a97a..214ada7b3 100644 --- a/js/src/runtime/attempt-runner.ts +++ b/js/src/runtime/attempt-runner.ts @@ -75,6 +75,8 @@ import { /** Configuration for an {@link AttemptRunner}. */ export interface AttemptRunnerOptions { + readonly cancellationSignal: AbortSignal; + readonly cancellationRetryIntervalMs: number; readonly completions: CompletionPipeline; readonly errorHandler: RiverErrorHandler | undefined; /** Client-level hooks, plugin hooks first. */ @@ -114,6 +116,9 @@ export class AttemptRunner { >(); #cancellationCheck: Promise | undefined; #cancellationCheckPending = false; + #cancellationGeneration = 0; + readonly #cancellationRetryIntervalMs: number; + readonly #cancellationSignal: AbortSignal; /** * For each claim in flight, the `attemptKey`s of cancellations that * need a durable check once the claim registers its attempts. @@ -136,6 +141,8 @@ export class AttemptRunner { readonly #workers: Workers; constructor(context: RuntimeContext, options: AttemptRunnerOptions) { + this.#cancellationRetryIntervalMs = options.cancellationRetryIntervalMs; + this.#cancellationSignal = options.cancellationSignal; this.#completions = options.completions; this.#context = context; this.#errorHandler = options.errorHandler; @@ -228,7 +235,7 @@ export class AttemptRunner { return; } try { - await this.recoverCancellations(signal); + await this.#checkCancellations(); } catch (error: unknown) { if (signal.aborted) return; this.#context.logger.error( @@ -239,6 +246,12 @@ export class AttemptRunner { } } + /** Reconcile a reconnect or overflow, including claims not yet registered. */ + requestCancellationCheck(): void { + this.#cancellationGeneration++; + void this.#checkCancellations(); + } + /** Cancel this process's attempts of a job just cancelled through this client. */ cancelLocal(job: JobRow): void { const attempt = this.#attempts.get( @@ -514,10 +527,13 @@ export class AttemptRunner { cancelled(id: bigint): boolean; end(): void; } { + const generation = this.#cancellationGeneration; const pending = new Set(); this.#claimCancellations.add(pending); return { - cancelled: (id) => pending.has(attemptKey(id, this.#context.clientId)), + cancelled: (id) => + generation !== this.#cancellationGeneration || + pending.has(attemptKey(id, this.#context.clientId)), end: () => { this.#claimCancellations.delete(pending); }, @@ -926,13 +942,22 @@ export class AttemptRunner { while (this.#cancellationCheckPending) { this.#cancellationCheckPending = false; try { - await this.recoverCancellations(this.#runSignal); + await this.recoverCancellations(this.#cancellationSignal); } catch (error: unknown) { - if (this.#runSignal.aborted) return; + if (this.#cancellationSignal.aborted) return; this.#context.logger.warn( "River could not confirm a cancellation notification", { error: describeError(error) } ); + try { + await this.#context.timer.delay( + this.#cancellationRetryIntervalMs, + this.#cancellationSignal + ); + } catch { + return; + } + this.#cancellationCheckPending = true; } } } finally { diff --git a/js/src/runtime/notification-pump.conformance.test.ts b/js/src/runtime/notification-pump.conformance.test.ts index 05b7c2c88..a1f9135a9 100644 --- a/js/src/runtime/notification-pump.conformance.test.ts +++ b/js/src/runtime/notification-pump.conformance.test.ts @@ -111,6 +111,7 @@ describe.each(["postgres", "sqlite"] as const)( }, } as unknown as AttemptRunner; const pump = new NotificationPump(context, { + signal: context.claimSignal, leaderResigned: () => { effects.push(["leaderResigned"]); }, diff --git a/js/src/runtime/notification-pump.ts b/js/src/runtime/notification-pump.ts index eae6930b1..ddbd12af7 100644 --- a/js/src/runtime/notification-pump.ts +++ b/js/src/runtime/notification-pump.ts @@ -2,14 +2,13 @@ * The notification pump: supervises the backend's notification streams, * waking queues on inserts, applying queue controls and cancellations, and * forwarding leadership resignation requests. Notifications are hints; - * polling remains the durable path. + * reconnects request durable cancellation checks. */ import type { RuntimeNotification } from "../driver.js"; import { LifecycleError } from "../errors.js"; import type { AttemptRunner } from "./attempt-runner.js"; import type { RuntimeContext } from "./context.js"; import { backOffAfterFailure } from "./context.js"; -import { describeError } from "./failures.js"; import { notificationCancellation, notificationLeaderResigned, @@ -20,6 +19,8 @@ import type { QueueProducer } from "./queue-producer.js"; /** Configuration for a {@link NotificationPump}. */ export interface NotificationPumpOptions { + /** Remains live until every producer has drained. */ + readonly signal: AbortSignal; /** Another client resigned leadership. */ readonly leaderResigned: () => void; readonly producer: QueueProducer; @@ -35,6 +36,7 @@ export class NotificationPump { readonly #producer: QueueProducer; readonly #resignLeadership: () => PromiseLike | undefined; readonly #runner: AttemptRunner; + readonly #signal: AbortSignal; constructor(context: RuntimeContext, options: NotificationPumpOptions) { this.#context = context; @@ -42,6 +44,7 @@ export class NotificationPump { this.#producer = options.producer; this.#resignLeadership = options.resignLeadership; this.#runner = options.runner; + this.#signal = options.signal; } /** @@ -147,24 +150,7 @@ export class NotificationPump { #recoverMissedNotifications(): void { this.#producer.wakeAll(); this.#producer.wakeControl(); - this.#recoverMissedCancellations(); - } - - /** - * Cancel attempts whose cancellation notice was lost while a stream - * reconnected. A failed check is logged; the attempt then finishes as if - * its notice never arrived, as it would before this recovery. - */ - #recoverMissedCancellations(): void { - const recovery = this.#runner - .recoverCancellations(this.#context.claimSignal) - .catch((error: unknown) => { - this.#context.logger.warn( - "River could not check running jobs for missed cancellations", - { error: describeError(error) } - ); - }); - this.#context.trackTask(recovery); + this.#runner.requestCancellationCheck(); } async #remoteCancellationLoop(): Promise { @@ -181,7 +167,7 @@ export class NotificationPump { }, () => { if (subscribed) { - this.#recoverMissedCancellations(); + this.#runner.requestCancellationCheck(); } else { subscribed = true; } @@ -202,7 +188,7 @@ export class NotificationPump { connected: () => void, resubscribes: () => boolean = () => true ): Promise { - const signal = this.#context.claimSignal; + const signal = this.#signal; let failures = 0; while (!signal.aborted) { let failure: unknown; diff --git a/producer.go b/producer.go index 070aa5f17..82d6650f5 100644 --- a/producer.go +++ b/producer.go @@ -112,8 +112,8 @@ type producerConfig struct { QueueEventCallback func(event *Event) // QueuePollInterval is the amount of time between periodic checks for - // queue setting changes and job cancellation requests. This is only used in - // poll-only mode (when no notifier is provided). + // queue setting changes and job cancellation requests in poll-only mode. + // It also sets the retry interval for failed cancellation checks. QueuePollInterval time.Duration // QueueReportInterval is the amount of time between periodic reports // of the queue status. @@ -389,7 +389,7 @@ drainControls: var ( controlSub *notifier.Subscription insertSub *notifier.Subscription - reconnectSub *notifier.ListenerReadySubscription + reconnectSub *notifier.RecoverySubscription ) if p.config.Notifier != nil { var err error @@ -425,7 +425,7 @@ drainControls: } // Recover cancellations whose NOTIFY was lost while disconnected. - reconnectSub = p.config.Notifier.RegisterListenerReadyFunc(p.TriggerCancelCheck) + reconnectSub = p.config.Notifier.RegisterRecoveryFunc(p.TriggerCancelCheck) } subroutineCtx, cancelSubroutines := context.WithCancelCause(context.WithoutCancel(fetchCtx)) @@ -492,7 +492,7 @@ drainControls: return nil } -// TriggerCancelCheck wakes the cancellation poller without blocking notification delivery. +// TriggerCancelCheck requests a cancellation check without blocking notification delivery. func (p *producer) TriggerCancelCheck() { p.cancelGeneration.Add(1) select { @@ -1046,21 +1046,29 @@ func (p *producer) handleWorkerDone(job *rivertype.JobRow) { p.jobResultCh <- job } -// Cancellation polling is independent of fetching: saturated and paused producers -// must still cancel jobs, including while draining during graceful shutdown. +// Cancellation checks run independently of fetching and remain live while draining. +// With notifications, the timer only retries failures; otherwise it polls regularly. func (p *producer) pollForJobCancellations(ctx context.Context, wg *sync.WaitGroup) { defer wg.Done() - ticker := time.NewTicker(p.config.QueuePollInterval) - defer ticker.Stop() + timer := time.NewTimer(p.config.QueuePollInterval) + defer timer.Stop() + if p.config.Notifier != nil { + timer.Stop() + } for { select { case <-ctx.Done(): return - case <-ticker.C: + case <-timer.C: case <-p.cancelPollCh: } - if err := p.pollForJobCancellationsOnce(ctx); err != nil { + timer.Stop() + err := p.pollForJobCancellationsOnce(ctx) + if p.config.Notifier == nil || err != nil { + timer.Reset(p.config.QueuePollInterval) + } + if err != nil { if !errors.Is(context.Cause(ctx), startstop.ErrStop) { p.Logger.ErrorContext(ctx, p.Name+": Error fetching job cancellation requests", slog.String("err", err.Error())) } diff --git a/producer_test.go b/producer_test.go index 46e920fe8..2424a9251 100644 --- a/producer_test.go +++ b/producer_test.go @@ -9,6 +9,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/require" @@ -216,7 +217,7 @@ func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { prod := baseservice.Init(archetype, &producer{ activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, cancelPollCh: make(chan struct{}, 1), - config: &producerConfig{MaxWorkers: 2, QueuePollInterval: 20 * time.Millisecond}, + config: &producerConfig{MaxWorkers: 2, Notifier: ¬ifier.Notifier{}, QueuePollInterval: 20 * time.Millisecond}, exec: exec, fetchLimiter: chanutil.NewDebouncedChan(ctx, time.Millisecond, true), jobResultCh: make(chan *rivertype.JobRow), @@ -274,6 +275,51 @@ func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { } } +func TestProducer_CancellationChecksWaitForHints(t *testing.T) { + t.Parallel() + + for _, failFirst := range []bool{false, true} { + t.Run(fmt.Sprintf("Retry=%v", failFirst), func(t *testing.T) { + t.Parallel() + + synctest.Test(t, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var queries atomic.Int32 + prod := baseservice.Init(riversharedtest.BaseServiceArchetype(t), &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: {}}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{Notifier: ¬ifier.Notifier{}, QueuePollInterval: time.Second}, + exec: &producerTestExecutor{jobGetCancelRequestedFunc: func(context.Context, *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + query := queries.Add(1) + if failFirst && query == 1 { + return nil, errors.New("temporary failure") + } + return nil, nil + }}, + }) + var wg sync.WaitGroup + wg.Add(1) + go prod.pollForJobCancellations(ctx, &wg) + time.Sleep(10 * time.Second) + require.Zero(t, queries.Load()) + + prod.TriggerCancelCheck() + synctest.Wait() + require.EqualValues(t, 1, queries.Load()) + time.Sleep(10 * time.Second) + want := 1 + if failFirst { + want = 2 + } + require.EqualValues(t, want, queries.Load()) + cancel() + wg.Wait() + }) + }) + } +} + func TestProducer_HandleControlNotification(t *testing.T) { t.Parallel() @@ -373,7 +419,7 @@ func TestProducer_JobCancelSurvivesNotifierReconnect(t *testing.T) { notif := notifier.New(archetype, mockListener) var initiallyReady testsignal.TestSignal[struct{}] initiallyReady.Init(t) - readySub := notif.RegisterListenerReadyFunc(func() { initiallyReady.Signal(struct{}{}) }) + readySub := notif.RegisterRecoveryFunc(func() { initiallyReady.Signal(struct{}{}) }) require.NoError(t, notif.Start(ctx)) initiallyReady.WaitOrTimeout() readySub.Unregister() @@ -1100,7 +1146,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin } }) - t.Run("CancellationPolling", func(t *testing.T) { + t.Run("CancellationWhilePausedSaturatedOrStopping", func(t *testing.T) { t.Parallel() for _, state := range []string{"Paused", "Running", "Stopping"} { @@ -1151,11 +1197,12 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin } // Use the executor directly so no local client shortcut can cancel - // the worker. No more jobs can be fetched in any of these states. + // the worker. Notifications or poll-only checks must work even though + // no more jobs can be fetched in any of these states. _, err := bundle.exec.JobCancel(ctx, &riverdriver.JobCancelParams{ ID: jobID, CancelAttemptedAt: time.Now().UTC(), - ControlTopic: "unsubscribed_control", + ControlTopic: string(notifier.NotificationTopicControl), Schema: producer.config.Schema, }) require.NoError(t, err) diff --git a/ruby/lib/client_runtime.rb b/ruby/lib/client_runtime.rb index 9ada74962..9092499fe 100644 --- a/ruby/lib/client_runtime.rb +++ b/ruby/lib/client_runtime.rb @@ -24,6 +24,8 @@ def initialize(started_at, started_monotonic, scheduled_at) attr_reader :periodic_jobs def initialize(client, driver, config) + @cancellation_generation = 0 + @cancellation_pending = false @client = client @condition = ConditionVariable.new @config = config @@ -33,7 +35,7 @@ def initialize(client, driver, config) @maintenance_condition = ConditionVariable.new @maintenance_generation = 0 @mutex = Mutex.new - @notifications_disabled = false + @notifications_disabled = config.poll_only || !driver.respond_to?(:notification_listener) @producer_threads = {} @queue_configs = config.queues.dup @queue_generations = Hash.new(0) @@ -282,6 +284,7 @@ def wake check_remote_cancellations(entries: {id => entry}) rescue => error @config.logger.error("River cancellation check failed: #{error.full_message}") + request_cancellation_check end end @@ -301,34 +304,37 @@ def wake end private def cancellation_loop + retrying = false loop do + @mutex.synchronize do + interval = @queue_configs.values.map { |config| config.resolved_fetch_poll_interval(@config) }.min || @config.fetch_poll_interval + deadline = (@config.poll_only || @notifications_disabled || retrying) ? monotonic_now + interval : nil + loop do + return if @stop_requested && @running.empty? + break if @cancellation_pending + + deadline ||= monotonic_now + interval if @notifications_disabled + remaining = deadline && deadline - monotonic_now + break if remaining && remaining <= 0 + + @condition.wait(@mutex, remaining) + end + @cancellation_pending = false + end + retrying = false begin check_remote_cancellations rescue => error + retrying = true @config.logger.error("River cancellation check failed: #{error.full_message}") end - - 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 + # Include registered attempts whose worker has not started yet: a reconnect + # can recover their cancellation between registration and begin_work. + entries ||= @mutex.synchronize { @running.dup } return if entries.empty? @driver.job_get_cancelled_ids(entries.keys).each do |id| @@ -558,7 +564,7 @@ def wake operation.call end - private def launch(row) + private def launch(row, generation = nil) gate = ::Queue.new thread = Thread.new do gate.pop @@ -571,7 +577,7 @@ def wake end @mutex.synchronize do - @running[row.id] = {check_before_work: @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) || (generation && generation != @cancellation_generation), queue: row.queue, thread: thread, working: false} end gate.push(true) end @@ -681,6 +687,7 @@ def wake unless listener listener = @driver.notification_listener @mutex.synchronize { @notifications_disabled = listener.nil? } + request_cancellation_check ready << true if ready ready = nil return unless listener @@ -765,11 +772,14 @@ def wake queue_state = @driver.queue_get(queue) publish_queue(queue_state.paused_at ? EVENT_QUEUE_PAUSED : EVENT_QUEUE_RESUMED, queue_state) if queue_state if capacity.positive? && !queue_state&.paused_at - @mutex.synchronize { @fetching_queues[queue] = [] } + generation = @mutex.synchronize do + @fetching_queues[queue] = [] + @cancellation_generation + end begin jobs = @driver.job_get_available(attempted_by: @config.id, max: capacity, queue: queue, **fetch_options) last_fetch = monotonic_now - jobs.each { |job| launch(job) } + jobs.each { |job| launch(job, generation) } next unless jobs.empty? ensure @mutex.synchronize { @fetching_queues.delete(queue) } @@ -803,6 +813,14 @@ def wake @mutex.synchronize { @subscriptions.dup }.each { |subscription| subscription.publish(event) } end + private def request_cancellation_check + @mutex.synchronize do + @cancellation_generation += 1 + @cancellation_pending = true + @condition.broadcast + end + end + private def queue_stopping?(queue) @mutex.synchronize { @stop_requested || @removed_queues[queue] } end @@ -823,11 +841,12 @@ def wake return unless message["job_id"].is_a?(Integer) @mutex.synchronize do + @cancellation_generation += 1 + @cancellation_pending = true + @condition.broadcast entry = @running[message["job_id"]] if entry && entry[:queue] == queue 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/sig/runtime.rbs b/ruby/sig/runtime.rbs index 1f2b8d73c..ba7212e48 100644 --- a/ruby/sig/runtime.rbs +++ b/ruby/sig/runtime.rbs @@ -360,6 +360,8 @@ module River def initialize: (Time, Float, Time) -> void end + @cancellation_generation: Integer + @cancellation_pending: bool @cancellation_thread: Thread? @client: Client @condition: ConditionVariable @@ -408,7 +410,7 @@ module River def begin_work: (Integer) -> untyped def cancellation_loop: () -> untyped - def check_remote_cancellations: () -> untyped + def check_remote_cancellations: (?entries: Hash[Integer, untyped]?) -> untyped def error_handler_cancel?: (Exception, Job) -> bool def execute: (JobRow) -> untyped def finish_failed: (JobRow, Job, Exception, JobTiming, ?cancelled: bool, ?worker: untyped) -> untyped @@ -418,7 +420,7 @@ module River def finish_work: (Integer) -> untyped def invoke_plugins: (Symbol, *untyped) -> Array[untyped] def invoke_worker: (untyped, Job) -> untyped - def launch: (JobRow) -> untyped + def launch: (JobRow, ?Integer?) -> untyped def maintenance_loop: () -> untyped def monotonic_now: () -> Float def next_retry: (JobRow, Exception, Time, ?worker: untyped) -> Time @@ -430,6 +432,7 @@ module River def periodic_jobs_changed: () -> untyped def producer_loop: (String, QueueConfig) -> untyped def publish: (Symbol, JobRow, JobTiming) -> untyped + def request_cancellation_check: () -> void def queue_stopping?: (String) -> bool def receive_notification: (untyped, untyped) -> untyped def remove_subscription: (Subscription) -> Subscription? diff --git a/ruby/spec/client_runtime_branch_spec.rb b/ruby/spec/client_runtime_branch_spec.rb index 39eaefcfe..3cea6e3ad 100644 --- a/ruby/spec/client_runtime_branch_spec.rb +++ b/ruby/spec/client_runtime_branch_spec.rb @@ -445,7 +445,7 @@ def execute(runtime, value) end value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { now } - value.define_singleton_method(:launch) { |_| @stop_requested = true if fetched_at.length == 2 } + value.define_singleton_method(:launch) { |_, _generation| @stop_requested = true if fetched_at.length == 2 } value.define_singleton_method(:wait) do |duration| waits << duration now += (waits.length == 1) ? 0.25 : duration @@ -469,7 +469,7 @@ def execute(runtime, value) end value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { now } - value.define_singleton_method(:launch) { |_| } + value.define_singleton_method(:launch) { |_, _generation| } value.define_singleton_method(:wait) { |_| @stop_requested = true } value.send(:producer_loop, "branch", River::QueueConfig.new(max_workers: 1, fetch_cooldown: 1, fetch_poll_interval: 1)) @@ -488,7 +488,7 @@ def execute(runtime, value) } value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { 100.0 } - value.define_singleton_method(:launch) do |job| + value.define_singleton_method(:launch) do |job, _generation| @running[job.id] = {queue: "branch", thread: Thread.current, working: true} end value.define_singleton_method(:wait) { |_| @stop_requested = true } @@ -528,7 +528,7 @@ def execute(runtime, value) end value = runtime(driver: driver) value.define_singleton_method(:monotonic_now) { now } - value.define_singleton_method(:launch) { |_| @stop_requested = true if fetched_at.length == 2 } + value.define_singleton_method(:launch) { |_, _generation| @stop_requested = true if fetched_at.length == 2 } value.define_singleton_method(:wait) do |duration| @stop_requested = true if paused_at paused_at = Time.now.utc @@ -813,7 +813,7 @@ def work(_job) end end - it "checks cancellation across queues in one batch, excluding inactive work" do + it "checks cancellation across queues, including registered work that has not started" do raised = [] fake_thread = Object.new fake_thread.define_singleton_method(:raise) { |error| raised << error } @@ -821,7 +821,7 @@ def work(_job) driver = Object.new driver.define_singleton_method(:job_get_cancelled_ids) { |ids| checked << ids - [3] + [2, 3] } value = runtime(driver: driver) value.instance_variable_set( @@ -836,7 +836,8 @@ def work(_job) value.send(:check_remote_cancellations) expect(raised).to eq([River::JobCancelError]) - expect(checked).to eq([[1, 3]]) + expect(checked).to eq([[1, 2, 3]]) + expect { value.send(:begin_work, 2) }.to raise_error(River::JobCancelError) end it "interrupts work that begins after stop was requested" do diff --git a/ruby/spec/notification_runtime_spec.rb b/ruby/spec/notification_runtime_spec.rb index 2e9d18b45..139c008fe 100644 --- a/ruby/spec/notification_runtime_spec.rb +++ b/ruby/spec/notification_runtime_spec.rb @@ -12,6 +12,41 @@ )) end + [false, true].each do |fail_first| + it "waits for cancellation hints and retries failures=#{fail_first}" do + driver.define_singleton_method(:notification_listener) { Object.new } + value = River::ClientRuntime.new(Object.new, driver, River::Config.new( + fetch_cooldown: 0.001, fetch_poll_interval: 0.01, leader_election_disabled: true, logger: Logger.new(output) + )) + queries = Queue.new + failures = 0 + driver.define_singleton_method(:job_get_cancelled_ids) do |_ids| + queries << true + if fail_first && failures.zero? + failures += 1 + raise "temporary read failure" + end + [] + end + value.instance_variable_set(:@running, 42 => {queue: "one", thread: Thread.current, working: true}) + checker = Thread.new { value.send(:cancellation_loop) } + sleep 0.05 + expect(queries).to be_empty + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) + Timeout.timeout(2) { queries.pop } + Timeout.timeout(2) { queries.pop } if fail_first + sleep 0.05 + expect(queries).to be_empty + ensure + value&.instance_variable_get(:@mutex)&.synchronize do + value.instance_variable_set(:@stop_requested, true) + value.instance_variable_get(:@running).clear + value.instance_variable_get(:@condition).broadcast + end + checker&.join + end + end + it "backs off repeated connection failures without indefinitely delaying shutdown" do listener = Object.new listener.define_singleton_method(:poll) { |_| raise "connection lost" } @@ -80,6 +115,35 @@ expect(replacement[:cancelled]).to be_nil end + it "retries a failed cancellation check made before the worker starts" do + raised = Queue.new + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + entry = {check_before_work: true, queue: "one", thread: thread, working: false} + runtime.instance_variable_set(:@running, 42 => entry) + queries = 0 + driver.define_singleton_method(:job_get_cancelled_ids) do |ids| + queries += 1 + raise "temporary read failure" if queries == 1 + + ids + end + + runtime.send(:begin_work, 42) + checker = Thread.new { runtime.send(:cancellation_loop) } + + expect(Timeout.timeout(2) { raised.pop }).to eq(River::JobCancelError) + expect(entry[:cancelled]).to be true + expect(output.string).to include("temporary read failure") + ensure + runtime.instance_variable_get(:@mutex).synchronize do + runtime.instance_variable_set(:@stop_requested, true) + runtime.instance_variable_get(:@running).clear + runtime.instance_variable_get(:@condition).broadcast + end + checker&.join + end + it "ignores invalid payloads without disrupting the receiver" do [nil, "{", "null", "[]", "{}"].each { |payload| runtime.send(:receive_notification, "river_insert", payload) } runtime.send(:receive_notification, "unknown", "{}") diff --git a/ruby/spec/runtime_draining_shared_examples.rb b/ruby/spec/runtime_draining_shared_examples.rb index 66e0cea34..863ce6372 100644 --- a/ruby/spec/runtime_draining_shared_examples.rb +++ b/ruby/spec/runtime_draining_shared_examples.rb @@ -108,10 +108,9 @@ 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. + # A peer notification must still be delivered while this client drains. River::Client.new(@driver).job_cancel(job.id) + Timeout.timeout(5) { poll_failed.pop } if poll_failure event = Timeout.timeout(3) { subscription.pop } expect(event).to have_attributes(kind: :job_cancelled, job: have_attributes(id: job.id, state: "cancelled")) diff --git a/rust/riverqueue/src/client/attempts.rs b/rust/riverqueue/src/client/attempts.rs index f51ffe204..5280d7569 100644 --- a/rust/riverqueue/src/client/attempts.rs +++ b/rust/riverqueue/src/client/attempts.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::Mutex; +use std::sync::atomic::Ordering; use tokio_util::sync::CancellationToken; @@ -27,12 +28,14 @@ pub(super) fn register_running_attempt( inner: &ClientInner, job_id: i64, cancellation: &CancellationToken, + generation: u64, ) { inner .running .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .insert(job_id, cancellation.clone()); - // Recheck after registration even if a notice arrived during the claim. - inner.cancellation_wakeup.notify_one(); + if inner.cancellation_generation.load(Ordering::Acquire) != generation { + inner.cancellation_wakeup.notify_one(); + } } diff --git a/rust/riverqueue/src/client/builder.rs b/rust/riverqueue/src/client/builder.rs index 058214cf1..22c6f3eb2 100644 --- a/rust/riverqueue/src/client/builder.rs +++ b/rust/riverqueue/src/client/builder.rs @@ -2,7 +2,6 @@ use std::collections::HashMap; use std::sync::atomic::AtomicBool; -#[cfg(test)] use std::sync::atomic::AtomicU64; use std::sync::{Arc, Mutex, RwLock}; use std::time::Duration; @@ -804,6 +803,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_generation: AtomicU64::new(0), cancellation_wakeup: Notify::new(), completion_sender: Mutex::new(None), database: self.database, diff --git a/rust/riverqueue/src/client/jobs.rs b/rust/riverqueue/src/client/jobs.rs index 94ed79cce..adf913f19 100644 --- a/rust/riverqueue/src/client/jobs.rs +++ b/rust/riverqueue/src/client/jobs.rs @@ -259,7 +259,7 @@ impl JobCancelRequest<'_> { // 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() { - inner.cancellation_wakeup.notify_one(); + inner.request_cancellation_check(); } Ok(row) } diff --git a/rust/riverqueue/src/client/mod.rs b/rust/riverqueue/src/client/mod.rs index 66a71958e..011d52f5f 100644 --- a/rust/riverqueue/src/client/mod.rs +++ b/rust/riverqueue/src/client/mod.rs @@ -101,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_generation: AtomicU64, cancellation_wakeup: Notify, completion_sender: Mutex>>, pub(crate) database: Database, @@ -201,6 +202,11 @@ impl ClientInner { !self.poll_only && self.database.delivers_notifications() } + fn request_cancellation_check(&self) { + self.cancellation_generation.fetch_add(1, Ordering::AcqRel); + self.cancellation_wakeup.notify_one(); + } + /// Whether this client receives notifications from other clients, which /// needs a backend listener and a client that isn't poll-only. When it /// doesn't, it wakes its own runtime directly after committing a change, diff --git a/rust/riverqueue/src/client/notifier.rs b/rust/riverqueue/src/client/notifier.rs index 07ab86e67..82d35a070 100644 --- a/rust/riverqueue/src/client/notifier.rs +++ b/rust/riverqueue/src/client/notifier.rs @@ -125,7 +125,7 @@ pub(super) fn dispatch_notification( match payload.action.as_str() { "cancel" => { if payload.job_id.is_some() { - inner.cancellation_wakeup.notify_one(); + inner.request_cancellation_check(); } } // Like Go's producer, a metadata change is handled at once @@ -221,6 +221,7 @@ async fn idles_without_listen_notify( } info!("River's database does not support LISTEN/NOTIFY; polling instead"); report_ready(ready); + inner.request_cancellation_check(); cancel.cancelled().await; Ok(true) } @@ -297,6 +298,7 @@ async fn listen_until_error( "timed out connecting and subscribing".to_owned(), ) })??; + inner.request_cancellation_check(); debug!("River notification listener healthy"); *attempt = 0; report_ready(ready); @@ -325,6 +327,7 @@ async fn listen_until_error( Ok(Ok(None)) => { warn!("River notification listener reconnected; waking producers"); wake_all_producers(queue_notifications); + inner.request_cancellation_check(); } Ok(Err(listener_error)) => return Err(listener_error.into()), Err(_) => { @@ -384,6 +387,7 @@ pub(super) async fn run_sqlite_notifications( .connect(&mut connection) .await .map_err(sqlite_backend_error)?; + inner.request_cancellation_check(); } for topic in [ crate::NOTIFICATION_TOPIC_CONTROL, diff --git a/rust/riverqueue/src/client/producer.rs b/rust/riverqueue/src/client/producer.rs index e6045da78..92afed2a7 100644 --- a/rust/riverqueue/src/client/producer.rs +++ b/rust/riverqueue/src/client/producer.rs @@ -2,6 +2,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::Duration; use chrono::Utc; @@ -54,8 +55,7 @@ pub(super) async fn run_dynamic_queues( mut changes: watch::Receiver, queues_ready: oneshot::Sender<()>, ) -> Result<(), Error> { - // Recover lost cancellation notices, including after the notifier stops - // during graceful shutdown, until every producer has drained. + // Cancellation checks remain independent of fetching until every producer drains. let _cancellation_poll = AbortOnDrop(tokio::spawn(poll_job_cancellations( Arc::clone(&inner), fetch_cancel.clone(), @@ -138,26 +138,21 @@ fn jittered_fetch_poll_interval(interval: Duration) -> Duration { crate::maintenance::random_duration(interval, interval.saturating_add(jitter)) } -/// Cancels this client's running attempts whose jobs have a cancellation -/// request, checking every [`QUEUE_CONFIG_POLL_INTERVAL`]. 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) { - let mut poll = tokio::time::interval(QUEUE_CONFIG_POLL_INTERVAL); - poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - // The first tick completes at once; running jobs start later. - poll.tick().await; +/// Checks cancellations on notification/reconnect hints, retrying failed reads. +/// Without a notifier, polls at the queue-control interval instead. +pub(super) async fn poll_job_cancellations(inner: Arc, stopping: CancellationToken) { + let mut retry = false; loop { tokio::select! { - _ = poll.tick() => {}, + () = tokio::time::sleep(QUEUE_CONFIG_POLL_INTERVAL), if retry || !inner.has_notifier() => {}, () = inner.cancellation_wakeup.notified() => {}, } - if let Err(poll_error) = check_job_cancellations(&inner).await + let result = check_job_cancellations(&inner).await; + retry = result.is_err(); + if let Err(check_error) = result && !stopping.is_cancelled() { - error!(error = %crate::error::Chain(&poll_error), "River failed to check for job cancellation requests"); + error!(error = %crate::error::Chain(&check_error), "River failed to check for job cancellation requests"); } } } @@ -849,6 +844,7 @@ pub(super) async fn run_queue( if available == 0 { continue; } + let cancellation_generation = inner.cancellation_generation.load(Ordering::Acquire); let fetched = match (&session, claims_through_session) { (Some(session), true) => { match claim_through_session( @@ -897,7 +893,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, row.id, &cancellation); + register_running_attempt(&inner, row.id, &cancellation, cancellation_generation); let task_inner = Arc::clone(&inner); let completion_sender = completion_sender.clone(); let session_row = attempts.session.is_some().then(|| row.clone()); diff --git a/rust/riverqueue/src/client/run.rs b/rust/riverqueue/src/client/run.rs index f8c42bac0..fb91a1e92 100644 --- a/rust/riverqueue/src/client/run.rs +++ b/rust/riverqueue/src/client/run.rs @@ -175,8 +175,8 @@ impl Service { /// When each part of a started client stops, matching River Go's order. /// /// A stop request cancels the client's fetch token. Claims, leadership with -/// its maintenance services, the notifier, and extension runtime services -/// all stop then, like Go's services started on the fetch context. Running +/// its maintenance services, and extension runtime services stop then. +/// The notifier remains available until every queue has drained. Running /// jobs keep going until they finish or `work_cancel` fires, and each /// producer keeps reporting to its extension session until its last attempt /// has left, from a token of its own that the stop request doesn't cancel. @@ -185,7 +185,9 @@ struct StopTokens { claims: CancellationToken, /// Leader election and leader-owned maintenance. leadership: CancellationToken, - /// The notifier and extension runtime services. + /// Notification delivery remains live until queues finish draining. + notifier: CancellationToken, + /// Extension runtime services. services: CancellationToken, } @@ -194,6 +196,7 @@ impl StopTokens { Self { claims: fetch_cancel.child_token(), leadership: fetch_cancel.child_token(), + notifier: CancellationToken::new(), services: fetch_cancel.child_token(), } } @@ -247,7 +250,14 @@ impl Supervisor { let Some((service, started_at)) = self.services.remove(&task_id) else { continue; }; - let stopping = self.fetch_cancel.is_cancelled(); + if service == Service::Queues { + self.stop.notifier.cancel(); + } + let stopping = if service == Service::Notifier { + self.stop.notifier.is_cancelled() + } else { + self.fetch_cancel.is_cancelled() + }; match outcome { // The completer ends once every producer has dropped its sender. Ok(()) if stopping || service == Service::Completer => {} @@ -374,7 +384,7 @@ impl Supervisor { )) } Service::Notifier => { - cancel = self.stop.services.child_token(); + cancel = self.stop.notifier.child_token(); let ready = self.notifier_ready.clone(); #[cfg(test)] if inner @@ -526,9 +536,9 @@ impl Stopper { /// Requests a soft stop. /// /// The client stops fetching new jobs and lets running jobs finish before - /// it stops. Leader election, maintenance, and the notification listener - /// stop at once, while each queue's producer keeps reporting its - /// running jobs until they finish. When the builder's `soft_stop_timeout` + /// it stops. Leader election and maintenance stop at once. The notification + /// listener keeps delivering cancellations, and each queue's producer keeps + /// reporting its running jobs until they finish. When the builder's `soft_stop_timeout` /// is set, jobs still running after that timeout are cancelled as if by /// [`Stopper::stop_and_cancel`]. pub fn stop(&self) { diff --git a/rust/riverqueue/src/client/tests.rs b/rust/riverqueue/src/client/tests.rs index 45e3ae72a..28bb36613 100644 --- a/rust/riverqueue/src/client/tests.rs +++ b/rust/riverqueue/src/client/tests.rs @@ -236,14 +236,23 @@ async fn go_notification_fixtures_dispatch_correctly() { } #[tokio::test] -async fn registration_wakes_cancellation_check() { +async fn registration_without_a_hint_does_not_query_cancellations() { #[cfg(feature = "postgres")] let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); #[cfg(not(feature = "postgres"))] let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); let client = Client::builder(pool).build().unwrap(); let cancellation = CancellationToken::new(); - register_running_attempt(&client.inner, 42, &cancellation); + register_running_attempt(&client.inner, 42, &cancellation, 0); + assert!( + tokio::time::timeout(Duration::ZERO, client.inner.cancellation_wakeup.notified(),) + .await + .is_err() + ); + // A hint consumed before registration still requires a fresh snapshot afterward. + client.inner.request_cancellation_check(); + client.inner.cancellation_wakeup.notified().await; + register_running_attempt(&client.inner, 43, &cancellation, 0); tokio::time::timeout( Duration::from_secs(1), client.inner.cancellation_wakeup.notified(), @@ -253,6 +262,61 @@ async fn registration_wakes_cancellation_check() { assert!(!cancellation.is_cancelled()); } +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn cancellation_checks_wait_for_hints_and_retry_failed_reads() { + let queries = Arc::new(AtomicU64::new(0)); + let observed = Arc::clone(&queries); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .before_acquire(move |_, _| { + let observed = Arc::clone(&observed); + Box::pin(async move { + observed.fetch_add(1, Ordering::AcqRel); + Ok(true) + }) + }) + .connect("sqlite::memory:") + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + let cancellation = CancellationToken::new(); + register_running_attempt(&client.inner, 42, &cancellation, 0); + let task = tokio::spawn(super::producer::poll_job_cancellations( + Arc::clone(&client.inner), + CancellationToken::new(), + )); + let before = queries.load(Ordering::Acquire); + tokio::time::sleep(Duration::from_millis(2100)).await; + assert_eq!(before, queries.load(Ordering::Acquire)); + client.inner.request_cancellation_check(); + tokio::time::timeout(Duration::from_secs(5), async { + while queries.load(Ordering::Acquire) == before { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap(); + // The first check fails because the table is absent. Acquiring its sole connection + // waits for that check to finish; repairing the database must be enough to recover. + let mut connection = pool.acquire().await.unwrap(); + sqlx::query("CREATE TABLE river_job (id INTEGER, state TEXT, metadata TEXT)") + .execute(&mut *connection) + .await + .unwrap(); + sqlx::query(r#"INSERT INTO river_job VALUES (42, 'running', '{"cancel_attempted_at":"2026-01-01T00:00:00Z"}')"#).execute(&mut *connection).await.unwrap(); + drop(connection); + tokio::time::timeout(Duration::from_secs(5), cancellation.cancelled()) + .await + .unwrap(); + let after = queries.load(Ordering::Acquire); + tokio::time::sleep(Duration::from_millis(2100)).await; + assert_eq!(after, queries.load(Ordering::Acquire)); + task.abort(); + let _ = task.await; + pool.close().await; +} + #[cfg(feature = "sqlite")] #[tokio::test] async fn delayed_cancellation_signals_recheck_retry_marker() { @@ -272,7 +336,12 @@ async fn delayed_cancellation_signals_recheck_retry_marker() { 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); + register_running_attempt( + &client.inner, + id, + &original, + client.inner.cancellation_generation.load(Ordering::Acquire), + ); 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(); @@ -282,7 +351,12 @@ async fn delayed_cancellation_signals_recheck_retry_marker() { .await .unwrap(); let replacement = CancellationToken::new(); - register_running_attempt(&client.inner, id, &replacement); + register_running_attempt( + &client.inner, + id, + &replacement, + client.inner.cancellation_generation.load(Ordering::Acquire), + ); dispatch_notification( &client.inner, &client.inner.queue_notifications, diff --git a/rust/riverqueue/tests/poll_only_cancellation.rs b/rust/riverqueue/tests/poll_only_cancellation.rs index a01234f70..457ff1fed 100644 --- a/rust/riverqueue/tests/poll_only_cancellation.rs +++ b/rust/riverqueue/tests/poll_only_cancellation.rs @@ -1,5 +1,5 @@ -//! Durable cancellation of running jobs, including when notifications are -//! disabled or a cancellation notice is lost, on every backend. +//! Durable cancellation of running jobs, with notifications enabled or disabled, +//! including during graceful shutdown, on every backend. //! //! Postgres tests run in a unique schema and fail rather than skip when //! `RIVER_RUST_DATABASE_URL` is unset; SQLite tests use a temporary file. @@ -11,8 +11,8 @@ mod support; use std::{convert::Infallible, sync::Arc, time::Duration}; use riverqueue::{ - __private::Database, Client, EventKind, Job, JobArgs, JobState, JobUpdateParams, QueueConfig, - WorkContext, WorkOutcome, Workers, + __private::Database, Client, EventKind, Job, JobArgs, JobState, QueueConfig, WorkContext, + WorkOutcome, Workers, }; use serde::{Deserialize, Serialize}; use tokio::sync::Semaphore; @@ -56,7 +56,7 @@ impl Backend { } } -/// Durable cancellation survives lost notifications and graceful shutdown. +/// Both notification delivery and poll-only checks remain live during graceful shutdown. async fn polls_for_remote_cancellation(backend: Backend) { for (notifications, while_stopping) in [(false, false), (false, true), (true, false), (true, true)] @@ -100,27 +100,7 @@ async fn polls_for_remote_cancellation(backend: Backend) { if while_stopping { run.stopper().stop(); } - 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(); - } + other.jobs().cancel(id).await.unwrap(); let event = tokio::time::timeout(TIMEOUT, events.recv()) .await From 5fbb860965e8acfd6f87455c711c60e771b2f470 Mon Sep 17 00:00:00 2001 From: Brandur Date: Tue, 27 Oct 2026 19:33:02 -0500 Subject: [PATCH 6/6] =?UTF-8?q?=F0=9F=A4=96=20Deliver=20attempt-scoped=20c?= =?UTF-8?q?ancellations=20without=20database=20reads?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Include the attempt number and attempted_at in cancellation notifications across Go, Java, JavaScript, Ruby, and Rust. Cancel matching active workers directly so a saturated pool cannot block delivery, and ignore notices for older executions even when an interrupted job reuses its attempt number. Keep durable reconciliation for legacy notices, reconnects, and claims that have not yet registered their active execution. Restrict Java cancellation recovery to batches of local active job IDs and use the configured polling interval for poll-only operation and read retries. Treat the cancellation marker's presence consistently across Postgres and SQLite. Cover blocked reads, stale execution identities, transaction boundaries, and shared notification payloads, including JavaScript's native stream. --- .../cmd/generatefixtures/protocol_values.go | 2 + .../src/main/java/com/riverqueue/Client.java | 2 +- .../main/java/com/riverqueue/Protocol.java | 37 ++++- .../main/java/com/riverqueue/RunHandle.java | 50 ++++--- .../resources/com/riverqueue/postgres.sql | 2 +- .../main/resources/com/riverqueue/sqlite.sql | 2 +- .../com/riverqueue/CancellationCheckTest.java | 127 ++++++++++++++++++ .../java/com/riverqueue/ProtocolTest.java | 10 +- js/driver/pg/src/driver.integration.test.ts | 63 +++++---- js/driver/pg/src/driver.ts | 3 +- js/driver/pg/src/sql/jobs.ts | 4 +- js/driver/pg/src/sql/notify.ts | 25 +++- js/driver/sqlite/src/driver.test.ts | 8 +- js/driver/sqlite/src/driver.ts | 2 +- js/etc/riverqueue.unstable-driver.api.md | 2 + js/src/driver.ts | 2 + js/src/runtime.test.ts | 85 ++++++++++++ js/src/runtime/attempt-runner.ts | 19 ++- js/src/runtime/notification-payloads.ts | 27 ++++ .../notification-pump.conformance.test.ts | 19 ++- js/src/runtime/notification-pump.ts | 16 ++- producer.go | 23 +++- producer_test.go | 37 ++++- .../internal/dbsqlc/river_job.sql.go | 4 +- riverdriver/riverdrivertest/listener.go | 17 ++- .../riverpgxv5/internal/dbsqlc/river_job.sql | 4 +- .../internal/dbsqlc/river_job.sql.go | 4 +- .../internal/dbsqlc/river_notification.sql | 7 +- .../internal/dbsqlc/river_notification.sql.go | 13 +- ruby/lib/client_runtime.rb | 13 +- ruby/lib/driver/runtime.rb | 2 +- ruby/spec/conformance_spec.rb | 10 +- .../insert_notification_shared_examples.rb | 8 +- ruby/spec/notification_runtime_spec.rb | 31 +++++ ruby/spec/runtime_draining_shared_examples.rb | 6 + rust/riverqueue/src/client/attempts.rs | 22 ++- rust/riverqueue/src/client/mod.rs | 3 +- rust/riverqueue/src/client/notifier.rs | 21 ++- rust/riverqueue/src/client/producer.rs | 13 +- rust/riverqueue/src/client/tests.rs | 111 +++++++++++---- rust/riverqueue/src/storage/postgres.rs | 4 +- rust/riverqueue/src/storage/sqlite.rs | 2 + 42 files changed, 714 insertions(+), 148 deletions(-) diff --git a/conformance/cmd/generatefixtures/protocol_values.go b/conformance/cmd/generatefixtures/protocol_values.go index 1fa2f71dd..7b4a9e40b 100644 --- a/conformance/cmd/generatefixtures/protocol_values.go +++ b/conformance/cmd/generatefixtures/protocol_values.go @@ -166,6 +166,8 @@ func makeProtocolNotifications(root string) ([]protocolNotification, error) { switch action { case "cancel": values["job_id"] = exampleJobID + values["attempt"] = 1 + values["attempted_at"] = "2026-01-02T03:04:05.678Z" case "metadata_changed": values["metadata"] = map[string]any{"owner": "candidate"} } diff --git a/java/river/src/main/java/com/riverqueue/Client.java b/java/river/src/main/java/com/riverqueue/Client.java index ceafd0d1b..9b930dd6f 100644 --- a/java/river/src/main/java/com/riverqueue/Client.java +++ b/java/river/src/main/java/com/riverqueue/Client.java @@ -656,7 +656,7 @@ public Job cancel(Connection connection, long id) { if (!rows.next()) return get(c, id); var job = read(rows); pilot.afterStateChange(c, driver(), job); - notify(c, Protocol.cancel(id, job.queue())); + notify(c, Protocol.cancel(job)); return job; } }); diff --git a/java/river/src/main/java/com/riverqueue/Protocol.java b/java/river/src/main/java/com/riverqueue/Protocol.java index f13bdaae5..6a3f57912 100644 --- a/java/river/src/main/java/com/riverqueue/Protocol.java +++ b/java/river/src/main/java/com/riverqueue/Protocol.java @@ -1,8 +1,8 @@ package com.riverqueue; +import java.time.Instant; 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. */ @@ -23,6 +23,28 @@ final class Protocol { private Protocol() {} + static Notification cancel(Job job) { + if (job.attemptedAt() == null) return cancel(job.id(), job.queue()); + return cancel(job.id(), job.queue(), job.attempt(), job.attemptedAt()); + } + + static Notification cancel(long id, String queue, int attempt, Instant attemptedAt) { + return new Notification( + TOPIC_CONTROL, + Json.encode( + Map.of( + "action", + "cancel", + "attempt", + attempt, + "attempted_at", + attemptedAt, + "job_id", + id, + "queue", + queue))); + } + static Notification cancel(long id, String queue) { return new Notification( TOPIC_CONTROL, Json.encode(Map.of("action", "cancel", "job_id", id, "queue", queue))); @@ -53,12 +75,14 @@ static Notification resigned(String id) { TOPIC_LEADERSHIP, Json.encode(Map.of("action", "resigned", "leader_id", id))); } + record CancellationNotice(long id, int attempt, Instant attemptedAt) {} + record Notification(String topic, String payload) {} /** Shared by both database listeners; dispatching itself requires no database I/O. */ record Dispatcher( String clientId, - LongConsumer cancellation, + Consumer cancellation, Consumer queues, Consumer leadership) { void dispatch(String topic, String payload) { @@ -70,9 +94,12 @@ void dispatch(String topic, String payload) { switch (topic) { case TOPIC_CONTROL -> { if (action.equals("cancel") && value.path("job_id").isIntegralNumber()) { - // 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()); + var attemptedAt = value.path("attempted_at"); + cancellation.accept( + new CancellationNotice( + value.path("job_id").asLong(), + value.path("attempt").asInt(0), + attemptedAt.isString() ? Instant.parse(attemptedAt.asString()) : null)); } else if ((action.equals("metadata_changed") || action.equals("pause") || action.equals("resume")) diff --git a/java/river/src/main/java/com/riverqueue/RunHandle.java b/java/river/src/main/java/com/riverqueue/RunHandle.java index c956ddc68..0242dbd99 100644 --- a/java/river/src/main/java/com/riverqueue/RunHandle.java +++ b/java/river/src/main/java/com/riverqueue/RunHandle.java @@ -87,10 +87,7 @@ public final class RunHandle implements AutoCloseable { id = river.id(); notifications = new Protocol.Dispatcher( - id, - ignored -> requestCancellationCheck(), - this::queueNotification, - this::leadershipNotification); + id, this::cancelNotification, this::queueNotification, this::leadershipNotification); queues.putAll(config.queues); periodic = new PeriodicJobBundle( @@ -356,6 +353,15 @@ private boolean wakeDueFetches() { return !ready.isEmpty(); } + private void cancelNotification(Protocol.CancellationNotice notice) { + var context = active.get(notice.id()); + if (context != null && notice.attempt() > 0 && notice.attemptedAt() != null) { + if (context.job().attempt() == notice.attempt() + && notice.attemptedAt().equals(context.job().attemptedAt())) + context.requestCancellation(WorkContext.Cancellation.REMOTE); + } else requestCancellationCheck(); + } + private void claim(String queue, int count) { long generation = cancellationGeneration.get(); List jobs = @@ -1426,27 +1432,35 @@ private void pollCancellations() { boolean retry = false; while (hasActiveWork()) { try { - if (cancellationPollOnly || retry) cancellationWake.tryAcquire(50, TimeUnit.MILLISECONDS); + if (cancellationPollOnly || retry) + cancellationWake.tryAcquire(config.pollInterval.toNanos(), TimeUnit.NANOSECONDS); else cancellationWake.acquire(); cancellationWake.drainPermits(); if (!hasActiveWork()) return; retry = false; var attempts = Map.copyOf(active); if (!attempts.isEmpty()) { - 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); + var ids = new ArrayList<>(attempts.keySet()); + for (int start = 0; start < ids.size(); start += 1000) { + var batch = ids.subList(start, Math.min(start + 1000, ids.size())); + river.transaction( + c -> { + try (var statement = + Sql.prepare( + c, + Sql.query(river.database(), "cancel_requested"), + Json.encode(batch)); + var rows = statement.executeQuery()) { + while (rows.next()) { + // A late query result must not cancel a subsequent attempt. + var context = attempts.get(rows.getLong(1)); + if (context != null) + context.requestCancellation(WorkContext.Cancellation.REMOTE); + } } - } - return null; - }); + return null; + }); + } } } catch (InterruptedException e) { Thread.currentThread().interrupt(); diff --git a/java/river/src/main/resources/com/riverqueue/postgres.sql b/java/river/src/main/resources/com/riverqueue/postgres.sql index 10e7df6f5..d6151487b 100644 --- a/java/river/src/main/resources/com/riverqueue/postgres.sql +++ b/java/river/src/main/resources/com/riverqueue/postgres.sql @@ -181,7 +181,7 @@ SELECT {columns} FROM {schema}river_job WHERE state='running' AND attempted_at>'cancel_attempted_at' IS NOT NULL +SELECT id FROM {schema}river_job WHERE id IN (SELECT value::bigint FROM jsonb_array_elements_text(?::jsonb)) AND state='running' AND jsonb_exists(metadata, 'cancel_attempted_at') -- name: notification_cursor SELECT coalesce(max(id),0) FROM {schema}river_notification diff --git a/java/river/src/main/resources/com/riverqueue/sqlite.sql b/java/river/src/main/resources/com/riverqueue/sqlite.sql index c859aa7a2..bcea6e2e2 100644 --- a/java/river/src/main/resources/com/riverqueue/sqlite.sql +++ b/java/river/src/main/resources/com/riverqueue/sqlite.sql @@ -198,7 +198,7 @@ AND queue NOT IN (SELECT value FROM json_each(?)) AND (? OR queue IN (SELECT val SELECT {columns} FROM {schema}river_job WHERE state='running' AND attempted_at? ORDER BY id LIMIT 1000 -- name: cancel_requested -SELECT id FROM river_job WHERE state='running' AND (CASE WHEN typeof(metadata)<>'text' OR json_valid(metadata) THEN json_extract(metadata,'$.cancel_attempted_at') END) IS NOT NULL +SELECT id FROM {schema}river_job WHERE id IN (SELECT CAST(value AS integer) FROM json_each(?)) AND state='running' AND (CASE WHEN typeof(metadata)<>'text' OR json_valid(metadata) THEN json_type(metadata,'$.cancel_attempted_at') END) IS NOT NULL -- name: notification_cursor SELECT coalesce(max(id),0) FROM {schema}river_notification diff --git a/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java b/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java index e42814197..87bd4498d 100644 --- a/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java +++ b/java/river/src/test/java/com/riverqueue/CancellationCheckTest.java @@ -8,12 +8,14 @@ import java.sql.Connection; import java.sql.SQLException; import java.time.Duration; +import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import javax.sql.DataSource; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; @@ -23,6 +25,39 @@ class CancellationCheckTest { @RegisterExtension final TestDatabase databases = new TestDatabase(); @TempDir Path directory; + @Test + void cancellationQueryOnlyReadsRequestedJobs() throws Exception { + var database = databases.open(directory.resolve("river.db")); + var client = Client.builder(database).build(); + var type = JobType.of("cancellation_query", String.class); + var first = client.insert(type, "first").job(); + var second = client.insert(type, "second").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='running', metadata='{\"cancel_attempted_at\":null}' WHERE id IN (?, ?)", + first.id(), + second.id())) { + statement.executeUpdate(); + } + try (var statement = + Sql.prepare( + connection, + Sql.query(database, "cancel_requested"), + Json.encode(List.of(first.id()))); + var rows = statement.executeQuery()) { + assertTrue(rows.next()); + assertEquals(first.id(), rows.getLong(1)); + assertFalse(rows.next()); + } + return null; + }); + } + @ParameterizedTest @ValueSource(booleans = {false, true}) @SuppressWarnings("try") // Resources intentionally scope background activity. @@ -108,4 +143,96 @@ void checksOnlyOnHintsAndRetriesFailures(boolean failFirst) throws Exception { } } } + + @Test + @SuppressWarnings("try") + void identifiedNoticeBypassesBlockedReadAndIgnoresOlderAttempts() throws Exception { + var database = databases.open(directory.resolve("river.db")); + var block = new AtomicBoolean(); + var readStarted = new CompletableFuture(); + var releaseRead = new CompletableFuture(); + var source = + (DataSource) + Proxy.newProxyInstance( + DataSource.class.getClassLoader(), + new Class[] {DataSource.class}, + (proxy, method, args) -> { + if (!method.getName().equals("getConnection")) + throw new UnsupportedOperationException(); + var connection = database.connection(); + return Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + (ignored, operation, arguments) -> { + if (operation.getName().equals("prepareStatement") + && arguments[0].equals(Sql.query(database, "cancel_requested")) + && block.compareAndSet(true, false)) { + readStarted.complete(null); + releaseRead.get(5, TimeUnit.SECONDS); + } + try { + return operation.invoke(connection, arguments); + } catch (InvocationTargetException error) { + throw error.getCause(); + } + }); + }); + var ready = new CompletableFuture(); + var entered = new CompletableFuture>(); + var cancelled = new CompletableFuture(); + var type = JobType.of("cancellation_delivery", String.class); + var client = + Client.builder(new Database(source, database.dialect()).withSchema(database.schema())) + .queue("default", 1) + .leadership(false) + .pollInterval(Duration.ofMillis(10)) + .stopTimeout(Duration.ofSeconds(5)) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .worker( + type, + context -> { + entered.complete(context); + context.awaitCancellation(); + cancelled.complete(null); + context.checkCancelled(); + }) + .build(); + try (var run = client.start()) { + try { + ready.get(5, TimeUnit.SECONDS); + var inserted = client.insert(type, "working").job(); + var context = entered.get(5, TimeUnit.SECONDS); + var job = context.job(); + block.set(true); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(inserted.id(), "default")); + return null; + }); + readStarted.get(5, TimeUnit.SECONDS); + client.transaction( + connection -> { + client.notify( + connection, + Protocol.cancel(job.id(), job.queue(), job.attempt() + 1, job.attemptedAt())); + client.notify( + connection, + Protocol.cancel( + job.id(), job.queue(), job.attempt(), job.attemptedAt().minusSeconds(1))); + return null; + }); + assertThrows( + java.util.concurrent.TimeoutException.class, + () -> cancelled.get(100, TimeUnit.MILLISECONDS)); + client.cancel(job.id()); + cancelled.get(2, TimeUnit.SECONDS); + } finally { + releaseRead.complete(null); + if (entered.isDone()) entered.join().requestCancellation(WorkContext.Cancellation.SHUTDOWN); + } + } + } } diff --git a/java/river/src/test/java/com/riverqueue/ProtocolTest.java b/java/river/src/test/java/com/riverqueue/ProtocolTest.java index ca198346e..994513338 100644 --- a/java/river/src/test/java/com/riverqueue/ProtocolTest.java +++ b/java/river/src/test/java/com/riverqueue/ProtocolTest.java @@ -83,7 +83,10 @@ Stream notificationDispatch() throws Exception { var leadership = new ArrayList(); var dispatcher = new Protocol.Dispatcher( - clientId, cancellations::add, queues::add, leadership::add); + clientId, + notice -> cancellations.add(notice.id()), + queues::add, + leadership::add); dispatcher.dispatch( fixture.required("topic").asString(), Json.encode(fixture.required("payload"))); @@ -130,7 +133,8 @@ void notificationTopicsRouteCancellationChecks() throws Exception { var queues = new ArrayList(); var leadership = new ArrayList(); var dispatcher = - new Protocol.Dispatcher("observer", cancellations::add, queues::add, leadership::add); + new Protocol.Dispatcher( + "observer", notice -> cancellations.add(notice.id()), queues::add, leadership::add); String payload = Json.encode(cancel.required("payload")); dispatcher.dispatch("unrelated", payload); dispatcher.dispatch(Protocol.TOPIC_LEADERSHIP, payload); @@ -183,7 +187,7 @@ private static WorkContext context(long id, JsonNode metadata) { Stream notificationEncoding() throws Exception { var notifications = Map.of( - "cancel", Protocol.cancel(42, "priority"), + "cancel", Protocol.cancel(42, "priority", 1, Instant.parse("2026-01-02T03:04:05.678Z")), "insert", Protocol.insert("priority"), "metadata_changed", Protocol.queueMetadataChanged("priority", Map.of("owner", "candidate")), diff --git a/js/driver/pg/src/driver.integration.test.ts b/js/driver/pg/src/driver.integration.test.ts index 084b0ae7f..9bcc9aa89 100644 --- a/js/driver/pg/src/driver.integration.test.ts +++ b/js/driver/pg/src/driver.integration.test.ts @@ -2495,32 +2495,43 @@ describe("PgDriver integration", () => { await iterator.return(undefined); }); - it("decodes exact bigint IDs from cancellation notifications", async () => { - const exactID = 9_007_199_254_741_001n; - await pool.query( - `INSERT INTO river_job (id, args, kind, max_attempts) - VALUES ($1::bigint, '{}'::jsonb, $2::text, 25)`, - [exactID.toString(10), `${filePrefix}_large_cancel_id`] - ); - const abort = new AbortController(); - const iterator = driver.jobCancellationSubscribe( - `${filePrefix}_cancel_owner`, - abort.signal - ); - const next = iterator.next(); - await waitForListenerPID(pool, "river_control"); - - await driver.jobCancel(exactID); - await expect(next).resolves.toEqual({ - done: false, - value: { - attemptedBy: `${filePrefix}_cancel_owner`, - id: exactID, - }, - }); - abort.abort(); - await iterator.return(undefined); - }); + it.each([false, true])( + "decodes exact bigint IDs from cancellation notifications with identity=%s", + async (identified) => { + const exactID = 9_007_199_254_741_001n; + const attemptedAt = Temporal.Instant.from("2026-01-02T03:04:05.6789Z"); + await pool.query( + `INSERT INTO river_job (id, args, kind, max_attempts, attempt, attempted_at, state) + VALUES ($1::bigint, '{}'::jsonb, $2::text, 25, $3::smallint, $4::timestamptz, $5::river_job_state)`, + [ + exactID.toString(10), + `${filePrefix}_large_cancel_id`, + identified ? 2 : 0, + identified ? attemptedAt.toString() : null, + identified ? "running" : "available", + ] + ); + const abort = new AbortController(); + const iterator = driver.jobCancellationSubscribe( + `${filePrefix}_cancel_owner`, + abort.signal + ); + const next = iterator.next(); + await waitForListenerPID(pool, "river_control"); + + await driver.jobCancel(exactID); + await expect(next).resolves.toEqual({ + done: false, + value: { + ...(identified ? { attempt: 2, attemptedAt } : {}), + attemptedBy: `${filePrefix}_cancel_owner`, + id: exactID, + }, + }); + abort.abort(); + await iterator.return(undefined); + } + ); it("notifies listeners of a committed insert notification", async () => { const abort = new AbortController(); diff --git a/js/driver/pg/src/driver.ts b/js/driver/pg/src/driver.ts index 3287ad65c..01ca76d25 100644 --- a/js/driver/pg/src/driver.ts +++ b/js/driver/pg/src/driver.ts @@ -5,6 +5,7 @@ import type { DriverInsertResult, DriverRecord, InsertDriverOptions, + JobCancellationNotice, JobClaimOptions, JobClaimParams, JobClaimResult, @@ -861,7 +862,7 @@ export class PgRuntime implements RuntimeDriver< attemptedBy: string, signal: AbortSignal, ready?: () => void - ): AsyncGenerator<{ attemptedBy: string; id: bigint }> { + ): AsyncGenerator { return notifySql.jobCancellationSubscribe( this.#db, attemptedBy, diff --git a/js/driver/pg/src/sql/jobs.ts b/js/driver/pg/src/sql/jobs.ts index f95c75659..e675eaa30 100644 --- a/js/driver/pg/src/sql/jobs.ts +++ b/js/driver/pg/src/sql/jobs.ts @@ -72,7 +72,7 @@ export async function jobCancelWithOptions( const sql = { with: ` WITH locked_job AS ( - SELECT id, queue, state, finalized_at + SELECT id, attempt, attempted_at, queue, state, finalized_at FROM ${jobTable} WHERE id = $1::bigint FOR UPDATE @@ -84,6 +84,8 @@ export async function jobCancelWithOptions( concat(coalesce($4::text, current_schema()), '.', $2::text), json_build_object( 'action', 'cancel', + 'attempt', attempt, + 'attempted_at', attempted_at, 'job_id', id, 'queue', queue )::text diff --git a/js/driver/pg/src/sql/notify.ts b/js/driver/pg/src/sql/notify.ts index ac966c5d5..408ec41c0 100644 --- a/js/driver/pg/src/sql/notify.ts +++ b/js/driver/pg/src/sql/notify.ts @@ -3,7 +3,10 @@ import { LinkedAbortSignal, quoteIdentifier } from "riverqueue/unstable-driver"; import { Buffer } from "node:buffer"; import type { Notification as NodePgNotification } from "pg"; import { ConfigurationError } from "riverqueue"; -import type { RuntimeNotification } from "riverqueue/unstable-driver"; +import type { + JobCancellationNotice, + RuntimeNotification, +} from "riverqueue/unstable-driver"; import type { PgDatabase, PgQueryable } from "../database.js"; import { POSTGRES_IDENTIFIER_MAX_BYTES } from "../database.js"; import { @@ -212,7 +215,7 @@ export async function* jobCancellationSubscribe( attemptedBy: string, signal: AbortSignal, ready?: () => void -): AsyncGenerator<{ attemptedBy: string; id: bigint }> { +): AsyncGenerator { for await (const notification of listen( db, ["river_control"], @@ -236,7 +239,23 @@ export async function* jobCancellationSubscribe( } try { const id = exactCancellationID(notification.payload, value.job_id); - if (id !== null) yield { attemptedBy, id }; + if (id !== null) { + if ( + "attempt" in value && + typeof value.attempt === "number" && + Number.isSafeInteger(value.attempt) && + value.attempt > 0 && + "attempted_at" in value && + typeof value.attempted_at === "string" + ) { + yield { + attempt: value.attempt, + attemptedAt: Temporal.Instant.from(value.attempted_at), + attemptedBy, + id, + }; + } else yield { attemptedBy, id }; + } } catch { // Ignore malformed notifications; polling remains the durable path. } diff --git a/js/driver/sqlite/src/driver.test.ts b/js/driver/sqlite/src/driver.test.ts index ed9c4b77a..3868059e2 100644 --- a/js/driver/sqlite/src/driver.test.ts +++ b/js/driver/sqlite/src/driver.test.ts @@ -336,7 +336,7 @@ describe("SqliteDriver", () => { database .prepare( - "UPDATE river_job SET state = 'running', attempt = 1 WHERE id = ?" + "UPDATE river_job SET state = 'running', attempt = 1, attempted_at = '2026-01-02T03:04:05.678Z' WHERE id = ?" ) .run(inserted.id); expect(await driver.jobDelete(inserted.id)).toMatchObject({ @@ -348,7 +348,7 @@ describe("SqliteDriver", () => { // Like River for Go, a retry publishes no insert notification. expect(control.map(({ topic }) => topic)).toEqual(["river_control"]); expect(control[0]?.payload).toBe( - `{"action":"cancel","job_id":${inserted.id},"queue":"default"}` + `{"action":"cancel","attempt":0,"attempted_at":null,"job_id":${inserted.id},"queue":"default"}` ); }); @@ -431,7 +431,7 @@ describe("SqliteDriver", () => { .job; database .prepare( - "UPDATE river_job SET state = 'running', attempt = 1 WHERE id = ?" + "UPDATE river_job SET state = 'running', attempt = 1, attempted_at = '2026-01-02T03:04:05.678Z' WHERE id = ?" ) .run(inserted.id); const cancels = () => @@ -460,7 +460,7 @@ describe("SqliteDriver", () => { }); expect(cancels()).toEqual([ { - payload: `{"action":"cancel","job_id":${inserted.id},"queue":"default"}`, + payload: `{"action":"cancel","attempt":1,"attempted_at":"2026-01-02T03:04:05.678Z","job_id":${inserted.id},"queue":"default"}`, topic: "river_control", }, ]); diff --git a/js/driver/sqlite/src/driver.ts b/js/driver/sqlite/src/driver.ts index ead745c98..f992a8906 100644 --- a/js/driver/sqlite/src/driver.ts +++ b/js/driver/sqlite/src/driver.ts @@ -623,7 +623,7 @@ export class SqliteRuntime implements RuntimeDriver< insertNotification( database, NOTIFICATION_TOPIC_CONTROL, - `{"action":"cancel","job_id":${updated.id},"queue":${JSON.stringify(updated.queue)}}` + `{"action":"cancel","attempt":${updated.attempt},"attempted_at":${JSON.stringify(updated.attemptedAt?.toString() ?? null)},"job_id":${updated.id},"queue":${JSON.stringify(updated.queue)}}` ); } return { diff --git a/js/etc/riverqueue.unstable-driver.api.md b/js/etc/riverqueue.unstable-driver.api.md index 9830babf3..7c55c41d6 100644 --- a/js/etc/riverqueue.unstable-driver.api.md +++ b/js/etc/riverqueue.unstable-driver.api.md @@ -541,6 +541,8 @@ export interface JobArgsTransformPlugin extends RiverPlugin { ```ts /** Remote cancellation of an attempt currently owned by this process. */ export interface JobCancellationNotice { + readonly attempt?: number; + readonly attemptedAt?: Temporal.Instant; readonly attemptedBy: string; readonly id: bigint; } diff --git a/js/src/driver.ts b/js/src/driver.ts index 7134cc221..98af925b7 100644 --- a/js/src/driver.ts +++ b/js/src/driver.ts @@ -303,6 +303,8 @@ export function jobCompletionKey( /** Remote cancellation of an attempt currently owned by this process. */ export interface JobCancellationNotice { + readonly attempt?: number; + readonly attemptedAt?: Temporal.Instant; readonly attemptedBy: string; readonly id: bigint; } diff --git a/js/src/runtime.test.ts b/js/src/runtime.test.ts index 7ddfbb19b..751914b61 100644 --- a/js/src/runtime.test.ts +++ b/js/src/runtime.test.ts @@ -4606,6 +4606,91 @@ describe("Client runtime database faults", () => { await run.stop(); }); + it("delivers identified cancellations during blocked reads and ignores older attempts", async () => { + const driver = new FakeRuntimeDriver(); + const started = Promise.withResolvers(); + const readStarted = Promise.withResolvers(); + const releaseRead = Promise.withResolvers(); + const releaseWorker = Promise.withResolvers(); + const deliver = Promise.withResolvers(); + const attemptedAt = Temporal.Instant.from("2026-01-02T03:04:05.6789Z"); + driver.claim = [{ ...fakeJob(), attempt: 2, attemptedAt }]; + Object.assign(driver, { + jobGetCancelRequested: () => { + readStarted.resolve(undefined); + return releaseRead.promise; + }, + }); + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + await deliver.promise; + // Legacy hint starts a read, simulating a saturated pool. + yield { + topic: "control", + payload: '{"action":"cancel","job_id":101,"queue":"default"}', + }; + await readStarted.promise; + const workerSignal = await started.promise; + for (const identity of [ + { attempt: 1, attempted_at: attemptedAt.toString() }, + { + attempt: 2, + attempted_at: attemptedAt.subtract({ seconds: 1 }).toString(), + }, + ]) { + yield { + topic: "control", + payload: JSON.stringify({ + action: "cancel", + job_id: 101, + queue: "default", + ...identity, + }), + }; + expect(workerSignal.aborted).toBe(false); + } + yield { + topic: "control", + payload: JSON.stringify({ + action: "cancel", + job_id: 101, + queue: "default", + attempt: 2, + attempted_at: attemptedAt.toString(), + }), + }; + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + queues: { default: { maxWorkers: 1 } }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + started.resolve(signal); + await releaseWorker.promise; + signal.throwIfAborted(); + } + ), + }); + const run = await client.start(); + try { + const signal = await started.promise; + deliver.resolve(undefined); + await waitUntil(() => signal.aborted); + expect(signal.reason).toMatchObject({ name: "JobCancelledError" }); + } finally { + deliver.resolve(undefined); + releaseRead.resolve([]); + releaseWorker.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + for (const failFirst of [false, true]) { it(`only checks cancellations on hints and retries failures=${failFirst}`, async () => { const driver = new FakeRuntimeDriver(); diff --git a/js/src/runtime/attempt-runner.ts b/js/src/runtime/attempt-runner.ts index 214ada7b3..c0bfee966 100644 --- a/js/src/runtime/attempt-runner.ts +++ b/js/src/runtime/attempt-runner.ts @@ -169,9 +169,24 @@ export class AttemptRunner { controller.abort(reason); } - /** Check a cancellation hint against durable state before aborting work. */ - cancelAttempt(id: bigint, attemptedBy: string): void { + /** Cancel a matching execution, or reconcile a legacy hint against durable state. */ + cancelAttempt( + id: bigint, + attemptedBy: string, + identity?: { attempt: number; attemptedAt: Temporal.Instant } + ): void { const key = attemptKey(id, attemptedBy); + const active = this.#attempts.get(key); + if (active !== undefined && identity !== undefined) { + if ( + active.row.attempt === identity.attempt && + active.row.attemptedAt?.epochNanoseconds === + identity.attemptedAt.epochNanoseconds + ) { + active.controller.abort(new JobCancelledError(id)); + } + return; + } for (const pending of this.#claimCancellations) pending.add(key); if (this.#attempts.has(key)) void this.#checkCancellations(); } diff --git a/js/src/runtime/notification-payloads.ts b/js/src/runtime/notification-payloads.ts index 080b64cd2..72c8f8874 100644 --- a/js/src/runtime/notification-payloads.ts +++ b/js/src/runtime/notification-payloads.ts @@ -19,6 +19,33 @@ export function notificationQueue(payload: string): string | null { return null; } +/** Optional execution identity lets cancellation bypass a blocked database read. */ +export function notificationCancellationAttempt( + payload: string +): { attempt: number; attemptedAt: Temporal.Instant } | undefined { + try { + const value: unknown = JSON.parse(payload); + if ( + value !== null && + typeof value === "object" && + "attempt" in value && + typeof value.attempt === "number" && + Number.isSafeInteger(value.attempt) && + value.attempt > 0 && + "attempted_at" in value && + typeof value.attempted_at === "string" + ) { + return { + attempt: value.attempt, + attemptedAt: Temporal.Instant.from(value.attempted_at), + }; + } + } catch { + // Old or malformed identities require durable reconciliation. + } + return undefined; +} + /** The job ID a control notification asks to cancel, if any. */ export function notificationCancellation(payload: string): bigint | null { if (!/"action"\s*:\s*"cancel"/.test(payload)) return null; diff --git a/js/src/runtime/notification-pump.conformance.test.ts b/js/src/runtime/notification-pump.conformance.test.ts index a1f9135a9..d657fa49c 100644 --- a/js/src/runtime/notification-pump.conformance.test.ts +++ b/js/src/runtime/notification-pump.conformance.test.ts @@ -106,8 +106,18 @@ describe.each(["postgres", "sqlite"] as const)( }, } as unknown as QueueProducer; const runner = { - cancelAttempt: (id: bigint, owner: string) => { - effects.push(["cancel", id, owner]); + cancelAttempt: ( + id: bigint, + owner: string, + identity: { attempt: number; attemptedAt: Temporal.Instant } + ) => { + effects.push([ + "cancel", + id, + owner, + identity.attempt, + identity.attemptedAt.toString(), + ]); }, } as unknown as AttemptRunner; const pump = new NotificationPump(context, { @@ -129,10 +139,7 @@ describe.each(["postgres", "sqlite"] as const)( // Semantic expectations are independent of the fixture's payload. // Checking all effects also rules out unrelated cancellations/wakeups. const expected: Record = { - cancel: [ - ["cancel", 42n, clientId], - ["refresh", "priority"], - ], + cancel: [["cancel", 42n, clientId, 1, "2026-01-02T03:04:05.678Z"]], insert: [["wake", "priority"]], metadata_changed: [["refresh", "priority"]], pause: [["refresh", "priority"]], diff --git a/js/src/runtime/notification-pump.ts b/js/src/runtime/notification-pump.ts index ddbd12af7..958a17205 100644 --- a/js/src/runtime/notification-pump.ts +++ b/js/src/runtime/notification-pump.ts @@ -11,6 +11,7 @@ import type { RuntimeContext } from "./context.js"; import { backOffAfterFailure } from "./context.js"; import { notificationCancellation, + notificationCancellationAttempt, notificationLeaderResigned, notificationQueue, notificationRequestsLeadershipResignation, @@ -94,7 +95,12 @@ export class NotificationPump { case "control": { const cancellation = notificationCancellation(notification.payload); if (cancellation !== null) { - this.#runner.cancelAttempt(cancellation, this.#context.clientId); + this.#runner.cancelAttempt( + cancellation, + this.#context.clientId, + notificationCancellationAttempt(notification.payload) + ); + return; } const queueName = notificationQueue(notification.payload); if (queueName === "*") { @@ -163,7 +169,13 @@ export class NotificationPump { "remote cancellation stream", (signal, ready) => subscribe(this.#context.clientId, signal, ready), (notice) => { - this.#runner.cancelAttempt(notice.id, notice.attemptedBy); + this.#runner.cancelAttempt( + notice.id, + notice.attemptedBy, + notice.attempt !== undefined && notice.attemptedAt !== undefined + ? { attempt: notice.attempt, attemptedAt: notice.attemptedAt } + : undefined + ); }, () => { if (subscribed) { diff --git a/producer.go b/producer.go index 82d6650f5..612621afa 100644 --- a/producer.go +++ b/producer.go @@ -549,10 +549,12 @@ const ( ) type controlEventPayload struct { - Action controlAction `json:"action"` - JobID int64 `json:"job_id,omitempty"` - Metadata json.RawMessage `json:"metadata,omitempty"` - Queue string `json:"queue"` + Action controlAction `json:"action"` + Attempt int `json:"attempt,omitempty"` + AttemptedAt *time.Time `json:"attempted_at,omitempty"` + JobID int64 `json:"job_id,omitempty"` + Metadata json.RawMessage `json:"metadata,omitempty"` + Queue string `json:"queue"` } type insertPayload struct { @@ -591,7 +593,18 @@ func (p *producer) handleControlNotification(fetchCtx, workCtx context.Context) ) return } - p.TriggerCancelCheck() + // New senders identify the execution, so delivery never needs a pool + // connection. Legacy notices and claims still being registered reconcile. + p.activeJobsMu.Lock() + executor := p.activeJobs[decoded.JobID] + p.activeJobsMu.Unlock() + if executor != nil && decoded.Attempt > 0 && decoded.AttemptedAt != nil { + if executor.JobRow.Attempt == decoded.Attempt && executor.JobRow.AttemptedAt != nil && executor.JobRow.AttemptedAt.Equal(*decoded.AttemptedAt) { + executor.Cancel(workCtx) + } + } else { + p.TriggerCancelCheck() + } default: p.Logger.DebugContext(workCtx, p.Name+": Received job control notification with unknown action", slog.String("action", string(decoded.Action)), diff --git a/producer_test.go b/producer_test.go index 2424a9251..b3ff82973 100644 --- a/producer_test.go +++ b/producer_test.go @@ -157,11 +157,13 @@ func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { blockedFetch bool cancelFromLookup bool queryError bool + remote bool }{ {name: "BlockedFetchCancelDelivery", blockedFetch: true}, {name: "BlockedFetchQueryResult", blockedFetch: true, cancelFromLookup: true}, {name: "CancelDelivery"}, {name: "QueryResult", cancelFromLookup: true}, + {name: "RemoteCancelDelivery", remote: true}, {name: "RetryAfterError", cancelFromLookup: true, queryError: true}, } { t.Run(testCase.name, func(t *testing.T) { @@ -185,7 +187,7 @@ func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { 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}}) + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelJob, JobRow: &rivertype.JobRow{ID: 1, Attempt: 1, AttemptedAt: new(time.Now().UTC())}}) var queryAttempts int exec := &producerTestExecutor{ jobGetAvailableFunc: func(context.Context, *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) { @@ -217,7 +219,7 @@ func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { prod := baseservice.Init(archetype, &producer{ activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, cancelPollCh: make(chan struct{}, 1), - config: &producerConfig{MaxWorkers: 2, Notifier: ¬ifier.Notifier{}, QueuePollInterval: 20 * time.Millisecond}, + config: &producerConfig{MaxWorkers: 2, Notifier: ¬ifier.Notifier{}, Queue: QueueDefault, QueuePollInterval: 20 * time.Millisecond}, exec: exec, fetchLimiter: chanutil.NewDebouncedChan(ctx, time.Millisecond, true), jobResultCh: make(chan *rivertype.JobRow), @@ -253,7 +255,9 @@ func TestProducer_CancellationCheckDoesNotBlock(t *testing.T) { require.Positive(t, signals.QueryStarted.WaitOrTimeout()) } - if !testCase.cancelFromLookup { + if testCase.remote { + prod.handleControlNotification(ctx, ctx)(notifier.NotificationTopicControl, fmt.Sprintf(`{"action":"cancel","attempt":1,"attempted_at":%q,"job_id":1,"queue":"default"}`, executor.JobRow.AttemptedAt.Format(time.RFC3339Nano))) + } else if !testCase.cancelFromLookup { client := &Client[struct{}]{producersByQueueName: map[string]*producer{QueueDefault: prod}} job := *executor.JobRow job.Queue = QueueDefault @@ -320,6 +324,33 @@ func TestProducer_CancellationChecksWaitForHints(t *testing.T) { } } +func TestProducer_CancellationNotificationAttemptIdentity(t *testing.T) { + t.Parallel() + + attemptedAt := time.Date(2026, 1, 2, 3, 4, 5, 678900000, time.UTC) + ctx, cancel := context.WithCancelCause(context.Background()) + defer cancel(nil) + archetype := riversharedtest.BaseServiceArchetype(t) + executor := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancel, JobRow: &rivertype.JobRow{Attempt: 2, AttemptedAt: &attemptedAt}}) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, + cancelPollCh: make(chan struct{}, 1), + config: &producerConfig{Queue: QueueDefault}, + }) + deliver := prod.handleControlNotification(context.Background(), context.Background()) + for _, payload := range []string{ + `{"action":"cancel","attempt":1,"attempted_at":"2026-01-02T03:04:05.6789Z","job_id":1,"queue":"default"}`, + `{"action":"cancel","attempt":2,"attempted_at":"2026-01-02T03:04:04.6789Z","job_id":1,"queue":"default"}`, + } { + deliver(notifier.NotificationTopicControl, payload) + require.NoError(t, ctx.Err()) + require.Empty(t, prod.cancelPollCh) + } + deliver(notifier.NotificationTopicControl, `{"action":"cancel","attempt":2,"attempted_at":"2026-01-02T03:04:05.6789Z","job_id":1,"queue":"default"}`) + require.ErrorIs(t, context.Cause(ctx), rivertype.ErrJobCancelledRemotely) + require.Empty(t, prod.cancelPollCh) +} + func TestProducer_HandleControlNotification(t *testing.T) { t.Parallel() diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 05b4fd854..3681ea7b6 100644 --- a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go @@ -16,7 +16,7 @@ import ( const jobCancel = `-- name: JobCancel :one WITH locked_job AS ( SELECT - id, queue, state, finalized_at + id, attempt, attempted_at, queue, state, finalized_at FROM /* TEMPLATE: schema */river_job WHERE river_job.id = $1 FOR UPDATE @@ -26,7 +26,7 @@ notification AS ( id, CASE WHEN $2::boolean THEN pg_notify( concat(coalesce($3::text, current_schema()), '.', $4::text), - json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text ) END FROM locked_job diff --git a/riverdriver/riverdrivertest/listener.go b/riverdriver/riverdrivertest/listener.go index 6b8df6de8..e87efe52c 100644 --- a/riverdriver/riverdrivertest/listener.go +++ b/riverdriver/riverdrivertest/listener.go @@ -2,6 +2,7 @@ package riverdrivertest import ( "context" + "encoding/json" "fmt" "testing" "time" @@ -111,8 +112,20 @@ func exerciseListener[TTx any](ctx context.Context, t *testing.T, driverWithPool notification := waitForNotification(ctx, t, listener) require.Equal(t, string(notifier.NotificationTopicControl), notification.Topic) - require.JSONEq(t, fmt.Sprintf(`{"action":"cancel","job_id":%d,"queue":%q}`, - committedJob.ID, committedJob.Queue), notification.Payload) + var payload struct { + Action string `json:"action"` + Attempt int `json:"attempt"` + AttemptedAt *time.Time `json:"attempted_at"` + JobID int64 `json:"job_id"` + Queue string `json:"queue"` + } + require.NoError(t, json.Unmarshal([]byte(notification.Payload), &payload)) + require.Equal(t, "cancel", payload.Action) + require.Equal(t, committedJob.Attempt, payload.Attempt) + require.NotNil(t, payload.AttemptedAt) + require.True(t, committedJob.AttemptedAt.Equal(*payload.AttemptedAt)) + require.Equal(t, committedJob.ID, payload.JobID) + require.Equal(t, committedJob.Queue, payload.Queue) requireNoNotification(ctx, t, listener) }) diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql index ac038c1c4..7b4b2ede6 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql @@ -40,7 +40,7 @@ CREATE TABLE river_job ( -- name: JobCancel :one WITH locked_job AS ( SELECT - id, queue, state, finalized_at + id, attempt, attempted_at, queue, state, finalized_at FROM /* TEMPLATE: schema */river_job WHERE river_job.id = @id FOR UPDATE @@ -50,7 +50,7 @@ notification AS ( id, CASE WHEN @notify::boolean THEN pg_notify( concat(coalesce(sqlc.narg('schema')::text, current_schema()), '.', @control_topic::text), - json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text ) END FROM locked_job diff --git a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go index 069001869..af3ef49e5 100644 --- a/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go +++ b/riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql.go @@ -16,7 +16,7 @@ import ( const jobCancel = `-- name: JobCancel :one WITH locked_job AS ( SELECT - id, queue, state, finalized_at + id, attempt, attempted_at, queue, state, finalized_at FROM /* TEMPLATE: schema */river_job WHERE river_job.id = $1 FOR UPDATE @@ -26,7 +26,7 @@ notification AS ( id, CASE WHEN $2::boolean THEN pg_notify( concat(coalesce($3::text, current_schema()), '.', $4::text), - json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text + json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text ) END FROM locked_job diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql index 4c78b2a9d..39b8ce561 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql +++ b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql @@ -30,10 +30,11 @@ FROM /* TEMPLATE: schema */river_notification; -- name: NotificationInsertJobCancel :exec INSERT INTO /* TEMPLATE: schema */river_notification (payload, topic) -VALUES ( - json_object('action', 'cancel', 'job_id', cast(@job_id AS integer), 'queue', cast(@queue AS text)), +SELECT + json_object('action', 'cancel', 'attempt', job.attempt, 'attempted_at', strftime('%Y-%m-%dT%H:%M:%fZ', job.attempted_at), 'job_id', job.id, 'queue', cast(@queue AS text)), @topic -); +FROM /* TEMPLATE: schema */river_job AS job +WHERE job.id = @job_id; -- name: NotificationInsertMany :exec INSERT INTO /* TEMPLATE: schema */river_notification ( diff --git a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go index a5cf11495..032918c56 100644 --- a/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go +++ b/riverdriver/riversqlite/internal/dbsqlc/river_notification.sql.go @@ -91,20 +91,21 @@ func (q *Queries) NotificationGetLastID(ctx context.Context, db DBTX) (int64, er const notificationInsertJobCancel = `-- name: NotificationInsertJobCancel :exec INSERT INTO /* TEMPLATE: schema */river_notification (payload, topic) -VALUES ( - json_object('action', 'cancel', 'job_id', cast(?1 AS integer), 'queue', cast(?2 AS text)), - ?3 -) +SELECT + json_object('action', 'cancel', 'attempt', job.attempt, 'attempted_at', strftime('%Y-%m-%dT%H:%M:%fZ', job.attempted_at), 'job_id', job.id, 'queue', cast(?1 AS text)), + ?2 +FROM /* TEMPLATE: schema */river_job AS job +WHERE job.id = ?3 ` type NotificationInsertJobCancelParams struct { - JobID int64 Queue string Topic string + JobID int64 } func (q *Queries) NotificationInsertJobCancel(ctx context.Context, db DBTX, arg *NotificationInsertJobCancelParams) error { - _, err := db.ExecContext(ctx, notificationInsertJobCancel, arg.JobID, arg.Queue, arg.Topic) + _, err := db.ExecContext(ctx, notificationInsertJobCancel, arg.Queue, arg.Topic, arg.JobID) return err } diff --git a/ruby/lib/client_runtime.rb b/ruby/lib/client_runtime.rb index 9092499fe..5ab995e8b 100644 --- a/ruby/lib/client_runtime.rb +++ b/ruby/lib/client_runtime.rb @@ -577,7 +577,7 @@ def wake end @mutex.synchronize do - @running[row.id] = {check_before_work: @fetching_queues.fetch(row.queue).include?(row.id) || (generation && generation != @cancellation_generation), queue: row.queue, thread: thread, working: false} + @running[row.id] = {attempt: row.attempt, attempted_at: row.attempted_at, check_before_work: @fetching_queues.fetch(row.queue).include?(row.id) || (generation && generation != @cancellation_generation), queue: row.queue, thread: thread, working: false} end gate.push(true) end @@ -841,10 +841,17 @@ def wake return unless message["job_id"].is_a?(Integer) @mutex.synchronize do + entry = @running[message["job_id"]] + if entry && message["attempt"].is_a?(Integer) && message["attempt"].positive? && message["attempted_at"].is_a?(String) + if entry[:queue] == queue && entry[:attempt] == message["attempt"] && entry[:attempted_at] == Time.iso8601(message["attempted_at"]) + entry[:cancelled] = true + entry[:thread].raise(JobCancelError) if entry[:working] + end + return + end @cancellation_generation += 1 @cancellation_pending = true @condition.broadcast - entry = @running[message["job_id"]] if entry && entry[:queue] == queue entry[:check_before_work] = true unless entry[:working] elsif (pending = @fetching_queues[queue]) @@ -869,7 +876,7 @@ def wake wake_maintenance if message["leader_id"].is_a?(String) && message["leader_id"] != @config.id end end - rescue JSON::ParserError, TypeError => error + rescue ArgumentError, JSON::ParserError, TypeError => error @config.logger.warn("River ignored invalid notification: #{error.message}") end diff --git a/ruby/lib/driver/runtime.rb b/ruby/lib/driver/runtime.rb index ed64584f1..afc90cf63 100644 --- a/ruby/lib/driver/runtime.rb +++ b/ruby/lib/driver/runtime.rb @@ -60,7 +60,7 @@ def job_cancel(id, now: Time.now.utc) row = job_get_by_id(updated_id || original_id) # Other engines interrupt running workers through this notification; # the durable marker remains Ruby's polling fallback. - runtime_notify("river_control", action: "cancel", job_id: row.id, queue: row.queue) if updated_id + runtime_notify("river_control", action: "cancel", attempt: row.attempt, attempted_at: row.attempted_at&.iso8601(6), job_id: row.id, queue: row.queue) if updated_id row end end diff --git a/ruby/spec/conformance_spec.rb b/ruby/spec/conformance_spec.rb index d08025afc..dc68fb3cf 100644 --- a/ruby/spec/conformance_spec.rb +++ b/ruby/spec/conformance_spec.rb @@ -163,7 +163,7 @@ class ConformanceRecord < ActiveRecord::Base row = @client.insert(River::JobArgsHash.new(:notify, {}), queue: "priority").job if name == "cancel" @driver.send(:runtime_execute, "UPDATE river_job SET id = 42 WHERE id = #{row.id}") - @driver.job_claim(id: 42, attempted_by: "conformance") + @driver.job_claim(id: 42, attempted_by: "conformance", allow_scheduled: true, now: Time.iso8601("2026-01-02T03:04:05.678Z")) @client.job_cancel(42) end when "metadata_changed", "pause", "resume" @@ -189,7 +189,13 @@ class ConformanceRecord < ActiveRecord::Base ((name == "insert") ? "insert" : "control") end expect(@driver.send(:runtime_value, notification, :topic)).to eq(protocol.fetch("topics").fetch(topic)) - expect(JSON.parse(@driver.send(:runtime_value, notification, :payload))).to eq(fixture.fetch("payload")) + payload = JSON.parse(@driver.send(:runtime_value, notification, :payload)) + expected = fixture.fetch("payload").dup + if name == "cancel" + payload["attempted_at"] = Time.iso8601(payload.fetch("attempted_at")) + expected["attempted_at"] = Time.iso8601(expected.fetch("attempted_at")) + end + expect(payload).to eq(expected) end end diff --git a/ruby/spec/insert_notification_shared_examples.rb b/ruby/spec/insert_notification_shared_examples.rb index 2ac9fce40..3805f0c45 100644 --- a/ruby/spec/insert_notification_shared_examples.rb +++ b/ruby/spec/insert_notification_shared_examples.rb @@ -73,12 +73,12 @@ def insert_notifications it "publishes Go-compatible cancellation only when the transaction commits" do client = River::Client.new(@driver) row = client.insert(River::JobArgsHash.new(:notify, {})).job - @driver.job_claim(id: row.id, attempted_by: "test") + claimed = @driver.job_claim(id: row.id, attempted_by: "test") @driver.transaction do client.job_cancel(row.id) expect(insert_notifications).to be_empty end - expect(insert_notifications).to eq([{"action" => "cancel", "job_id" => row.id, "queue" => "default"}]) + expect(insert_notifications).to eq([{"action" => "cancel", "attempt" => claimed.attempt, "attempted_at" => claimed.attempted_at.iso8601(6), "job_id" => row.id, "queue" => "default"}]) end it "does not notify rolled-back, missing, or already-finalized cancellations" do @@ -107,14 +107,14 @@ def cancellation_notifications it "makes cancellation and its outbox notification visible together on commit" do client = River::Client.new(@driver) row = client.insert(River::JobArgsHash.new(:notify, {})).job - @driver.job_claim(id: row.id, attempted_by: "test") + claimed = @driver.job_claim(id: row.id, attempted_by: "test") @driver.transaction do client.job_cancel(row.id) expect(cancellation_notifications.length).to eq(1) observer = Thread.new { cancellation_notifications } expect(Timeout.timeout(5) { observer.value }).to be_empty end - expect(cancellation_notifications).to eq([{"action" => "cancel", "job_id" => row.id, "queue" => "default"}]) + expect(cancellation_notifications).to eq([{"action" => "cancel", "attempt" => claimed.attempt, "attempted_at" => claimed.attempted_at.iso8601(6), "job_id" => row.id, "queue" => "default"}]) end it "rolls back both cancellation and the outbox row" do diff --git a/ruby/spec/notification_runtime_spec.rb b/ruby/spec/notification_runtime_spec.rb index 139c008fe..58043d8f3 100644 --- a/ruby/spec/notification_runtime_spec.rb +++ b/ruby/spec/notification_runtime_spec.rb @@ -47,6 +47,34 @@ end end + it "delivers an identified cancellation while a read is blocked and ignores older attempts" do + value = runtime + queried = Queue.new + release = Queue.new + raised = Queue.new + driver.define_singleton_method(:job_get_cancelled_ids) do |_ids| + queried << true + release.pop + [] + end + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + attempted_at = Time.iso8601("2026-01-02T03:04:05.6789Z") + entry = {attempt: 2, attempted_at: attempted_at, queue: "one", thread: thread, working: true} + value.instance_variable_set(:@running, 42 => entry) + checker = Thread.new { value.send(:check_remote_cancellations) } + Timeout.timeout(2) { queried.pop } + [{attempt: 1, attempted_at: attempted_at}, {attempt: 2, attempted_at: attempted_at - 1}].each do |identity| + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", job_id: 42, queue: "one", attempt: identity[:attempt], attempted_at: identity[:attempted_at].iso8601(6))) + expect(raised).to be_empty + end + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", job_id: 42, queue: "one", attempt: 2, attempted_at: attempted_at.iso8601(6))) + expect(Timeout.timeout(2) { raised.pop }).to eq(River::JobCancelError) + ensure + release&.push(true) + checker&.join + end + it "backs off repeated connection failures without indefinitely delaying shutdown" do listener = Object.new listener.define_singleton_method(:poll) { |_| raise "connection lost" } @@ -146,6 +174,9 @@ it "ignores invalid payloads without disrupting the receiver" do [nil, "{", "null", "[]", "{}"].each { |payload| runtime.send(:receive_notification, "river_insert", payload) } + runtime.instance_variable_set(:@running, 42 => {attempt: 1, queue: "one"}) + runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", attempt: 1, attempted_at: "invalid", job_id: 42, queue: "one")) + runtime.instance_variable_set(:@running, {}) runtime.send(:receive_notification, "unknown", "{}") runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: "42")) runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) diff --git a/ruby/spec/runtime_draining_shared_examples.rb b/ruby/spec/runtime_draining_shared_examples.rb index 863ce6372..587f58cf1 100644 --- a/ruby/spec/runtime_draining_shared_examples.rb +++ b/ruby/spec/runtime_draining_shared_examples.rb @@ -86,6 +86,12 @@ end super(ids) end + if poll_failure + # Legacy senders require a durable check, including its retry after a failed read. + @driver.define_singleton_method(:runtime_notify) do |topic, payload| + super(topic, payload.except(:attempt, :attempted_at)) + end + end runtime.define_singleton_method(:queue_stopping?) do |queue| stopping = super(queue) if stopping && !signalled diff --git a/rust/riverqueue/src/client/attempts.rs b/rust/riverqueue/src/client/attempts.rs index 5280d7569..b211a3a8a 100644 --- a/rust/riverqueue/src/client/attempts.rs +++ b/rust/riverqueue/src/client/attempts.rs @@ -8,8 +8,15 @@ use tokio_util::sync::CancellationToken; use crate::client::ClientInner; +#[derive(Clone, Debug, PartialEq)] +pub(super) struct RunningAttempt { + pub(super) attempt: i32, + pub(super) attempted_at: Option>, + pub(super) cancellation: CancellationToken, +} + pub(super) fn remove_running_attempt( - running: &Mutex>, + running: &Mutex>, job_id: i64, cancellation: &CancellationToken, ) { @@ -18,7 +25,7 @@ pub(super) fn remove_running_attempt( .unwrap_or_else(std::sync::PoisonError::into_inner); if running .get(&job_id) - .is_some_and(|active| active == cancellation) + .is_some_and(|active| active.cancellation == *cancellation) { running.remove(&job_id); } @@ -29,12 +36,21 @@ pub(super) fn register_running_attempt( job_id: i64, cancellation: &CancellationToken, generation: u64, + attempt: i32, + attempted_at: Option>, ) { inner .running .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(job_id, cancellation.clone()); + .insert( + job_id, + RunningAttempt { + attempt, + attempted_at, + cancellation: cancellation.clone(), + }, + ); if inner.cancellation_generation.load(Ordering::Acquire) != generation { inner.cancellation_wakeup.notify_one(); } diff --git a/rust/riverqueue/src/client/mod.rs b/rust/riverqueue/src/client/mod.rs index 011d52f5f..6b4c813ff 100644 --- a/rust/riverqueue/src/client/mod.rs +++ b/rust/riverqueue/src/client/mod.rs @@ -58,7 +58,6 @@ use std::{ use chrono::{DateTime, Utc}; use tokio::sync::{Notify, broadcast, mpsc, watch}; -use tokio_util::sync::CancellationToken; use self::completer::CompletionUpdate; pub(crate) use self::notify_limiter::InsertNotifyLimiter; @@ -146,7 +145,7 @@ pub(crate) struct ClientInner { queue_notifications: broadcast::Sender, queues: RwLock>, pub(crate) retry_policy: Arc, - running: Mutex>, + running: Mutex>, #[cfg(feature = "postgres")] pub(crate) schema: SchemaName, soft_stop_timeout: Option, diff --git a/rust/riverqueue/src/client/notifier.rs b/rust/riverqueue/src/client/notifier.rs index 82d35a070..d8cacad27 100644 --- a/rust/riverqueue/src/client/notifier.rs +++ b/rust/riverqueue/src/client/notifier.rs @@ -39,6 +39,8 @@ const LISTENER_PING_INTERVAL: Duration = Duration::from_secs(5); #[derive(Deserialize)] pub(super) struct ControlNotification { pub(super) action: String, + pub(super) attempt: Option, + pub(super) attempted_at: Option>, pub(super) job_id: Option, pub(super) queue: Option, } @@ -124,7 +126,24 @@ pub(super) fn dispatch_notification( }; match payload.action.as_str() { "cancel" => { - if payload.job_id.is_some() { + if let Some(id) = payload.job_id { + if let (Some(attempt), Some(attempted_at)) = ( + payload.attempt.filter(|attempt| *attempt > 0), + payload.attempted_at, + ) { + let running = inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(active) = running.get(&id) { + if active.attempt == attempt + && active.attempted_at == Some(attempted_at) + { + active.cancellation.cancel(); + } + return; + } + } inner.request_cancellation_check(); } } diff --git a/rust/riverqueue/src/client/producer.rs b/rust/riverqueue/src/client/producer.rs index 92afed2a7..9d3f3be68 100644 --- a/rust/riverqueue/src/client/producer.rs +++ b/rust/riverqueue/src/client/producer.rs @@ -175,8 +175,8 @@ pub(super) async fn check_job_cancellations(inner: &ClientInner) -> Result<(), E 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(); + if let Some(attempt) = attempts.get(&id) { + attempt.cancellation.cancel(); } } } @@ -893,7 +893,14 @@ 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, row.id, &cancellation, cancellation_generation); + register_running_attempt( + &inner, + row.id, + &cancellation, + cancellation_generation, + row.attempt, + row.attempted_at, + ); let task_inner = Arc::clone(&inner); let completion_sender = completion_sender.clone(); let session_row = attempts.session.is_some().then(|| row.clone()); diff --git a/rust/riverqueue/src/client/tests.rs b/rust/riverqueue/src/client/tests.rs index 28bb36613..e09504450 100644 --- a/rust/riverqueue/src/client/tests.rs +++ b/rust/riverqueue/src/client/tests.rs @@ -1,7 +1,8 @@ use serde::Deserialize; use serde_json::Map; +use tokio_util::sync::CancellationToken; -use super::attempts::{register_running_attempt, remove_running_attempt}; +use super::attempts::{RunningAttempt, register_running_attempt, remove_running_attempt}; #[cfg(feature = "postgres")] use super::completer::CompletionBatcher; #[cfg(feature = "sqlite")] @@ -94,12 +95,26 @@ fn completion_cleanup_preserves_newer_attempt() { let job_id = 42; let first = CancellationToken::new(); let second = CancellationToken::new(); - let running = Mutex::new(HashMap::from([(job_id, first.clone())])); + let running = Mutex::new(HashMap::from([( + job_id, + RunningAttempt { + attempt: 1, + attempted_at: None, + cancellation: first.clone(), + }, + )])); running .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(job_id, second.clone()); + .insert( + job_id, + RunningAttempt { + attempt: 2, + attempted_at: None, + cancellation: second.clone(), + }, + ); remove_running_attempt(&running, job_id, &first); assert_eq!( @@ -107,7 +122,11 @@ fn completion_cleanup_preserves_newer_attempt() { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .get(&job_id), - Some(&second) + Some(&RunningAttempt { + attempt: 2, + attempted_at: None, + cancellation: second.clone() + }) ); remove_running_attempt(&running, job_id, &second); @@ -171,7 +190,14 @@ async fn go_notification_fixtures_dispatch_correctly() { let cancellation = CancellationToken::new(); let unrelated = CancellationToken::new(); for (job_id, token) in [(42, &cancellation), (43, &unrelated)] { - inner.running.lock().unwrap().insert(job_id, token.clone()); + inner.running.lock().unwrap().insert( + job_id, + RunningAttempt { + attempt: 1, + attempted_at: Some("2026-01-02T03:04:05.678Z".parse().unwrap()), + cancellation: token.clone(), + }, + ); } dispatch_notification( @@ -204,30 +230,20 @@ async fn go_notification_fixtures_dispatch_correctly() { ), _ => unreachable!(), } + assert_eq!(cancellation.is_cancelled(), name == "cancel", "{name}"); assert!( - !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", + .is_err(), "{name}" ); assert!(!unrelated.is_cancelled(), "{name}"); assert!( - matches!( - producer.try_recv(), - Err(broadcast::error::TryRecvError::Empty) - ), + producer.try_recv().is_err(), "{name}: unexpected producer notification" ); assert!( - matches!( - leadership.try_recv(), - Err(broadcast::error::TryRecvError::Empty) - ), + leadership.try_recv().is_err(), "{name}: unexpected leadership notification" ); } @@ -235,6 +251,40 @@ async fn go_notification_fixtures_dispatch_correctly() { pool.close().await; } +#[tokio::test] +async fn identified_cancellation_needs_no_database_and_ignores_older_attempts() { + #[cfg(feature = "postgres")] + let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); + #[cfg(not(feature = "postgres"))] + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); + let cancellation = CancellationToken::new(); + let attempted_at: chrono::DateTime = "2026-01-02T03:04:05.6789Z".parse().unwrap(); + register_running_attempt(&client.inner, 42, &cancellation, 0, 2, Some(attempted_at)); + for (attempt, at) in [ + (1, attempted_at), + (2, attempted_at - chrono::Duration::seconds(1)), + (2, attempted_at), + ] { + let payload = serde_json::json!({"action":"cancel", "attempt":attempt, "attempted_at":at, "job_id":42,"queue":"default"}); + dispatch_notification( + &client.inner, + &client.inner.queue_notifications, + crate::NOTIFICATION_TOPIC_CONTROL, + &payload.to_string(), + ); + assert_eq!( + cancellation.is_cancelled(), + attempt == 2 && at == attempted_at + ); + } + assert!( + tokio::time::timeout(Duration::ZERO, client.inner.cancellation_wakeup.notified()) + .await + .is_err() + ); +} + #[tokio::test] async fn registration_without_a_hint_does_not_query_cancellations() { #[cfg(feature = "postgres")] @@ -243,7 +293,7 @@ async fn registration_without_a_hint_does_not_query_cancellations() { let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); let client = Client::builder(pool).build().unwrap(); let cancellation = CancellationToken::new(); - register_running_attempt(&client.inner, 42, &cancellation, 0); + register_running_attempt(&client.inner, 42, &cancellation, 0, 1, None); assert!( tokio::time::timeout(Duration::ZERO, client.inner.cancellation_wakeup.notified(),) .await @@ -252,7 +302,7 @@ async fn registration_without_a_hint_does_not_query_cancellations() { // A hint consumed before registration still requires a fresh snapshot afterward. client.inner.request_cancellation_check(); client.inner.cancellation_wakeup.notified().await; - register_running_attempt(&client.inner, 43, &cancellation, 0); + register_running_attempt(&client.inner, 43, &cancellation, 0, 1, None); tokio::time::timeout( Duration::from_secs(1), client.inner.cancellation_wakeup.notified(), @@ -281,7 +331,7 @@ async fn cancellation_checks_wait_for_hints_and_retry_failed_reads() { .unwrap(); let client = Client::builder(pool.clone()).build().unwrap(); let cancellation = CancellationToken::new(); - register_running_attempt(&client.inner, 42, &cancellation, 0); + register_running_attempt(&client.inner, 42, &cancellation, 0, 1, None); let task = tokio::spawn(super::producer::poll_job_cancellations( Arc::clone(&client.inner), CancellationToken::new(), @@ -341,6 +391,8 @@ async fn delayed_cancellation_signals_recheck_retry_marker() { id, &original, client.inner.cancellation_generation.load(Ordering::Acquire), + 1, + None, ); client.jobs().cancel(id).await.unwrap(); sqlx::query("UPDATE river_job SET state = 'cancelled', finalized_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = ?").bind(id).execute(&pool).await.unwrap(); @@ -356,6 +408,8 @@ async fn delayed_cancellation_signals_recheck_retry_marker() { id, &replacement, client.inner.cancellation_generation.load(Ordering::Acquire), + 1, + None, ); dispatch_notification( &client.inner, @@ -737,11 +791,14 @@ async fn completer_abandons_its_backlog_when_a_batch_fails_during_shutdown() { let now = std::time::Instant::now(); for job_id in 0..i64::try_from(updates).unwrap() { let cancellation = CancellationToken::new(); - inner - .running - .lock() - .unwrap() - .insert(job_id, cancellation.clone()); + inner.running.lock().unwrap().insert( + job_id, + RunningAttempt { + attempt: 0, + attempted_at: None, + cancellation: cancellation.clone(), + }, + ); sender .try_send(CompletionUpdate { attempt: None, diff --git a/rust/riverqueue/src/storage/postgres.rs b/rust/riverqueue/src/storage/postgres.rs index 08b0b7fe0..4673aac8a 100644 --- a/rust/riverqueue/src/storage/postgres.rs +++ b/rust/riverqueue/src/storage/postgres.rs @@ -45,9 +45,9 @@ impl Backend for PostgresBackend<'_> { let table = self.schema.qualify("river_job"); let sql = format!( "WITH locked AS (\ - SELECT id, queue, state, finalized_at FROM {table} WHERE id = $1 FOR UPDATE\ + SELECT id, attempt, attempted_at, queue, state, finalized_at FROM {table} WHERE id = $1 FOR UPDATE\ ), notified AS (\ - SELECT id, CASE WHEN $5::boolean THEN pg_notify(concat(coalesce($2::text, current_schema()), '.', $3::text), json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text) END \ + SELECT id, CASE WHEN $5::boolean THEN pg_notify(concat(coalesce($2::text, current_schema()), '.', $3::text), json_build_object('action', 'cancel', 'attempt', attempt, 'attempted_at', attempted_at, 'job_id', id, 'queue', queue)::text) END \ FROM locked WHERE state NOT IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NULL\ ), updated AS (\ UPDATE {table} AS job SET \ diff --git a/rust/riverqueue/src/storage/sqlite.rs b/rust/riverqueue/src/storage/sqlite.rs index 57694909e..1f01ba3c4 100644 --- a/rust/riverqueue/src/storage/sqlite.rs +++ b/rust/riverqueue/src/storage/sqlite.rs @@ -70,6 +70,8 @@ impl Backend for SqliteBackend<'_> { }; let payload = serde_json::json!({ "action": "cancel", + "attempt": row.attempt, + "attempted_at": row.attempted_at, "job_id": id, "queue": row.queue, })