From 358a79126ce45ae788c92121d3db881a317538da Mon Sep 17 00:00:00 2001 From: Chris George Date: Sun, 23 Aug 2026 14:32:18 -0700 Subject: [PATCH 1/3] 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 b347dea0d..1c7ad7176 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). 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 250adaeae30003904aef6ad933d9eb305423627d Mon Sep 17 00:00:00 2001 From: Chris George Date: Sun, 23 Aug 2026 14:53:09 -0700 Subject: [PATCH 2/3] 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 1c7ad7176..07e05cd9b 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 6775f580b0a8c8603709768f0a9dc95b2b2b3982 Mon Sep 17 00:00:00 2001 From: Chris George Date: Tue, 20 Oct 2026 07:35:49 -0700 Subject: [PATCH 3/3] 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 07e05cd9b..b12bb392f 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). 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: