Basic stuck job detection

Here, try to make some inroads on a feature we've been talking about for
a while: detection of stuck jobs.

Unfortunately in Go it's quite easy to accidentally park a job by using
a `select` on a channel that won't return and forgetting a separate
branch for `<-ctx.Done()` so that it won't respect job timeouts either.

Here, add in some basic detection for that case. Eventually we'd like to
give users some options for what to do in case jobs become stuck, but
here we do only the simplest things for now: log when we detect a stuck
job and count the number of stuck jobs in a producer's stats loop.

In the future we may want to have some additional intelligence like
having producers move stuck jobs to a separate bucket up to a certain
limit before crashing (the next best option because it's not possible to
manually kill goroutines).

Author: brandur

internal/jobexecutor/job_executor.go

@@ -104,20 +104,23 @@ func (r *jobExecutorResult) ErrorStr() string {
 type JobExecutor struct {
 	baseservice.BaseService
 
-	CancelFunc               context.CancelCauseFunc
-	ClientJobTimeout         time.Duration
-	Completer                jobcompleter.JobCompleter
-	ClientRetryPolicy        ClientRetryPolicy
-	DefaultClientRetryPolicy ClientRetryPolicy
-	ErrorHandler             ErrorHandler
-	HookLookupByJob          *hooklookup.JobHookLookup
-	HookLookupGlobal         hooklookup.HookLookupInterface
-	InformProducerDoneFunc   func(jobRow *rivertype.JobRow)
-	JobRow                   *rivertype.JobRow
-	MiddlewareLookupGlobal   middlewarelookup.MiddlewareLookupInterface
-	SchedulerInterval        time.Duration
-	WorkerMiddleware         []rivertype.WorkerMiddleware
-	WorkUnit                 workunit.WorkUnit
+	CancelFunc                context.CancelCauseFunc
+	ClientJobTimeout          time.Duration
+	Completer                 jobcompleter.JobCompleter
+	ClientRetryPolicy         ClientRetryPolicy
+	DefaultClientRetryPolicy  ClientRetryPolicy
+	ErrorHandler              ErrorHandler
+	HookLookupByJob           *hooklookup.JobHookLookup
+	HookLookupGlobal          hooklookup.HookLookupInterface
+	InformProducerDoneFunc    func(jobRow *rivertype.JobRow)
+	InformProducerStuckFunc   func()
+	InformProducerUnstuckFunc func()
+	JobRow                    *rivertype.JobRow
+	MiddlewareLookupGlobal    middlewarelookup.MiddlewareLookupInterface
+	SchedulerInterval         time.Duration
+	StuckThresholdOverride    time.Duration
+	WorkerMiddleware          []rivertype.WorkerMiddleware
+	WorkUnit                  workunit.WorkUnit
 
 	// Meant to be used from within the job executor only.
 	start time.Time
@@ -171,6 +174,49 @@ func (e *JobExecutor) execute(ctx context.Context) (res *jobExecutorResult) {
 	metadataUpdates := make(map[string]any)
 	ctx = context.WithValue(ctx, ContextKeyMetadataUpdates, metadataUpdates)
 
+	// Watches for jobs that may have become stuck. i.e. They've run longer than
+	// their job timeout (plus a small margin) and don't appear to be responding
+	// to context cancellation (unfortunately, quite an easy error to make in
+	// Go).
+	//
+	// Currently we don't do anything if we notice a job is stuck. Knowing about
+	// stuck jobs is just used for informational purposes in the producer in
+	// generating periodic stats.
+	if e.ClientJobTimeout > 0 {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+
+		go func() {
+			const stuckThresholdDefault = 5 * time.Second
+
+			select {
+			case <-ctx.Done():
+				// cancellation or execution finished
+
+			case <-time.After(e.ClientJobTimeout + cmp.Or(e.StuckThresholdOverride, stuckThresholdDefault)):
+				e.InformProducerStuckFunc()
+
+				e.Logger.WarnContext(ctx, e.Name+": Job appears to be stuck",
+					slog.Int64("job_id", e.JobRow.ID),
+					slog.String("kind", e.JobRow.Kind),
+					slog.Duration("timeout", e.ClientJobTimeout),
+				)
+
+				// In case the executor ever becomes unstuck, inform the
+				// producer. However, if we got all the way here there's a good
+				// chance this will never happen (the worker is really stuck and
+				// will never return).
+				defer e.InformProducerUnstuckFunc()
+
+				defer e.Logger.InfoContext(ctx, e.Name+": Job became unstuck",
+					slog.Duration("duration", time.Since(e.start)),
+					slog.Int64("job_id", e.JobRow.ID),
+					slog.String("kind", e.JobRow.Kind),
+				)
+			}
+		}()
+	}
+
 	defer func() {
 		if recovery := recover(); recovery != nil {
 			e.Logger.ErrorContext(ctx, e.Name+": panic recovery; possible bug with Worker",

internal/jobexecutor/job_executor_test.go

@@ -184,18 +184,20 @@ func TestJobExecutor_Execute(t *testing.T) {
 		t.Cleanup(func() { cancel(nil) })
 
 		executor := baseservice.Init(archetype, &JobExecutor{
-			CancelFunc:               cancel,
-			ClientRetryPolicy:        &retrypolicytest.RetryPolicyNoJitter{},
-			Completer:                bundle.completer,
-			DefaultClientRetryPolicy: &retrypolicytest.RetryPolicyNoJitter{},
-			ErrorHandler:             bundle.errorHandler,
-			HookLookupByJob:          hooklookup.NewJobHookLookup(),
-			HookLookupGlobal:         hooklookup.NewHookLookup(nil),
-			InformProducerDoneFunc:   func(job *rivertype.JobRow) {},
-			JobRow:                   bundle.jobRow,
-			MiddlewareLookupGlobal:   middlewarelookup.NewMiddlewareLookup(nil),
-			SchedulerInterval:        riverinternaltest.SchedulerShortInterval,
-			WorkUnit:                 workUnitFactory.MakeUnit(bundle.jobRow),
+			CancelFunc:                cancel,
+			ClientRetryPolicy:         &retrypolicytest.RetryPolicyNoJitter{},
+			Completer:                 bundle.completer,
+			DefaultClientRetryPolicy:  &retrypolicytest.RetryPolicyNoJitter{},
+			ErrorHandler:              bundle.errorHandler,
+			HookLookupByJob:           hooklookup.NewJobHookLookup(),
+			HookLookupGlobal:          hooklookup.NewHookLookup(nil),
+			InformProducerDoneFunc:    func(job *rivertype.JobRow) {},
+			InformProducerStuckFunc:   func() {},
+			InformProducerUnstuckFunc: func() {},
+			JobRow:                    bundle.jobRow,
+			MiddlewareLookupGlobal:    middlewarelookup.NewMiddlewareLookup(nil),
+			SchedulerInterval:         riverinternaltest.SchedulerShortInterval,
+			WorkUnit:                  workUnitFactory.MakeUnit(bundle.jobRow),
 		})
 
 		return executor, bundle
@@ -696,6 +698,36 @@ func TestJobExecutor_Execute(t *testing.T) {
 		})
 	})
 
+	t.Run("StuckDetection", func(t *testing.T) {
+		t.Parallel()
+
+		executor, bundle := setup(t)
+
+		executor.ClientJobTimeout = 5 * time.Millisecond
+		executor.StuckThresholdOverride = 1 * time.Nanosecond // must be greater than 0 to take effect
+
+		var (
+			informProducerStuckReceived   = make(chan struct{})
+			informProducerUnstuckReceived = make(chan struct{})
+		)
+		executor.InformProducerStuckFunc = func() {
+			close(informProducerStuckReceived)
+		}
+		executor.InformProducerUnstuckFunc = func() {
+			close(informProducerUnstuckReceived)
+		}
+
+		executor.WorkUnit = newWorkUnitFactoryWithCustomRetry(func() error {
+			riversharedtest.WaitOrTimeout(t, informProducerStuckReceived)
+			return nil
+		}, nil).MakeUnit(bundle.jobRow)
+
+		executor.Execute(ctx)
+		_ = riversharedtest.WaitOrTimeout(t, bundle.updateCh)
+
+		riversharedtest.WaitOrTimeout(t, informProducerUnstuckReceived)
+	})
+
 	t.Run("Panic", func(t *testing.T) {
 		t.Parallel()
 

producer.go

@@ -209,6 +209,7 @@ type producer struct {
 	// An atomic count of the number of jobs actively being worked on. This is
 	// written to by the main goroutine, but read by the dispatcher.
 	numJobsActive atomic.Int32
+	numJobsStuck  atomic.Int32
 
 	numJobsRan atomic.Uint64
 	paused     bool
@@ -771,20 +772,26 @@ func (p *producer) heartbeatLogLoop(ctx context.Context, wg *sync.WaitGroup) {
 	ticker := time.NewTicker(5 * time.Second)
 	defer ticker.Stop()
 	type jobCount struct {
-		ran    uint64
 		active int
+		ran    uint64
+		stuck  int
 	}
 	var prevCount jobCount
 	for {
 		select {
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-			curCount := jobCount{ran: p.numJobsRan.Load(), active: int(p.numJobsActive.Load())}
+			curCount := jobCount{
+				active: int(p.numJobsActive.Load()),
+				ran:    p.numJobsRan.Load(),
+				stuck:  int(p.numJobsStuck.Load()),
+			}
 			if curCount != prevCount {
 				p.Logger.InfoContext(ctx, p.Name+": Producer job counts",
 					slog.Uint64("num_completed_jobs", curCount.ran),
 					slog.Int("num_jobs_running", curCount.active),
+					slog.Int("num_jobs_stuck", curCount.stuck),
 					slog.String("queue", p.config.Queue),
 				)
 			}
@@ -806,19 +813,21 @@ func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.
 		jobCtx, jobCancel := context.WithCancelCause(workCtx)
 
 		executor := baseservice.Init(&p.Archetype, &jobexecutor.JobExecutor{
-			CancelFunc:               jobCancel,
-			ClientJobTimeout:         p.jobTimeout,
-			ClientRetryPolicy:        p.retryPolicy,
-			Completer:                p.completer,
-			DefaultClientRetryPolicy: &DefaultClientRetryPolicy{},
-			ErrorHandler:             p.errorHandler,
-			HookLookupByJob:          p.config.HookLookupByJob,
-			HookLookupGlobal:         p.config.HookLookupGlobal,
-			MiddlewareLookupGlobal:   p.config.MiddlewareLookupGlobal,
-			InformProducerDoneFunc:   p.handleWorkerDone,
-			JobRow:                   job,
-			SchedulerInterval:        p.config.SchedulerInterval,
-			WorkUnit:                 workUnit,
+			CancelFunc:                jobCancel,
+			ClientJobTimeout:          p.jobTimeout,
+			ClientRetryPolicy:         p.retryPolicy,
+			Completer:                 p.completer,
+			DefaultClientRetryPolicy:  &DefaultClientRetryPolicy{},
+			ErrorHandler:              p.errorHandler,
+			HookLookupByJob:           p.config.HookLookupByJob,
+			HookLookupGlobal:          p.config.HookLookupGlobal,
+			MiddlewareLookupGlobal:    p.config.MiddlewareLookupGlobal,
+			InformProducerDoneFunc:    p.handleWorkerDone,
+			InformProducerStuckFunc:   func() { p.numJobsStuck.Add(1) },
+			InformProducerUnstuckFunc: func() { p.numJobsStuck.Add(-1) },
+			JobRow:                    job,
+			SchedulerInterval:         p.config.SchedulerInterval,
+			WorkUnit:                  workUnit,
 		})
 		p.addActiveJob(job.ID, executor)