fix memory leak of job cancellation contexts

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed a memory leak caused by not always cancelling the context used to enable jobs to be cancelled remotely. [PR #243](https://github.com/riverqueue/river/pull/243).
+
 ## [0.0.23] - 2024-02-29
 
 ### Added

job_executor.go

@@ -17,6 +17,10 @@ import (
 	"github.com/riverqueue/river/rivertype"
 )
 
+// Error used in CancelFunc in cases where the job was not cancelled for
+// purposes of resource cleanup. Should never be user visible.
+var errExecutorDefaultCancel = errors.New("context cancelled as executor finished")
+
 // UnknownJobKindError is returned when a Client fetches and attempts to
 // work a job that has not been registered on the Client's Workers bundle (using
 // AddWorker).
@@ -138,6 +142,9 @@ func (e *jobExecutor) Cancel() {
 }
 
 func (e *jobExecutor) Execute(ctx context.Context) {
+	// Ensure that the context is cancelled no matter what, or it will leak:
+	defer e.CancelFunc(errExecutorDefaultCancel)
+
 	e.start = e.TimeNowUTC()
 	e.stats = &jobstats.JobStatistics{
 		QueueWaitDuration: e.start.Sub(e.JobRow.ScheduledAt),

job_executor_test.go

@@ -174,7 +174,12 @@ func TestJobExecutor_Execute(t *testing.T) {
 			jobRow:            job,
 		}
 
+		// allocate this context just so we can set the CancelFunc:
+		_, cancel := context.WithCancelCause(ctx)
+		t.Cleanup(func() { cancel(nil) })
+
 		executor := baseservice.Init(archetype, &jobExecutor{
+			CancelFunc:             cancel,
 			ClientRetryPolicy:      &retryPolicyNoJitter{},
 			Completer:              bundle.completer,
 			ErrorHandler:           bundle.errorHandler,
@@ -618,6 +623,22 @@ func TestJobExecutor_Execute(t *testing.T) {
 		require.True(t, bundle.errorHandler.HandlePanicCalled)
 	})
 
+	t.Run("CancelFuncCleanedUpEvenWithoutCancel", func(t *testing.T) {
+		t.Parallel()
+
+		executor, bundle := setup(t)
+
+		executor.WorkUnit = newWorkUnitFactoryWithCustomRetry(func() error { return nil }, nil).MakeUnit(bundle.jobRow)
+
+		workCtx, cancelFunc := context.WithCancelCause(ctx)
+		executor.CancelFunc = cancelFunc
+
+		executor.Execute(workCtx)
+		executor.Completer.Wait()
+
+		require.ErrorIs(t, context.Cause(workCtx), errExecutorDefaultCancel)
+	})
+
 	runCancelTest := func(t *testing.T, returnErr error) *rivertype.JobRow { //nolint:thelper
 		executor, bundle := setup(t)
 
@@ -640,6 +661,7 @@ func TestJobExecutor_Execute(t *testing.T) {
 
 		workCtx, cancelFunc := context.WithCancelCause(ctx)
 		executor.CancelFunc = cancelFunc
+		t.Cleanup(func() { cancelFunc(nil) })
 
 		executor.Execute(workCtx)
 		executor.Completer.Wait()

producer.go

@@ -358,6 +358,7 @@ func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype.
 			workUnit = workInfo.workUnitFactory.MakeUnit(job)
 		}
 
+		// jobCancel will always be called by the executor to prevent leaks.
 		jobCtx, jobCancel := context.WithCancelCause(workCtx)
 
 		executor := baseservice.Init(&p.Archetype, &jobExecutor{