Add pilot hooks for job deletion

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- Job deletion through `Client.JobDelete`, `Client.JobDeleteMany`, their `Tx` variants, and the job cleaner now goes through optional pilot hooks, so a pilot can take part in deletes as it already can in cancels and retries. Behavior with the standard pilot is unchanged. [PR #1500](https://github.com/riverqueue/river/pull/1500).
+
 ### Fixed
 
 - Fixed `AddWorkerSafely` leaving a worker's primary kind and earlier aliases registered when a later kind alias conflicts with an existing registration. Failed registrations now leave the worker registry unchanged. [PR #1440](https://github.com/riverqueue/river/pull/1440).

client.go

@@ -1003,6 +1003,7 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client
 				CancelledJobRetentionPeriod: config.CancelledJobRetentionPeriod,
 				CompletedJobRetentionPeriod: config.CompletedJobRetentionPeriod,
 				DiscardedJobRetentionPeriod: config.DiscardedJobRetentionPeriod,
+				Pilot:                       client.pilot,
 				QueuesExcluded:              client.pilot.JobCleanerQueuesExcluded(),
 				Schema:                      config.Schema,
 				Timeout:                     config.JobCleanerTimeout,
@@ -1601,10 +1602,7 @@ func (c *Client[TTx]) jobCancel(ctx context.Context, exec riverdriver.Executor,
 // deleted row if it was deleted. Jobs in the running state are not deleted,
 // instead returning rivertype.ErrJobRunning.
 func (c *Client[TTx]) JobDelete(ctx context.Context, id int64) (*rivertype.JobRow, error) {
-	return c.driver.GetExecutor().JobDelete(ctx, &riverdriver.JobDeleteParams{
-		ID:     id,
-		Schema: c.config.Schema,
-	})
+	return c.jobDelete(ctx, c.driver.GetExecutor(), id)
 }
 
 // JobDeleteTx deletes the job with the given ID from the database, returning the
@@ -1614,10 +1612,23 @@ func (c *Client[TTx]) JobDelete(ctx context.Context, id int64) (*rivertype.JobRo
 // until the transaction commits, and if the transaction rolls back, so too is
 // the deleted job.
 func (c *Client[TTx]) JobDeleteTx(ctx context.Context, tx TTx, id int64) (*rivertype.JobRow, error) {
-	return c.driver.UnwrapExecutor(tx).JobDelete(ctx, &riverdriver.JobDeleteParams{
+	return c.jobDelete(ctx, c.driver.UnwrapExecutor(tx), id)
+}
+
+func (c *Client[TTx]) jobDelete(ctx context.Context, exec riverdriver.Executor, id int64) (*rivertype.JobRow, error) {
+	params := &riverdriver.JobDeleteParams{
 		ID:     id,
 		Schema: c.config.Schema,
-	})
+	}
+
+	if pilot, ok := c.pilot.(riverpilot.PilotJobDeleter); ok {
+		return pilot.JobDelete(ctx, exec, params)
+	}
+
+	// Temporary compatibility fallback for Pilot implementations from before
+	// PilotJobDeleter. Once Pilot embeds PilotJobDeleter, replace the assertion
+	// above and this fallback with a direct call to c.pilot.JobDelete.
+	return exec.JobDelete(ctx, params)
 }
 
 // JobGet fetches a single job by its ID. Returns the up-to-date JobRow for the
@@ -2484,7 +2495,18 @@ func (c *Client[TTx]) jobDeleteMany(ctx context.Context, exec riverdriver.Execut
 		return nil, err
 	}
 
-	jobs, err := exec.JobDeleteMany(ctx, (*riverdriver.JobDeleteManyParams)(listParams))
+	driverParams := (*riverdriver.JobDeleteManyParams)(listParams)
+
+	var jobs []*rivertype.JobRow
+	if pilot, ok := c.pilot.(riverpilot.PilotJobDeleter); ok {
+		jobs, err = pilot.JobDeleteMany(ctx, exec, driverParams)
+	} else {
+		// Temporary compatibility fallback for Pilot implementations from before
+		// PilotJobDeleter. Once Pilot embeds PilotJobDeleter, replace the
+		// assertion above and this fallback with a direct call to
+		// c.pilot.JobDeleteMany.
+		jobs, err = exec.JobDeleteMany(ctx, driverParams)
+	}
 	if err != nil {
 		return nil, err
 	}

client_pilot_test.go

@@ -26,6 +26,8 @@ type pilotSpy struct {
 
 	jobCancelCalls                atomic.Int64
 	jobCleanerQueuesExcludedCalls atomic.Int64
+	jobDeleteCalls                atomic.Int64
+	jobDeleteManyCalls            atomic.Int64
 	jobInsertManyCalls            atomic.Int64
 	jobRetryCalls                 atomic.Int64
 	pilotInitCalls                atomic.Int64
@@ -69,6 +71,16 @@ func (p *pilotSpy) JobCleanerQueuesExcluded() []string {
 	return p.StandardPilot.JobCleanerQueuesExcluded()
 }
 
+func (p *pilotSpy) JobDelete(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobDeleteParams) (*rivertype.JobRow, error) {
+	p.jobDeleteCalls.Add(1)
+	return p.StandardPilot.JobDelete(ctx, exec, params)
+}
+
+func (p *pilotSpy) JobDeleteMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobDeleteManyParams) ([]*rivertype.JobRow, error) {
+	p.jobDeleteManyCalls.Add(1)
+	return p.StandardPilot.JobDeleteMany(ctx, exec, params)
+}
+
 func (p *pilotSpy) JobGetAvailable(ctx context.Context, exec riverdriver.Executor, state riverpilot.ProducerState, params *riverdriver.JobGetAvailableParams) (*riverdriver.JobGetAvailableResult, error) {
 	p.testSignals.JobGetAvailable.Signal(struct{}{})
 	return p.StandardPilot.JobGetAvailable(ctx, exec, state, params)
@@ -130,6 +142,12 @@ func (p *pilotSpy) QueueMetadataChanged(ctx context.Context, exec riverdriver.Ex
 	return p.StandardPilot.QueueMetadataChanged(ctx, exec, params)
 }
 
+// pilotWithoutJobDeleter hides PilotJobDeleter to simulate a pilot built
+// against a River version from before it existed.
+type pilotWithoutJobDeleter struct {
+	riverpilot.Pilot
+}
+
 func Test_Client_PilotUsage(t *testing.T) {
 	t.Parallel()
 
@@ -224,6 +242,100 @@ func Test_Client_PilotUsage(t *testing.T) {
 		})
 	})
 
+	t.Run("JobDeleteManyTxUsesPilot", func(t *testing.T) {
+		t.Parallel()
+
+		client, pilot := setup(t, nil)
+
+		insertRes, err := client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+
+		withClientTx(t, client, func(tx pgx.Tx) {
+			deleteRes, err := client.JobDeleteManyTx(ctx, tx, NewJobDeleteManyParams().IDs(insertRes.Job.ID))
+			require.NoError(t, err)
+			require.Len(t, deleteRes.Jobs, 1)
+			require.Equal(t, int64(1), pilot.jobDeleteManyCalls.Load())
+		})
+	})
+
+	t.Run("JobDeleteManyUsesPilot", func(t *testing.T) {
+		t.Parallel()
+
+		client, pilot := setup(t, nil)
+
+		insertRes, err := client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+
+		deleteRes, err := client.JobDeleteMany(ctx, NewJobDeleteManyParams().IDs(insertRes.Job.ID))
+		require.NoError(t, err)
+		require.Len(t, deleteRes.Jobs, 1)
+		require.Equal(t, int64(1), pilot.jobDeleteManyCalls.Load())
+	})
+
+	t.Run("JobDeleteTxUsesPilot", func(t *testing.T) {
+		t.Parallel()
+
+		client, pilot := setup(t, nil)
+
+		insertRes, err := client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+
+		withClientTx(t, client, func(tx pgx.Tx) {
+			_, err = client.JobDeleteTx(ctx, tx, insertRes.Job.ID)
+			require.NoError(t, err)
+			require.Equal(t, int64(1), pilot.jobDeleteCalls.Load())
+		})
+	})
+
+	t.Run("JobDeleteUsesExecutorFallbackForPilotWithoutJobDeleter", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			dbPool       = riversharedtest.DBPool(ctx, t)
+			driver       = riverpgxv5.New(dbPool)
+			schema       = riverdbtest.TestSchema(ctx, t, driver, nil)
+			pilot        = &pilotWithoutJobDeleter{Pilot: &pilotSpy{}}
+			pluginDriver = newDriverWithPlugin(t, dbPool)
+		)
+		pluginDriver.pilot = pilot
+
+		_, implementsJobDeleter := riverpilot.Pilot(pilot).(riverpilot.PilotJobDeleter)
+		require.False(t, implementsJobDeleter)
+
+		client, err := NewClient(pluginDriver, newTestConfig(t, schema))
+		require.NoError(t, err)
+
+		insertRes1, err := client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+		insertRes2, err := client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+
+		_, err = client.JobDelete(ctx, insertRes1.Job.ID)
+		require.NoError(t, err)
+
+		deleteRes, err := client.JobDeleteMany(ctx, NewJobDeleteManyParams().IDs(insertRes2.Job.ID))
+		require.NoError(t, err)
+		require.Len(t, deleteRes.Jobs, 1)
+
+		_, err = client.JobGet(ctx, insertRes1.Job.ID)
+		require.ErrorIs(t, err, ErrNotFound)
+		_, err = client.JobGet(ctx, insertRes2.Job.ID)
+		require.ErrorIs(t, err, ErrNotFound)
+	})
+
+	t.Run("JobDeleteUsesPilot", func(t *testing.T) {
+		t.Parallel()
+
+		client, pilot := setup(t, nil)
+
+		insertRes, err := client.Insert(ctx, noOpArgs{}, nil)
+		require.NoError(t, err)
+
+		_, err = client.JobDelete(ctx, insertRes.Job.ID)
+		require.NoError(t, err)
+		require.Equal(t, int64(1), pilot.jobDeleteCalls.Load())
+	})
+
 	t.Run("JobInsertManyUsesPilot", func(t *testing.T) {
 		t.Parallel()
 

internal/maintenance/job_cleaner.go

@@ -11,6 +11,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"
@@ -54,6 +55,9 @@ type JobCleanerConfig struct {
 	// Interval is the amount of time to wait between runs of the cleaner.
 	Interval time.Duration
 
+	// Pilot controls driver-level behavior that can be customized by plugins.
+	Pilot riverpilot.Pilot
+
 	// QueuesExcluded are queues that'll be excluded from cleaning.
 	QueuesExcluded []string
 
@@ -80,6 +84,9 @@ func (c *JobCleanerConfig) mustValidate() *JobCleanerConfig {
 	if c.Interval <= 0 {
 		panic("JobCleanerConfig.Interval must be above zero")
 	}
+	if c.Pilot == nil {
+		panic("JobCleanerConfig.Pilot must be set")
+	}
 	if c.Timeout <= 0 {
 		panic("JobCleanerConfig.Timeout must be above zero")
 	}
@@ -110,6 +117,10 @@ type JobCleaner struct {
 
 func NewJobCleaner(archetype *baseservice.Archetype, config *JobCleanerConfig, exec riverdriver.Executor) *JobCleaner {
 	batchSizes := config.WithDefaults()
+	pilot := config.Pilot
+	if pilot == nil {
+		pilot = &riverpilot.StandardPilot{}
+	}
 
 	return baseservice.Init(archetype, &JobCleaner{
 		Config: (&JobCleanerConfig{
@@ -119,6 +130,7 @@ func NewJobCleaner(archetype *baseservice.Archetype, config *JobCleanerConfig, e
 			DiscardedJobRetentionPeriod: cmp.Or(config.DiscardedJobRetentionPeriod, riversharedmaintenance.DiscardedJobRetentionPeriodDefault),
 			QueuesExcluded:              config.QueuesExcluded,
 			Interval:                    cmp.Or(config.Interval, riversharedmaintenance.JobCleanerIntervalDefault),
+			Pilot:                       pilot,
 			Schema:                      config.Schema,
 			Timeout:                     cmp.Or(config.Timeout, riversharedmaintenance.JobCleanerTimeoutDefault),
 		}).mustValidate(),
@@ -176,6 +188,17 @@ func (s *JobCleaner) batchSize() int {
 	return s.Config.Default
 }
 
+func (s *JobCleaner) deleteBefore(ctx context.Context, params *riverdriver.JobDeleteBeforeParams) (int, error) {
+	if pilot, ok := s.Config.Pilot.(riverpilot.PilotJobDeleter); ok {
+		return pilot.JobDeleteBefore(ctx, s.exec, params)
+	}
+
+	// Temporary compatibility fallback for Pilot implementations from before
+	// PilotJobDeleter. Once Pilot embeds PilotJobDeleter, replace the assertion
+	// above and this fallback with a direct call to s.Config.Pilot.JobDeleteBefore.
+	return s.exec.JobDeleteBefore(ctx, params)
+}
+
 type jobCleanerRunOnceResult struct {
 	NumJobsDeleted int
 }
@@ -193,7 +216,7 @@ func (s *JobCleaner) runOnce(ctx context.Context) (*jobCleanerRunOnceResult, err
 				return 0, nil
 			}
 
-			numDeleted, err := s.exec.JobDeleteBefore(ctx, &riverdriver.JobDeleteBeforeParams{
+			numDeleted, err := s.deleteBefore(ctx, &riverdriver.JobDeleteBeforeParams{
 				CancelledDoDelete:           s.Config.CancelledJobRetentionPeriod != -1,
 				CancelledFinalizedAtHorizon: time.Now().Add(-s.Config.CancelledJobRetentionPeriod),
 				CompletedDoDelete:           s.Config.CompletedJobRetentionPeriod != -1,

internal/maintenance/job_cleaner_test.go

@@ -2,6 +2,7 @@ package maintenance
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -10,13 +11,31 @@ 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"
 	"github.com/riverqueue/river/rivershared/testfactory"
 	"github.com/riverqueue/river/rivertype"
 )
 
+type jobCleanerPilotSpy struct {
+	riverpilot.StandardPilot
+
+	jobDeleteBeforeCalls  atomic.Int64
+	jobDeleteBeforeParams *riverdriver.JobDeleteBeforeParams
+}
+
+func (p *jobCleanerPilotSpy) JobDeleteBefore(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobDeleteBeforeParams) (int, error) {
+	p.jobDeleteBeforeCalls.Add(1)
+	p.jobDeleteBeforeParams = params
+	return p.StandardPilot.JobDeleteBefore(ctx, exec, params)
+}
+
+type jobCleanerPilotWithoutJobDeleter struct {
+	riverpilot.Pilot
+}
+
 func TestJobCleaner(t *testing.T) {
 	t.Parallel()
 
@@ -436,4 +455,44 @@ func TestJobCleaner(t *testing.T) {
 			}
 		}
 	})
+
+	t.Run("UsesPilot", func(t *testing.T) {
+		t.Parallel()
+
+		cleaner, bundle := setup(t)
+		pilot := &jobCleanerPilotSpy{}
+		cleaner.Config.Pilot = pilot
+
+		job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateCancelled), FinalizedAt: new(bundle.cancelledDeleteHorizon.Add(-1 * time.Hour))})
+
+		res, err := cleaner.runOnce(ctx)
+		require.NoError(t, err)
+		require.Equal(t, 1, res.NumJobsDeleted)
+
+		require.Equal(t, int64(1), pilot.jobDeleteBeforeCalls.Load())
+		require.Equal(t, cleaner.Config.Default, pilot.jobDeleteBeforeParams.Max)
+		require.True(t, pilot.jobDeleteBeforeParams.CancelledDoDelete)
+
+		_, err = bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID, Schema: cleaner.Config.Schema})
+		require.ErrorIs(t, err, rivertype.ErrNotFound)
+	})
+
+	t.Run("UsesExecutorFallbackForPilotWithoutJobDeleter", func(t *testing.T) {
+		t.Parallel()
+
+		cleaner, bundle := setup(t)
+		cleaner.Config.Pilot = &jobCleanerPilotWithoutJobDeleter{Pilot: cleaner.Config.Pilot}
+
+		_, implementsJobDeleter := cleaner.Config.Pilot.(riverpilot.PilotJobDeleter)
+		require.False(t, implementsJobDeleter)
+
+		job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{State: new(rivertype.JobStateCancelled), FinalizedAt: new(bundle.cancelledDeleteHorizon.Add(-1 * time.Hour))})
+
+		res, err := cleaner.runOnce(ctx)
+		require.NoError(t, err)
+		require.Equal(t, 1, res.NumJobsDeleted)
+
+		_, err = bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID, Schema: cleaner.Config.Schema})
+		require.ErrorIs(t, err, rivertype.ErrNotFound)
+	})
 }