From 358334baadcae3c453ad6b8d11ba2f3e23da1a15 Mon Sep 17 00:00:00 2001 From: Chris George Date: Sun, 23 Aug 2026 14:32:18 -0700 Subject: [PATCH 1/4] 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 b0724cece..d54268254 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- A `JobCancel()` sent while a client's notifier is disconnected/reconnecting is no longer lost. Previously, the cancellation's `NOTIFY` could commit while nothing was `LISTEN`ing (Postgres `NOTIFY` is fire-and-forget), leaving the running job unaware of the cancellation until `JobRescuer`'s much longer stuck-job sweep. Each client now reconciles its currently running jobs against `cancel_attempted_at` by primary key every time its notifier (re)establishes healthy listening, closing the gap with no added steady-state query load. [PR #1361](https://github.com/riverqueue/river/pull/1361), as discussed in [#1358](https://github.com/riverqueue/river/issues/1358). - Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440). - Fixed the notifier retrying without backoff and flooding logs when the context passed to `Client.Start` reaches its deadline. [PR #1495](https://github.com/riverqueue/river/pull/1495). - Fixed `JobSnooze` accepting negative durations despite documenting that they panic. [PR #1529](https://github.com/riverqueue/river/pull/1529). diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index e860cb17b..b53358ddf 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 76ad0b820..9848e8ba1 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 3532d19832a4ffcf7c258f958d52fb27e02a86b6 Mon Sep 17 00:00:00 2001 From: Chris George Date: Sun, 23 Aug 2026 14:53:09 -0700 Subject: [PATCH 2/4] Deliver cancellation locally when the client has no notifier --- CHANGELOG.md | 1 + client.go | 73 +++++++++++++++++++++++++++--- client_test.go | 118 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 187 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d54268254..c8785e1a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- `JobCancel()` against a job running in the same process is now delivered immediately for clients with no notifier at all, including an explicit `Config.PollOnly` on an otherwise listener-capable driver. Previously, the local delivery path only activated for drivers that can't support `LISTEN`/`NOTIFY` at all, so `PollOnly: true` on the Postgres or SQLite drivers fell all the way back to `JobRescuer`'s stuck-job sweep for a job it could have cancelled synchronously in-process. [PR #1361](https://github.com/riverqueue/river/pull/1361). - 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). - Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440). - Fixed the notifier retrying without backoff and flooding logs when the context passed to `Client.Start` reaches its deadline. [PR #1495](https://github.com/riverqueue/river/pull/1495). diff --git a/client.go b/client.go index e06af4b0c..63cca3e69 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 } @@ -2982,6 +2993,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 472bac298..60d930c2e 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 2cea719d19048b4176713ed704dd0f2209fb7aa3 Mon Sep 17 00:00:00 2001 From: Chris George Date: Tue, 20 Oct 2026 07:35:49 -0700 Subject: [PATCH 3/4] 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. --- CHANGELOG.md | 5 +- 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 + 13 files changed, 676 insertions(+), 278 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c8785e1a9..55be01b14 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,8 +13,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed -- `JobCancel()` against a job running in the same process is now delivered immediately for clients with no notifier at all, including an explicit `Config.PollOnly` on an otherwise listener-capable driver. Previously, the local delivery path only activated for drivers that can't support `LISTEN`/`NOTIFY` at all, so `PollOnly: true` on the Postgres or SQLite drivers fell all the way back to `JobRescuer`'s stuck-job sweep for a job it could have cancelled synchronously in-process. [PR #1361](https://github.com/riverqueue/river/pull/1361). -- 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). +- Recover `JobCancel()` notifications lost during notifier disconnects by reconciling running jobs after reconnect. [PR #1361](https://github.com/riverqueue/river/pull/1361). +- Deliver `JobCancel()` locally in `PollOnly` mode without blocking unstarted or stopping clients. [PR #1361](https://github.com/riverqueue/river/pull/1361). +- Clear previous cancellation metadata when retrying a job. [PR #1361](https://github.com/riverqueue/river/pull/1361). - Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440). - Fixed the notifier retrying without backoff and flooding logs when the context passed to `Client.Start` reaches its deadline. [PR #1495](https://github.com/riverqueue/river/pull/1495). - Fixed `JobSnooze` accepting negative durations despite documenting that they panic. [PR #1529](https://github.com/riverqueue/river/pull/1529). diff --git a/client.go b/client.go index 63cca3e69..1686411d2 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, @@ -2993,55 +2984,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 60d930c2e..3f61ff946 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 b53358ddf..98f1d7801 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 9848e8ba1..0d138ca61 100644 --- a/internal/notifier/notifier_test.go +++ b/internal/notifier/notifier_test.go @@ -19,6 +19,7 @@ import ( "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/startstoptest" + "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivershared/util/maputil" "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" @@ -281,6 +282,57 @@ func TestNotifier(t *testing.T) { require.Equal(t, TopicAndPayload{testTopic1, "msg1"}, riversharedtest.WaitOrTimeout(t, notifyChan)) }) + t.Run("RegisterListenerReadyFunc", func(t *testing.T) { + t.Parallel() + + notifier, _ := setup(t, nil) + notifier.testDisableSleep = true + + var errorNum int + + listenerMock := NewListenerMock(notifier.listener) + listenerMock.waitForNotificationFunc = func(ctx context.Context) (*riverdriver.Notification, error) { + // Returns an error the first time, but then works after. + errorNum++ + switch errorNum { + case 1: + return nil, errors.New("error during wait") + default: + return listenerMock.Listener.WaitForNotification(ctx) + } + } + notifier.listener = listenerMock + + var ready testsignal.TestSignal[struct{}] + ready.Init(t) + sub := notifier.RegisterListenerReadyFunc(func() { ready.Signal(struct{}{}) }) + t.Cleanup(sub.Unregister) + + start(t, notifier) + + // Fires once on the initial connect. + ready.WaitOrTimeout() + + // And again once the notifier reconnects after the injected failure. + require.EqualError(t, notifier.testSignals.BackoffError.WaitOrTimeout(), "error during wait") + ready.WaitOrTimeout() + + notifier.Stop() + sub.Unregister() + + // WaitForNotification is entered after listener-ready callbacks finish. + var waiting testsignal.TestSignal[struct{}] + waiting.Init(t) + listenerMock.waitForNotificationFunc = func(ctx context.Context) (*riverdriver.Notification, error) { + waiting.Signal(struct{}{}) + return listenerMock.Listener.WaitForNotification(ctx) + } + + require.NoError(t, notifier.Start(ctx)) + waiting.WaitOrTimeout() + ready.RequireEmpty() + }) + t.Run("SingleTopicMultipleSubscribers", func(t *testing.T) { t.Parallel() @@ -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 a4de886f8..5080dc421 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 6e651cbfe..e48c4adf7 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 0ff9958d2..7787761f6 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 c4f80bdc7aff952f634eda030aeba2988f88b759 Mon Sep 17 00:00:00 2001 From: Brandur Date: Sat, 24 Oct 2026 12:20:20 -0500 Subject: [PATCH 4/4] Keep cancellation recovery active while workers drain 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 | 71 ++- client_test.go | 386 +++++++++++++- internal/notifier/notifier.go | 13 +- .../src/main/java/com/riverqueue/Client.java | 6 +- .../main/java/com/riverqueue/Protocol.java | 15 +- .../src/main/java/com/riverqueue/Workers.java | 65 ++- .../resources/com/riverqueue/postgres.sql | 8 +- .../main/resources/com/riverqueue/sqlite.sql | 8 +- .../java/com/riverqueue/ProtocolTest.java | 36 +- .../test/java/com/riverqueue/RetryTest.java | 112 ++++ .../test/java/com/riverqueue/WorkersTest.java | 179 +++++++ 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 | 222 ++++---- producer_test.go | 486 +++++++++++++----- .../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, 2541 insertions(+), 751 deletions(-) diff --git a/client.go b/client.go index 1686411d2..3a117a053 100644 --- a/client.go +++ b/client.go @@ -1231,6 +1231,16 @@ func (c *Client[TTx]) Start(ctx context.Context) error { fetchCtx := withClient(fetchCtx, c) workCtx = withClient(workCtx, c) + // Keep cancellation notifications and reconnects live until producers + // finish draining. The notifier is stopped with the other services below. + if c.notifier != nil { + if err := c.notifier.StartWithConnectContext(context.WithoutCancel(fetchCtx), fetchCtx); err != nil { + workCancel(err) + stopServicesOnError() + return err + } + } + if err := startstop.StartAll(fetchCtx, c.services...); err != nil { workCancel(err) stopServicesOnError() @@ -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 @@ -2408,14 +2414,15 @@ func (c *Client[TTx]) producerAdd(queueName string, queueConfig QueueConfig) (*p } func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) error { - c.producersMu.Lock() - defer c.producersMu.Unlock() - + c.producersMu.RLock() producer, ok := c.producersByQueueName[queueName] + c.producersMu.RUnlock() if !ok { return &QueueNotFoundError{Name: queueName} } + // QueueBundle serializes additions and removals. Keep the registry readable + // while draining so workers can still receive local cancellation requests. shouldStop, stopped, finalizeStop := producer.StopInit() if shouldStop { select { @@ -2427,7 +2434,9 @@ func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) erro } } + c.producersMu.Lock() delete(c.producersByQueueName, queueName) + c.producersMu.Unlock() return nil } @@ -2815,7 +2824,7 @@ func (c *Client[TTx]) QueuePause(ctx context.Context, name string, opts *QueuePa return err } - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) return nil } @@ -2885,7 +2894,7 @@ func (c *Client[TTx]) QueueResume(ctx context.Context, name string, opts *QueueP return err } - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) return nil } @@ -2945,7 +2954,7 @@ func (c *Client[TTx]) QueueUpdate(ctx context.Context, name string, params *Queu return nil, err } - c.notifyProducerWithoutListenerQueueControlEvent(name, controlEvent) + c.notifyProducerWithoutListenerQueueControlEvent(ctx, name, controlEvent) return queue, nil } @@ -2967,38 +2976,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 } @@ -3008,10 +3014,25 @@ func (c *Client[TTx]) notifyProducerOfCancelWithoutNotifier(ctx context.Context, default: } - // cancelCh is also drained while a database fetch is in progress. + producer.activeJobsMu.Lock() + executor := producer.activeJobs[job.ID] + if executor != nil { + sameAttemptedAt := executor.JobRow.AttemptedAt == job.AttemptedAt || + (executor.JobRow.AttemptedAt != nil && job.AttemptedAt != nil && executor.JobRow.AttemptedAt.Equal(*job.AttemptedAt)) + if executor.JobRow.Attempt != job.Attempt || !sameAttemptedAt { + executor = nil + } + } + producer.activeJobsMu.Unlock() + + // Keep the matching execution attached even if delivery waits on the buffer. select { - case producer.cancelCh <- controlEvent.JobID: - producer.testSignals.QueueControlEventTriggered.Signal(controlEvent) + case producer.cancelCh <- producerCancelRequest{executor: executor, jobID: job.ID}: + producer.testSignals.QueueControlEventTriggered.Signal(&controlEventPayload{ + Action: controlActionCancel, + JobID: job.ID, + Queue: job.Queue, + }) case <-ctx.Done(): case <-state.done: } diff --git a/client_test.go b/client_test.go index 3f61ff946..f0df01026 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() @@ -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() @@ -2871,12 +3021,12 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { setup := func(t *testing.T) *testBundle { t.Helper() - targetProducer := &producer{cancelCh: make(chan int64, 1)} + targetProducer := &producer{cancelCh: make(chan producerCancelRequest, 1)} ctx, cancel := context.WithCancel(context.Background()) ctx, _, started, stopped := targetProducer.StartInit(ctx) - targetProducer.cancelState.Store(&producerCancelState{done: ctx.Done()}) + targetProducer.runState.Store(&producerRunState{done: ctx.Done()}) started() - targetProducer.cancelCh <- 123 + targetProducer.cancelCh <- producerCancelRequest{jobID: 123} stop := sync.OnceFunc(func() { cancel() stopped() @@ -2892,6 +3042,50 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { } } + for _, testCase := range []struct { + name string + replaceBeforeDelivery bool + reuseAttempt bool + }{ + {name: "BufferedCancellationKeepsOriginalExecutor"}, + {name: "NewAttemptBeforeDeliveryIsReconciled", replaceBeforeDelivery: true}, + {name: "NewAttemptedAtBeforeDeliveryIsReconciled", replaceBeforeDelivery: true, reuseAttempt: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + <-bundle.producer.cancelCh + archetype := riversharedtest.BaseServiceArchetype(t) + originalCtx, cancelOriginal := context.WithCancelCause(context.Background()) + defer cancelOriginal(nil) + replacementCtx, cancelReplacement := context.WithCancelCause(context.Background()) + defer cancelReplacement(nil) + attemptedAt := time.Now().UTC() + job := &rivertype.JobRow{ID: 456, Attempt: 1, AttemptedAt: &attemptedAt, Queue: QueueDefault} + original := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelOriginal, JobRow: job}) + replacementJob := *job + replacementJob.AttemptedAt = new(attemptedAt.Add(time.Second)) + if !testCase.reuseAttempt { + replacementJob.Attempt++ + } + replacement := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelReplacement, JobRow: &replacementJob}) + bundle.producer.activeJobs = map[int64]*jobexecutor.JobExecutor{job.ID: original} + if testCase.replaceBeforeDelivery { + bundle.producer.activeJobs[job.ID] = replacement + } + + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), job) + bundle.producer.activeJobs[job.ID] = replacement + request := <-bundle.producer.cancelCh + require.Equal(t, !testCase.replaceBeforeDelivery, request.cancel(context.Background())) + require.NoError(t, replacementCtx.Err()) + if !testCase.replaceBeforeDelivery { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + } + }) + } + t.Run("FullBufferContextCancelled", func(t *testing.T) { t.Parallel() @@ -2902,7 +3096,7 @@ 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.notifyProducerOfCancelWithoutNotifier(ctx, &rivertype.JobRow{Queue: QueueDefault}) returned.Signal(struct{}{}) }() @@ -2917,7 +3111,7 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { var returned testsignal.TestSignal[struct{}] returned.Init(t) go func() { - bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), QueueDefault, &controlEventPayload{}) + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{Queue: QueueDefault}) returned.Signal(struct{}{}) }() bundle.stopped() @@ -2932,13 +3126,9 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { bundle := setup(t) <-bundle.producer.cancelCh - bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), QueueDefault, &controlEventPayload{ - Action: controlActionCancel, - JobID: 456, - Queue: QueueDefault, - }) + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{ID: 456, Queue: QueueDefault}) - require.Equal(t, int64(456), <-bundle.producer.cancelCh) + require.Equal(t, int64(456), (<-bundle.producer.cancelCh).jobID) }) t.Run("StopInProgress", func(t *testing.T) { @@ -2958,7 +3148,7 @@ 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.notifyProducerOfCancelWithoutNotifier(ctx, &rivertype.JobRow{Queue: QueueDefault}) returned.Signal(struct{}{}) }() @@ -2967,6 +3157,126 @@ func Test_Client_NotifyProducerOfCancelWithoutNotifier(t *testing.T) { }) } +func Test_Client_NotifyProducerWithoutListenerQueueControlEvent(t *testing.T) { + t.Parallel() + + type testBundle struct { + client *Client[pgx.Tx] + fetchStopped <-chan struct{} + producer *producer + stopFunc func() + } + + setup := func(t *testing.T) *testBundle { + t.Helper() + + ctx, cancel := context.WithCancel(context.Background()) + runCtx, cancelRun := context.WithCancel(context.Background()) + targetProducer := &producer{ + cancelCh: make(chan producerCancelRequest, 1), + queueControlCh: make(chan *controlEventPayload, 1), + } + ctx, _, started, stopped := targetProducer.StartInit(ctx) + targetProducer.runState.Store(&producerRunState{done: runCtx.Done(), fetchDone: ctx.Done()}) + started() + targetProducer.queueControlCh <- &controlEventPayload{Action: controlActionPause} + stopFunc := sync.OnceFunc(func() { + cancel() + cancelRun() + stopped() + }) + t.Cleanup(stopFunc) + + return &testBundle{ + client: &Client[pgx.Tx]{ + driver: &DriverPollOnly{}, + producersByQueueName: map[string]*producer{QueueDefault: targetProducer}, + }, + fetchStopped: ctx.Done(), + producer: targetProducer, + stopFunc: stopFunc, + } + } + + t.Run("FullBufferContextCancelled", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerWithoutListenerQueueControlEvent(ctx, QueueDefault, &controlEventPayload{}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.queueControlCh, 1) + }) + + t.Run("ProducerNotStarted", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + bundle.producer.runState.Store(nil) + + var returned testsignal.TestSignal[struct{}] + returned.Init(t) + go func() { + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, &controlEventPayload{}) + returned.Signal(struct{}{}) + }() + + returned.WaitOrTimeout() + require.Len(t, bundle.producer.queueControlCh, 1) + }) + + t.Run("QueueRemoval", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + var controlsReturned testsignal.TestSignal[struct{}] + var queueRemoved testsignal.TestSignal[error] + controlsReturned.Init(t) + queueRemoved.Init(t) + go func() { + queueRemoved.Signal(bundle.client.producerRemove(context.Background(), QueueDefault)) + }() + riversharedtest.WaitOrTimeout(t, bundle.fetchStopped) + + // The queue is still registered while draining, but queue controls must + // not wait on its full buffer or hold the registry lock needed by removal. + go func() { + for _, action := range []controlAction{controlActionPause, controlActionResume, controlActionMetadataChanged} { + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, &controlEventPayload{Action: action}) + } + controlsReturned.Signal(struct{}{}) + }() + controlsReturned.WaitOrTimeout() + require.Len(t, bundle.producer.queueControlCh, 1) + + // Cancellation delivery remains live until the running attempts drain. + bundle.client.notifyProducerOfCancelWithoutNotifier(context.Background(), &rivertype.JobRow{ID: 123, Queue: QueueDefault}) + require.Equal(t, int64(123), (<-bundle.producer.cancelCh).jobID) + bundle.stopFunc() + require.NoError(t, queueRemoved.WaitOrTimeout()) + require.Empty(t, bundle.client.producersByQueueName) + }) + + t.Run("RunningProducer", func(t *testing.T) { + t.Parallel() + + bundle := setup(t) + <-bundle.producer.queueControlCh + controlEvent := &controlEventPayload{Action: controlActionResume, Queue: QueueDefault} + bundle.client.notifyProducerWithoutListenerQueueControlEvent(context.Background(), QueueDefault, controlEvent) + + require.Same(t, controlEvent, <-bundle.producer.queueControlCh) + }) +} + func Test_Client_Stop_Common(t *testing.T) { t.Parallel() @@ -8934,6 +9244,47 @@ func Test_Client_Start_Error(t *testing.T) { }) } +func Test_Client_Start_NotifierCancellation(t *testing.T) { + t.Parallel() + + ctx := context.Background() + dbPool := riversharedtest.DBPool(ctx, t) + driver := riverpgxv5.New(dbPool) + schema := riverdbtest.TestSchema(ctx, t, driver, nil) + listener := newProducerTestListenerMock(driver.GetListener(&riverdriver.GetListenenerParams{Schema: schema})) + var connecting, releaseConnect testsignal.TestSignal[struct{}] + var startReturned testsignal.TestSignal[error] + connecting.Init(t) + releaseConnect.Init(t) + startReturned.Init(t) + listener.connectFunc = func(ctx context.Context) error { + connecting.Signal(struct{}{}) + select { + case <-ctx.Done(): + return ctx.Err() + case <-releaseConnect.WaitC(): + return errors.New("test released blocked connection") + } + } + client, err := NewClient(&DriverWithListener{Driver: driver, listener: listener}, newTestConfig(t, schema)) + require.NoError(t, err) + + startCtx, cancel := context.WithCancel(ctx) + defer cancel() + go func() { startReturned.Signal(client.Start(startCtx)) }() + t.Cleanup(func() { + releaseConnect.Signal(struct{}{}) + stopCtx, stopCancel := context.WithTimeout(ctx, 5*time.Second) + defer stopCancel() + require.NoError(t, client.Stop(stopCtx)) + }) + connecting.WaitOrTimeout() + cancel() + // The connection must honor startup cancellation, even though the notifier + // stays alive after a successful start until workers finish draining. + require.NoError(t, startReturned.WaitOrTimeout()) +} + func Test_Client_YugabyteQueueControl(t *testing.T) { t.Parallel() @@ -10825,6 +11176,17 @@ func NewDriverPollOnly(dbPool *pgxpool.Pool) *DriverPollOnly { func (d *DriverPollOnly) SupportsListener() bool { return false } +// DriverWithListener lets tests inject faults into a real listener. +type DriverWithListener struct { + *riverpgxv5.Driver + + listener riverdriver.Listener +} + +func (d *DriverWithListener) GetListener(_ *riverdriver.GetListenenerParams) riverdriver.Listener { + return d.listener +} + // DriverWithoutListenNotify simulates a driver without any listen/notify // support at all. An example of this is SQLite through `riversqlite` where // SQLite doesn't support listen/notify or anything like it. We could use the diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index 98f1d7801..93d3d34bc 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 b3a63512f..a867856d2 100644 --- a/java/river/src/main/java/com/riverqueue/Client.java +++ b/java/river/src/main/java/com/riverqueue/Client.java @@ -781,11 +781,7 @@ public Job retry(Connection connection, long id) { var before = locked(c, id); try (var statement = Sql.prepare( - c, - Sql.query(database, "retry"), - database.timestamp(clock.instant()), - id, - database.timestamp(clock.instant())); + c, Sql.query(database, "retry"), database.timestamp(clock.instant()), id); var rows = statement.executeQuery()) { boolean updated = rows.next(); var job = updated ? read(rows) : before; diff --git a/java/river/src/main/java/com/riverqueue/Protocol.java b/java/river/src/main/java/com/riverqueue/Protocol.java index 1e0a1da43..f13bdaae5 100644 --- a/java/river/src/main/java/com/riverqueue/Protocol.java +++ b/java/river/src/main/java/com/riverqueue/Protocol.java @@ -2,6 +2,7 @@ import java.util.Map; import java.util.function.Consumer; +import java.util.function.LongConsumer; import java.util.regex.Pattern; /** Stored and transmitted values shared with the other River implementations. */ @@ -57,8 +58,7 @@ record Notification(String topic, String payload) {} /** Shared by both database listeners; dispatching itself requires no database I/O. */ record Dispatcher( String clientId, - Map> active, - Map pendingCancellation, + LongConsumer cancellation, Consumer queues, Consumer leadership) { void dispatch(String topic, String payload) { @@ -70,14 +70,9 @@ void dispatch(String topic, String payload) { switch (topic) { case TOPIC_CONTROL -> { if (action.equals("cancel") && value.path("job_id").isIntegralNumber()) { - long jobId = value.path("job_id").asLong(); - // Publish pending cancellation first so an attempt registered concurrently sees it. - pendingCancellation.put(jobId, System.nanoTime()); - var context = active.get(jobId); - if (context != null) { - context.requestCancellation(WorkContext.Cancellation.REMOTE); - pendingCancellation.remove(jobId); - } + // Notices identify a job, not an attempt. Confirm its durable marker + // before cancelling so a delayed notice cannot undo a retry. + cancellation.accept(value.path("job_id").asLong()); } else if ((action.equals("metadata_changed") || action.equals("pause") || action.equals("resume")) diff --git a/java/river/src/main/java/com/riverqueue/Workers.java b/java/river/src/main/java/com/riverqueue/Workers.java index 567d515e1..7dc8e0dbc 100644 --- a/java/river/src/main/java/com/riverqueue/Workers.java +++ b/java/river/src/main/java/com/riverqueue/Workers.java @@ -31,7 +31,6 @@ public final class Workers implements AutoCloseable { private final Map> active = new ConcurrentHashMap<>(); private final Map activeThreads = new ConcurrentHashMap<>(); - private final Map pendingCancellation = new ConcurrentHashMap<>(); private final AtomicBoolean maintenanceBusy = new AtomicBoolean(); private volatile Instant nextReindex = Instant.now() @@ -40,6 +39,7 @@ public final class Workers implements AutoCloseable { .plusDays(1) .atStartOfDay() .toInstant(java.time.ZoneOffset.UTC); + private final Semaphore cancellationWake = new Semaphore(0); private final Builder config; private final RetryPolicy defaultRetryPolicy = RetryPolicy.defaults(); private final java.util.concurrent.BlockingQueue completions = @@ -81,7 +81,10 @@ private Workers(Builder config) { id = config.id == null ? "java-" + UUID.randomUUID() : config.id; notifications = new Protocol.Dispatcher( - id, active, pendingCancellation, this::queueNotification, this::leadershipNotification); + id, + ignored -> cancellationWake.release(), + this::queueNotification, + this::leadershipNotification); queues.putAll(config.queues); periodic = new PeriodicJobBundle( @@ -99,6 +102,7 @@ private Workers(Builder config) { unsubscribeInsert = river.onInsertCommit(wake::release); if (!config.pollOnly) executor.submit(this::listen); executor.submit(this::completeBatches); + executor.submit(this::pollCancellations); executor.submit(this::supervise); executor.submit(this::dispatch); if (config.leadership) executor.submit(this::elect); @@ -244,7 +248,6 @@ public Subscription subscribe(Consumer subscriber, EventKind... kinds) { private void dispatch() { long nextPoll = 0; long nextMaintenance = 0; - long nextCancelPoll = 0; while (running.get()) { try { long now = System.nanoTime(); @@ -272,25 +275,6 @@ private void dispatch() { queues.keySet().forEach(this::heartbeat); nextMaintenance = now + config.serviceInterval.toNanos(); } - if (now >= nextCancelPoll) { - river.transaction( - c -> { - try (var statement = - Sql.prepare(c, Sql.query(river.database(), "cancel_requested")); - var rows = statement.executeQuery()) { - while (rows.next()) { - var context = active.get(rows.getLong(1)); - if (context != null) - context.requestCancellation(WorkContext.Cancellation.REMOTE); - } - } - return null; - }); - pendingCancellation - .entrySet() - .removeIf(entry -> now - entry.getValue() > Duration.ofMinutes(1).toNanos()); - nextCancelPoll = now + TimeUnit.MILLISECONDS.toNanos(50); - } if (wake.tryAcquire(10, TimeUnit.MILLISECONDS)) wake.release(); } catch (InterruptedException e) { Thread.currentThread().interrupt(); @@ -443,8 +427,6 @@ private void launch(Job row, WorkContext context, WorkContext.Step active.put(row.id(), context); if (row.metadata().has("cancel_attempted_at")) context.requestCancellation(WorkContext.Cancellation.REMOTE); - if (pendingCancellation.remove(row.id()) != null) - context.requestCancellation(WorkContext.Cancellation.REMOTE); if (cancelOnStop) context.requestCancellation(WorkContext.Cancellation.SHUTDOWN); executor.submit( () -> { @@ -1394,6 +1376,41 @@ private void observe(String event) { } } + private void pollCancellations() { + // Fetches may block, and graceful stop ends dispatch before active work. + // Keep cancellation checks independent until every claimed attempt drains. + while (running.get() || dispatcherStopped.getCount() > 0 || !active.isEmpty()) { + try { + cancellationWake.drainPermits(); + var attempts = Map.copyOf(active); + if (!attempts.isEmpty()) { + river.transaction( + c -> { + try (var statement = + Sql.prepare(c, Sql.query(river.database(), "cancel_requested")); + var rows = statement.executeQuery()) { + while (rows.next()) { + // A late query result must not cancel a subsequent attempt. + var context = attempts.get(rows.getLong(1)); + if (context != null) + context.requestCancellation(WorkContext.Cancellation.REMOTE); + } + } + return null; + }); + } + } catch (Exception e) { + report(e); + } + try { + cancellationWake.tryAcquire(50, TimeUnit.MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + } + } + private void publish(Event event) { for (var subscriber : subscribers) try { diff --git a/java/river/src/main/resources/com/riverqueue/postgres.sql b/java/river/src/main/resources/com/riverqueue/postgres.sql index d5a92e3ca..10e7df6f5 100644 --- a/java/river/src/main/resources/com/riverqueue/postgres.sql +++ b/java/river/src/main/resources/com/riverqueue/postgres.sql @@ -56,9 +56,11 @@ DELETE FROM {schema}river_job WHERE id=? AND state != 'running' RETURNING {colum -- name: delete_filtered DELETE FROM {schema}river_job WHERE {where} AND id=? AND state != 'running' RETURNING {columns} -- name: retry -UPDATE {schema}river_job SET state='available', finalized_at=NULL, -max_attempts=CASE WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=? -WHERE id=? AND state != 'running' AND NOT (state='available' AND scheduled_at < ?) RETURNING {columns} +WITH retry_params AS (SELECT CAST(? AS timestamptz) AS now) +UPDATE {schema}river_job SET state='available', finalized_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, +metadata=CASE WHEN jsonb_typeof(metadata)='object' THEN metadata - 'cancel_attempted_at' ELSE metadata END, +max_attempts=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END +WHERE id=? AND state != 'running' AND ((jsonb_typeof(metadata)='object' AND jsonb_exists(metadata, 'cancel_attempted_at')) OR NOT (state='available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {columns} -- name: reset_jobs DELETE FROM {schema}river_job -- name: reset_queues diff --git a/java/river/src/main/resources/com/riverqueue/sqlite.sql b/java/river/src/main/resources/com/riverqueue/sqlite.sql index f38cfd267..2ccbff7de 100644 --- a/java/river/src/main/resources/com/riverqueue/sqlite.sql +++ b/java/river/src/main/resources/com/riverqueue/sqlite.sql @@ -55,9 +55,11 @@ DELETE FROM {schema}river_job WHERE id=? AND state != 'running' RETURNING {colum -- name: delete_filtered DELETE FROM {schema}river_job WHERE {where} AND id=? AND state != 'running' RETURNING {columns} -- name: retry -UPDATE {schema}river_job SET state='available', finalized_at=NULL, -max_attempts=CASE WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=? -WHERE id=? AND state != 'running' AND NOT (state='available' AND scheduled_at < ?) RETURNING {columns} +WITH retry_params AS (SELECT CAST(? AS text) AS now) +UPDATE {schema}river_job SET state='available', finalized_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, +metadata=CASE WHEN typeof(metadata)='text' AND NOT json_valid(metadata) THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, +max_attempts=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt=max_attempts THEN max_attempts+1 ELSE max_attempts END, scheduled_at=CASE WHEN state='available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END +WHERE id=? AND state != 'running' AND ((CASE WHEN typeof(metadata)='text' AND NOT json_valid(metadata) THEN NULL ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL OR NOT (state='available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {columns} -- name: reset_jobs DELETE FROM {schema}river_job -- name: reset_queues diff --git a/java/river/src/test/java/com/riverqueue/ProtocolTest.java b/java/river/src/test/java/com/riverqueue/ProtocolTest.java index eb0424058..116f0b250 100644 --- a/java/river/src/test/java/com/riverqueue/ProtocolTest.java +++ b/java/river/src/test/java/com/riverqueue/ProtocolTest.java @@ -5,7 +5,6 @@ import java.time.Instant; import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Random; @@ -79,27 +78,17 @@ Stream notificationDispatch() throws Exception { DynamicTest.dynamicTest( name + " received by " + clientId, () -> { - var target = context(42); - var unrelated = context(43); + var cancellations = new ArrayList(); var queues = new ArrayList(); var leadership = new ArrayList(); - var pending = new HashMap(); var dispatcher = new Protocol.Dispatcher( - clientId, - Map.of(42L, target, 43L, unrelated), - pending, - queues::add, - leadership::add); + clientId, cancellations::add, queues::add, leadership::add); dispatcher.dispatch( fixture.required("topic").asString(), Json.encode(fixture.required("payload"))); - assertEquals( - name.equals("cancel") ? WorkContext.Cancellation.REMOTE : null, - target.cancellation()); - assertNull(unrelated.cancellation(), "Notification must not cancel another job"); - assertTrue(pending.isEmpty(), "A registered attempt consumes its cancellation"); + assertEquals(name.equals("cancel") ? List.of(42L) : List.of(), cancellations); switch (name) { case "insert", "metadata_changed", "pause", "resume" -> { assertEquals(List.of(new Protocol.QueueNotice("priority", name)), queues); @@ -130,44 +119,35 @@ Stream notificationDispatch() throws Exception { } @Test - void notificationTopicsAndPendingCancellation() throws Exception { + void notificationTopicsRouteCancellationChecks() throws Exception { var fixture = Conformance.fixture("protocol_values.json").required("notifications"); var cancel = java.util.stream.StreamSupport.stream(fixture.spliterator(), false) .filter(value -> value.required("name").asString().equals("cancel")) .findFirst() .orElseThrow(); - var unrelated = context(43); - var pending = new HashMap(); + var cancellations = new ArrayList(); var queues = new ArrayList(); var leadership = new ArrayList(); var dispatcher = - new Protocol.Dispatcher( - "observer", Map.of(43L, unrelated), pending, queues::add, leadership::add); + new Protocol.Dispatcher("observer", cancellations::add, queues::add, leadership::add); String payload = Json.encode(cancel.required("payload")); dispatcher.dispatch("unrelated", payload); dispatcher.dispatch(Protocol.TOPIC_LEADERSHIP, payload); - assertTrue(pending.isEmpty()); assertTrue(queues.isEmpty()); assertTrue(leadership.isEmpty()); dispatcher.dispatch(Protocol.TOPIC_INSERT, payload); assertEquals(List.of(new Protocol.QueueNotice("priority", "insert")), queues); - assertTrue( - pending.isEmpty(), "A control payload on the insert topic must not cancel an attempt"); queues.clear(); + assertTrue(cancellations.isEmpty()); dispatcher.dispatch(cancel.required("topic").asString(), payload); - assertEquals(java.util.Set.of(42L), pending.keySet()); - assertNull(unrelated.cancellation()); + assertEquals(List.of(42L), cancellations); assertTrue(queues.isEmpty()); assertTrue(leadership.isEmpty()); } - private static WorkContext context(long id) { - return context(id, Json.object()); - } - private static WorkContext context(long id, JsonNode metadata) { var now = Instant.parse("2026-01-02T03:04:05Z"); var job = diff --git a/java/river/src/test/java/com/riverqueue/RetryTest.java b/java/river/src/test/java/com/riverqueue/RetryTest.java index 52385d523..77577ec99 100644 --- a/java/river/src/test/java/com/riverqueue/RetryTest.java +++ b/java/river/src/test/java/com/riverqueue/RetryTest.java @@ -9,6 +9,7 @@ import java.time.Instant; import java.time.ZoneOffset; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; @@ -27,6 +28,55 @@ class RetryTest { @RegisterExtension final TestDatabase databases = new TestDatabase(); @TempDir Path directory; + @Test + void retryClearsCancellation() { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry cancelled").job(); + var cancelled = client.cancel(inserted.id()); + assertTrue(cancelled.metadata().has("cancel_attempted_at")); + var retried = client.retry(inserted.id()); + assertEquals(Job.State.AVAILABLE, retried.state()); + assertFalse(retried.metadata().has("cancel_attempted_at")); + assertFalse(client.get(inserted.id()).metadata().has("cancel_attempted_at")); + } + + @Test + void retryClearsCancellationWithoutReorderingAvailableJob() { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database, Clock.fixed(NOW, ZoneOffset.UTC)); + var inserted = + client + .insert( + TYPE, + "already available", + InsertOptions.builder() + .metadata(Map.of("cancel_attempted_at", NOW.toString(), "preserved", true)) + .scheduledAt(NOW.minusSeconds(3600)) + .build()) + .job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='available', attempt=max_attempts WHERE id=?", + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + return null; + }); + var retried = client.retry(inserted.id()); + assertEquals(inserted.scheduledAt(), retried.scheduledAt()); + assertEquals(inserted.maxAttempts(), retried.maxAttempts()); + var expectedMetadata = (tools.jackson.databind.node.ObjectNode) inserted.metadata().deepCopy(); + expectedMetadata.remove("cancel_attempted_at"); + assertEquals(expectedMetadata, retried.metadata()); + assertEquals(retried, client.get(inserted.id())); + } + @ParameterizedTest @EnumSource(Job.State.class) void retryNotifiesOnlyWhenTheRowChanges(Job.State state) { @@ -67,6 +117,36 @@ void retryNotifiesOnlyWhenTheRowChanges(Job.State state) { } } + @ParameterizedTest + @ValueSource( + strings = {"null", "42", "false", "\"cancel_attempted_at\"", "[\"cancel_attempted_at\"]"}) + void retryPreservesNonObjectMetadata(String metadata) { + var database = databases.open(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry metadata").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE " + + database.prefix() + + "river_job SET state='retryable', metadata=" + + (database.dialect() == Database.Dialect.POSTGRES + ? "CAST(? AS jsonb)" + : "jsonb(?)") + + " WHERE id=?", + metadata, + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + return null; + }); + + assertEquals(Job.State.AVAILABLE, client.retry(inserted.id()).state()); + assertEquals(Json.parse(metadata), client.get(inserted.id()).metadata()); + } + @ParameterizedTest @ValueSource(strings = {"local", "peer", "transaction"}) void retryWakesWorkersAfterCommit(String mode) throws Exception { @@ -216,4 +296,36 @@ void sqliteNotificationFailureRollsBackTheRetry() { job.metadata().path(Protocol.METADATA_OUTPUT).asString()); assertEquals(0, notifications.get()); } + + @Test + void sqliteRetryQueryPreservesMalformedMetadata() { + var database = TestDatabase.sqlite(directory.resolve("river.db")); + var client = new Client(database); + var inserted = client.insert(TYPE, "retry malformed metadata").job(); + client.transaction( + connection -> { + try (var statement = + Sql.prepare( + connection, + "UPDATE river_job SET state='retryable', metadata='[not json' WHERE id=?", + inserted.id())) { + assertEquals(1, statement.executeUpdate()); + } + // The public retry validates the locked row first. Exercise the SQL + // directly so cleanup cannot turn a decode error into a SQL error. + try (var statement = + Sql.prepare( + connection, + Sql.query(database, "retry"), + database.timestamp(NOW), + inserted.id()); + var rows = statement.executeQuery()) { + assertTrue(rows.next()); + assertEquals("[not json", rows.getString("metadata")); + } + return null; + }); + var error = assertThrows(IllegalArgumentException.class, () -> client.get(inserted.id())); + assertTrue(error.getMessage().contains("metadata")); + } } diff --git a/java/river/src/test/java/com/riverqueue/WorkersTest.java b/java/river/src/test/java/com/riverqueue/WorkersTest.java index fc3a7c56d..cc1f36e32 100644 --- a/java/river/src/test/java/com/riverqueue/WorkersTest.java +++ b/java/river/src/test/java/com/riverqueue/WorkersTest.java @@ -71,6 +71,120 @@ void builderRequiresServiceIntervalsBelowQueueRetention(String value, boolean va .getMessage()); } + @Test + void cancellationContinuesWhileFetchingIsBlocked() throws Exception { + var entered = new CompletableFuture(); + var fetching = new CompletableFuture(); + var release = new CompletableFuture(); + var cancelled = new CompletableFuture(); + var client = + client() + .withPilot( + new Pilot() { + @Override + public int workerLimit(String queue, int configured) { + if (entered.isDone()) { + fetching.complete(null); + release.join(); + } + return configured; + } + }); + var workers = + workers(client) + .pollInterval(Duration.ofMillis(10)) + .add( + TYPE, + context -> { + entered.complete(null); + context.awaitCancellation(); + cancelled.complete(context.cancellation()); + context.checkCancelled(); + }) + .start(); + try { + var inserted = client.insert(TYPE, "cancel while fetching").job(); + fetching.get(5, TimeUnit.SECONDS); + client.cancel(inserted.id()); + assertEquals(WorkContext.Cancellation.REMOTE, cancelled.get(5, TimeUnit.SECONDS)); + } finally { + release.complete(null); + workers.stopAndCancel(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void cancellationContinuesAfterGracefulStop(boolean pollOnly) throws Exception { + var client = client(); + var entered = new CompletableFuture>(); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .add( + TYPE, + context -> { + entered.complete(context); + context.awaitCancellation(); + context.checkCancelled(); + }) + .start(); + try (var subscription = workers.subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel while draining").job(); + entered.get(5, TimeUnit.SECONDS); + // The timeout leaves active work draining after the dispatcher and + // listener have stopped. A peer can still cancel that work. + assertThrows(RiverException.class, workers::stop); + client.cancel(inserted.id()); + assertEquals(Job.State.CANCELLED, outcome.get(5, TimeUnit.SECONDS).job().state()); + workers.stop(); + } finally { + workers.stopAndCancel(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void cancellationDuringArgumentDecodingIsReconciled(boolean pollOnly) throws Exception { + var decoding = new CompletableFuture(); + var release = new CompletableFuture(); + var client = + client() + .withPilot( + new Pilot() { + @Override + public tools.jackson.databind.JsonNode decode( + Job job) { + decoding.complete(null); + release.join(); + return job.args(); + } + }); + var outcome = new CompletableFuture(); + var workers = + workers(client) + .pollOnly(pollOnly) + .add( + TYPE, + context -> { + context.awaitCancellation(); + context.checkCancelled(); + }) + .start(); + try (var subscription = workers.subscribe(outcome::complete)) { + var inserted = client.insert(TYPE, "cancel before registration").job(); + decoding.get(5, TimeUnit.SECONDS); + client.cancel(inserted.id()); + release.complete(null); + assertEquals(Job.State.CANCELLED, outcome.get(5, TimeUnit.SECONDS).job().state()); + workers.stop(); + } finally { + release.complete(null); + workers.stopAndCancel(); + } + } + @ParameterizedTest @CsvSource({"REMOTE,false", "REMOTE,true", "TIMEOUT,false", "TIMEOUT,true"}) void cancellationRespectsTheHandlersOutcome(WorkContext.Cancellation cause, boolean interrupted) @@ -694,6 +808,71 @@ void invalidRetryPoliciesFallBack(String policy) throws Exception { } } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void retriedJobIgnoresStaleCancellation(boolean delayedUntilRunning) throws Exception { + var client = client(); + var ready = new CompletableFuture(); + var paused = new LinkedBlockingQueue(); + var entered = new CompletableFuture(); + var finish = new CompletableFuture(); + if (!delayedUntilRunning) finish.complete(null); + var outcome = new CompletableFuture(); + var errors = new LinkedBlockingQueue(); + var workers = + workers(client) + .pollOnly(false) + .errorHandler(errors::add) + .observe( + event -> { + if (event.equals("listen_ready")) ready.complete(null); + }) + .add( + TYPE, + context -> { + entered.complete(null); + finish.join(); + context.checkCancelled(); + }) + .start(); + try (var pauseSubscription = workers.subscribe(paused::add, Workers.EventKind.QUEUE_PAUSED); + var outcomeSubscription = + workers.subscribe(outcome::complete, Workers.EventKind.JOB_COMPLETED)) { + ready.get(5, TimeUnit.SECONDS); + var inserted = + client + .insert( + TYPE, + "retry after an idle cancellation", + InsertOptions.builder().scheduledAt(Instant.now().plusSeconds(3600)).build()) + .job(); + client.cancel(inserted.id()); + client.queues().pause("default"); + // The pause follows cancellation on the same listener, so receiving it + // proves the unmatched cancellation was dispatched before retry. + assertNotNull(paused.poll(5, TimeUnit.SECONDS)); + client.retry(inserted.id()); + client.queues().resume("default"); + if (delayedUntilRunning) { + entered.get(5, TimeUnit.SECONDS); + client.transaction( + connection -> { + client.notify(connection, Protocol.cancel(inserted.id(), "default")); + return null; + }); + client.queues().pause("default"); + assertNotNull(paused.poll(5, TimeUnit.SECONDS)); + finish.complete(null); + } + assertEquals(inserted.id(), outcome.get(5, TimeUnit.SECONDS).job().id()); + workers.stop(); + } finally { + finish.complete(null); + workers.stopAndCancel(); + } + assertTrue(errors.isEmpty(), () -> "Unexpected worker errors: " + errors); + } + @ParameterizedTest @CsvSource({"false,1,5000", "true,1,5000", "false,3600,1", "true,3600,1"}) void shortDelaysWakeFetchingBeforeThePollInterval( diff --git a/js/driver/pg/src/driver.integration.test.ts b/js/driver/pg/src/driver.integration.test.ts index 25ecd0f93..4ef3e1c25 100644 --- a/js/driver/pg/src/driver.integration.test.ts +++ b/js/driver/pg/src/driver.integration.test.ts @@ -531,11 +531,23 @@ describe("PgDriver integration", () => { await expect(driver.jobGetCancelRequested([])).resolves.toEqual([]); }); + it("ignores cancellation names stored in non-object metadata", async () => { + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_cancel_array`) + ); + await pool.query( + "UPDATE river_job SET state = 'running', metadata = $1::jsonb WHERE id = $2", + ['["cancel_attempted_at"]', job.id.toString()] + ); + await expect(driver.jobGetCancelRequested([job.id])).resolves.toEqual([]); + }); + it("retries non-running jobs but leaves running jobs unchanged", async () => { const scheduledAt = Temporal.Instant.from("2026-08-30T15:00:00Z"); const now = Temporal.Instant.from("2026-08-30T12:00:00.123456Z"); const retryable = await driver.jobInsert( insertParams(`${filePrefix}_retry`, { + metadata: { cancel_attempted_at: now.toString(), preserved: true }, scheduledAt, state: "retryable", }) @@ -558,10 +570,33 @@ describe("PgDriver integration", () => { }); expect(retried!.state).toBe("available"); + expect(retried!.metadata).toEqual({ preserved: true }); expect(retried!.scheduledAt.toString()).toBe("2026-08-30T12:00:00.123456Z"); expect(unchanged!.state).toBe("running"); }); + it("clears cancellation without reordering an already available retry", async () => { + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_retry_available`, { + metadata: { cancel_attempted_at: now.toString(), preserved: true }, + scheduledAt: now.subtract({ hours: 1 }), + state: "available", + }) + ); + await pool.query( + "UPDATE river_job SET attempt = max_attempts WHERE id = $1", + [job.id.toString()] + ); + const expectedMetadata = { ...job.metadata }; + delete expectedMetadata.cancel_attempted_at; + const retried = await driver.jobRetryWithOptions({ id: job.id, now }); + expect(retried?.metadata).toEqual(expectedMetadata); + expect(retried?.scheduledAt.toString()).toBe(job.scheduledAt.toString()); + expect(retried?.maxAttempts).toBe(job.maxAttempts); + expect((await driver.jobGet(job.id))?.metadata).toEqual(expectedMetadata); + }); + it("returns the committed row to the loser of a cancel or retry race", async () => { const inserted = await driver.jobInsert( insertParams(`${filePrefix}_race`, { @@ -633,6 +668,29 @@ describe("PgDriver integration", () => { ); }); + it.each(["null", "42", '"cancel_attempted_at"', '["cancel_attempted_at"]'])( + "retries jobs without changing non-object metadata: %s", + async (metadata) => { + const { job } = await driver.jobInsert( + insertParams(`${filePrefix}_retry_metadata`, { state: "retryable" }) + ); + await pool.query( + "UPDATE river_job SET metadata = $1::jsonb WHERE id = $2::bigint", + [metadata, job.id.toString()] + ); + + expect(await driver.jobRetry(job.id)).toMatchObject({ + id: job.id, + state: "available", + }); + const persisted = await pool.query<{ metadata: unknown }>( + "SELECT metadata FROM river_job WHERE id = $1::bigint", + [job.id.toString()] + ); + expect(persisted.rows[0]?.metadata).toEqual(JSON.parse(metadata)); + } + ); + it("acts on rows River can't fully read by ID and lists them like Go", async () => { const kind = `${filePrefix}_poison`; const uniqueKey = new Uint8Array(32).fill(13); diff --git a/js/driver/pg/src/sql/jobs.ts b/js/driver/pg/src/sql/jobs.ts index 37b91ecb6..f95c75659 100644 --- a/js/driver/pg/src/sql/jobs.ts +++ b/js/driver/pg/src/sql/jobs.ts @@ -294,6 +294,7 @@ export async function jobGetCancelRequested( SELECT id::text AS id FROM ${db.table("river_job")} WHERE id = any($1::bigint[]) + AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' ORDER BY id @@ -822,19 +823,25 @@ export async function jobRetryWithOptions( SET state = 'available', max_attempts = CASE + WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN river_job.max_attempts WHEN river_job.attempt = river_job.max_attempts THEN river_job.max_attempts + 1 ELSE river_job.max_attempts END, - finalized_at = NULL, - scheduled_at = coalesce($2::timestamptz, now()) + finalized_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN river_job.finalized_at ELSE NULL END, + metadata = CASE + WHEN jsonb_typeof(river_job.metadata) = 'object' + THEN river_job.metadata - 'cancel_attempted_at' + ELSE river_job.metadata + END, + scheduled_at = CASE WHEN river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) THEN river_job.scheduled_at ELSE coalesce($2::timestamptz, now()) END FROM job_to_update WHERE river_job.id = job_to_update.id AND river_job.state != 'running' - AND NOT ( + AND ((jsonb_typeof(river_job.metadata) = 'object' AND river_job.metadata ? 'cancel_attempted_at') OR NOT ( river_job.state = 'available' AND river_job.scheduled_at < coalesce($2::timestamptz, now()) - ) + )) RETURNING river_job.* ) `, diff --git a/js/driver/sqlite/src/compatibility.test.ts b/js/driver/sqlite/src/compatibility.test.ts index e73c52129..1d413bbf1 100644 --- a/js/driver/sqlite/src/compatibility.test.ts +++ b/js/driver/sqlite/src/compatibility.test.ts @@ -633,7 +633,7 @@ describe("SqliteDriver compatibility with Go's riversqlite", () => { errors: null, errors_null: 1n, finalized_at: null, - metadata: cancelled, + metadata: { m: "x" }, scheduled_at: "2026-01-02 06:06:07.123", state: "available", }); diff --git a/js/driver/sqlite/src/driver.test.ts b/js/driver/sqlite/src/driver.test.ts index 3506895e0..70d0f92b8 100644 --- a/js/driver/sqlite/src/driver.test.ts +++ b/js/driver/sqlite/src/driver.test.ts @@ -326,10 +326,12 @@ describe("SqliteDriver", () => { const retried = await driver.jobRetryDetailed(inserted.id, { now }); expect(retried.status).toBe("retried"); - if (retried.status === "retried") + if (retried.status === "retried") { expect(retried.job.state).toBe("available"); + expect(retried.job.metadata).not.toHaveProperty("cancel_attempted_at"); + } expect((await driver.jobRetryDetailed(inserted.id, { now })).status).toBe( - "unchanged" + "retried" ); database @@ -350,6 +352,79 @@ describe("SqliteDriver", () => { ); }); + test("increments exhausted attempts at the exact retry timestamp", async () => { + const { database, driver } = await setup(); + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert({ + args: {}, + kind: "retry_boundary", + maxAttempts: 1, + scheduledAt: now, + state: "available", + }); + database + .prepare("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .run(job.id); + const retried = await driver.jobRetryDetailed(job.id, { now }); + expect(retried).toMatchObject({ + status: "retried", + job: { maxAttempts: 2 }, + }); + }); + + test("clears cancellation without reordering an already available retry", async () => { + const { database, driver } = await setup(); + const now = Temporal.Instant.from("2026-08-30T12:00:00Z"); + const { job } = await driver.jobInsert({ + args: {}, + kind: "retry_available", + metadata: { cancel_attempted_at: now.toString(), preserved: true }, + scheduledAt: now.subtract({ hours: 1 }), + state: "available", + }); + database + .prepare("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .run(job.id); + const expectedMetadata = { ...job.metadata }; + delete expectedMetadata.cancel_attempted_at; + const retried = await driver.jobRetryDetailed(job.id, { now }); + expect(retried.status).toBe("retried"); + if (retried.status !== "retried") + throw new Error("expected a metadata update"); + expect(retried.job.metadata).toEqual(expectedMetadata); + expect(retried.job.scheduledAt.toString()).toBe(job.scheduledAt.toString()); + expect(retried.job.maxAttempts).toBe(job.maxAttempts); + expect((await driver.jobGet(job.id))?.metadata).toEqual(expectedMetadata); + }); + + test.each(["[not json", "null", '["cancel_attempted_at"]'])( + "retries jobs without changing malformed metadata: %s", + async (metadata) => { + const { database, driver } = await setup(); + const { job } = await driver.jobInsert({ + args: {}, + finalizedAt: Temporal.Now.instant(), + kind: "retry_metadata", + state: "discarded", + }); + database + .prepare("UPDATE river_job SET metadata = ? WHERE id = ?") + .run(metadata, job.id); + + expect(await driver.jobRetryDetailed(job.id)).toMatchObject({ + job: { id: job.id, state: "available" }, + status: "retried", + }); + expect( + database + .prepare( + "SELECT CASE WHEN typeof(metadata) = 'text' THEN metadata ELSE json(metadata) END AS metadata FROM river_job WHERE id = ?" + ) + .get(job.id)?.metadata + ).toBe(metadata); + } + ); + test("notifies a running job's cancellation in the cancelling transaction, like Go", async () => { const { database, driver } = await setup(); const inserted = (await driver.jobInsert({ args: {}, kind: "control_job" })) diff --git a/js/driver/sqlite/src/driver.ts b/js/driver/sqlite/src/driver.ts index 98a0891c5..ead745c98 100644 --- a/js/driver/sqlite/src/driver.ts +++ b/js/driver/sqlite/src/driver.ts @@ -1438,21 +1438,29 @@ export class SqliteRuntime implements RuntimeDriver< const raw = resultStatement( database, ` + WITH retry_params AS (SELECT ? AS now) UPDATE river_job SET state = 'available', max_attempts = CASE + WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = ? + finalized_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, + metadata = CASE + WHEN ${invalidJsonTextSql("metadata")} THEN metadata + ELSE jsonb_remove(metadata, '$.cancel_attempted_at') + END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END WHERE id = ? AND state != 'running' - AND (state != 'available' OR scheduled_at > ?) + AND ((CASE WHEN ${invalidJsonTextSql("metadata")} THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR NOT (state = 'available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING ${JOB_COLUMNS} ` - ).get(timestamp, id, timestamp); + ).get(timestamp, id); const updated = raw === undefined ? null : decodeJobRowPartial(raw).job; const job = updated ?? getJob(database, id, { partial: true }); if (job === null) return { status: "not_found" }; diff --git a/js/src/pilot-runtime.test.ts b/js/src/pilot-runtime.test.ts index 1d0372194..de64ed451 100644 --- a/js/src/pilot-runtime.test.ts +++ b/js/src/pilot-runtime.test.ts @@ -1067,11 +1067,15 @@ describe("pilot services and teardown", () => { }); const run = await client.start(); - const first = await timer.waitFor(({ kind }) => kind === "delay"); + const first = await timer.waitFor( + ({ kind, ms }) => kind === "delay" && ms !== 2_000 + ); expect(runs).toBe(1); await timer.advance(first.ms); await waitUntil(() => runs === 2); - const second = await timer.waitFor(({ kind }) => kind === "delay"); + const second = await timer.waitFor( + ({ kind, ms }) => kind === "delay" && ms !== 2_000 + ); // Backoff grows while the service keeps failing quickly. expect(second.ms).toBeGreaterThan(first.ms); await timer.advance(second.ms); @@ -1079,7 +1083,9 @@ describe("pilot services and teardown", () => { // A failure after a healthy minute starts backoff over. await timer.advance(60_000); longRun.reject(new Error("service failed after a while")); - const third = await timer.waitFor(({ kind }) => kind === "delay"); + const third = await timer.waitFor( + ({ kind, ms }) => kind === "delay" && ms !== 2_000 + ); expect(third.ms).toBe(first.ms); await timer.advance(third.ms); await waitUntil(() => runs === 4); diff --git a/js/src/runtime.test.ts b/js/src/runtime.test.ts index c44c05346..781973d07 100644 --- a/js/src/runtime.test.ts +++ b/js/src/runtime.test.ts @@ -618,6 +618,7 @@ describe("Client runtime", () => { clientId: "runtime-test", completionBatchSize: 1, leaderElectionDisabled: true, + queueControlPollInterval: { milliseconds: 1234 }, queues: { default: { maxWorkers: 1, pollInterval: { milliseconds: 10_000 } }, }, @@ -633,7 +634,9 @@ describe("Client runtime", () => { expect(driver.completionFailures).toBe(0); expect(driver.completionCalls).toBe(3); expect(driver.completions[0]?.kind).toBe("complete"); - expect(timer.delays.slice(-2)).toEqual([1_000, 2_000]); + expect(timer.delays.filter((delay) => delay !== 1234).slice(-2)).toEqual([ + 1_000, 2_000, + ]); expect(timer.activeTimeouts).toBe(0); expect(run.state).toBe("stopped"); }); @@ -858,6 +861,7 @@ describe("Client runtime", () => { onMetric: (metric) => void metrics.push(metric), }, leaderElectionDisabled: true, + queueControlPollInterval: { milliseconds: 1234 }, queues: { default: { maxWorkers: 1, pollInterval: { milliseconds: 10_000 } }, }, @@ -875,7 +879,9 @@ describe("Client runtime", () => { name: "job_completion_requeued", }); // Every failed attempt backs off, including the last before a requeue. - expect(timer.delays).toEqual([1_000, 2_000, 4_000, 1_000, 2_000]); + expect(timer.delays.filter((delay) => delay !== 1234)).toEqual([ + 1_000, 2_000, 4_000, 1_000, 2_000, + ]); expect(events).toEqual(["job_started", "job_completed"]); }); @@ -4330,7 +4336,9 @@ describe("Client runtime database faults", () => { await run.stop(); expect(run.state).toBe("stopped"); - expect(timer.delays).toEqual([250, 500, 1_000]); + expect(timer.delays.filter((delay) => delay !== 2_000)).toEqual([ + 250, 500, 1_000, + ]); expect( logs.map(({ attributes, level }) => [ level, @@ -4377,7 +4385,7 @@ describe("Client runtime database faults", () => { await waitUntil(() => events.includes("queue_paused")); expect(run.diagnostics.queues.default?.paused).toBe(true); expect(run.state).toBe("running"); - expect(timer.delays).toEqual([250, 500]); + expect(timer.delays.filter((delay) => delay !== 1)).toEqual([250, 500]); expect(logs.map(({ message }) => message)).toEqual([ "River queue control poll failed; retrying after backoff", "River queue control poll failed; retrying after backoff", @@ -4388,6 +4396,14 @@ describe("Client runtime database faults", () => { it("cancels a job whose cancellation arrived while its claim was in flight", async () => { const driver = new FakeRuntimeDriver(); driver.claim = [fakeJob()]; + driver.listRows = [ + { + ...fakeJob(), + state: "running", + attemptedBy: ["runtime-test"], + metadata: { cancel_attempted_at: null }, + }, + ]; const release = Promise.withResolvers(); driver.claimGate = release.promise; const handled = Promise.withResolvers(); @@ -4516,7 +4532,7 @@ describe("Client runtime database faults", () => { () => driver.completions.length === 1 && subscriptions === 3 ); expect(run.state).toBe("running"); - expect(timer.delays).toEqual([250, 250]); + expect(timer.delays.filter((delay) => delay !== 2_000)).toEqual([250, 250]); expect(logs.map(({ level, message }) => [level, message])).toEqual([ [ "warn", @@ -4551,6 +4567,14 @@ describe("Client runtime database faults", () => { }, }); driver.claim = [fakeJob()]; + driver.listRows = [ + { + ...fakeJob(), + state: "running", + attemptedBy: ["runtime-test"], + metadata: { cancel_attempted_at: null }, + }, + ]; const logs: LogEntry[] = []; const timer = new FakeTimer(); const definition = defineJob({ kind: "test" }); @@ -4582,6 +4606,264 @@ describe("Client runtime database faults", () => { await run.stop(); }); + for (const stream of ["runtime", "remote", "local"] as const) { + for (const duringClaim of [false, true]) { + it(`revalidates a delayed ${stream} cancellation ${duringClaim ? "during a retry claim" : "against a running retry"}`, async () => { + const driver = new FakeRuntimeDriver(); + const deliver = Promise.withResolvers(); + const handled = Promise.withResolvers(); + const releaseClaim = Promise.withResolvers(); + const finish = Promise.withResolvers(); + const started = Promise.withResolvers(); + let checks = 0; + Object.assign(driver, { + jobGetCancelRequested: () => { + checks++; + return []; + }, + }); + driver.claim = [{ ...fakeJob(), attempt: 2 }]; + if (duringClaim) driver.claimGate = releaseClaim.promise; + if (stream === "runtime") { + driver.notificationStream = async function* (_topics, signal, ready) { + ready(); + await deliver.promise; + yield { + topic: "control", + payload: '{"action":"cancel","job_id":101,"queue":"default"}', + }; + handled.resolve(undefined); + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { once: true }) + ); + }; + } else if (stream === "remote") { + Object.defineProperty(driver, "runtimeNotificationSubscribe", { + value: undefined, + }); + Object.defineProperty(driver, "jobCancellationSubscribe", { + value: async function* (attemptedBy: string, signal: AbortSignal) { + await deliver.promise; + yield { attemptedBy, id: 101n }; + handled.resolve(undefined); + await new Promise((resolve) => + signal.addEventListener("abort", () => resolve(), { + once: true, + }) + ); + }, + }); + } + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + queues: { + default: { maxWorkers: 1, pollInterval: { milliseconds: 10_000 } }, + }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + started.resolve(signal); + await finish.promise; + signal.throwIfAborted(); + } + ), + }); + const run = await client.start(); + try { + await waitUntil(() => driver.claimed.length === 1); + if (!duringClaim) await started.promise; + if (stream === "local") { + // The cancelled row returned by a delayed commit belongs to attempt A. + driver.cancelled = { ...fakeJob(), attempt: 1, state: "running" }; + await client.jobs.cancel(101n); + } else { + deliver.resolve(undefined); + await handled.promise; + } + releaseClaim.resolve(undefined); + const signal = await started.promise; + await waitUntil(() => checks > 0); + expect(signal.aborted).toBe(false); + finish.resolve(undefined); + await waitUntil(() => driver.completions.length === 1); + expect(driver.completions[0]?.kind).toBe("complete"); + } finally { + deliver.resolve(undefined); + releaseClaim.resolve(undefined); + finish.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + } + } + + for (const batched of [false, true]) { + for (const phase of ["running", "claim", "draining"] as const) { + it(`recovers a lost cancellation while ${phase}, using ${batched ? "batches" : "rows"}`, async () => { + const driver = new FakeRuntimeDriver(); + const claim = Promise.withResolvers(); + const started = Promise.withResolvers(); + let marked = false; + let failures = 0; + const read = (ids: readonly bigint[]) => { + if (!marked) return []; + if (failures++ === 0) + throw new Error("temporary cancellation read failure"); + return ids; + }; + if (batched) { + Object.assign(driver, { jobGetCancelRequested: read }); + } else { + Object.assign(driver, { + jobGet: (id: bigint) => + read([id]).length === 0 + ? null + : { + ...fakeJob(), + id, + attemptedBy: ["runtime-test"], + state: "running", + metadata: { cancel_attempted_at: null }, + }, + }); + } + driver.claim = [fakeJob()]; + if (phase === "claim") driver.claimGate = claim.promise; + let handlerSignal: AbortSignal | undefined; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + logger: recordingLogger([]), + queues: { default: { maxWorkers: 1 } }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + handlerSignal = signal; + started.resolve(undefined); + await new Promise((resolve) => { + if (signal.aborted) resolve(); + else + signal.addEventListener("abort", () => resolve(), { + once: true, + }); + }); + signal.throwIfAborted(); + } + ), + }); + overrideRuntimeTiming(client, { + random: () => 0.5, + timer: new FakeTimer(), + }); + const run = await client.start(); + try { + if (phase === "claim") { + await waitUntil(() => driver.claimRequests.length > 0); + marked = true; + claim.resolve(undefined); + } + await started.promise; + const stopping = phase === "draining" ? run.stop() : undefined; + marked = true; + await waitUntil(() => driver.completions.length === 1); + expect(handlerSignal?.reason).toMatchObject({ + name: "JobCancelledError", + }); + expect(driver.completions[0]?.kind).toBe("cancel"); + expect(failures).toBeGreaterThanOrEqual(2); + await stopping; + await run.stop(); + } finally { + claim.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + } + } + + for (const batched of [false, true]) { + it(`does not apply a stale ${batched ? "batch" : "row"} cancellation result to a new attempt`, async () => { + const driver = new FakeRuntimeDriver(); + const firstRead = Promise.withResolvers(); + const readStarted = Promise.withResolvers(); + const finishFirst = Promise.withResolvers(); + const finishSecond = Promise.withResolvers(); + const signals: AbortSignal[] = []; + let reads = 0; + const read = async () => { + reads++; + if (reads !== 1) return false; + readStarted.resolve(undefined); + await firstRead.promise; + return true; + }; + if (batched) { + Object.assign(driver, { + jobGetCancelRequested: async () => ((await read()) ? [7n] : []), + }); + } else { + Object.assign(driver, { + jobGet: async () => + (await read()) + ? { + ...fakeJob(), + attemptedBy: ["runtime-test"], + state: "running", + metadata: { cancel_attempted_at: null }, + } + : null, + }); + } + driver.claim = [fakeJob()]; + const client = new Client(driver, { + clientId: "runtime-test", + completionBatchSize: 1, + leaderElectionDisabled: true, + queues: { + default: { + fetchCooldown: { milliseconds: 1 }, + maxWorkers: 1, + pollInterval: { milliseconds: 1 }, + }, + }, + workers: new Workers().add( + defineJob({ kind: "test" }), + async ({ signal }) => { + signals.push(signal); + await (signals.length === 1 + ? finishFirst.promise + : finishSecond.promise); + } + ), + }); + overrideRuntimeTiming(client, { + random: () => 0.5, + timer: new FakeTimer(), + }); + const run = await client.start(); + try { + await readStarted.promise; + finishFirst.resolve(undefined); + await waitUntil(() => driver.completions.length === 1); + driver.claim = [fakeJob()]; + await waitUntil(() => signals.length === 2); + firstRead.resolve(undefined); + await waitUntil(() => reads >= 2); + expect(signals[1]?.aborted).toBe(false); + finishSecond.resolve(undefined); + await run.stop(); + } finally { + firstRead.resolve(undefined); + finishFirst.resolve(undefined); + finishSecond.resolve(undefined); + await run.stop({ mode: "cancel" }); + } + }); + } + for (const [stream, read] of [ ["remote cancellation", "each row"], ["runtime notification", "each row"], diff --git a/js/src/runtime.ts b/js/src/runtime.ts index 1be2c83ad..a8ce8d425 100644 --- a/js/src/runtime.ts +++ b/js/src/runtime.ts @@ -215,7 +215,7 @@ export class RuntimeController { * still be cancelled while its queue drains. */ readonly #cancellationPollAbort = new AbortController(); - /** How often a runtime without notifications checks for cancellations. */ + /** How often the runtime checks for durable cancellation requests. */ readonly #cancellationPollIntervalMs: number; readonly #claimAbort = new AbortController(); readonly #client: Client; @@ -664,18 +664,15 @@ export class RuntimeController { } async #initialize(): Promise { - // Like River for Go's client, learn whether the database delivers - // notifications before anything starts. Without them, running jobs - // learn of cancellations by polling until every producer has drained. + // Notifications deliver cancellations promptly. Keep a durable fallback + // until every producer drains, including during reconnects and shutdown. const listen = await this.#listens(); - if (!listen) { - this.#trackTask( - this.#runner.pollCancellations( - this.#cancellationPollIntervalMs, - this.#cancellationPollAbort.signal - ) - ); - } + this.#trackTask( + this.#runner.pollCancellations( + this.#cancellationPollIntervalMs, + this.#cancellationPollAbort.signal + ) + ); // Like River for Go's notifier, notifications are listening before any // queue claims, and a failure to listen fails the start. await this.#notifications.start(listen); diff --git a/js/src/runtime/attempt-runner.ts b/js/src/runtime/attempt-runner.ts index cca2ea5a5..29423aee0 100644 --- a/js/src/runtime/attempt-runner.ts +++ b/js/src/runtime/attempt-runner.ts @@ -108,10 +108,15 @@ const CANCEL_POLL_TIMEOUT_MS = 10_000; export class AttemptRunner { readonly #attemptExecutor = new AttemptExecutor(); /** In-flight attempts by `attemptKey`. */ - readonly #attempts = new Map(); + readonly #attempts = new Map< + string, + { readonly controller: AbortController; readonly row: JobRow } + >(); + #cancellationCheck: Promise | undefined; + #cancellationCheckPending = false; /** * For each claim in flight, the `attemptKey`s of cancellations that - * found no attempt running here. + * need a durable check once the claim registers its attempts. */ readonly #claimCancellations = new Set>(); readonly #completions: CompletionPipeline; @@ -153,18 +158,15 @@ export class AttemptRunner { /** Abort every in-flight attempt with `reason`. */ abortAttempts(reason: unknown): void { - for (const controller of this.#attempts.values()) controller.abort(reason); + for (const { controller } of this.#attempts.values()) + controller.abort(reason); } - /** Cancel the attempt of job `id` that `attemptedBy` owns, if it runs here. */ + /** Check a cancellation hint against durable state before aborting work. */ cancelAttempt(id: bigint, attemptedBy: string): void { const key = attemptKey(id, attemptedBy); - const attempt = this.#attempts.get(key); - if (attempt !== undefined) { - attempt.abort(new JobCancelledError(id)); - return; - } for (const pending of this.#claimCancellations) pending.add(key); + if (this.#attempts.has(key)) void this.#checkCancellations(); } /** @@ -178,30 +180,40 @@ export class AttemptRunner { const clientId = this.#context.clientId; const driver = this.#context.driver; const requested = driver.jobGetCancelRequested?.bind(driver); - const ids = this.#ownAttemptIds(); + const attempts = this.#ownAttempts(); if (requested === undefined) { // A backend without the batched read checks each job's row. - for (const id of ids) { + for (const [id, controller] of attempts) { signal.throwIfAborted(); - const job = await raceWithAbort(driver.jobGet(id), signal); - if ( - job?.state === "running" && - job.attemptedBy.at(-1) === clientId && - job.metadata.cancel_attempted_at !== undefined - ) { - this.cancelAttempt(id, clientId); + const timeout = this.#context.timer.timeout( + CANCEL_POLL_TIMEOUT_MS, + () => new Error("River's cancellation check timed out") + ); + const link = new LinkedAbortSignal([signal, timeout.signal]); + try { + const job = await raceWithAbort(driver.jobGet(id), link.signal); + if ( + job?.state === "running" && + job.attemptedBy.at(-1) === clientId && + Object.hasOwn(job.metadata, "cancel_attempted_at") + ) { + controller.abort(new JobCancelledError(id)); + } + } finally { + link[Symbol.dispose](); + timeout.dispose(); } } return; } - await this.#cancelRequested(requested, ids, signal); + await this.#cancelRequested(requested, attempts, signal); } /** * Cancel this client's running attempts whose jobs have a cancellation - * request, checking every `intervalMs` until `signal` aborts, like River - * for Go's producers without a notifier. Each check reads the attempts' - * jobs in batches of 1,000 IDs, each bounded by ten seconds, and reads + * request, checking every `intervalMs` until `signal` aborts. Each check + * reads the attempts' jobs in batches of 1,000 IDs (or one row at a time + * without a batched read), each bounded by ten seconds, and reads * nothing while no attempt runs. A failed check is logged and the next * one tries again. */ @@ -209,9 +221,6 @@ export class AttemptRunner { intervalMs: number, signal: AbortSignal ): Promise { - const driver = this.#context.driver; - const requested = driver.jobGetCancelRequested?.bind(driver); - if (requested === undefined) return; for (;;) { try { await this.#context.timer.delay(intervalMs, signal); @@ -219,7 +228,7 @@ export class AttemptRunner { return; } try { - await this.#cancelRequested(requested, this.#ownAttemptIds(), signal); + await this.recoverCancellations(signal); } catch (error: unknown) { if (signal.aborted) return; this.#context.logger.error( @@ -232,17 +241,19 @@ export class AttemptRunner { /** Cancel this process's attempts of a job just cancelled through this client. */ cancelLocal(job: JobRow): void { - const keyPrefix = `${job.id.toString(10)}:`; - for (const [key, controller] of this.#attempts) { - if (key.startsWith(keyPrefix)) { - controller.abort(new JobCancelledError(job.id)); - } - } - // The job may be in a claim still in flight here. - const key = attemptKey(job.id, this.#context.clientId); - if (!this.#attempts.has(key)) { - for (const pending of this.#claimCancellations) pending.add(key); + const attempt = this.#attempts.get( + attemptKey(job.id, this.#context.clientId) + ); + if ( + attempt !== undefined && + attempt.row.attempt === job.attempt && + attempt.row.attemptedAt?.epochNanoseconds === + job.attemptedAt?.epochNanoseconds + ) { + attempt.controller.abort(new JobCancelledError(job.id)); + return; } + this.cancelAttempt(job.id, this.#context.clientId); } /** Diagnostics reported by each work executor in use. */ @@ -348,21 +359,21 @@ export class AttemptRunner { * an unknown kind, its attempt fails before hooks or middleware run, so the * error handler sees it and the retry policy retries or discards it. * - * `cancelled` starts the attempt already cancelled, for a job whose - * cancellation arrived while the claim that took it was in flight. + * `checkCancellation` confirms a notice received while the claim was in flight + * before starting work. A stale notice cannot cancel a retried attempt. */ async run( row: JobRow, releaseCapacity: () => void, decodeError?: Error, - cancelled = false + checkCancellation = false ): Promise { const registration = workerRegistration(this.#workers, row.kind); const attemptController = new AbortController(); const timeoutController = new AbortController(); const key = attemptKey(row.id, this.#context.clientId); - this.#attempts.set(key, attemptController); - if (cancelled) attemptController.abort(new JobCancelledError(row.id)); + this.#attempts.set(key, { controller: attemptController, row }); + if (checkCancellation) await this.#checkCancellations(); // Like Go, a remote cancellation or shutdown is tracked apart from the // attempt's own timeout: a remote cancellation that arrives after the // timeout still decides how the attempt ends. @@ -479,7 +490,8 @@ export class AttemptRunner { timers.dispose(); this.#peers?.abandon(metadataState); metadataState.active = false; - this.#attempts.delete(key); + if (this.#attempts.get(key)?.controller === attemptController) + this.#attempts.delete(key); } } @@ -495,9 +507,8 @@ export class AttemptRunner { /** * Keep cancellations of jobs not worked here that arrive until `end()`, - * for one claim, like River for Go's producer does for one fetch: an - * attempt the claim starts for such a job starts cancelled. Other - * cancellations are discarded when the claim ends. + * for one claim. Confirm them against durable state before starting work. + * Other hints are discarded when the claim ends. */ watchCancellations(): { cancelled(id: bigint): boolean; @@ -871,15 +882,15 @@ export class AttemptRunner { } /** - * Cancel the attempts among `ids` whose jobs have a cancellation request, + * Cancel the snapshotted attempts whose jobs have a cancellation request, * reading them in batches of 1,000 IDs, each bounded by ten seconds. */ async #cancelRequested( requested: NonNullable, - ids: readonly bigint[], + attempts: ReadonlyMap, signal: AbortSignal ): Promise { - const clientId = this.#context.clientId; + const ids = [...attempts.keys()]; for (let start = 0; start < ids.length; start += CANCEL_POLL_BATCH_SIZE) { const timeout = this.#context.timer.timeout( CANCEL_POLL_TIMEOUT_MS, @@ -893,7 +904,11 @@ export class AttemptRunner { }), link.signal ); - for (const id of cancelled) this.cancelAttempt(id, clientId); + // Never look up a fresh attempt by ID after the query: a retry may + // already be running with the same ID and client owner. + for (const id of cancelled) { + attempts.get(id)?.abort(new JobCancelledError(id)); + } } finally { link[Symbol.dispose](); timeout.dispose(); @@ -901,12 +916,44 @@ export class AttemptRunner { } } - /** The job IDs of the attempts this client runs itself. */ - #ownAttemptIds(): bigint[] { + /** Coalesce notification checks while preserving notices received during a read. */ + #checkCancellations(): Promise { + this.#cancellationCheckPending = true; + if (this.#cancellationCheck !== undefined) return this.#cancellationCheck; + const check = (async () => { + try { + while (this.#cancellationCheckPending) { + this.#cancellationCheckPending = false; + try { + await this.recoverCancellations(this.#runSignal); + } catch (error: unknown) { + if (this.#runSignal.aborted) return; + this.#context.logger.warn( + "River could not confirm a cancellation notification", + { error: describeError(error) } + ); + } + } + } finally { + this.#cancellationCheck = undefined; + } + })(); + this.#cancellationCheck = check; + this.#context.trackTask(check); + return check; + } + + /** Snapshot the executions themselves, so late results cannot cancel retries. */ + #ownAttempts(): Map { const suffix = `:${this.#context.clientId}`; - return [...this.#attempts.keys()] - .filter((key) => key.endsWith(suffix)) - .map((key) => BigInt(key.slice(0, -suffix.length))); + return new Map( + [...this.#attempts] + .filter(([key]) => key.endsWith(suffix)) + .map(([key, { controller }]) => [ + BigInt(key.slice(0, -suffix.length)), + controller, + ]) + ); } async #reportStuck( diff --git a/producer.go b/producer.go index 4b00d16a6..a5f1dcfd0 100644 --- a/producer.go +++ b/producer.go @@ -45,18 +45,20 @@ const ( // Test-only properties. type producerTestSignals struct { - CancelHandledDuringFetch testsignal.TestSignal[int64] // notifies when a cancellation is handled during a fetch - DeletedExpiredQueueRecords testsignal.TestSignal[struct{}] // notifies when the producer deletes expired queue records - ExecutorShutdownStarted testsignal.TestSignal[struct{}] // notifies when the producer starts draining active jobs - JobFetchTriggered testsignal.TestSignal[struct{}] // notifies when the producer's fetch limiter is triggered via triggerJobFetch - MetadataChanged testsignal.TestSignal[struct{}] // notifies when the producer detects a metadata change - Paused testsignal.TestSignal[struct{}] // notifies when the producer is paused - PolledQueueConfig testsignal.TestSignal[struct{}] // notifies when the producer polls for queue settings - QueueControlEventTriggered testsignal.TestSignal[*controlEventPayload] // notifies when a queue control event is triggered via triggerQueueControlEvent - ReportedProducerStatus testsignal.TestSignal[struct{}] // notifies when the producer reports its own status - ReportedQueueStatus testsignal.TestSignal[struct{}] // notifies when the producer reports queue status - Resumed testsignal.TestSignal[struct{}] // notifies when the producer is resumed - StartedExecutors testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass + CancelHandledDuringFetch testsignal.TestSignal[int64] // notifies when a cancellation is handled during a fetch + DeletedExpiredQueueRecords testsignal.TestSignal[struct{}] // notifies when the producer deletes expired queue records + ExecutorShutdownStarted testsignal.TestSignal[struct{}] // notifies when the producer starts draining active jobs + JobFetchTriggered testsignal.TestSignal[struct{}] // notifies when the producer's fetch limiter is triggered via triggerJobFetch + MetadataChanged testsignal.TestSignal[struct{}] // notifies when the producer detects a metadata change + Paused testsignal.TestSignal[struct{}] // notifies when the producer is paused + PolledQueueConfig testsignal.TestSignal[struct{}] // notifies when the producer polls for queue settings + QueueControlEventTriggered testsignal.TestSignal[*controlEventPayload] // notifies when a queue control event is triggered via triggerQueueControlEvent + Reconciled testsignal.TestSignal[struct{}] // notifies when durable cancellation state was applied + ReconnectHandledDuringFetch testsignal.TestSignal[struct{}] // notifies when a reconnect is handled during a fetch + ReportedProducerStatus testsignal.TestSignal[struct{}] // notifies when the producer reports its own status + ReportedQueueStatus testsignal.TestSignal[struct{}] // notifies when the producer reports queue status + Resumed testsignal.TestSignal[struct{}] // notifies when the producer is resumed + StartedExecutors testsignal.TestSignal[struct{}] // notifies when runOnce finishes a pass } func (ts *producerTestSignals) Init(tb testutil.TestingTB) { @@ -67,6 +69,7 @@ func (ts *producerTestSignals) Init(tb testutil.TestingTB) { ts.Paused.Init(tb) ts.PolledQueueConfig.Init(tb) ts.QueueControlEventTriggered.Init(tb) + ts.ReconnectHandledDuringFetch.Init(tb) ts.ReportedQueueStatus.Init(tb) ts.ReportedProducerStatus.Init(tb) ts.Resumed.Init(tb) @@ -217,8 +220,7 @@ type producer struct { // Receives job IDs from notifier, polling, and local cancellation calls. // Only the main goroutine reads it. - cancelCh chan int64 - cancelState atomic.Pointer[producerCancelState] + cancelCh chan producerCancelRequest // Set to true when the producer thinks it should trigger another fetch as // soon as slots are available. This is written and read by the main @@ -246,6 +248,7 @@ type producer struct { reconnectCh chan struct{} retryPolicy ClientRetryPolicy + runState atomic.Pointer[producerRunState] testSignals producerTestSignals } @@ -264,7 +267,7 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi producer := baseservice.Init(archetype, &producer{ activeJobs: make(map[int64]*jobexecutor.JobExecutor), - cancelCh: make(chan int64, 1000), + cancelCh: make(chan producerCancelRequest, 1000), completer: config.Completer, config: config.mustValidate(), exec: exec, @@ -313,6 +316,18 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } + // Controls left over from the previous run must not override current queue + // settings or cancel a retried attempt. Delivery can race with shutdown. +drainControls: + for { + select { + case <-p.cancelCh: + case <-p.queueControlCh: + default: + break drainControls + } + } + // Workers can be registered after the client is constructed, so capture // kinds at startup. Keep an empty registry non-nil to avoid fetching all kinds. p.fetchKind = nil @@ -407,7 +422,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return err } - controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(workCtx)) + controlSub, err = p.config.Notifier.Listen(fetchCtx, notifier.NotificationTopicControl, p.handleControlNotification(fetchCtx, workCtx)) if err != nil { stopped() if strings.HasSuffix(err.Error(), "conn closed") || errors.Is(err, context.Canceled) { @@ -426,12 +441,12 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } 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() { @@ -475,8 +490,6 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { } p.fetchAndRunLoop(fetchCtx, workCtx) - p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) - p.executorShutdownLoop(workCtx) p.Logger.DebugContext(workCtx, p.Name+": Shutdown loop exited, awaiting subroutines", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) cancelSubroutines(fmt.Errorf("producer stopped: %w", startstop.ErrStop)) @@ -506,9 +519,25 @@ func (p *producer) TriggerJobFetch() { // listen/notify. This is used by clients using drivers that don't support // listeners to wake a producer immediately after a queue control event was // known to be performed so the producer doesn't have to wait on polling. -func (p *producer) TriggerQueueControlEvent(controlEvent *controlEventPayload) { - p.queueControlCh <- controlEvent - p.testSignals.QueueControlEventTriggered.Signal(controlEvent) +func (p *producer) TriggerQueueControlEvent(ctx context.Context, controlEvent *controlEventPayload) { + state := p.runState.Load() + if state == nil || controlEvent == nil { + return + } + + // Queue controls are only consumed while fetching. Unlike cancellations, + // they must stop being delivered as soon as the producer starts draining. + select { + case <-state.fetchDone: + return + default: + } + select { + case p.queueControlCh <- controlEvent: + p.testSignals.QueueControlEventTriggered.Signal(controlEvent) + case <-ctx.Done(): + case <-state.fetchDone: + } } type controlAction string @@ -531,7 +560,7 @@ type insertPayload struct { Queue string `json:"queue"` } -func (p *producer) handleControlNotification(workCtx context.Context) func(notifier.NotificationTopic, string) { +func (p *producer) handleControlNotification(fetchCtx, workCtx context.Context) func(notifier.NotificationTopic, string) { return func(topic notifier.NotificationTopic, payload string) { var decoded controlEventPayload if err := json.Unmarshal([]byte(payload), &decoded); err != nil { @@ -541,12 +570,15 @@ func (p *producer) handleControlNotification(workCtx context.Context) func(notif switch decoded.Action { case controlActionMetadataChanged, controlActionPause, controlActionResume: + if fetchCtx.Err() != nil { + return + } if decoded.Queue != rivercommon.AllQueuesString && decoded.Queue != p.config.Queue { p.Logger.DebugContext(workCtx, p.Name+": Queue control notification for other queue", slog.String("action", string(decoded.Action))) return } select { - case <-workCtx.Done(): + case <-fetchCtx.Done(): case p.queueControlCh <- &decoded: default: p.Logger.WarnContext(workCtx, p.Name+": Queue control notification dropped due to full buffer", slog.String("action", string(decoded.Action))) @@ -562,7 +594,7 @@ func (p *producer) handleControlNotification(workCtx context.Context) func(notif } select { case <-workCtx.Done(): - case p.cancelCh <- decoded.JobID: + case p.cancelCh <- producerCancelRequest{jobID: decoded.JobID}: default: p.Logger.WarnContext(workCtx, p.Name+": Job cancel notification dropped due to full buffer", slog.Int64("job_id", decoded.JobID)) } @@ -582,7 +614,8 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { p.fetchLimiter.Call() fetchResultCh := make(chan producerFetchResult) - reconcileCtx, cancelReconcile := context.WithCancel(fetchCtx) + // Reconciliation must remain live while workers drain after fetching stops. + reconcileCtx, cancelReconcile := context.WithCancel(context.WithoutCancel(fetchCtx)) var reconcileWG sync.WaitGroup defer func() { cancelReconcile() @@ -633,6 +666,7 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { } else { reconcileAttempt = 0 p.applyReconcileResult(workCtx, result) + p.testSignals.Reconciled.Signal(struct{}{}) } if reconcilePending { reconcilePending = false @@ -640,17 +674,15 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { } } +fetchLoop: for { select { case <-fetchCtx.Done(): - return + break fetchLoop case msg := <-p.queueControlCh: switch msg.Action { case controlActionCancel: - // This path is only expected to take effect in poll-only mode, and - // only works for the case of a single process. Multi-process setups - // will have to wait for the next poll event for a cancel to take effect. - p.maybeCancelJob(workCtx, msg.JobID) + startReconcile() case controlActionMetadataChanged: p.Logger.DebugContext(workCtx, p.Name+": Queue metadata changed", slog.String("queue", p.config.Queue), slog.String("queue_in_message", msg.Queue)) p.testSignals.MetadataChanged.Signal(struct{}{}) @@ -684,8 +716,10 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { default: p.Logger.DebugContext(workCtx, p.Name+": Unknown queue control action", "action", msg.Action) } - case jobID := <-p.cancelCh: - p.maybeCancelJob(workCtx, jobID) + case request := <-p.cancelCh: + if !request.cancel(workCtx) { + startReconcile() + } case <-p.reconnectCh: startReconcile() case result := <-reconcileResultCh: @@ -696,7 +730,7 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { // the fetchLimiter is also ready to fire: select { case <-fetchCtx.Done(): - return + break fetchLoop default: } case result := <-p.jobResultCh: @@ -711,6 +745,9 @@ func (p *producer) fetchAndRunLoop(fetchCtx, workCtx context.Context) { } } } + + p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) + p.executorShutdownLoop(workCtx, reconcileResultCh, handleReconcileResult, startReconcile) } // Loops every FetchPollInterval to check for jobs. This is meant as a back up @@ -760,9 +797,7 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } - // A fetched job may not have an executor yet. Keep unmatched cancellations - // only for this fetch so unrelated IDs cannot affect later fetches. - pendingCancelJobs := make(map[int64]struct{}) + var reconcileAfterFetch bool go p.dispatchWork(workCtx, limit, fetchResultCh) for { @@ -771,7 +806,10 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr if result.err != nil { p.Logger.ErrorContext(workCtx, p.Name+": Error fetching jobs", slog.String("err", result.err.Error()), slog.String("queue", p.config.Queue)) } else if numLocked := len(result.jobs); numLocked > 0 { - p.startNewExecutors(workCtx, result.jobs, result.decodeErrors, pendingCancelJobs) + p.startNewExecutors(workCtx, result.jobs, result.decodeErrors) + if reconcileAfterFetch { + startReconcile() + } if numLocked == limit { // Fetch returned the maximum number of jobs that were requested, @@ -782,21 +820,27 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } return case <-p.reconnectCh: + // Claimed jobs may still be absent from activeJobs until this fetch + // returns. Recheck them after starting their executors too. + reconcileAfterFetch = true startReconcile() + p.testSignals.ReconnectHandledDuringFetch.Signal(struct{}{}) case result := <-reconcileResultCh: handleReconcileResult(result) case result := <-p.jobResultCh: p.removeActiveJob(result) - case jobID := <-p.cancelCh: - if !p.maybeCancelJob(workCtx, jobID) { - pendingCancelJobs[jobID] = struct{}{} + case request := <-p.cancelCh: + if !request.cancel(workCtx) { + // The claim may not be registered yet. Check durable state once it is. + reconcileAfterFetch = true + startReconcile() } - p.testSignals.CancelHandledDuringFetch.Signal(jobID) + p.testSignals.CancelHandledDuringFetch.Signal(request.jobID) } } } -func (p *producer) executorShutdownLoop(ctx context.Context) { +func (p *producer) executorShutdownLoop(ctx context.Context, reconcileResultCh <-chan producerReconcileResult, handleReconcileResult func(producerReconcileResult), startReconcile func()) { p.testSignals.ExecutorShutdownStarted.Signal(struct{}{}) // No more jobs will be fetched or executed. However, we must wait for all @@ -810,10 +854,16 @@ func (p *producer) executorShutdownLoop(ctx context.Context) { } select { - case jobID := <-p.cancelCh: - p.maybeCancelJob(ctx, jobID) + case request := <-p.cancelCh: + if !request.cancel(ctx) { + startReconcile() + } case result := <-p.jobResultCh: p.removeActiveJob(result) + case <-p.reconnectCh: + startReconcile() + case result := <-reconcileResultCh: + handleReconcileResult(result) } } } @@ -877,24 +927,12 @@ func (p *producer) addActiveJob(id int64, executor *jobexecutor.JobExecutor) { // applyReconcileResult runs on the producer goroutine and ignores replaced executions. func (p *producer) applyReconcileResult(ctx context.Context, result producerReconcileResult) { - for _, job := range result.jobs { + for _, id := range result.cancelRequestedIDs { p.activeJobsMu.Lock() - executor := p.activeJobs[job.ID] + executor := p.activeJobs[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) + if executor != nil && executor == result.executors[id] { + executor.Cancel(ctx) } } } @@ -940,17 +978,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() @@ -966,15 +993,15 @@ func (p *producer) handleNotifierReconnect(ctx context.Context, resultCh chan<- 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{ + cancelRequestedIDs, err := timeoututil.WithTimeoutV(ctx, 10*time.Second, p.Name+".handleNotifierReconnect", func(ctx context.Context) ([]int64, error) { + return p.exec.JobGetCancelRequested(ctx, &riverdriver.JobGetCancelRequestedParams{ ID: ids, Schema: p.config.Schema, }) }) select { case <-ctx.Done(): - case resultCh <- producerReconcileResult{err: err, executors: executors, jobs: jobs}: + case resultCh <- producerReconcileResult{cancelRequestedIDs: cancelRequestedIDs, err: err, executors: executors}: } }() } @@ -1120,7 +1147,7 @@ func (p *producer) heartbeatLogLoop(ctx context.Context, wg *sync.WaitGroup) { // couldn't be decoded isn't worked, but it still gets an executor that fails // its attempt with the decode error, so that it's retried or discarded // through the normal error handling path instead of being left running. -func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.JobRow, decodeErrors map[int64]error, pendingCancelJobs map[int64]struct{}) { +func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.JobRow, decodeErrors map[int64]error) { defaultClientRetryPolicy := retrypolicy.NewDefault(p.Time) for _, job := range jobs { @@ -1161,10 +1188,6 @@ func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype. WorkUnit: workUnit, }) p.addActiveJob(job.ID, executor) - if _, ok := pendingCancelJobs[job.ID]; ok { - // Cancel before execution so the worker starts with a cancelled context. - executor.Cancel(workCtx) - } go executor.Execute(jobCtx) } @@ -1206,8 +1229,9 @@ func (p *producer) pollForJobCancellations(ctx context.Context, wg *sync.WaitGro func (p *producer) pollForJobCancellationsOnce(ctx context.Context) error { p.activeJobsMu.Lock() - jobIDs := slices.Collect(maps.Keys(p.activeJobs)) + executors := maps.Clone(p.activeJobs) p.activeJobsMu.Unlock() + jobIDs := slices.Collect(maps.Keys(executors)) // Bound query size (including SQLite's parameter count) and skip database // access entirely when there are no active jobs. @@ -1223,10 +1247,10 @@ func (p *producer) pollForJobCancellationsOnce(ctx context.Context) error { } for _, id := range cancelIDs { - select { - case <-ctx.Done(): - return ctx.Err() - case p.cancelCh <- id: + // Signal the snapshotted execution: a retry may already be running + // with the same ID by the time the query returns. + if executor := executors[id]; executor != nil { + executor.Cancel(ctx) } } } @@ -1387,8 +1411,19 @@ func (p *producer) reportQueueStatusOnce(ctx context.Context) { p.testSignals.ReportedQueueStatus.Signal(struct{}{}) } -type producerCancelState struct { - done <-chan struct{} +// A local request retains the executor that matched the cancelled database row. +// Requests without that identity must be confirmed against durable state. +type producerCancelRequest struct { + executor *jobexecutor.JobExecutor + jobID int64 +} + +func (r producerCancelRequest) cancel(ctx context.Context) bool { + if r.executor == nil { + return false + } + r.executor.Cancel(ctx) + return true } type producerFetchResult struct { @@ -1398,9 +1433,14 @@ type producerFetchResult struct { } type producerReconcileResult struct { - err error - executors map[int64]*jobexecutor.JobExecutor - jobs []*rivertype.JobRow + cancelRequestedIDs []int64 + err error + executors map[int64]*jobexecutor.JobExecutor +} + +type producerRunState struct { + done <-chan struct{} + fetchDone <-chan struct{} } type errorHandlerAdapter struct { diff --git a/producer_test.go b/producer_test.go index c63e04325..4c51e50dc 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' @@ -152,17 +152,14 @@ func TestProducer_ApplyReconcileResult(t *testing.T) { t.Parallel() for _, testCase := range []struct { - name string - attempt int - metadata string - replaced bool - wantCancel bool + name string + cancelRequested bool + 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}, + {name: "Cancelled", cancelRequested: true, wantCancel: true}, + {name: "NoCancellation"}, + {name: "ReplacedExecutor", cancelRequested: true, replaced: true}, } { t.Run(testCase.name, func(t *testing.T) { t.Parallel() @@ -177,9 +174,13 @@ func TestProducer_ApplyReconcileResult(t *testing.T) { } prod := baseservice.Init(archetype, &producer{activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}}) + var cancelRequestedIDs []int64 + if testCase.cancelRequested { + cancelRequestedIDs = []int64{1} + } prod.applyReconcileResult(context.Background(), producerReconcileResult{ - executors: map[int64]*jobexecutor.JobExecutor{1: snapshotExecutor}, - jobs: []*rivertype.JobRow{{ID: 1, Attempt: testCase.attempt, Metadata: []byte(testCase.metadata)}}, + cancelRequestedIDs: cancelRequestedIDs, + executors: map[int64]*jobexecutor.JobExecutor{1: snapshotExecutor}, }) if testCase.wantCancel { require.ErrorIs(t, context.Cause(ctx), rivertype.ErrJobCancelledRemotely) @@ -190,142 +191,231 @@ func TestProducer_ApplyReconcileResult(t *testing.T) { } } -// 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{ + cancelCh: make(chan producerCancelRequest, 1), + config: &producerConfig{Queue: QueueDefault}, + queueControlCh: make(chan *controlEventPayload, 1), + }) + fetchCtx, cancelFetch := context.WithCancel(context.Background()) + handleNotification := prod.handleControlNotification(fetchCtx, context.Background()) + cancelFetch() + payload, err := json.Marshal(&controlEventPayload{Action: action, JobID: 123, Queue: QueueDefault}) + require.NoError(t, err) + + handleNotification(notifier.NotificationTopicControl, string(payload)) + + require.Empty(t, prod.queueControlCh) + if action == controlActionCancel { + require.Len(t, prod.cancelCh, 1) + require.Equal(t, int64(123), (<-prod.cancelCh).jobID) + } else { + require.Empty(t, prod.cancelCh) } - 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: queuePollIntervalDefault, + 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.ReconnectHandledDuringFetch.Init(t) + prod.pilot = &afterJobGetAvailablePilot{ + Pilot: pilot, + afterJobGetAvailableFunc: func(res *riverdriver.JobGetAvailableResult, err error) { + if err == nil && res != nil && len(res.Jobs) > 0 { + fetchReturned.Signal(res.Jobs[0].ID) + <-releaseFetch.WaitC() + } + }, + } + } + prod.testSignals.ExecutorShutdownStarted.Init(t) + fetchCtx, fetchCancel := context.WithCancel(ctx) + defer fetchCancel() + workCtx, workCancel := context.WithCancel(ctx) + + require.NoError(t, prod.StartWorkContext(fetchCtx, workCtx)) + // Cancel workers before waiting for producer shutdown. + t.Cleanup(prod.Stop) + t.Cleanup(workCancel) + t.Cleanup(closeReleaseReconnect) + t.Cleanup(releaseFetchFunc) + + var jobID int64 + if testCase.fetchPending { + jobID = fetchReturned.WaitOrTimeout() + } else { + jobID = jobStarted.WaitOrTimeout() + } + require.Positive(t, jobID) + if testCase.stop { + fetchCancel() + if !testCase.fetchPending { + prod.testSignals.ExecutorShutdownStarted.WaitOrTimeout() + } + } - // 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.ReconnectHandledDuringFetch.WaitOrTimeout() + releaseFetchFunc() + } - update := riversharedtest.WaitOrTimeout(t, jobUpdates) - require.Equal(t, rivertype.JobStateCancelled, update[0].Job.State, - "job should have observed cancellation after the notifier reconnected and reconciled active jobs against cancel_attempted_at") + 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 +580,52 @@ func TestProducer_MetricEmitHook(t *testing.T) { }) } +func TestProducer_PollForJobCancellations(t *testing.T) { + t.Parallel() + + for _, replaced := range []bool{false, true} { + t.Run(fmt.Sprintf("ReplacedExecutor_%t", replaced), func(t *testing.T) { + t.Parallel() + + archetype := riversharedtest.BaseServiceArchetype(t) + originalCtx, cancelOriginal := context.WithCancelCause(context.Background()) + defer cancelOriginal(nil) + replacementCtx, cancelReplacement := context.WithCancelCause(context.Background()) + defer cancelReplacement(nil) + original := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelOriginal, JobRow: &rivertype.JobRow{ID: 1}}) + replacement := baseservice.Init(archetype, &jobexecutor.JobExecutor{CancelFunc: cancelReplacement, JobRow: &rivertype.JobRow{ID: 1}}) + prod := baseservice.Init(archetype, &producer{ + activeJobs: map[int64]*jobexecutor.JobExecutor{1: original}, + cancelCh: make(chan producerCancelRequest, 1), + config: &producerConfig{}, + }) + prod.exec = &producerTestExecutor{ + jobGetCancelRequestedFunc: func(ctx context.Context, params *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { + require.Equal(t, []int64{1}, params.ID) + if replaced { + // The old attempt finishes and is retried before the query returns. + cancelOriginal(nil) + prod.activeJobs[1] = replacement + } + return []int64{1}, nil + }, + } + + require.NoError(t, prod.pollForJobCancellationsOnce(context.Background())) + // Process any queued cancellations as the producer loop would. + select { + case request := <-prod.cancelCh: + request.cancel(context.Background()) + default: + } + require.NoError(t, replacementCtx.Err()) + if !replaced { + require.ErrorIs(t, context.Cause(originalCtx), rivertype.ErrJobCancelledRemotely) + } + }) + } +} + func TestProducer_PollOnly(t *testing.T) { t.Parallel() @@ -586,7 +722,7 @@ func TestProducer_ReconcileDoesNotBlock(t *testing.T) { <-releaseFetch return &riverdriver.JobGetAvailableResult{}, nil }, - jobGetByIDManyFunc: func(ctx context.Context, _ *riverdriver.JobGetByIDManyParams) ([]*rivertype.JobRow, error) { + jobGetCancelRequestedFunc: func(ctx context.Context, _ *riverdriver.JobGetCancelRequestedParams) ([]int64, error) { deadline, ok := ctx.Deadline() if ok { signals.QueryStarted.Signal(time.Until(deadline)) @@ -598,7 +734,7 @@ func TestProducer_ReconcileDoesNotBlock(t *testing.T) { 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 + return []int64{1}, nil } <-ctx.Done() return nil, ctx.Err() @@ -609,7 +745,7 @@ func TestProducer_ReconcileDoesNotBlock(t *testing.T) { require.NoError(t, err) prod := baseservice.Init(archetype, &producer{ activeJobs: map[int64]*jobexecutor.JobExecutor{1: executor}, - cancelCh: make(chan int64, 1), + cancelCh: make(chan producerCancelRequest, 1), config: &producerConfig{MaxWorkers: 2}, exec: exec, fetchLimiter: chanutil.NewDebouncedChan(ctx, time.Millisecond, true), @@ -642,7 +778,7 @@ func TestProducer_ReconcileDoesNotBlock(t *testing.T) { } if !testCase.cancelFromLookup { - prod.cancelCh <- 1 + prod.cancelCh <- producerCancelRequest{executor: executor, jobID: 1} } select { case <-jobCtx.Done(): @@ -759,7 +895,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin } } - mustInsert := func(ctx context.Context, t *testing.T, producer *producer, bundle *testBundle, args JobArgs) { + mustInsert := func(ctx context.Context, t *testing.T, producer *producer, bundle *testBundle, args JobArgs) int64 { t.Helper() insertParams, err := insertParamsFromConfigArgsAndOptions(bundle.archetype, bundle.config, args, &InsertOpts{ @@ -778,11 +914,12 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin insertParams.ScheduledAt = &bundle.timeBeforeStart } - _, err = bundle.exec.JobInsertFastMany(ctx, &riverdriver.JobInsertFastManyParams{ + inserted, err := bundle.exec.JobInsertFastMany(ctx, &riverdriver.JobInsertFastManyParams{ Jobs: []*riverdriver.JobInsertFastParams{(*riverdriver.JobInsertFastParams)(insertParams)}, Schema: producer.config.Schema, }) require.NoError(t, err) + return inserted[0].Job.ID } startProducer := func(t *testing.T, fetchCtx, workCtx context.Context, producer *producer) { @@ -839,6 +976,45 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin producer.testSignals.ReportedQueueStatus.WaitOrTimeout() }) + t.Run("RestartDiscardsStaleControls", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.config.QueuePollInterval = time.Hour + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, _ *Job[noOpArgs]) error { return ctx.Err() })) + startProducer(t, ctx, ctx, producer) + producer.Stop() + + // Model controls buffered just as the previous run stopped fetching. + producer.queueControlCh <- &controlEventPayload{Action: controlActionPause, Queue: bundle.queue} + producer.queueControlCh <- &controlEventPayload{Action: controlActionMetadataChanged, Metadata: []byte(`{"old":true}`), Queue: bundle.queue} + _, err := bundle.exec.QueueUpdate(ctx, &riverdriver.QueueUpdateParams{ + Metadata: []byte(`{"current":true}`), + MetadataDoUpdate: true, + Name: bundle.queue, + Schema: producer.config.Schema, + }) + require.NoError(t, err) + + id := mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + _, err = bundle.exec.JobCancel(ctx, &riverdriver.JobCancelParams{ + CancelAttemptedAt: time.Now(), + ControlTopic: string(notifier.NotificationTopicControl), + ID: id, + Schema: producer.config.Schema, + }) + require.NoError(t, err) + _, err = bundle.exec.JobRetry(ctx, &riverdriver.JobRetryParams{ID: id, Schema: producer.config.Schema}) + require.NoError(t, err) + producer.cancelCh <- producerCancelRequest{jobID: id} + + startProducer(t, ctx, ctx, producer) + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + producer.testSignals.MetadataChanged.RequireEmpty() + producer.testSignals.Paused.RequireEmpty() + }) + t.Run("UnknownJobKind", func(t *testing.T) { t.Parallel() @@ -1021,6 +1197,33 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRetryable, update.Job.State) }) + t.Run("CancellationNoticeWithoutMarkerDoesNotCancelJob", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.testSignals.Reconciled.Init(t) + var entered testsignal.TestSignal[struct{}] + var release testsignal.TestSignal[struct{}] + entered.Init(t) + release.Init(t) + AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, _ *Job[noOpArgs]) error { + entered.Signal(struct{}{}) + <-release.WaitC() + return ctx.Err() + })) + id := mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + startProducer(t, ctx, ctx, producer) + t.Cleanup(func() { release.Signal(struct{}{}) }) + entered.WaitOrTimeout() + + // An old notification can arrive after retry cleared the durable marker. + producer.cancelCh <- producerCancelRequest{jobID: id} + producer.testSignals.Reconciled.WaitOrTimeout() + release.Signal(struct{}{}) + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + }) + t.Run("CancellationWhileFetchResultPendingReachesJob", func(t *testing.T) { t.Parallel() @@ -1054,6 +1257,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin testutil.JobArgsReflectKind[JobArgs] } AddWorker(bundle.workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + <-ctx.Done() workerErr.Signal(ctx.Err()) return ctx.Err() })) @@ -1073,7 +1277,7 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateRunning, jobAfterCancel.State) // Send directly to avoid depending on notifier delivery timing. - producer.cancelCh <- jobID + producer.cancelCh <- producerCancelRequest{jobID: jobID} require.Equal(t, jobID, producer.testSignals.CancelHandledDuringFetch.WaitOrTimeout()) releaseFetch.Signal(struct{}{}) diff --git a/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go b/riverdriver/riverdatabasesql/internal/dbsqlc/river_job.sql.go index 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 5080dc421..0b05ee242 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 e48c4adf7..13271c6a8 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 7787761f6..11c26f388 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 1fefd422a..ff7eb13ac 100644 --- a/riverdriver/riversqlite/river_sqlite_driver.go +++ b/riverdriver/riversqlite/river_sqlite_driver.go @@ -859,7 +859,7 @@ func (e *Executor) JobRetry(ctx context.Context, params *riverdriver.JobRetryPar return nil, interpretError(err) } - return jobRowFromInternal(job) + return jobRowFromInternal(jobFromReturningRow((*dbsqlc.JobGetAvailableRow)(job))) }) } diff --git a/ruby/lib/client_runtime.rb b/ruby/lib/client_runtime.rb index ac064ca8e..30d702ec8 100644 --- a/ruby/lib/client_runtime.rb +++ b/ruby/lib/client_runtime.rb @@ -120,7 +120,8 @@ def healthy? @mutex.synchronize do @stop_requested || (@producer_threads.all? { |name, thread| @removed_queues[name] || thread.alive? } && (!@maintenance_thread || @maintenance_thread.alive?) && - (!@notification_thread || @notifications_disabled || @notification_thread.alive?)) + (!@notification_thread || @notifications_disabled || @notification_thread.alive?) && + (!@cancellation_thread || @cancellation_thread.alive?)) end end @@ -197,6 +198,9 @@ def start @started = true @stop_requested = false @stopped = false + thread = Thread.new { cancellation_loop } + @cancellation_thread = thread + @threads << thread end begin @@ -243,6 +247,7 @@ def stop(cancel: false, wait: true) @stopped = true @threads.clear @producer_threads.clear + @cancellation_thread = nil @maintenance_thread = nil @notification_thread = nil @leader = false @@ -271,6 +276,15 @@ def wake end private def begin_work(id) + entry = @mutex.synchronize { @running.fetch(id) } + if entry[:check_before_work] + begin + check_remote_cancellations(entries: {id => entry}) + rescue => error + @config.logger.error("River cancellation check failed: #{error.full_message}") + end + end + interruption = @mutex.synchronize do entry = @running.fetch(id) if @stop_requested @@ -286,12 +300,45 @@ def wake raise interruption if interruption end - private def check_remote_cancellations(queue) - entries = @mutex.synchronize { @running.select { |_id, entry| entry[:queue] == queue && entry[:working] } } + private def cancellation_loop + loop do + begin + check_remote_cancellations + rescue => error + @config.logger.error("River cancellation check failed: #{error.full_message}") + end + + finished = @mutex.synchronize do + if @stop_requested && @running.empty? + true + else + interval = @queue_configs.values.map { |config| config.resolved_fetch_poll_interval(@config) }.min || @config.fetch_poll_interval + # Continue polling while producers drain, independently of fetches + # and their cooldowns. Completions and queue changes wake us early. + @condition.wait(@mutex, interval) unless @running.any? { |_id, entry| entry[:check_cancellation] } + false + end + end + break if finished + end + end + + private def check_remote_cancellations(entries: nil) + entries ||= @mutex.synchronize do + @running.select { |_id, entry| entry[:working] || entry[:check_cancellation] }.tap do |snapshot| + snapshot.each_value { |entry| entry[:check_cancellation] = false } + end + end return if entries.empty? @driver.job_get_cancelled_ids(entries.keys).each do |id| - entries.fetch(id)[:thread].raise(JobCancelError) + @mutex.synchronize do + entry = entries.fetch(id) + next unless @running[id].equal?(entry) + + entry[:cancelled] = true + entry[:thread].raise(JobCancelError) if entry[:working] + end end end @@ -524,7 +571,7 @@ def wake end @mutex.synchronize do - @running[row.id] = {cancelled: @fetching_queues.fetch(row.queue).include?(row.id), queue: row.queue, thread: thread, working: false} + @running[row.id] = {check_before_work: @fetching_queues.fetch(row.queue).include?(row.id), queue: row.queue, thread: thread, working: false} end gate.push(true) end @@ -698,7 +745,6 @@ def wake draining = queue_stopping?(queue) break if draining && running_count(queue).zero? - check_remote_cancellations(queue) if draining wait_for_running_jobs(queue, poll_interval) next @@ -779,8 +825,9 @@ def wake @mutex.synchronize do entry = @running[message["job_id"]] if entry && entry[:queue] == queue - entry[:cancelled] = true - entry[:thread].raise(JobCancelError) if entry[:working] + entry[:check_before_work] = true unless entry[:working] + entry[:check_cancellation] = true + @condition.broadcast elsif (pending = @fetching_queues[queue]) # A claim may have committed before its workers are registered. pending << message["job_id"] diff --git a/ruby/lib/driver/runtime.rb b/ruby/lib/driver/runtime.rb index f353230ca..a4ab6d395 100644 --- a/ruby/lib/driver/runtime.rb +++ b/ruby/lib/driver/runtime.rb @@ -242,27 +242,29 @@ def job_rescue_stuck(horizon:, retry_policy:, now: Time.now.utc, max: 1_000, log end def job_retry(id, now: Time.now.utc) + already_available = "state = 'available' AND scheduled_at < #{runtime_time(now)}" + cancellation_requested = runtime_postgres? ? "(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')" : "(CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL" updated_id, job = transaction do # Check the counter in the write itself so a concurrent claim cannot # advance it between validation and the update. updated_id = runtime_returning_ids(<<~SQL).first UPDATE river_job SET state = 'available', - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - metadata = #{runtime_postgres? ? "metadata - 'cancel_attempted_at'" : "jsonb_remove(metadata, '$.cancel_attempted_at')"}, - finalized_at = NULL, - scheduled_at = #{runtime_time(now)} + max_attempts = CASE WHEN #{already_available} THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, + metadata = #{runtime_postgres? ? "CASE WHEN jsonb_typeof(metadata) = 'object' THEN metadata - 'cancel_attempted_at' ELSE metadata END" : "CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END"}, + finalized_at = CASE WHEN #{already_available} THEN finalized_at ELSE NULL END, + scheduled_at = CASE WHEN #{already_available} THEN scheduled_at ELSE #{runtime_time(now)} END WHERE id = #{Integer(id)} AND state != 'running' - AND (state != 'available' OR scheduled_at > #{runtime_time(now)}) - AND attempt < #{River::MAX_ATTEMPTS_LIMIT} + AND (#{cancellation_requested} OR NOT (#{already_available})) + AND ((#{already_available}) OR attempt < #{River::MAX_ATTEMPTS_LIMIT}) RETURNING id SQL job = job_get_by_id(updated_id || id) runtime_notify("river_insert", queue: job.queue) if updated_id && job [updated_id, job] end - if !updated_id && job && job.state != "running" && (job.state != "available" || job.scheduled_at > now) && job.attempt >= River::MAX_ATTEMPTS_LIMIT + if !updated_id && job && job.state != "running" && (job.state != "available" || job.scheduled_at >= now) && job.attempt >= River::MAX_ATTEMPTS_LIMIT raise ArgumentError, "cannot retry a job with #{River::MAX_ATTEMPTS_LIMIT} or more attempts" end job @@ -489,7 +491,7 @@ def request_resign end private def runtime_cancel_attempted - runtime_postgres? ? "metadata ? 'cancel_attempted_at'" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" + runtime_postgres? ? "(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')" : "(CASE WHEN json_valid(metadata, 10) THEN metadata -> 'cancel_attempted_at' END) IS NOT NULL" end private def runtime_claim_jobs(predicate, max, attempted_by, now) diff --git a/ruby/sig/runtime.rbs b/ruby/sig/runtime.rbs index ab94b3a7e..6b7edbcc5 100644 --- a/ruby/sig/runtime.rbs +++ b/ruby/sig/runtime.rbs @@ -358,6 +358,7 @@ module River def initialize: (Time, Float, Time) -> void end + @cancellation_thread: Thread? @client: Client @condition: ConditionVariable @config: Config @@ -404,7 +405,8 @@ module River private def begin_work: (Integer) -> untyped - def check_remote_cancellations: (String) -> untyped + def cancellation_loop: () -> untyped + def check_remote_cancellations: () -> untyped def error_handler_cancel?: (Exception, Job) -> bool def execute: (JobRow) -> untyped def finish_failed: (JobRow, Job, Exception, JobTiming, ?cancelled: bool, ?worker: untyped) -> untyped diff --git a/ruby/spec/client_runtime_branch_spec.rb b/ruby/spec/client_runtime_branch_spec.rb index 999be512c..03394a6af 100644 --- a/ruby/spec/client_runtime_branch_spec.rb +++ b/ruby/spec/client_runtime_branch_spec.rb @@ -243,7 +243,7 @@ def execute(runtime, value) value.periodic_jobs.add(periodic) - expect(value.instance_variable_get(:@threads)).to eq([thread]) + expect(value.instance_variable_get(:@threads)).to eq([value.instance_variable_get(:@cancellation_thread), thread]) expect(entered).to be_empty ensure release << true if release @@ -329,7 +329,10 @@ def execute(runtime, value) value.start - expect(value.instance_variable_get(:@threads)).to be_empty + expect(value.instance_variable_get(:@producer_threads)).to be_empty + ensure + driver.define_singleton_method(:leader_release) { |_| } + value&.stop end it "allows queue changes during startup without launching a producer twice" do @@ -344,9 +347,12 @@ def execute(runtime, value) value.start value.send(:start_producer, "dynamic", River::QueueConfig.new(max_workers: 1)) - expect(value.instance_variable_get(:@threads).length).to eq(2) + expect(value.instance_variable_get(:@producer_threads).length).to eq(2) value.queue_remove("branch") value.queue_remove("dynamic") + ensure + driver.define_singleton_method(:leader_release) { |_| } + value&.stop end it "joins only work belonging to a queue as it is removed" do @@ -471,15 +477,10 @@ def execute(runtime, value) expect(fetched_at).to eq([100.0]) end - it "keeps observing active jobs when shutdown interrupts the fetch cooldown" do - checked = [] + it "drains active jobs when shutdown interrupts the fetch cooldown" do fetched = 0 job = row driver = Object.new - driver.define_singleton_method(:job_get_cancelled_ids) { |ids| - checked << ids - [] - } driver.define_singleton_method(:queue_get) { |_| nil } driver.define_singleton_method(:job_get_available) { |**| fetched += 1 @@ -496,7 +497,22 @@ def execute(runtime, value) value.send(:producer_loop, "branch", River::QueueConfig.new(max_workers: 2, fetch_cooldown: 1, fetch_poll_interval: 1)) expect(fetched).to eq(1) - expect(checked).to eq([[1], [1]]) + expect(value.instance_variable_get(:@running)).to be_empty + end + + it "drains active jobs when a fetch fails during shutdown" do + driver = Object.new + value = runtime(driver: driver) + value.instance_variable_set(:@running, {1 => {queue: "branch", working: true}}) + driver.define_singleton_method(:queue_get) do |_queue| + value.instance_variable_set(:@stop_requested, true) + raise "fetch failed during shutdown" + end + value.define_singleton_method(:wait_for_running_jobs) { |_, _| @running.clear } + + value.send(:producer_loop, "branch", River::QueueConfig.new(max_workers: 2)) + + expect(value.instance_variable_get(:@running)).to be_empty end it "rechecks queue pause state after waiting for its fetch cooldown" do @@ -728,7 +744,58 @@ def work(_job) expect { execute(value, row) }.not_to raise_error end - it "checks cancellation in one batch, excluding inactive work and other queues" do + [false, true].each do |stopping| + it "cancels work during a blocked fetch#{" and graceful stop" if stopping}" do + started, fetching, release = Queue.new, Queue.new, Queue.new + worker = Object.new + worker.define_singleton_method(:work) do |_job| + started << true + Queue.new.pop + end + job = row + marked = false + claims = 0 + driver = Object.new + driver.define_singleton_method(:queue_upsert) { |_name| } + driver.define_singleton_method(:queue_get) { |_name| nil } + driver.define_singleton_method(:job_get_available) do |**| + claims += 1 + if claims == 1 + [job] + else + fetching << true + release.pop + [] + end + end + driver.define_singleton_method(:job_get_cancelled_ids) { |ids| marked ? ids : [] } + driver.define_singleton_method(:job_set_state_if_running) do |**attributes| + job.state = attributes.fetch(:state) + job + end + value = described_class.new(Object.new, driver, + config(worker: worker, queues: {branch: 2}).with( + fetch_cooldown: 0.001, fetch_poll_interval: 0.01, leader_election_disabled: true + )) + subscription = value.subscribe([:job_cancelled]) + value.start + Timeout.timeout(2) { + started.pop + fetching.pop + } + value.stop(wait: false) if stopping + marked = true + + event = Timeout.timeout(2) { subscription.pop } + expect(event.job.state).to eq(River::JOB_STATE_CANCELLED) + ensure + value&.stop(cancel: true, wait: false) + release << true + value&.stop + end + end + + it "checks cancellation across queues in one batch, excluding inactive work" do raised = [] fake_thread = Object.new fake_thread.define_singleton_method(:raise) { |error| raised << error } @@ -748,10 +815,10 @@ def work(_job) } ) - value.send(:check_remote_cancellations, "branch") + value.send(:check_remote_cancellations) expect(raised).to eq([River::JobCancelError]) - expect(checked).to eq([[3]]) + expect(checked).to eq([[1, 3]]) end it "interrupts work that begins after stop was requested" do diff --git a/ruby/spec/driver_runtime_feature_spec.rb b/ruby/spec/driver_runtime_feature_spec.rb index 3c793561f..4771894d3 100644 --- a/ruby/spec/driver_runtime_feature_spec.rb +++ b/ruby/spec/driver_runtime_feature_spec.rb @@ -121,7 +121,7 @@ expect(postgres_driver.send(:runtime_state, River::JOB_STATE_AVAILABLE)).to end_with("::river_job_state") expect(postgres_driver.send(:runtime_time, Time.utc(2026, 1, 2))).to end_with("::timestamptz") expect(postgres_driver.send(:runtime_merge_metadata, {"ready" => true})).to include("metadata ||") - expect(postgres_driver.send(:runtime_cancel_attempted)).to eq("metadata ? 'cancel_attempted_at'") + expect(postgres_driver.send(:runtime_cancel_attempted)).to eq("(jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at')") expect(postgres_driver.send(:runtime_update_value, :attempted_by, ["one"])).to include("ARRAY[", "::text[]") expect(postgres_driver.send(:runtime_update_value, :errors, [7])).to include("ARRAY[", "::jsonb[]") expect(postgres_driver.send(:runtime_queue_columns)).to eq("name, created_at, metadata, paused_at, updated_at") diff --git a/ruby/spec/driver_runtime_shared_examples.rb b/ruby/spec/driver_runtime_shared_examples.rb index 058f54ce4..4218f86df 100644 --- a/ruby/spec/driver_runtime_shared_examples.rb +++ b/ruby/spec/driver_runtime_shared_examples.rb @@ -293,6 +293,13 @@ def work(_job) expect(driver.job_get_by_id(row.id).state).to eq("running") end + it "ignores cancellation names stored in non-object metadata" do + job = insert_job(1) + metadata = driver.send(:runtime_json, ["cancel_attempted_at"]) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = #{metadata} WHERE id = #{job.id}") + expect(driver.job_get_cancelled_ids([job.id])).to eq([]) + end + it "looks up cancellation markers for a batch, excluding missing and uncancelled jobs" do cancelled = insert_job(1) unmarked = insert_job(2) @@ -483,6 +490,51 @@ def insert_job(value, **options) expect(driver.job_retry(-1)).to be_nil end + it "increments exhausted attempts at the exact retry timestamp" do + job = insert_job(1, scheduled_at: now, max_attempts: 1) + driver.send(:runtime_execute, "UPDATE river_job SET attempt = max_attempts WHERE id = #{job.id}") + expect(driver.job_retry(job.id, now: now)).to have_attributes(max_attempts: 2, scheduled_at: job.scheduled_at) + end + + it "clears cancellation metadata without reordering an already available retry" do + job = insert_job(1, scheduled_at: now - 60, max_attempts: 1, + metadata: {"cancel_attempted_at" => now.iso8601, "preserved" => true}) + driver.send(:runtime_execute, "UPDATE river_job SET attempt = max_attempts WHERE id = #{job.id}") + + retried = driver.job_retry(job.id, now: now) + expect(retried).to have_attributes(state: "available", attempt: 1, max_attempts: 1, + scheduled_at: job.scheduled_at, metadata: job.metadata.except("cancel_attempted_at")) + expect(driver.job_get_by_id(job.id).metadata).to eq(retried.metadata) + end + + it "preserves malformed SQLite metadata and reports its decoding error", database: :sqlite do + skip "requires SQLite" if driver.send(:runtime_postgres?) + + job = insert_job(1, scheduled_at: now + 60) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = '[not json' WHERE id = #{job.id}") + + expect { driver.job_retry(job.id, now: now) }.to raise_error(River::JobRowDecodeError, /metadata/) + row = driver.send(:runtime_query_rows, "SELECT metadata FROM river_job WHERE id = #{job.id}").first + expect(driver.send(:runtime_value, row, :metadata)).to eq("[not json") + end + + [nil, 42, false, "cancel_attempted_at", ["cancel_attempted_at"]].each do |metadata| + it "preserves non-object metadata #{metadata.inspect} when retrying" do + job = insert_job(1, scheduled_at: now + 60) + encoded = driver.send(:runtime_json, JSON.generate(metadata)) + driver.send(:runtime_execute, "UPDATE river_job SET metadata = #{encoded} WHERE id = #{job.id}") + + if metadata.nil? + expect(driver.job_retry(job.id, now: now)).to have_attributes(state: "available") + else + expect { driver.job_retry(job.id, now: now) }.to raise_error(River::JobRowDecodeError) + end + column = driver.send(:runtime_postgres?) ? "metadata::text" : "json(metadata)" + row = driver.send(:runtime_query_rows, "SELECT #{column} AS metadata FROM river_job WHERE id = #{job.id}").first + expect(JSON.parse(driver.send(:runtime_value, row, :metadata))).to eq(metadata) + end + end + it "rejects manual retries after the last representable attempt without changing the job" do job = insert_job(1, max_attempts: 32_767) original = client.job_update(job.id, attempt: 32_767, finalized_at: now, state: :discarded) diff --git a/ruby/spec/notification_runtime_spec.rb b/ruby/spec/notification_runtime_spec.rb index d3e6e6d57..2e9d18b45 100644 --- a/ruby/spec/notification_runtime_spec.rb +++ b/ruby/spec/notification_runtime_spec.rb @@ -30,6 +30,7 @@ end it "cancels an attempt before its worker starts but ignores another queue's job" do + driver.define_singleton_method(:job_get_cancelled_ids) { |ids| ids } entry = {queue: "one", thread: Thread.current, working: false} runtime.instance_variable_set(:@running, 42 => entry) runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "two", job_id: 42)) @@ -40,6 +41,45 @@ expect(entry[:working]).to be false end + [false, true].each do |working| + it "ignores a delayed cancellation notice after retry#{" while working" if working}" do + raised = [] + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + entry = {queue: "one", thread: thread, working: working} + runtime.instance_variable_set(:@running, 42 => entry) + driver.define_singleton_method(:job_get_cancelled_ids) { |_ids| [] } + + runtime.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) + runtime.send(:check_remote_cancellations) + runtime.send(:begin_work, 42) unless working + + expect(raised).to be_empty + expect(entry[:cancelled]).to be_nil + expect(entry[:working]).to be true + end + end + + it "does not interrupt a replacement registered during a cancellation check" do + raised = [] + thread = Object.new + thread.define_singleton_method(:raise) { |error| raised << error } + original = {queue: "one", thread: thread, working: true} + replacement = original.dup + value = runtime + value.instance_variable_set(:@running, 42 => original) + driver.define_singleton_method(:job_get_cancelled_ids) do |ids| + value.instance_variable_set(:@running, 42 => replacement) + ids + end + + value.send(:receive_notification, "river_control", JSON.generate(action: "cancel", queue: "one", job_id: 42)) + value.send(:check_remote_cancellations) + + expect(raised).to be_empty + expect(replacement[:cancelled]).to be_nil + end + it "ignores invalid payloads without disrupting the receiver" do [nil, "{", "null", "[]", "{}"].each { |payload| runtime.send(:receive_notification, "river_insert", payload) } runtime.send(:receive_notification, "unknown", "{}") diff --git a/ruby/spec/runtime_draining_shared_examples.rb b/ruby/spec/runtime_draining_shared_examples.rb index 879b6b742..66e0cea34 100644 --- a/ruby/spec/runtime_draining_shared_examples.rb +++ b/ruby/spec/runtime_draining_shared_examples.rb @@ -77,9 +77,11 @@ runtime = client.instance_variable_get(:@runtime) signalled = false failures = 0 + poll_failed = Queue.new @driver.define_singleton_method(:job_get_cancelled_ids) do |ids| if poll_failure && signalled && failures.zero? failures += 1 + poll_failed << true raise "temporary cancellation query failure" end super(ids) @@ -106,6 +108,7 @@ nil end Timeout.timeout(5) { draining.pop } + Timeout.timeout(5) { poll_failed.pop } if poll_failure # Persist through another client so only database polling can deliver it. River::Client.new(@driver).job_cancel(job.id) diff --git a/rust/riverqueue/src/client/attempts.rs b/rust/riverqueue/src/client/attempts.rs index 7483935d5..f51ffe204 100644 --- a/rust/riverqueue/src/client/attempts.rs +++ b/rust/riverqueue/src/client/attempts.rs @@ -2,11 +2,10 @@ use std::collections::HashMap; use std::sync::Mutex; -use std::sync::atomic::{AtomicU64, Ordering}; use tokio_util::sync::CancellationToken; -use crate::client::{ClientInner, PENDING_CANCELLATION_LIMIT, PENDING_CANCELLATION_RETENTION}; +use crate::client::ClientInner; pub(super) fn remove_running_attempt( running: &Mutex>, @@ -25,99 +24,15 @@ pub(super) fn remove_running_attempt( } pub(super) fn register_running_attempt( - running: &Mutex>, - pending_cancellations: &Mutex>, + inner: &ClientInner, job_id: i64, cancellation: &CancellationToken, ) { - // Keep the locks in this order here and in `signal_running_attempt` so a - // cancellation cannot fall between checking the active map and recording - // a just-fetched attempt. - let mut running = running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - running.insert(job_id, cancellation.clone()); - let should_cancel = pending_cancellations + inner + .running .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .remove(&job_id) - .is_some(); - drop(running); - if should_cancel { - cancellation.cancel(); - } -} - -pub(super) fn signal_running_attempt( - running: &Mutex>, - pending_cancellations: &Mutex>, - fetch_registration_windows: &AtomicU64, - job_id: i64, -) { - let running = running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(cancellation) = running.get(&job_id).cloned() { - drop(running); - cancellation.cancel(); - return; - } - if fetch_registration_windows.load(Ordering::SeqCst) == 0 { - return; - } - - let now = std::time::Instant::now(); - let mut pending = pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - pending.retain(|_, received_at| { - now.saturating_duration_since(*received_at) <= PENDING_CANCELLATION_RETENTION - }); - if pending.len() >= PENDING_CANCELLATION_LIMIT - && let Some(oldest_job_id) = pending - .iter() - .min_by_key(|(_, received_at)| **received_at) - .map(|(job_id, _)| *job_id) - { - pending.remove(&oldest_job_id); - } - pending.insert(job_id, now); -} - -pub(super) struct FetchRegistrationGuard<'a> { - pub(super) inner: &'a ClientInner, -} - -impl<'a> FetchRegistrationGuard<'a> { - pub(super) fn new(inner: &'a ClientInner) -> Self { - inner - .fetch_registration_windows - .fetch_add(1, Ordering::SeqCst); - Self { inner } - } -} - -impl Drop for FetchRegistrationGuard<'_> { - fn drop(&mut self) { - // Synchronize the last-window transition with - // `signal_running_attempt`, which holds this lock while deciding - // whether to retain an unmatched cancellation. - let _running = self - .inner - .running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if self - .inner - .fetch_registration_windows - .fetch_sub(1, Ordering::SeqCst) - == 1 - { - self.inner - .pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clear(); - } - } + .insert(job_id, cancellation.clone()); + // Recheck after registration even if a notice arrived during the claim. + inner.cancellation_wakeup.notify_one(); } diff --git a/rust/riverqueue/src/client/builder.rs b/rust/riverqueue/src/client/builder.rs index 0c378c36d..058214cf1 100644 --- a/rust/riverqueue/src/client/builder.rs +++ b/rust/riverqueue/src/client/builder.rs @@ -1,12 +1,14 @@ //! Client configuration and construction. use std::collections::HashMap; -use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::sync::atomic::AtomicBool; +#[cfg(test)] +use std::sync::atomic::AtomicU64; use std::sync::{Arc, Mutex, RwLock}; use std::time::Duration; use serde_json::{Map, Value}; -use tokio::sync::{broadcast, watch}; +use tokio::sync::{Notify, broadcast, watch}; use crate::__private::Pilot; #[cfg(feature = "postgres")] @@ -802,6 +804,7 @@ impl ClientBuilder { inner: Arc::new(ClientInner { allow_legacy_job_kinds: self.allow_legacy_job_kinds, allow_unregistered_job_kinds: self.allow_unregistered_job_kinds, + cancellation_wakeup: Notify::new(), completion_sender: Mutex::new(None), database: self.database, default_max_attempts: self.default_max_attempts, @@ -809,7 +812,6 @@ impl ClientBuilder { events, fetch_cooldown: self.fetch_cooldown, fetch_kinds, - fetch_registration_windows: AtomicU64::new(0), hooks: self.hooks, id: self.id, insert_middleware: self.insert_middleware, @@ -823,7 +825,6 @@ impl ClientBuilder { #[cfg(test)] notifier_start_panics: AtomicU64::new(0), peer_owners: Mutex::new(HashMap::new()), - pending_cancellations: Mutex::new(HashMap::new()), periodic_jobs, pilot: self.pilot, poll_only: self.poll_only, diff --git a/rust/riverqueue/src/client/jobs.rs b/rust/riverqueue/src/client/jobs.rs index ac8482870..94ed79cce 100644 --- a/rust/riverqueue/src/client/jobs.rs +++ b/rust/riverqueue/src/client/jobs.rs @@ -9,7 +9,6 @@ use std::{ use serde_json::Map; use crate::__private::DatabaseConnection as PilotDatabaseConnection; -use crate::client::attempts::signal_running_attempt; use crate::client::request::{Target, request_type}; use crate::database::DatabaseTransactionExecutor; use crate::storage::Access; @@ -257,17 +256,10 @@ impl JobCancelRequest<'_> { let row = session.storage(inner).job_cancel(self.id).await?; session.commit().await?; // Without a listener (no backend listener, or a poll-only client), - // wake this client's running attempt directly, like Go's - // `notifyProducerWithoutListenerQueueControlEvent`. Other clients - // observe the committed notification through their own listener or - // outbox poll. + // wake the durable check. The post-commit signal can be delayed until + // a retry is running, so it must not cancel by job ID alone. if own_transaction && !inner.listens_for_notifications() { - signal_running_attempt( - &inner.running, - &inner.pending_cancellations, - &inner.fetch_registration_windows, - self.id, - ); + inner.cancellation_wakeup.notify_one(); } Ok(row) } diff --git a/rust/riverqueue/src/client/mod.rs b/rust/riverqueue/src/client/mod.rs index 68b517853..66a71958e 100644 --- a/rust/riverqueue/src/client/mod.rs +++ b/rust/riverqueue/src/client/mod.rs @@ -57,7 +57,7 @@ use std::{ }; use chrono::{DateTime, Utc}; -use tokio::sync::{broadcast, mpsc, watch}; +use tokio::sync::{Notify, broadcast, mpsc, watch}; use tokio_util::sync::CancellationToken; use self::completer::CompletionUpdate; @@ -79,8 +79,6 @@ use crate::{ pub(crate) const ATTEMPTED_BY_MAX: i32 = 100; const EVENT_BUFFER_CAPACITY: usize = 10_000; -const PENDING_CANCELLATION_LIMIT: usize = 10_000; -const PENDING_CANCELLATION_RETENTION: Duration = Duration::from_mins(1); // Large queues otherwise become limited by a single Postgres claim round trip. // Concurrent `SKIP LOCKED` claims safely divide the available worker slots. const PARALLEL_FETCH_MINIMUM: usize = 1_000; @@ -103,6 +101,7 @@ const PRODUCER_STALE_RETENTION: Duration = Duration::from_mins(5); pub(crate) struct ClientInner { allow_legacy_job_kinds: bool, allow_unregistered_job_kinds: bool, + cancellation_wakeup: Notify, completion_sender: Mutex>>, pub(crate) database: Database, default_max_attempts: i32, @@ -115,7 +114,6 @@ pub(crate) struct ClientInner { /// Kinds claims are restricted to, including aliases, or `None` to claim /// every kind. pub(crate) fetch_kinds: Option>, - fetch_registration_windows: AtomicU64, pub(crate) hooks: Vec>, pub(crate) id: String, insert_middleware: Vec>, @@ -139,7 +137,6 @@ pub(crate) struct ClientInner { notifier_start_panics: AtomicU64, /// Peer jobs owned by running attempts, mapped to their ledger. peer_owners: Mutex>, - pending_cancellations: Mutex>, pub(crate) periodic_jobs: PeriodicJobs, pub(crate) pilot: Arc, poll_only: bool, diff --git a/rust/riverqueue/src/client/notifier.rs b/rust/riverqueue/src/client/notifier.rs index be9aa976e..07ab86e67 100644 --- a/rust/riverqueue/src/client/notifier.rs +++ b/rust/riverqueue/src/client/notifier.rs @@ -22,7 +22,6 @@ use tracing::{debug, error, warn}; use crate::Error; use crate::client::ClientInner; -use crate::client::attempts::signal_running_attempt; use crate::client::backoff::exponential_backoff; #[cfg(feature = "sqlite")] use crate::client::sqlite_backend_error; @@ -125,13 +124,8 @@ pub(super) fn dispatch_notification( }; match payload.action.as_str() { "cancel" => { - if let Some(job_id) = payload.job_id { - signal_running_attempt( - &inner.running, - &inner.pending_cancellations, - &inner.fetch_registration_windows, - job_id, - ); + if payload.job_id.is_some() { + inner.cancellation_wakeup.notify_one(); } } // Like Go's producer, a metadata change is handled at once diff --git a/rust/riverqueue/src/client/producer.rs b/rust/riverqueue/src/client/producer.rs index fa68924f9..e6045da78 100644 --- a/rust/riverqueue/src/client/producer.rs +++ b/rust/riverqueue/src/client/producer.rs @@ -16,9 +16,7 @@ use tokio_util::sync::CancellationToken; use tracing::{debug, error, warn}; use crate::__private::DatabaseConnection as PilotDatabaseConnection; -use crate::client::attempts::{ - FetchRegistrationGuard, register_running_attempt, signal_running_attempt, -}; +use crate::client::attempts::register_running_attempt; use crate::client::backoff::{RestartBackoff, exponential_backoff}; use crate::client::completer::CompletionUpdate; use crate::client::executor::{AbortOnDrop, execute_job}; @@ -56,10 +54,8 @@ pub(super) async fn run_dynamic_queues( mut changes: watch::Receiver, queues_ready: oneshot::Sender<()>, ) -> Result<(), Error> { - // Without a notifier, running jobs learn of cancellation requests by - // polling until every producer has drained. Whether the client has one is - // known only once a Postgres server's capabilities are detected, so the - // poll checks each time. + // Recover lost cancellation notices, including after the notifier stops + // during graceful shutdown, until every producer has drained. let _cancellation_poll = AbortOnDrop(tokio::spawn(poll_job_cancellations( Arc::clone(&inner), fetch_cancel.clone(), @@ -143,10 +139,9 @@ fn jittered_fetch_poll_interval(interval: Duration) -> Duration { } /// Cancels this client's running attempts whose jobs have a cancellation -/// request, checking every [`QUEUE_CONFIG_POLL_INTERVAL`] while the client -/// has no notifier, like River Go's producers without one. That's checked -/// each time, since a Postgres server's lack of `LISTEN`/`NOTIFY` is -/// detected only once the client runs. +/// request, checking every [`QUEUE_CONFIG_POLL_INTERVAL`]. Notifications +/// deliver cancellations promptly; polling also recovers notices lost during +/// reconnects or while a fetched attempt was still being registered. /// It keeps checking after `stopping` fires, so jobs can still be cancelled /// while producers drain; the caller aborts it once they have. async fn poll_job_cancellations(inner: Arc, stopping: CancellationToken) { @@ -155,46 +150,42 @@ async fn poll_job_cancellations(inner: Arc, stopping: CancellationT // The first tick completes at once; running jobs start later. poll.tick().await; loop { - poll.tick().await; - if inner.has_notifier() { - continue; + tokio::select! { + _ = poll.tick() => {}, + () = inner.cancellation_wakeup.notified() => {}, } - let mut ids = inner - .running - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .keys() - .copied() - .collect::>(); - ids.sort_unstable(); - for batch in ids.chunks(JOB_CANCEL_POLL_BATCH_SIZE) { - let requested = tokio::time::timeout( - JOB_CANCEL_POLL_TIMEOUT, - crate::storage::job_cancel_requested(&inner, batch), - ) - .await - .map_err(|_| Error::runtime_context("job cancellation poll", "timed out")) - .and_then(|requested| requested); - match requested { - Ok(requested) => { - for id in requested { - signal_running_attempt( - &inner.running, - &inner.pending_cancellations, - &inner.fetch_registration_windows, - id, - ); - } - } - Err(poll_error) => { - if !stopping.is_cancelled() { - error!(error = %crate::error::Chain(&poll_error), "River failed to check for job cancellation requests"); - } - break; - } + if let Err(poll_error) = check_job_cancellations(&inner).await + && !stopping.is_cancelled() + { + error!(error = %crate::error::Chain(&poll_error), "River failed to check for job cancellation requests"); + } + } +} + +pub(super) async fn check_job_cancellations(inner: &ClientInner) -> Result<(), Error> { + let attempts = inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let mut ids = attempts.keys().copied().collect::>(); + ids.sort_unstable(); + for batch in ids.chunks(JOB_CANCEL_POLL_BATCH_SIZE) { + let requested = tokio::time::timeout( + JOB_CANCEL_POLL_TIMEOUT, + crate::storage::job_cancel_requested(inner, batch), + ) + .await + .map_err(|_| Error::runtime_context("job cancellation poll", "timed out"))??; + for id in requested { + // A query can outlive an attempt and its retry. Only signal the + // execution that was present in its snapshot. + if let Some(cancellation) = attempts.get(&id) { + cancellation.cancel(); } } } + Ok(()) } type ProducerOutcome = (String, u64, CancellationToken, Result<(), Error>); @@ -858,7 +849,6 @@ pub(super) async fn run_queue( if available == 0 { continue; } - let registration_guard = FetchRegistrationGuard::new(&inner); let fetched = match (&session, claims_through_session) { (Some(session), true) => { match claim_through_session( @@ -907,12 +897,7 @@ pub(super) async fn run_queue( for (row, decode_error) in claimed { let hard_cancel = work_cancel.child_token(); let cancellation = hard_cancel.child_token(); - register_running_attempt( - &inner.running, - &inner.pending_cancellations, - row.id, - &cancellation, - ); + register_running_attempt(&inner, row.id, &cancellation); let task_inner = Arc::clone(&inner); let completion_sender = completion_sender.clone(); let session_row = attempts.session.is_some().then(|| row.clone()); @@ -928,7 +913,6 @@ pub(super) async fn run_queue( .await; }); } - drop(registration_guard); while let Some(joined) = attempts.tasks.try_join_next_with_id() { attempts.exited(joined.map(|(task_id, ())| task_id), false); } diff --git a/rust/riverqueue/src/client/tests.rs b/rust/riverqueue/src/client/tests.rs index d52f83f76..45e3ae72a 100644 --- a/rust/riverqueue/src/client/tests.rs +++ b/rust/riverqueue/src/client/tests.rs @@ -1,7 +1,7 @@ use serde::Deserialize; use serde_json::Map; -use super::attempts::{register_running_attempt, remove_running_attempt, signal_running_attempt}; +use super::attempts::{register_running_attempt, remove_running_attempt}; #[cfg(feature = "postgres")] use super::completer::CompletionBatcher; #[cfg(feature = "sqlite")] @@ -171,12 +171,7 @@ async fn go_notification_fixtures_dispatch_correctly() { let cancellation = CancellationToken::new(); let unrelated = CancellationToken::new(); for (job_id, token) in [(42, &cancellation), (43, &unrelated)] { - register_running_attempt( - &inner.running, - &inner.pending_cancellations, - job_id, - token, - ); + inner.running.lock().unwrap().insert(job_id, token.clone()); } dispatch_notification( @@ -209,11 +204,18 @@ async fn go_notification_fixtures_dispatch_correctly() { ), _ => unreachable!(), } - assert_eq!(cancellation.is_cancelled(), name == "cancel", "{name}"); assert!( - !unrelated.is_cancelled(), - "{name}: cancelled an unrelated job" + !cancellation.is_cancelled(), + "{name}: notice cancelled without a durable check" ); + assert_eq!( + tokio::time::timeout(Duration::ZERO, inner.cancellation_wakeup.notified()) + .await + .is_ok(), + name == "cancel", + "{name}" + ); + assert!(!unrelated.is_cancelled(), "{name}"); assert!( matches!( producer.try_recv(), @@ -233,50 +235,73 @@ async fn go_notification_fixtures_dispatch_correctly() { pool.close().await; } -#[test] -fn pending_cancellation_reaches_fetched_attempt() { - let job_id = 42; +#[tokio::test] +async fn registration_wakes_cancellation_check() { + #[cfg(feature = "postgres")] + let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); + #[cfg(not(feature = "postgres"))] + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); let cancellation = CancellationToken::new(); - let fetch_registration_windows = AtomicU64::new(1); - let pending_cancellations = Mutex::new(HashMap::new()); - let running = Mutex::new(HashMap::new()); - - signal_running_attempt( - &running, - &pending_cancellations, - &fetch_registration_windows, - job_id, - ); - register_running_attempt(&running, &pending_cancellations, job_id, &cancellation); - - assert!(cancellation.is_cancelled()); - assert!( - pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() - ); + register_running_attempt(&client.inner, 42, &cancellation); + tokio::time::timeout( + Duration::from_secs(1), + client.inner.cancellation_wakeup.notified(), + ) + .await + .unwrap(); + assert!(!cancellation.is_cancelled()); } -#[test] -fn unmatched_cancellation_is_not_retained_without_fetch() { - let fetch_registration_windows = AtomicU64::new(0); - let pending_cancellations = Mutex::new(HashMap::new()); - let running = Mutex::new(HashMap::new()); - - signal_running_attempt( - &running, - &pending_cancellations, - &fetch_registration_windows, - 42, - ); - - assert!( - pending_cancellations - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn delayed_cancellation_signals_recheck_retry_marker() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + riverqueue_migrate::SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()) + .without_notifications() + .build() + .unwrap(); + let id: i64 = sqlx::query_scalar("INSERT INTO river_job (kind, args, state, attempt, max_attempts, metadata) VALUES ('test', jsonb('{}'), 'running', 1, 2, jsonb('{}')) RETURNING id") + .fetch_one(&pool).await.unwrap(); + let original = CancellationToken::new(); + register_running_attempt(&client.inner, id, &original); + client.jobs().cancel(id).await.unwrap(); + sqlx::query("UPDATE river_job SET state = 'cancelled', finalized_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = ?").bind(id).execute(&pool).await.unwrap(); + client.jobs().retry(id).await.unwrap(); + sqlx::query("UPDATE river_job SET state = 'running', attempt = 2 WHERE id = ?") + .bind(id) + .execute(&pool) + .await + .unwrap(); + let replacement = CancellationToken::new(); + register_running_attempt(&client.inner, id, &replacement); + dispatch_notification( + &client.inner, + &client.inner.queue_notifications, + crate::NOTIFICATION_TOPIC_CONTROL, + &format!(r#"{{"action":"cancel","job_id":{id},"queue":"default"}}"#), ); + super::producer::check_job_cancellations(&client.inner) + .await + .unwrap(); + assert!(!replacement.is_cancelled()); + // The local post-commit signal is also only a hint, and remains useful + // when the current attempt actually has a durable cancellation request. + client.jobs().cancel(id).await.unwrap(); + assert!(!replacement.is_cancelled()); + super::producer::check_job_cancellations(&client.inner) + .await + .unwrap(); + assert!(replacement.is_cancelled()); + pool.close().await; } fn retry_row(error_count: usize) -> JobRow { diff --git a/rust/riverqueue/src/database/sqlite.rs b/rust/riverqueue/src/database/sqlite.rs index 4e8b26e21..40dfe3444 100644 --- a/rust/riverqueue/src/database/sqlite.rs +++ b/rust/riverqueue/src/database/sqlite.rs @@ -935,25 +935,30 @@ pub(crate) async fn retry( ) -> Result, BackendError> { let sql = format!( r#" + WITH retry_params AS (SELECT ? AS now) UPDATE river_job SET state = 'available', max_attempts = CASE + WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, - finalized_at = NULL, - scheduled_at = ? + finalized_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN finalized_at ELSE NULL END, + metadata = CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) + THEN metadata ELSE jsonb_remove(metadata, '$.cancel_attempted_at') END, + scheduled_at = CASE WHEN state = 'available' AND scheduled_at < (SELECT now FROM retry_params) THEN scheduled_at ELSE (SELECT now FROM retry_params) END WHERE id = ? AND state != 'running' - AND (state != 'available' OR scheduled_at > ?) + AND ((CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN NULL + ELSE json_type(metadata, '$.cancel_attempted_at') END) IS NOT NULL + OR NOT (state = 'available' AND scheduled_at < (SELECT now FROM retry_params))) RETURNING {JOB_COLUMNS} "# ); sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) .bind(sqlite_time(now)) .bind(id) - .bind(sqlite_time(now)) .fetch_optional(&mut *connection) .await? .map(JobRecord::into_job) @@ -2523,6 +2528,46 @@ mod tests { assert_eq!(row.attempted_by, ["three", "four", "five"]); } + #[tokio::test] + async fn retry_at_exact_schedule_increments_max_attempts() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Available, now).await; + sqlx::query("UPDATE river_job SET attempt = max_attempts WHERE id = ?") + .bind(job.id) + .execute(&mut *connection) + .await + .unwrap(); + let retried = retry(&mut connection, job.id, now).await.unwrap().unwrap(); + assert_eq!(retried.max_attempts, job.max_attempts + 1); + assert_eq!(retried.scheduled_at, job.scheduled_at); + } + + #[tokio::test] + async fn retry_preserves_malformed_metadata_for_row_validation() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Retryable, now).await; + sqlx::query("UPDATE river_job SET metadata = '[not json' WHERE id = ?") + .bind(job.id) + .execute(&mut *connection) + .await + .unwrap(); + + let error = retry(&mut connection, job.id, now).await.unwrap_err(); + assert!( + matches!(error, BackendError::InvalidRow(ref message) if message.contains("metadata")) + ); + let metadata: String = sqlx::query_scalar("SELECT metadata FROM river_job WHERE id = ?") + .bind(job.id) + .fetch_one(&mut *connection) + .await + .unwrap(); + assert_eq!(metadata, "[not json"); + } + /// Inserts a job for the undecodable row tests. async fn insert_test_job( connection: &mut SqliteConnection, diff --git a/rust/riverqueue/src/storage/postgres.rs b/rust/riverqueue/src/storage/postgres.rs index 316eb3587..08b0b7fe0 100644 --- a/rust/riverqueue/src/storage/postgres.rs +++ b/rust/riverqueue/src/storage/postgres.rs @@ -80,7 +80,7 @@ impl Backend for PostgresBackend<'_> { let table = self.schema.qualify("river_job"); let sql = format!( "SELECT id FROM {table} \ - WHERE id = any($1) AND metadata ? 'cancel_attempted_at' AND state = 'running' \ + WHERE id = any($1) AND jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at' AND state = 'running' \ ORDER BY id" ); Ok(sqlx::query_scalar(AssertSqlSafe(sql)) @@ -238,10 +238,11 @@ impl Backend for PostgresBackend<'_> { let sql = format!( "WITH locked AS (SELECT id FROM {table} WHERE id = $1 FOR UPDATE), \ updated AS (UPDATE {table} AS job SET state = 'available', \ - max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, \ - finalized_at = NULL, scheduled_at = now() \ + max_attempts = CASE WHEN job.state = 'available' AND job.scheduled_at < now() THEN max_attempts WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, \ + finalized_at = CASE WHEN job.state = 'available' AND job.scheduled_at < now() THEN finalized_at ELSE NULL END, metadata = CASE WHEN jsonb_typeof(metadata) = 'object' \ + THEN metadata - 'cancel_attempted_at' ELSE metadata END, scheduled_at = CASE WHEN job.state = 'available' AND job.scheduled_at < now() THEN scheduled_at ELSE now() END \ FROM locked WHERE job.id = locked.id AND job.state != 'running' \ - AND NOT (job.state = 'available' AND job.scheduled_at < now()) RETURNING job.*) \ + AND ((jsonb_typeof(metadata) = 'object' AND metadata ? 'cancel_attempted_at') OR NOT (job.state = 'available' AND job.scheduled_at < now())) RETURNING job.*) \ SELECT {}, false AS unique_skipped_as_duplicate FROM updated AS job \ UNION ALL SELECT {}, false AS unique_skipped_as_duplicate \ FROM (SELECT * FROM {table} WHERE id = $1 FOR UPDATE) AS job \ diff --git a/rust/riverqueue/tests/client_handles.rs b/rust/riverqueue/tests/client_handles.rs index 8f00521a1..a3b8ccc51 100644 --- a/rust/riverqueue/tests/client_handles.rs +++ b/rust/riverqueue/tests/client_handles.rs @@ -546,6 +546,7 @@ macro_rules! scenarios { .unwrap(); let retried = jobs.retry(cancelled).tx(&mut tx).await.unwrap(); assert_eq!(retried.state, JobState::Available); + assert!(!retried.metadata.contains_key("cancel_attempted_at")); tx.commit().await.unwrap(); assert_eq!( @@ -557,6 +558,35 @@ macro_rules! scenarios { fixture.cleanup().await; } + #[tokio::test(flavor = "multi_thread")] + async fn retry_clears_cancellation_without_reordering_an_available_job() { + let fixture = Fixture::new().await; + let metadata: riverqueue::JobMetadata = + r#"{"cancel_attempted_at":"2026-01-01T00:00:00Z","preserved":true}"# + .parse() + .unwrap(); + let id = fixture + .client + .insert(args("already available")) + .opts(InsertOpts::default().with_metadata(metadata)) + .await + .unwrap() + .id(); + let jobs = fixture.client.jobs(); + let before = jobs.get(id).await.unwrap(); + assert_eq!(before.state, JobState::Available); + + let retried = jobs.retry(id).await.unwrap(); + assert!(!retried.metadata.contains_key("cancel_attempted_at")); + assert_eq!(retried.metadata.get::("preserved").unwrap(), Some(true)); + assert_eq!(retried.scheduled_at, before.scheduled_at); + assert_eq!(retried.max_attempts, before.max_attempts); + assert!( + !jobs.get(id).await.unwrap().metadata.contains_key("cancel_attempted_at") + ); + fixture.cleanup().await; + } + #[tokio::test(flavor = "multi_thread")] async fn job_requests_run_on_the_pool_without_a_transaction() { let fixture = Fixture::new().await; diff --git a/rust/riverqueue/tests/poll_only_cancellation.rs b/rust/riverqueue/tests/poll_only_cancellation.rs index 76b6928d1..a01234f70 100644 --- a/rust/riverqueue/tests/poll_only_cancellation.rs +++ b/rust/riverqueue/tests/poll_only_cancellation.rs @@ -1,5 +1,5 @@ -//! Cancellation of running jobs in clients without notifications, on every -//! backend. +//! Durable cancellation of running jobs, including when notifications are +//! disabled or a cancellation notice is lost, on every backend. //! //! Postgres tests run in a unique schema and fail rather than skip when //! `RIVER_RUST_DATABASE_URL` is unset; SQLite tests use a temporary file. @@ -11,8 +11,8 @@ mod support; use std::{convert::Infallible, sync::Arc, time::Duration}; use riverqueue::{ - __private::Database, Client, EventKind, Job, JobArgs, JobState, QueueConfig, WorkContext, - WorkOutcome, Workers, + __private::Database, Client, EventKind, Job, JobArgs, JobState, JobUpdateParams, QueueConfig, + WorkContext, WorkOutcome, Workers, }; use serde::{Deserialize, Serialize}; use tokio::sync::Semaphore; @@ -56,11 +56,11 @@ impl Backend { } } -/// A client without notifications cancels its running job once it polls the -/// cancellation another client requested, including while it's stopping and -/// waiting for that job. +/// Durable cancellation survives lost notifications and graceful shutdown. async fn polls_for_remote_cancellation(backend: Backend) { - for while_stopping in [false, true] { + for (notifications, while_stopping) in + [(false, false), (false, true), (true, false), (true, true)] + { let started = Arc::new(Semaphore::new(0)); let mut workers = Workers::new(); let worker_started = Arc::clone(&started); @@ -75,13 +75,17 @@ async fn polls_for_remote_cancellation(backend: Backend) { } }) .unwrap(); - let client = Client::builder(backend.database()) + let builder = Client::builder(backend.database()) .queue(riverqueue::QUEUE_DEFAULT, QueueConfig::new(1)) .without_leader_election() - .without_notifications() - .workers(workers) - .build() - .unwrap(); + .workers(workers); + let client = if notifications { + builder + } else { + builder.without_notifications() + } + .build() + .unwrap(); let other = Client::builder(backend.database()).build().unwrap(); let mut events = client.subscribe(&[EventKind::JobCancelled]).unwrap(); let mut run = client.start().unwrap(); @@ -96,7 +100,27 @@ async fn polls_for_remote_cancellation(backend: Backend) { if while_stopping { run.stopper().stop(); } - other.jobs().cancel(id).await.unwrap(); + if notifications { + // Persist a cancellation without its notification, as though that + // notice was lost while the listener reconnected. + other + .jobs() + .update( + id, + JobUpdateParams::default().metadata( + [( + "cancel_attempted_at".to_owned(), + chrono::Utc::now().to_rfc3339().into(), + )] + .into_iter() + .collect(), + ), + ) + .await + .unwrap(); + } else { + other.jobs().cancel(id).await.unwrap(); + } let event = tokio::time::timeout(TIMEOUT, events.recv()) .await