Add Pilot.JobRescueMany for active job rescue

This one's presented as an alternative to #1256 with its own alternative
Pro implementation and which may provide some superior properties around
potential bloat for table bloat in `river_job`.

Author: brandur

client.go

@@ -942,6 +942,7 @@ 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,
 				WorkUnitFactoryFunc: func(kind string) workunit.WorkUnitFactory {

internal/maintenance/job_rescuer.go

@@ -14,6 +14,7 @@ import (
 	"github.com/riverqueue/river/riverdriver"
 	"github.com/riverqueue/river/rivershared/baseservice"
 	"github.com/riverqueue/river/rivershared/circuitbreaker"
+	"github.com/riverqueue/river/rivershared/riverpilot"
 	"github.com/riverqueue/river/rivershared/riversharedmaintenance"
 	"github.com/riverqueue/river/rivershared/startstop"
 	"github.com/riverqueue/river/rivershared/testsignal"
@@ -50,6 +51,9 @@ type JobRescuerConfig struct {
 	// Interval is the amount of time to wait between runs of the rescuer.
 	Interval time.Duration
 
+	// Pilot controls driver-level behavior that can be customized by plugins.
+	Pilot riverpilot.Pilot
+
 	// RescueAfter is the amount of time for a job to be active before it is
 	// considered stuck and should be rescued.
 	RescueAfter time.Duration
@@ -70,6 +74,9 @@ func (c *JobRescuerConfig) mustValidate() *JobRescuerConfig {
 	if c.Interval <= 0 {
 		panic("RescuerConfig.Interval must be above zero")
 	}
+	if c.Pilot == nil {
+		panic("RescuerConfig.Pilot must be set")
+	}
 	if c.RescueAfter <= 0 {
 		panic("RescuerConfig.JobDuration must be above zero")
 	}
@@ -103,12 +110,17 @@ type JobRescuer struct {
 
 func NewRescuer(archetype *baseservice.Archetype, config *JobRescuerConfig, exec riverdriver.Executor) *JobRescuer {
 	batchSizes := config.WithDefaults()
+	pilot := config.Pilot
+	if pilot == nil {
+		pilot = &riverpilot.StandardPilot{}
+	}
 
 	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,
@@ -253,7 +265,7 @@ func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error)
 		}
 
 		if len(rescueManyParams.ID) > 0 {
-			_, err = s.exec.JobRescueMany(ctx, &rescueManyParams)
+			_, err = s.Config.Pilot.JobRescueMany(ctx, s.exec, &rescueManyParams)
 			if err != nil {
 				return nil, fmt.Errorf("error rescuing stuck jobs: %w", err)
 			}

internal/maintenance/job_rescuer_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"math"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -14,6 +15,7 @@ import (
 	"github.com/riverqueue/river/riverdbtest"
 	"github.com/riverqueue/river/riverdriver"
 	"github.com/riverqueue/river/riverdriver/riverpgxv5"
+	"github.com/riverqueue/river/rivershared/riverpilot"
 	"github.com/riverqueue/river/rivershared/riversharedmaintenance"
 	"github.com/riverqueue/river/rivershared/riversharedtest"
 	"github.com/riverqueue/river/rivershared/startstoptest"
@@ -57,6 +59,19 @@ func (p *SimpleClientRetryPolicy) NextRetry(job *rivertype.JobRow) time.Time {
 	return job.AttemptedAt.Add(timeutil.SecondsAsDuration(retrySeconds))
 }
 
+type jobRescueManyPilotSpy struct {
+	riverpilot.StandardPilot
+
+	calls  atomic.Int64
+	params *riverdriver.JobRescueManyParams
+}
+
+func (p *jobRescueManyPilotSpy) JobRescueMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRescueManyParams) (*struct{}, error) {
+	p.calls.Add(1)
+	p.params = params
+	return p.StandardPilot.JobRescueMany(ctx, exec, params)
+}
+
 func TestJobRescuer(t *testing.T) {
 	t.Parallel()
 
@@ -304,6 +319,22 @@ func TestJobRescuer(t *testing.T) {
 		riversharedtest.WaitOrTimeout(t, stopped)
 	})
 
+	t.Run("UsesPilotJobRescueMany", func(t *testing.T) {
+		t.Parallel()
+
+		rescuer, bundle := setup(t)
+		pilot := &jobRescueManyPilotSpy{}
+		rescuer.Config.Pilot = pilot
+
+		job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{Kind: ptrutil.Ptr(rescuerJobKind), State: ptrutil.Ptr(rivertype.JobStateRunning), AttemptedAt: ptrutil.Ptr(bundle.rescueHorizon.Add(-1 * time.Hour)), MaxAttempts: ptrutil.Ptr(5)})
+
+		_, err := rescuer.runOnce(ctx)
+		require.NoError(t, err)
+
+		require.Equal(t, int64(1), pilot.calls.Load())
+		require.Equal(t, []int64{job.ID}, pilot.params.ID)
+	})
+
 	t.Run("CanRunMultipleTimes", func(t *testing.T) {
 		t.Parallel()
 

metadata_test.go

@@ -4,8 +4,9 @@ import (
 	"context"
 	"testing"
 
-	"github.com/riverqueue/river/internal/jobexecutor"
 	"github.com/stretchr/testify/require"
+
+	"github.com/riverqueue/river/internal/jobexecutor"
 )
 
 func TestMetadataSet(t *testing.T) {

rivershared/riverpilot/pilot.go

@@ -41,6 +41,8 @@ type Pilot interface {
 		params *riverdriver.JobInsertFastManyParams,
 	) ([]*riverdriver.JobInsertFastResult, error)
 
+	JobRescueMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRescueManyParams) (*struct{}, error)
+
 	JobRetry(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRetryParams) (*rivertype.JobRow, error)
 
 	JobSetStateIfRunningMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error)

rivershared/riverpilot/standard_pilot.go

@@ -39,6 +39,10 @@ func (p *StandardPilot) JobInsertMany(
 	return exec.JobInsertFastMany(ctx, params)
 }
 
+func (p *StandardPilot) JobRescueMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRescueManyParams) (*struct{}, error) {
+	return exec.JobRescueMany(ctx, params)
+}
+
 func (p *StandardPilot) JobRetry(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobRetryParams) (*rivertype.JobRow, error) {
 	return exec.JobRetry(ctx, params)
 }