From 46d000ff5c5d937cb676723ebfbb285e8c4fcf23 Mon Sep 17 00:00:00 2001 From: Brandur Date: Tue, 21 Jul 2026 15:10:34 -0500 Subject: [PATCH 1/2] Track unhealthy heartbeat + pause job fetching when unhealthy This one's driven by another issue flagged out of the active rescue system: as currently implemented, the producer's heartbeat produces an error if it fails, but there's no additional feedback. So if a producer were to chronically start failing its heartbeats, it'd be possible for the `JobRescuer` (were its rescue queries to keep working in this case, which is not certain) to start incorrectly rescuing jobs as producers started to look inactive. This could then result in duplicate jobs. This change doesn't fully address that problem, but partly remediates it by having each producer track whether its heartbeat is unhealthy. If found to be unhealthy, the producer stops fetching new jobs because if a simple producer update is failing, presumably the rest of the database is also experiencing significant trouble. In this case jobs running in an unhealthy producer could still be rescued, but at least they wouldn't start work again as their producer is unhealthy, thus avoiding a potentially runaway feedback loop. --- producer.go | 77 +++++++++++++++++++++++++++++++++--------- producer_test.go | 88 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 16 deletions(-) diff --git a/producer.go b/producer.go index 88e30420..fb44f861 100644 --- a/producer.go +++ b/producer.go @@ -191,16 +191,17 @@ type producer struct { // Jobs which are currently being worked. Only used by main goroutine. activeJobs map[int64]*jobexecutor.JobExecutor - completer jobcompleter.JobCompleter - config *producerConfig - id atomic.Int64 // atomic because it's written at startup and read during shutdown - exec riverdriver.Executor - errorHandler jobexecutor.ErrorHandler - fetchLimiter *chanutil.DebouncedChan - metricEmitHooks []rivertype.HookMetricEmit // memoized hooks of type HookMetricEmit for reuse in dispatchWork - state riverpilot.ProducerState - pilot riverpilot.Pilot - workers *Workers + completer jobcompleter.JobCompleter + config *producerConfig + id atomic.Int64 // atomic because it's written at startup and read during shutdown + exec riverdriver.Executor + errorHandler jobexecutor.ErrorHandler + fetchLimiter *chanutil.DebouncedChan + heartbeatUnhealthy atomic.Bool // atomic because it's written by the status loop and read by the fetch loop + metricEmitHooks []rivertype.HookMetricEmit // memoized hooks of type HookMetricEmit for reuse in dispatchWork + state riverpilot.ProducerState + pilot riverpilot.Pilot + workers *Workers // Receives job IDs to cancel. Written by notifier goroutine, only read from // main goroutine. @@ -348,6 +349,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return err } p.id.Store(id) + p.heartbeatUnhealthy.Store(false) p.fetchLimiter = chanutil.NewDebouncedChan(fetchCtx, p.config.FetchCooldown, true) @@ -639,7 +641,7 @@ func (p *producer) jitteredFetchPollInterval() time.Duration { func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan producerFetchResult) { var limit int - if p.paused { + if p.paused || p.heartbeatUnhealthy.Load() { limit = 0 } else { limit = p.maxJobsToFetch() @@ -810,8 +812,10 @@ func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit { } func (p *producer) dispatchWork(workCtx context.Context, count int, fetchResultCh chan<- producerFetchResult) { - // When a queue is paused, innerFetchLoop dispatches with count zero so it can - // continue servicing state changes without attempting to lock jobs or emit metrics. + // When fetching is disabled because the queue is paused or producer + // heartbeats are unhealthy, innerFetchLoop dispatches with count zero so it + // can continue servicing state changes without attempting to lock jobs or + // emit metrics. if count <= 0 { fetchResultCh <- producerFetchResult{} return @@ -1064,23 +1068,64 @@ func (p *producer) reportProducerStatusOnce(ctx context.Context) { defer cancel() p.Logger.DebugContext(ctx, p.Name+": Reporting producer status", slog.Int64("id", p.id.Load()), slog.String("queue", p.config.Queue)) - err := p.pilot.ProducerKeepAlive(ctx, p.exec, &riverdriver.ProducerKeepAliveParams{ + params := &riverdriver.ProducerKeepAliveParams{ ID: p.id.Load(), QueueName: p.config.Queue, Schema: p.config.Schema, StaleUpdatedAtHorizon: p.Time.Now().Add(-p.config.StaleProducerRetentionPeriod), - }) + } + + const ( + producerReportMaxAttempts = 2 + producerReportRetryInterval = time.Second + ) + + var ( + attempts int + err error + ) + // Retry once within this report's timeout so that a transient database error + // doesn't stop all fetching until the next minute-long report interval. + for attempts < producerReportMaxAttempts { + attempts++ + err = p.pilot.ProducerKeepAlive(ctx, p.exec, params) + if err == nil || ctx.Err() != nil { + break + } + + if attempts < producerReportMaxAttempts { + p.Logger.WarnContext(ctx, p.Name+": Producer status update failed; retrying", + slog.Int("attempt", attempts), + slog.Int64("id", p.id.Load()), + slog.String("queue", p.config.Queue), + slog.Duration("retry_interval", producerReportRetryInterval), + slog.String("err", err.Error()), + ) + serviceutil.CancellableSleep(ctx, producerReportRetryInterval) + } + } if err != nil && errors.Is(context.Cause(ctx), startstop.ErrStop) { return } if err != nil { - p.Logger.ErrorContext(ctx, p.Name+": Producer status update, error updating in database", + p.heartbeatUnhealthy.Store(true) + p.Logger.ErrorContext(ctx, p.Name+": Producer status update failed; pausing job fetching", + slog.Int("attempts", attempts), slog.Int64("id", p.id.Load()), slog.String("queue", p.config.Queue), slog.String("err", err.Error()), ) return } + + if p.heartbeatUnhealthy.Swap(false) { + p.Logger.InfoContext(ctx, p.Name+": Producer status updates recovered; resuming job fetching", + slog.Int64("id", p.id.Load()), + slog.String("queue", p.config.Queue), + ) + p.fetchLimiter.Call() + } + p.testSignals.ReportedProducerStatus.Signal(struct{}{}) } diff --git a/producer_test.go b/producer_test.go index a23262f1..ee5a6fea 100644 --- a/producer_test.go +++ b/producer_test.go @@ -3,8 +3,10 @@ package river import ( "context" "encoding/json" + "errors" "fmt" "slices" + "sync/atomic" "testing" "time" @@ -25,6 +27,7 @@ import ( "github.com/riverqueue/river/rivershared/riversharedtest" "github.com/riverqueue/river/rivershared/startstoptest" "github.com/riverqueue/river/rivershared/testfactory" + "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivershared/util/ptrutil" "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/testutil" @@ -44,6 +47,34 @@ func (l *countingPluginLookup) ByKind(kind pluginlookup.PluginKind) []any { return l.PluginLookupInterface.ByKind(kind) } +type producerHeartbeatPilot struct { + riverpilot.StandardPilot + + failKeepAlive atomic.Bool + failNextKeepAlive atomic.Bool + jobGetAvailableCalls atomic.Int64 + jobGetAvailableSignal testsignal.TestSignal[struct{}] + keepAliveCalls atomic.Int64 +} + +func (p *producerHeartbeatPilot) Init(tb testutil.TestingTB) { + p.jobGetAvailableSignal.Init(tb) +} + +func (p *producerHeartbeatPilot) JobGetAvailable(ctx context.Context, exec riverdriver.Executor, state riverpilot.ProducerState, params *riverdriver.JobGetAvailableParams) ([]*rivertype.JobRow, error) { + p.jobGetAvailableCalls.Add(1) + p.jobGetAvailableSignal.Signal(struct{}{}) + return p.StandardPilot.JobGetAvailable(ctx, exec, state, params) +} + +func (p *producerHeartbeatPilot) ProducerKeepAlive(ctx context.Context, exec riverdriver.Executor, params *riverdriver.ProducerKeepAliveParams) error { + p.keepAliveCalls.Add(1) + if p.failKeepAlive.Load() || p.failNextKeepAlive.Swap(false) { + return errors.New("producer heartbeat failed") + } + return p.StandardPilot.ProducerKeepAlive(ctx, exec, params) +} + func Test_Producer_CanSafelyCompleteJobsWhileFetchingNewOnes(t *testing.T) { // We have encountered previous data races with the list of active jobs on // Producer because we need to know the count of active jobs in order to @@ -513,6 +544,63 @@ func testProducer(t *testing.T, makeProducer func(ctx context.Context, t *testin require.Equal(t, rivertype.JobStateCompleted, update.Job.State) }) + t.Run("HeartbeatFailureRetriesBeforeStoppingFetches", func(t *testing.T) { + t.Parallel() + + producer, _ := setup(t) + + pilot := &producerHeartbeatPilot{} + pilot.Init(t) + pilot.failNextKeepAlive.Store(true) + producer.pilot = pilot + + producer.reportProducerStatusOnce(ctx) + + require.Equal(t, int64(2), pilot.keepAliveCalls.Load()) + require.False(t, producer.heartbeatUnhealthy.Load()) + }) + + t.Run("HeartbeatFailureStopsFetchingUntilRecovery", func(t *testing.T) { + t.Parallel() + + producer, bundle := setup(t) + producer.config.FetchPollInterval = time.Hour + producer.config.ProducerReportInterval = time.Hour + + pilot := &producerHeartbeatPilot{} + pilot.Init(t) + producer.pilot = pilot + + AddWorker(bundle.workers, &noOpWorker{}) + startProducer(t, ctx, ctx, producer) + + // Wait for the initial empty fetch so there are no fetches in flight when + // the producer becomes unhealthy. + pilot.jobGetAvailableSignal.WaitOrTimeout() + + pilot.failKeepAlive.Store(true) + producer.reportProducerStatusOnce(ctx) + require.True(t, producer.heartbeatUnhealthy.Load()) + + jobGetAvailableCalls := pilot.jobGetAvailableCalls.Load() + mustInsert(ctx, t, producer, bundle, &noOpArgs{}) + producer.TriggerJobFetch() + + select { + case update := <-bundle.jobUpdates: + t.Fatalf("Unexpected job update while producer heartbeat was unhealthy: %+v", update) + case <-time.After(200 * time.Millisecond): + } + require.Equal(t, jobGetAvailableCalls, pilot.jobGetAvailableCalls.Load()) + + pilot.failKeepAlive.Store(false) + producer.reportProducerStatusOnce(ctx) + require.False(t, producer.heartbeatUnhealthy.Load()) + + update := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Equal(t, rivertype.JobStateCompleted, update.Job.State) + }) + t.Run("RegistersQueueStatus", func(t *testing.T) { t.Parallel() From d57228c08821e9a7e9d432cb076e6972ddf123ca Mon Sep 17 00:00:00 2001 From: Brandur Date: Tue, 21 Jul 2026 15:40:12 -0500 Subject: [PATCH 2/2] Pause `JobRescuer` when any producer is unhealthy This takes the functionality that tracks producer unhealthy heartbeats a step further by having `JobRescuer` pause in case any of its producers are unhealthy. The end goal is the same: don't accidentally rescue jobs where a producer has fallen out of sync because it can't write its heartbeats to the database. As before, this takes us further in the direction of resilience, but it's still not a perfect patch. The `JobRescuer` is only running on the client currently elected leader, so it requires that producers on that client be unhealthy for the pause to be initiated. If a producer on a client elsewhere is unhealthy, it won't be picked up. However, this should be okay in most situations because if one producer is unable to write heartbeats, there's a reasonable chance that all producers are having trouble doing it. --- client.go | 21 +++++-- client_test.go | 16 +++++ internal/maintenance/job_rescuer.go | 76 ++++++++++++++++++++---- internal/maintenance/job_rescuer_test.go | 55 +++++++++++++++++ 4 files changed, 153 insertions(+), 15 deletions(-) diff --git a/client.go b/client.go index 96d7b3eb..577216a3 100644 --- a/client.go +++ b/client.go @@ -960,10 +960,11 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client { jobRescuer := maintenance.NewRescuer(archetype, &maintenance.JobRescuerConfig{ - ClientRetryPolicy: config.RetryPolicy, - Pilot: client.pilot, - RescueAfter: config.RescueStuckJobsAfter, - Schema: config.Schema, + ClientRetryPolicy: config.RetryPolicy, + Pilot: client.pilot, + ProducersHealthyFunc: client.producersHealthy, + RescueAfter: config.RescueStuckJobsAfter, + Schema: config.Schema, WorkUnitFactoryFunc: func(kind string) workunit.WorkUnitFactory { if workerInfo, ok := config.Workers.workersMap[kind]; ok { return workerInfo.workUnitFactory @@ -2324,6 +2325,18 @@ func (c *Client[TTx]) producerRemove(ctx context.Context, queueName string) erro return nil } +func (c *Client[TTx]) producersHealthy() bool { + c.producersMu.RLock() + defer c.producersMu.RUnlock() + + for _, producer := range c.producersByQueueName { + if producer.heartbeatUnhealthy.Load() { + return false + } + } + return true +} + var nameRegex = regexp.MustCompile(`^(?:[a-z0-9])+(?:[_|\-]?[a-z0-9]+)*$`) func validateQueueName(queueName string) error { diff --git a/client_test.go b/client_test.go index 478ff6eb..d7898dfa 100644 --- a/client_test.go +++ b/client_test.go @@ -5859,6 +5859,22 @@ func Test_Client_Maintenance(t *testing.T) { requireJobHasState(jobNotYetStuck3.ID, jobNotYetStuck3.State) }) + t.Run("JobRescuerProducerHealth", func(t *testing.T) { + t.Parallel() + + client, _ := setup(t, newTestConfig(t, "")) + jobRescuer := maintenance.GetService[*maintenance.JobRescuer](client.queueMaintainer) + producer := client.producersByQueueName[QueueDefault] + + require.True(t, jobRescuer.Config.ProducersHealthyFunc()) + + producer.heartbeatUnhealthy.Store(true) + require.False(t, jobRescuer.Config.ProducersHealthyFunc()) + + producer.heartbeatUnhealthy.Store(false) + require.True(t, jobRescuer.Config.ProducersHealthyFunc()) + }) + t.Run("JobScheduler", func(t *testing.T) { t.Parallel() diff --git a/internal/maintenance/job_rescuer.go b/internal/maintenance/job_rescuer.go index 488c3371..bc7b5375 100644 --- a/internal/maintenance/job_rescuer.go +++ b/internal/maintenance/job_rescuer.go @@ -33,11 +33,13 @@ const ( // JobRescuerTestSignals are internal signals used exclusively in tests. type JobRescuerTestSignals struct { FetchedBatch testsignal.TestSignal[struct{}] // notifies when runOnce has fetched a batch of jobs + SkippedRun testsignal.TestSignal[struct{}] // notifies when runOnce was skipped because a producer was unhealthy UpdatedBatch testsignal.TestSignal[struct{}] // notifies when runOnce has updated rescued jobs from a batch } func (ts *JobRescuerTestSignals) Init(tb testutil.TestingTB) { ts.FetchedBatch.Init(tb) + ts.SkippedRun.Init(tb) ts.UpdatedBatch.Init(tb) } @@ -54,6 +56,11 @@ type JobRescuerConfig struct { // Pilot controls driver-level behavior that can be customized by plugins. Pilot riverpilot.Pilot + // ProducersHealthyFunc returns whether all producers in this client are + // healthy. Job rescue is paused when it returns false. This is an in-process + // safeguard; producers in other clients aren't visible through it. + ProducersHealthyFunc func() bool + // RescueAfter is the amount of time for a job to be active before it is // considered stuck and should be rescued. RescueAfter time.Duration @@ -99,6 +106,9 @@ type JobRescuer struct { exec riverdriver.Executor + // Tracks the logging transition into and out of a producer health pause. + pausedForUnhealthyProducers bool + // Circuit breaker that tracks consecutive timeout failures from the central // query. The query starts by using the full/default batch size, but after // this breaker trips (after N consecutive timeouts occur in a row), it @@ -114,16 +124,21 @@ func NewRescuer(archetype *baseservice.Archetype, config *JobRescuerConfig, exec if pilot == nil { pilot = &riverpilot.StandardPilot{} } + producersHealthyFunc := config.ProducersHealthyFunc + if producersHealthyFunc == nil { + producersHealthyFunc = func() bool { return true } + } return baseservice.Init(archetype, &JobRescuer{ Config: (&JobRescuerConfig{ - BatchSizes: batchSizes, - ClientRetryPolicy: config.ClientRetryPolicy, - Interval: cmp.Or(config.Interval, JobRescuerIntervalDefault), - Pilot: pilot, - RescueAfter: cmp.Or(config.RescueAfter, JobRescuerRescueAfterDefault), - Schema: config.Schema, - WorkUnitFactoryFunc: config.WorkUnitFactoryFunc, + BatchSizes: batchSizes, + ClientRetryPolicy: config.ClientRetryPolicy, + Interval: cmp.Or(config.Interval, JobRescuerIntervalDefault), + Pilot: pilot, + ProducersHealthyFunc: producersHealthyFunc, + RescueAfter: cmp.Or(config.RescueAfter, JobRescuerRescueAfterDefault), + Schema: config.Schema, + WorkUnitFactoryFunc: config.WorkUnitFactoryFunc, }).mustValidate(), exec: exec, reducedBatchSizeBreaker: riversharedmaintenance.ReducedBatchSizeBreaker(batchSizes), @@ -190,10 +205,31 @@ type metadataWithCancelAttemptedAt struct { CancelAttemptedAt time.Time `json:"cancel_attempted_at"` } +func (s *JobRescuer) producersHealthy(ctx context.Context) bool { + if s.Config.ProducersHealthyFunc() { + if s.pausedForUnhealthyProducers { + s.pausedForUnhealthyProducers = false + s.Logger.InfoContext(ctx, s.Name+": Producers healthy; resuming job rescue") + } + return true + } + + if !s.pausedForUnhealthyProducers { + s.pausedForUnhealthyProducers = true + s.Logger.WarnContext(ctx, s.Name+": Producer unhealthy; pausing job rescue") + } + s.TestSignals.SkippedRun.Signal(struct{}{}) + return false +} + func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error) { res := &rescuerRunOnceResult{} for { + if !s.producersHealthy(ctx) { + return res, nil + } + stuckJobs, err := s.getStuckJobs(ctx) if err != nil { if errors.Is(err, context.DeadlineExceeded) { @@ -205,9 +241,18 @@ func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error) s.reducedBatchSizeBreaker.ResetIfNotOpen() + // Producer health may have changed while the query was in flight. Check + // again before preparing or applying any job state transitions. + if !s.producersHealthy(ctx) { + return res, nil + } + s.TestSignals.FetchedBatch.Signal(struct{}{}) - now := time.Now().UTC() + var ( + batchRes rescuerRunOnceResult + now = time.Now().UTC() + ) rescueManyParams := riverdriver.JobRescueManyParams{ ID: make([]int64, 0, len(stuckJobs)), @@ -243,7 +288,7 @@ func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error) } if !metadata.CancelAttemptedAt.IsZero() { - res.NumJobsCancelled++ + batchRes.NumJobsCancelled++ addRescueParam(rivertype.JobStateCancelled, &now, job.ScheduledAt) // reused previous scheduled value continue } @@ -252,24 +297,33 @@ func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error) switch retryDecision { case jobRetryDecisionDiscard: - res.NumJobsDiscarded++ + batchRes.NumJobsDiscarded++ addRescueParam(rivertype.JobStateDiscarded, &now, job.ScheduledAt) // reused previous scheduled value case jobRetryDecisionIgnore: // job not timed out yet due to kind-specific timeout value; ignore case jobRetryDecisionRetry: - res.NumJobsRetried++ + batchRes.NumJobsRetried++ addRescueParam(rivertype.JobStateRetryable, nil, retryAt) } } if len(rescueManyParams.ID) > 0 { + // Recheck as close as possible to the state transition. Counts for this + // batch aren't added to the result unless its update is applied. + if !s.producersHealthy(ctx) { + return res, nil + } + _, err = s.Config.Pilot.JobRescueMany(ctx, s.exec, &rescueManyParams) if err != nil { return nil, fmt.Errorf("error rescuing stuck jobs: %w", err) } } + res.NumJobsCancelled += batchRes.NumJobsCancelled + res.NumJobsDiscarded += batchRes.NumJobsDiscarded + res.NumJobsRetried += batchRes.NumJobsRetried s.TestSignals.UpdatedBatch.Signal(struct{}{}) diff --git a/internal/maintenance/job_rescuer_test.go b/internal/maintenance/job_rescuer_test.go index a6fcdfe1..64fe3ec8 100644 --- a/internal/maintenance/job_rescuer_test.go +++ b/internal/maintenance/job_rescuer_test.go @@ -141,6 +141,61 @@ func TestJobRescuer(t *testing.T) { require.Equal(t, JobRescuerIntervalDefault, cleaner.Config.Interval) }) + t.Run("PausesWhileProducersUnhealthy", func(t *testing.T) { + t.Parallel() + + rescuer, bundle := setup(t) + var producersHealthy atomic.Bool + rescuer.Config.ProducersHealthyFunc = producersHealthy.Load + + job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{Kind: ptrutil.Ptr(rescuerJobKind), State: ptrutil.Ptr(rivertype.JobStateRunning), AttemptedAt: ptrutil.Ptr(bundle.rescueHorizon.Add(-time.Minute)), MaxAttempts: ptrutil.Ptr(5)}) + + res, err := rescuer.runOnce(ctx) + require.NoError(t, err) + require.Equal(t, &rescuerRunOnceResult{}, res) + rescuer.TestSignals.SkippedRun.WaitOrTimeout() + require.True(t, rescuer.pausedForUnhealthyProducers) + + jobAfter, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, jobAfter.State) + + producersHealthy.Store(true) + + res, err = rescuer.runOnce(ctx) + require.NoError(t, err) + require.Equal(t, int64(1), res.NumJobsRetried) + require.False(t, rescuer.pausedForUnhealthyProducers) + + jobAfter, err = bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRetryable, jobAfter.State) + }) + + t.Run("RechecksHealthBeforeRescue", func(t *testing.T) { + t.Parallel() + + rescuer, bundle := setup(t) + var healthChecks atomic.Int64 + rescuer.Config.ProducersHealthyFunc = func() bool { + // Stay healthy through the pre-fetch and post-fetch checks, then become + // unhealthy immediately before the job state update. + return healthChecks.Add(1) < 3 + } + + job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{Kind: ptrutil.Ptr(rescuerJobKind), State: ptrutil.Ptr(rivertype.JobStateRunning), AttemptedAt: ptrutil.Ptr(bundle.rescueHorizon.Add(-time.Minute)), MaxAttempts: ptrutil.Ptr(5)}) + + res, err := rescuer.runOnce(ctx) + require.NoError(t, err) + require.Equal(t, &rescuerRunOnceResult{}, res) + rescuer.TestSignals.SkippedRun.WaitOrTimeout() + rescuer.TestSignals.UpdatedBatch.RequireEmpty() + + jobAfter, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, jobAfter.State) + }) + t.Run("StartStopStress", func(t *testing.T) { t.Parallel()