Track unhealthy heartbeat + pause job fetching when unhealthy

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 {

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

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{}{})
 

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