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() 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()