From da2bc55b6bc59a211cca9ac16d112a14948175b7 Mon Sep 17 00:00:00 2001 From: Brandur Date: Sat, 13 Jun 2026 08:31:02 -0500 Subject: [PATCH] Add `Config.StopAbandonTimeout` to abandon jobs still running after stop timeouts Here, add a new `Config.StopAbandonTimeout` on top of the existing `SoftStopTimeout` whose job it is to recover badly behaving job as much as possible before coming to a full stop. Currently, if a client is stopping and is running jobs that don't respond to context cancellation, those jobs end up getting left in a `running` state, which means that they won't be recoverable again until they're rescued an hour later. `StopAbandonTimeout` engages after soft stop, and has each producer initiate abandonment of its running jobs, which means setting them to an error state (or discarded in case a job is out of retries). Because they're errored, they'll get to run immediately the next time a client starts up. Ideally, users don't need to depend on this functionality since the "correct" behavior would be to make sure that all jobs are able to respond to context cancellation, so we make this new feature optional. --- CHANGELOG.md | 1 + client.go | 136 +++++++--- client_test.go | 176 +++++++++++- ..._graceful_shutdown_stop_and_cancel_test.go | 8 +- example_graceful_shutdown_test.go | 6 +- internal/jobexecutor/job_executor.go | 6 + producer.go | 253 ++++++++++++++++-- producer_test.go | 253 ++++++++++++++++++ 8 files changed, 773 insertions(+), 66 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 18d32635..7009e69b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Added `Config.StopAbandonTimeout` to bound how long a client waits after job contexts are cancelled during shutdown. When the timeout elapses, jobs that have not begun finalizing are abandoned, recorded as failed, and set as errored (or discarded if retries are exhausted). Worker goroutines cannot be forcibly terminated and continue running until the process exits. [PR #1289](https://github.com/riverqueue/river/pull/1289). - Added `EventKindJobInterrupted`, emitted when a running job is interrupted because its client is shutting down, the job was cancelled, and has been made immediately available to be worked again. [PR #1290](https://github.com/riverqueue/river/pull/1290). ### Changed diff --git a/client.go b/client.go index bf07bfc0..6e8ea820 100644 --- a/client.go +++ b/client.go @@ -378,11 +378,9 @@ type Config struct { Schema string // SoftStopTimeout is the maximum amount of time that the client will wait - // for running jobs to finish during a stop before their contexts are - // cancelled. After the timeout elapses, the client escalates to a hard stop - // by cancelling the context of all running jobs. This applies regardless of - // how stop is initiated — whether by calling Stop, StopAndCancel, or by - // cancelling the context passed to Start. + // for running jobs to finish during a graceful stop before entering soft + // stop by cancelling job contexts. This applies when stop is initiated by + // calling Stop or by cancelling the context passed to Start. // // In combination with signal.NotifyContext on the context passed to Start, // this can simplify graceful stop to: @@ -393,16 +391,33 @@ type Config struct { // if err := client.Start(ctx); err != nil { ... } // <-client.Stopped() // - // The signal cancels the Start context, which initiates a soft stop. If + // The signal cancels the Start context, which initiates a graceful stop. If // running jobs haven't finished after SoftStopTimeout, their contexts are - // automatically cancelled to trigger a hard stop. + // cancelled. // - // StopAndCancel bypasses the timeout entirely and cancels job contexts - // immediately. + // StopAndCancel cancels job contexts immediately instead of waiting for + // SoftStopTimeout. // // Defaults to no timeout (wait indefinitely for jobs to finish). SoftStopTimeout time.Duration + // StopAbandonTimeout is the maximum amount of time that the client will wait + // after job contexts are cancelled during a shutdown "soft stop" before it + // abandons jobs still running (i.e. those which did not respond to context + // cancellation). Abandoned jobs are recorded as failed and made immediately + // available for retry, or discarded if they have no attempts remaining. + // + // Worker goroutines cannot be forcibly terminated, and may continue running + // after their jobs are abandoned. Job implementations should therefore be + // prepared for duplicate execution. + // + // The timer starts only after a soft stop has begun by cancelling job + // contexts, like after SoftStopTimeout elapses, StopAndCancel is called, or + // the Start context is cancelled without SoftStopTimeout configured. + // + // Defaults to no timeout (job abandonment disabled). + StopAbandonTimeout time.Duration + // SkipJobKindValidation causes the job kind format validation check to be // skipped. This is available as an interim stopgap for users that have // invalid job kind names, but would rather disable the check rather than @@ -536,6 +551,7 @@ func (c *Config) WithDefaults() *Config { RetryPolicy: retryPolicy, Schema: c.Schema, SoftStopTimeout: c.SoftStopTimeout, + StopAbandonTimeout: c.StopAbandonTimeout, SkipJobKindValidation: c.SkipJobKindValidation, SkipUnknownJobCheck: c.SkipUnknownJobCheck, Test: c.Test, @@ -566,6 +582,9 @@ func (c *Config) validate() error { if c.FetchPollInterval < c.FetchCooldown { return fmt.Errorf("FetchPollInterval cannot be shorter than FetchCooldown (%s)", c.FetchCooldown) } + if c.StopAbandonTimeout < 0 { + return errors.New("StopAbandonTimeout cannot be less than zero") + } if len(c.ID) > 100 { return errors.New("ID cannot be longer than 100 characters") } @@ -601,6 +620,9 @@ func (c *Config) validate() error { if c.Schema != "" && !postgresSchemaNameRE.MatchString(c.Schema) { return errors.New("Schema name can only contain letters, numbers, and underscores, and must start with a letter or underscore") } + if c.SoftStopTimeout < 0 { + return errors.New("SoftStopTimeout cannot be less than zero") + } for queue, queueConfig := range c.Queues { if err := queueConfig.validate(queue, c.FetchCooldown, c.FetchPollInterval); err != nil { @@ -1069,10 +1091,12 @@ func NewClient[TTx any](driver riverdriver.Driver[TTx], config *Config) (*Client // A graceful shutdown stops fetching new jobs but allows any previously fetched // jobs to complete. This can be initiated with the Stop method. // -// A more abrupt shutdown can be achieved by either cancelling the provided -// context or by calling StopAndCancel. This will not only stop fetching new -// jobs, but will also cancel the context for any currently-running jobs. If -// using StopAndCancel, there's no need to also call Stop. +// A soft stop cancels job contexts after fetching has stopped. It can be +// initiated by calling StopAndCancel, by cancelling the provided context when +// SoftStopTimeout is not configured, or by waiting for SoftStopTimeout to elapse +// during graceful stop. If StopAbandonTimeout is configured, jobs still running +// after that timeout will be abandoned and recorded as failed. If using +// StopAndCancel, there's no need to also call Stop. func (c *Client[TTx]) Start(ctx context.Context) error { fetchCtx, shouldStart, started, stopped := c.baseStartStop.StartInit(ctx) if !shouldStart { @@ -1086,9 +1110,13 @@ func (c *Client[TTx]) Start(ctx context.Context) error { // sure to take a channel reference before finishing stopped. c.stopped = c.baseStartStop.StoppedUnsafe() - producersAsServices := func() []startstop.Service { + producers := func() []*producer { + return maputil.Values(c.producersByQueueName) + } + + producersAsServices := func(producers []*producer) []startstop.Service { return sliceutil.Map( - maputil.Values(c.producersByQueueName), + producers, func(p *producer) startstop.Service { return p }, ) } @@ -1142,8 +1170,8 @@ func (c *Client[TTx]) Start(ctx context.Context) error { // We use separate contexts for fetching and working to allow for a // graceful stop. When SoftStopTimeout is configured, the work context // is detached from the start context so that cancelling the start - // context initiates a soft stop (with timeout escalation) rather than - // an immediate hard stop. When SoftStopTimeout is not configured, the + // context initiates a graceful stop (with timeout escalation) rather + // than an immediate soft stop. When SoftStopTimeout is not configured, the // work context inherits from the start context to preserve the // existing behavior where cancelling the start context is equivalent // to StopAndCancel. @@ -1166,7 +1194,7 @@ func (c *Client[TTx]) Start(ctx context.Context) error { for _, producer := range c.producersByQueueName { if err := producer.StartWorkContext(fetchCtx, workCtx); err != nil { workCancel(err) - startstop.StopAllParallel(producersAsServices()...) + startstop.StopAllParallel(producersAsServices(producers())...) stopServicesOnError() return err } @@ -1188,7 +1216,7 @@ func (c *Client[TTx]) Start(ctx context.Context) error { // Generate producer services while c.queues.startStopMu.Lock() is still // held. This is used for WaitAllStarted below, but don't use it elsewhere // because new producers may have been added while the client is running. - producerServices := producersAsServices() + producerServices := producersAsServices(producers()) go func() { // Wait for all subservices to start up before signaling our own start. @@ -1215,22 +1243,57 @@ func (c *Client[TTx]) Start(ctx context.Context) error { c.queues.startStopMu.Lock() defer c.queues.startStopMu.Unlock() + producerList := producers() + + abandonTimerCtx, abandonTimerCancel := context.WithCancel(context.WithoutCancel(ctx)) + defer abandonTimerCancel() + + startAbandonTimer := sync.OnceFunc(func() { + if c.config.StopAbandonTimeout <= 0 { + return + } + + go func() { + timer := time.NewTimer(c.config.StopAbandonTimeout) + defer timer.Stop() + + select { + case <-timer.C: + c.baseService.Logger.WarnContext(ctx, c.baseService.Name+": StopAbandonTimeout elapsed; abandoning remaining jobs", slog.Duration("stop_abandon_timeout", c.config.StopAbandonTimeout)) + for _, producer := range producerList { + producer.abandon() + } + case <-abandonTimerCtx.Done(): + } + }() + }) + + workCtx := c.queues.workCtx + go func() { + select { + case <-workCtx.Done(): + startAbandonTimer() + case <-abandonTimerCtx.Done(): + } + }() + // If SoftStopTimeout is configured, start a timer that will cancel - // the work context (escalating to a hard stop) if producers don't - // finish in time. StopAndCancel also calls workCancel, in which case - // this timer is a harmless no-op because the context is already done. + // the work context if producers don't finish in time. Once the work + // context is cancelled, the optional abandonment timer starts. if c.config.SoftStopTimeout > 0 { softStopTimer := time.AfterFunc(c.config.SoftStopTimeout, func() { c.baseService.Logger.WarnContext(ctx, c.baseService.Name+": Soft stop timeout; cancelling remaining job contexts", slog.Duration("soft_stop_timeout", c.config.SoftStopTimeout)) c.workCancel(rivercommon.ErrStop) + startAbandonTimer() }) defer softStopTimer.Stop() } // On stop, have the producers stop fetching first of all. c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": Stopping producers") - startstop.StopAllParallel(producersAsServices()...) + startstop.StopAllParallel(producersAsServices(producerList)...) c.baseService.Logger.DebugContext(ctx, c.baseService.Name+": All producers stopped") + abandonTimerCancel() c.workCancel(rivercommon.ErrStop) @@ -1259,12 +1322,17 @@ func (c *Client[TTx]) Start(ctx context.Context) error { // complete before exiting. If the provided context is done before shutdown has // completed, Stop will return immediately with the context's error. // -// If SoftStopTimeout is configured, running job contexts will be automatically -// cancelled after the timeout elapses, escalating to a hard stop. This also -// applies when stop is initiated by cancelling the context passed to Start. +// If SoftStopTimeout is configured, jobs still running after the timeout +// elapses have their contexts cancelled. +// +// If StopAbandonTimeout is configured, jobs still running after SoftStopTimeout +// and StopAbandonTimeout have elapsed (i.e. waited for jobs to stop gracefully +// before cancelling, then waited again for them to stop on cancel) are abandoned +// and recorded as failed so they can be retried immediately. This also applies +// when stop is initiated by cancelling the context passed to Start. // -// There's no need to call this method if a hard stop has already been initiated -// by cancelling the context passed to Start or by calling StopAndCancel. +// There's no need to call this method if shutdown has already been initiated by +// cancelling the context passed to Start or by calling StopAndCancel. func (c *Client[TTx]) Stop(ctx context.Context) error { shouldStop, stopped, finalizeStop := c.baseStartStop.StopInit() if !shouldStop { @@ -1283,10 +1351,11 @@ func (c *Client[TTx]) Stop(ctx context.Context) error { // StopAndCancel shuts down the client and cancels all work in progress. It is a // more aggressive stop than Stop because the contexts for any in-progress jobs -// are cancelled. However, it still waits for jobs to complete before returning, -// even though their contexts are cancelled. If the provided context is done -// before shutdown has completed, StopAndCancel will return immediately with the -// context's error. +// are cancelled immediately. If StopAbandonTimeout is configured, jobs that +// still remain running after the timeout are abandoned; otherwise, StopAndCancel +// waits for jobs to complete even though their contexts are cancelled. If the +// provided context is done before shutdown has completed, StopAndCancel will +// return immediately with the context's error. // // This can also be initiated by cancelling the context passed to Start. There is // no need to call this method if the context passed to Start is cancelled @@ -1298,7 +1367,7 @@ func (c *Client[TTx]) Stop(ctx context.Context) error { // graceful stop semantics without requiring manual orchestration of Stop and // StopAndCancel. func (c *Client[TTx]) StopAndCancel(ctx context.Context) error { - c.baseService.Logger.InfoContext(ctx, c.baseService.Name+": Hard stop started; cancelling all work") + c.baseService.Logger.InfoContext(ctx, c.baseService.Name+": Soft stop started; cancelling all work") c.workCancel(rivercommon.ErrStop) shouldStop, stopped, finalizeStop := c.baseStartStop.StopInit() @@ -2298,6 +2367,7 @@ func (c *Client[TTx]) producerAdd(queueName string, queueConfig QueueConfig) (*p JobStuckCount: &c.stuckJobCount, JobStuckThreshold: c.config.JobStuckThreshold, JobTimeout: c.config.JobTimeout, + JobUpdateCallback: c.subscriptionManager.distributeJobUpdates, MaxWorkers: queueConfig.MaxWorkers, Notifier: c.notifier, Queue: queueName, diff --git a/client_test.go b/client_test.go index cda5eb6b..a6c82a50 100644 --- a/client_test.go +++ b/client_test.go @@ -2742,6 +2742,138 @@ func Test_Client_StopAndCancel(t *testing.T) { }) } +func Test_Client_StopAbandonTimeout(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type JobArgs struct { + testutil.JobArgsReflectKind[JobArgs] + } + + setup := func(t *testing.T, configFunc func(config *Config), insertOpts *InsertOpts) (*Client[pgx.Tx], *rivertype.JobRow, <-chan *Event, chan struct{}, chan struct{}, func()) { + t.Helper() + + config := newTestConfig(t, "") + configFunc(config) + + jobContextDoneChan := make(chan struct{}) + jobReleasedChan := make(chan struct{}) + jobStartedChan := make(chan struct{}) + releaseJobChan := make(chan struct{}) + releaseJob := sync.OnceFunc(func() { close(releaseJobChan) }) + + AddWorker(config.Workers, WorkFunc(func(ctx context.Context, job *Job[JobArgs]) error { + close(jobStartedChan) + <-ctx.Done() + close(jobContextDoneChan) + <-releaseJobChan + close(jobReleasedChan) + return nil + })) + + client := runNewTestClient(ctx, t, config) + subscribeChan := subscribe(t, client) + t.Cleanup(releaseJob) + + insertRes, err := client.Insert(ctx, JobArgs{}, insertOpts) + require.NoError(t, err) + + riversharedtest.WaitOrTimeout(t, jobStartedChan) + + return client, insertRes.Job, subscribeChan, jobContextDoneChan, jobReleasedChan, releaseJob + } + + requireAbandonEvent := func(t *testing.T, subscribeChan <-chan *Event, jobID int64, state rivertype.JobState) { + t.Helper() + + event := riversharedtest.WaitOrTimeout(t, subscribeChan) + require.NotNil(t, event) + require.Equal(t, EventKindJobFailed, event.Kind) + require.Equal(t, jobID, event.Job.ID) + require.Equal(t, state, event.Job.State) + require.NotNil(t, event.JobStats) + + _, ok := <-subscribeChan + require.False(t, ok, "expected exactly one job event") + } + + requireAbandoned := func(t *testing.T, client *Client[pgx.Tx], jobID int64, state rivertype.JobState) { + t.Helper() + + jobAfter, err := client.JobGet(ctx, jobID) + require.NoError(t, err) + require.Equal(t, state, jobAfter.State) + require.Len(t, jobAfter.Errors, 1) + require.Equal(t, producerJobAbandonedError, jobAfter.Errors[0].Error) + require.Equal(t, 1, jobAfter.Errors[0].Attempt) + require.Empty(t, jobAfter.Errors[0].Trace) + if state == rivertype.JobStateDiscarded { + require.NotNil(t, jobAfter.FinalizedAt) + } else { + require.Nil(t, jobAfter.FinalizedAt) + } + } + + t.Run("AfterSoftStopTimeout", func(t *testing.T) { + t.Parallel() + + client, job, subscribeChan, jobContextDoneChan, jobReleasedChan, releaseJob := setup(t, func(config *Config) { + config.StopAbandonTimeout = 100 * time.Millisecond + config.SoftStopTimeout = 100 * time.Millisecond + }, nil) + + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + require.NoError(t, client.Stop(stopCtx)) + + riversharedtest.WaitOrTimeout(t, jobContextDoneChan) + requireAbandoned(t, client, job.ID, rivertype.JobStateAvailable) + requireAbandonEvent(t, subscribeChan, job.ID, rivertype.JobStateAvailable) + + releaseJob() + riversharedtest.WaitOrTimeout(t, jobReleasedChan) + }) + + t.Run("AfterStopAndCancel", func(t *testing.T) { + t.Parallel() + + client, job, subscribeChan, jobContextDoneChan, jobReleasedChan, releaseJob := setup(t, func(config *Config) { + config.StopAbandonTimeout = 100 * time.Millisecond + }, nil) + + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + require.NoError(t, client.StopAndCancel(stopCtx)) + + riversharedtest.WaitOrTimeout(t, jobContextDoneChan) + requireAbandoned(t, client, job.ID, rivertype.JobStateAvailable) + requireAbandonEvent(t, subscribeChan, job.ID, rivertype.JobStateAvailable) + + releaseJob() + riversharedtest.WaitOrTimeout(t, jobReleasedChan) + }) + + t.Run("AtMaxAttempts", func(t *testing.T) { + t.Parallel() + + client, job, subscribeChan, jobContextDoneChan, jobReleasedChan, releaseJob := setup(t, func(config *Config) { + config.StopAbandonTimeout = 100 * time.Millisecond + }, &InsertOpts{MaxAttempts: 1}) + + stopCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + require.NoError(t, client.StopAndCancel(stopCtx)) + + riversharedtest.WaitOrTimeout(t, jobContextDoneChan) + requireAbandoned(t, client, job.ID, rivertype.JobStateDiscarded) + requireAbandonEvent(t, subscribeChan, job.ID, rivertype.JobStateDiscarded) + + releaseJob() + riversharedtest.WaitOrTimeout(t, jobReleasedChan) + }) +} + func Test_Client_SoftStopTimeout(t *testing.T) { t.Parallel() @@ -2751,7 +2883,7 @@ func Test_Client_SoftStopTimeout(t *testing.T) { testutil.JobArgsReflectKind[JobArgs] } - t.Run("EscalatesToHardStopAfterTimeout", func(t *testing.T) { + t.Run("CancelsJobsAfterTimeout", func(t *testing.T) { t.Parallel() config := newTestConfig(t, "") @@ -2772,8 +2904,8 @@ func Test_Client_SoftStopTimeout(t *testing.T) { riversharedtest.WaitOrTimeout(t, jobStartedChan) - // Stop initiates a soft stop. The job won't finish on its own, but - // SoftStopTimeout should escalate to a hard stop after 100ms. + // Stop initiates a graceful stop. The job won't finish on its own, but + // SoftStopTimeout should cancel its context after 100ms. require.NoError(t, client.Stop(ctx)) // Verify the job's context was indeed cancelled. @@ -2878,7 +3010,7 @@ func Test_Client_SoftStopTimeout(t *testing.T) { require.NoError(t, client.Stop(ctx)) }) - t.Run("ContextCancellationEscalatesAfterTimeout", func(t *testing.T) { + t.Run("StartContextCancellationCancelsJobsAfterTimeout", func(t *testing.T) { t.Parallel() config := newTestConfig(t, "") @@ -2915,8 +3047,8 @@ func Test_Client_SoftStopTimeout(t *testing.T) { riversharedtest.WaitOrTimeout(t, jobStartedChan) - // Cancel the start context. This should initiate a soft stop, then - // escalate to hard stop after SoftStopTimeout. + // Cancel the start context. This should initiate a graceful stop, then + // cancel job contexts after SoftStopTimeout. startCtxCancel() riversharedtest.WaitOrTimeout(t, client.Stopped()) @@ -8780,6 +8912,22 @@ func Test_NewClient_Validations(t *testing.T) { }, wantErr: fmt.Errorf("FetchPollInterval cannot be shorter than FetchCooldown (%s)", 20*time.Millisecond), }, + { + name: "StopAbandonTimeout cannot be negative", + configFunc: func(config *Config) { + config.StopAbandonTimeout = -1 + }, + wantErr: errors.New("StopAbandonTimeout cannot be less than zero"), + }, + { + name: "StopAbandonTimeout may be overridden", + configFunc: func(config *Config) { + config.StopAbandonTimeout = 23 * time.Second + }, + validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper + require.Equal(t, 23*time.Second, client.config.StopAbandonTimeout) + }, + }, { name: "FetchPollInterval cannot be less than MinFetchPollInterval", configFunc: func(config *Config) { config.FetchPollInterval = time.Millisecond - 1 }, @@ -8982,6 +9130,22 @@ func Test_NewClient_Validations(t *testing.T) { }, wantErr: errors.New("Schema name can only contain letters, numbers, and underscores, and must start with a letter or underscore"), }, + { + name: "SoftStopTimeout cannot be negative", + configFunc: func(config *Config) { + config.SoftStopTimeout = -1 + }, + wantErr: errors.New("SoftStopTimeout cannot be less than zero"), + }, + { + name: "SoftStopTimeout may be overridden", + configFunc: func(config *Config) { + config.SoftStopTimeout = 23 * time.Second + }, + validateResult: func(t *testing.T, client *Client[pgx.Tx]) { //nolint:thelper + require.Equal(t, 23*time.Second, client.config.SoftStopTimeout) + }, + }, { name: "Queues can be nil when Workers is also nil", configFunc: func(config *Config) { diff --git a/example_graceful_shutdown_stop_and_cancel_test.go b/example_graceful_shutdown_stop_and_cancel_test.go index 30217c2d..86aea931 100644 --- a/example_graceful_shutdown_stop_and_cancel_test.go +++ b/example_graceful_shutdown_stop_and_cancel_test.go @@ -20,8 +20,8 @@ import ( // Example_gracefulShutdownStopCancel demonstrates graceful stop with explicit // fallback to StopAndCancel. When a SIGINT/SIGTERM arrives, Stop initiates a -// soft stop. If running jobs don't finish before the soft stop context expires, -// StopAndCancel cancels their contexts (hard stop). This example is intended to +// graceful stop. If running jobs don't finish before the graceful stop context +// expires, StopAndCancel cancels their contexts. This example is intended to // demonstrate advanced use of StopAndCancel. Generally, prefer the method shown // in Example_gracefulShutdown over the one here. func Example_gracefulShutdownStopAndCancel() { @@ -59,8 +59,8 @@ func Example_gracefulShutdownStopAndCancel() { } // Use signal.NotifyContext to detect SIGINT/SIGTERM, but don't pass the - // signal context to Start. Cancelling the Start context cancels running job - // contexts immediately, which is equivalent to StopAndCancel. + // signal context to Start. Cancelling the Start context would cancel running + // job contexts immediately, which is equivalent to StopAndCancel. signalCtx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM) defer stop() diff --git a/example_graceful_shutdown_test.go b/example_graceful_shutdown_test.go index bf99b8ac..72cb4d60 100644 --- a/example_graceful_shutdown_test.go +++ b/example_graceful_shutdown_test.go @@ -43,8 +43,8 @@ func (w *WaitsForCancelOnlyWorker) Work(ctx context.Context, job *river.Job[Wait // Example_gracefulShutdown demonstrates graceful stop using SoftStopTimeout. // When a SIGINT/SIGTERM arrives, the start context is cancelled, which -// initiates a soft stop. If running jobs don't finish within the configured -// SoftStopTimeout, their contexts are automatically cancelled (hard stop). +// initiates a graceful stop. If running jobs don't finish within the configured +// SoftStopTimeout, their contexts are automatically cancelled. func Example_gracefulShutdown() { ctx := context.Background() @@ -77,7 +77,7 @@ func Example_gracefulShutdown() { } // Use signal.NotifyContext to cancel the start context on SIGINT/SIGTERM. - // When the signal fires, the client initiates a soft stop. If running jobs + // When the signal fires, the client initiates a graceful stop. If running jobs // don't finish within SoftStopTimeout, their contexts are cancelled. signalCtx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM) defer stop() diff --git a/internal/jobexecutor/job_executor.go b/internal/jobexecutor/job_executor.go index adaa6265..12244345 100644 --- a/internal/jobexecutor/job_executor.go +++ b/internal/jobexecutor/job_executor.go @@ -121,6 +121,7 @@ type JobExecutor struct { Unstuck func() } SchedulerInterval time.Duration + ShouldReportResultFunc func() bool StuckThresholdOverride time.Duration WorkerMiddleware []rivertype.WorkerMiddleware WorkUnit workunit.WorkUnit @@ -158,6 +159,11 @@ func (e *JobExecutor) Execute(ctx context.Context) { res.Err = context.Cause(ctx) } + if e.ShouldReportResultFunc != nil && !e.ShouldReportResultFunc() { + e.ProducerCallbacks.JobDone(e.JobRow) + return + } + var multiJobErrors withJobsAndErrorsByID if res.Err != nil { multiJobErrors, _ = res.Err.(withJobsAndErrorsByID) diff --git a/producer.go b/producer.go index da11246e..7af2bbc1 100644 --- a/producer.go +++ b/producer.go @@ -15,6 +15,7 @@ import ( "github.com/riverqueue/river/internal/jobcompleter" "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/jobstats" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/retrypolicy" @@ -35,6 +36,7 @@ import ( ) const ( + producerJobAbandonedError = "job abandoned because River client StopAbandonTimeout elapsed" producerReportIntervalDefault = 30 * time.Second queuePollIntervalDefault = 2 * time.Second queueReportIntervalDefault = 10 * time.Minute @@ -89,6 +91,7 @@ type producerConfig struct { JobStuckCount *atomic.Int32 JobStuckThreshold time.Duration JobTimeout time.Duration + JobUpdateCallback func(ctx context.Context, updates []jobcompleter.CompleterJobUpdated) MaxWorkers int // Notifier is a notifier for subscribing to new job inserts and job @@ -190,8 +193,9 @@ type producer struct { baseservice.BaseService startstop.BaseStartStop - // Jobs which are currently being worked. Only used by main goroutine. - activeJobs map[int64]*jobexecutor.JobExecutor + // Jobs which are currently being worked. The map itself is only used by the + // main goroutine, while each entry coordinates finalization with its worker. + activeJobs map[int64]*producerActiveJob completer jobcompleter.JobCompleter config *producerConfig @@ -199,6 +203,8 @@ type producer struct { exec riverdriver.Executor errorHandler jobexecutor.ErrorHandler fetchLimiter *chanutil.DebouncedChan + abandonCh chan struct{} // signals that remaining running jobs should be abandoned and set to errored + abandonOnce *sync.Once // closes abandonCh exactly once metricEmitHooks []rivertype.HookMetricEmit // memoized hooks of type HookMetricEmit for reuse in dispatchWork state riverpilot.ProducerState pilot riverpilot.Pilot @@ -215,7 +221,7 @@ type producer struct { // Receives completed jobs from workers. Written by completed workers, only // read from main goroutine. - jobResultCh chan *rivertype.JobRow + jobResultCh chan *producerJobResult jobTimeout time.Duration @@ -247,13 +253,15 @@ func newProducer(archetype *baseservice.Archetype, exec riverdriver.Executor, pi } producer := baseservice.Init(archetype, &producer{ - activeJobs: make(map[int64]*jobexecutor.JobExecutor), + activeJobs: make(map[int64]*producerActiveJob), cancelCh: make(chan int64, 1000), completer: config.Completer, config: config.mustValidate(), exec: exec, errorHandler: errorHandler, - jobResultCh: make(chan *rivertype.JobRow, config.MaxWorkers), + abandonCh: make(chan struct{}), + abandonOnce: &sync.Once{}, + jobResultCh: make(chan *producerJobResult, config.MaxWorkers), jobTimeout: config.JobTimeout, pilot: pilot, queueControlCh: make(chan *controlEventPayload, 100), @@ -296,6 +304,9 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return nil } + p.abandonCh = make(chan struct{}) + p.abandonOnce = &sync.Once{} + isExpectedShutdownError := func(err error) bool { return errors.Is(err, startstop.ErrStop) || strings.HasSuffix(err.Error(), "conn closed") || fetchCtx.Err() != nil } @@ -429,7 +440,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { p.fetchAndRunLoop(fetchCtx, workCtx) p.Logger.DebugContext(workCtx, p.Name+": Entering shutdown loop", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) - p.executorShutdownLoop() + p.executorShutdownLoop(context.WithoutCancel(fetchCtx)) p.Logger.DebugContext(workCtx, p.Name+": Shutdown loop exited, awaiting subroutines", slog.String("queue", p.config.Queue), slog.Int64("id", p.id.Load())) cancelSubroutines(fmt.Errorf("producer stopped: %w", startstop.ErrStop)) @@ -484,6 +495,68 @@ type insertPayload struct { Queue string `json:"queue"` } +type producerActiveJob struct { + abandoned atomic.Bool + executor *jobexecutor.JobExecutor + finalizationMu sync.Mutex + workerFinalizationStarted bool +} + +func newProducerActiveJob(executor *jobexecutor.JobExecutor) *producerActiveJob { + activeJob := &producerActiveJob{executor: executor} + executor.Completer = &producerJobCompleter{ + activeJob: activeJob, + JobCompleter: executor.Completer, + } + executor.ShouldReportResultFunc = func() bool { return !activeJob.isAbandoned() } + return activeJob +} + +func (j *producerActiveJob) isAbandoned() bool { + return j.abandoned.Load() +} + +func (j *producerActiveJob) jobSetStateIfRunning(ctx context.Context, completer jobcompleter.JobCompleter, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + j.finalizationMu.Lock() + defer j.finalizationMu.Unlock() + + if j.abandoned.Load() { + return nil + } + + j.workerFinalizationStarted = true + return completer.JobSetStateIfRunning(ctx, stats, params) +} + +func (j *producerActiveJob) tryAbandon() bool { + if !j.finalizationMu.TryLock() { + return false + } + defer j.finalizationMu.Unlock() + + if j.abandoned.Load() || j.workerFinalizationStarted { + return false + } + + j.abandoned.Store(true) + return true +} + +type producerJobCompleter struct { + jobcompleter.JobCompleter + + activeJob *producerActiveJob +} + +func (c *producerJobCompleter) JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + return c.activeJob.jobSetStateIfRunning(ctx, c.JobCompleter, stats, params) +} + +type producerJobResult struct { + executor *jobexecutor.JobExecutor + job *rivertype.JobRow +} + func (p *producer) handleControlNotification(workCtx context.Context) func(notifier.NotificationTopic, string) { return func(topic notifier.NotificationTopic, payload string) { var decoded controlEventPayload @@ -677,12 +750,30 @@ func (p *producer) innerFetchLoop(workCtx context.Context, fetchResultCh chan pr } } -func (p *producer) executorShutdownLoop() { +func (p *producer) executorShutdownLoop(ctx context.Context) { // No more jobs will be fetched or executed. However, we must wait for all // in-progress jobs to complete. + abandonCh := p.abandonCh for len(p.activeJobs) != 0 { - result := <-p.jobResultCh - p.removeActiveJob(result) + select { + case result := <-p.jobResultCh: + p.removeActiveJob(result) + case <-abandonCh: + abandonCh = nil + p.drainJobResults() + p.abandonActiveJobs(ctx) + } + } +} + +func (p *producer) drainJobResults() { + for { + select { + case result := <-p.jobResultCh: + p.removeActiveJob(result) + default: + return + } } } @@ -738,17 +829,136 @@ func (p *producer) finalizeShutdown(ctx context.Context) { func (p *producer) addActiveJob(id int64, executor *jobexecutor.JobExecutor) { p.numJobsActive.Add(1) - p.activeJobs[id] = executor + p.activeJobs[id] = newProducerActiveJob(executor) +} + +func (p *producer) abandon() { + p.abandonOnce.Do(func() { close(p.abandonCh) }) } -func (p *producer) removeActiveJob(job *rivertype.JobRow) { - executor := p.activeJobs[job.ID] - delete(p.activeJobs, job.ID) - if executor == nil || executor.TryCloseSlot() { +func (p *producer) abandonActiveJobs(ctx context.Context) { + if len(p.activeJobs) == 0 { + return + } + + p.abandon() + + abandonedActiveJobs := make(map[int64]*producerActiveJob, len(p.activeJobs)) + for id, activeJob := range p.activeJobs { + if activeJob.tryAbandon() { + abandonedActiveJobs[id] = activeJob + } + } + + now := p.Time.Now() + params := &riverdriver.JobSetStateIfRunningManyParams{ + Attempt: make([]*int, 0, len(abandonedActiveJobs)), + ErrData: make([][]byte, 0, len(abandonedActiveJobs)), + FinalizedAt: make([]*time.Time, 0, len(abandonedActiveJobs)), + ID: make([]int64, 0, len(abandonedActiveJobs)), + MetadataDoMerge: make([]bool, 0, len(abandonedActiveJobs)), + MetadataUpdates: make([][]byte, 0, len(abandonedActiveJobs)), + Now: &now, + ScheduledAt: make([]*time.Time, 0, len(abandonedActiveJobs)), + Schema: p.config.Schema, + State: make([]rivertype.JobState, 0, len(abandonedActiveJobs)), + } + + for _, activeJob := range abandonedActiveJobs { + job := activeJob.executor.JobRow + errData, err := json.Marshal(rivertype.AttemptError{ + At: now, + Attempt: job.Attempt, + Error: producerJobAbandonedError, + }) + if err != nil { + panic(fmt.Errorf("error serializing job abandonment error: %w", err)) + } + + var setStateParams *riverdriver.JobSetStateIfRunningParams + if job.Attempt >= job.MaxAttempts { + setStateParams = riverdriver.JobSetStateDiscarded(job.ID, now, errData, nil) + } else { + setStateParams = riverdriver.JobSetStateErrorAvailable(job.ID, now, errData, nil) + } + + params.Attempt = append(params.Attempt, setStateParams.Attempt) + params.ErrData = append(params.ErrData, setStateParams.ErrData) + params.FinalizedAt = append(params.FinalizedAt, setStateParams.FinalizedAt) + params.ID = append(params.ID, setStateParams.ID) + params.MetadataDoMerge = append(params.MetadataDoMerge, setStateParams.MetadataDoMerge) + params.MetadataUpdates = append(params.MetadataUpdates, setStateParams.MetadataUpdates) + params.ScheduledAt = append(params.ScheduledAt, setStateParams.ScheduledAt) + params.State = append(params.State, setStateParams.State) + } + + if len(abandonedActiveJobs) > 0 { + timeoutCtx, cancel := context.WithTimeout(ctx, rivercommon.HotOperationTimeout) + defer cancel() + + completeStart := p.Time.Now() + jobs, err := p.pilot.JobSetStateIfRunningMany(timeoutCtx, p.exec, params) + completeDuration := p.Time.Now().Sub(completeStart) + if err != nil { + p.Logger.ErrorContext(ctx, p.Name+": Error setting abandoned jobs to errored", slog.String("err", err.Error()), slog.Int("num_jobs", len(params.ID)), slog.String("queue", p.config.Queue)) + } else { + abandonedJobs := make([]*rivertype.JobRow, 0, len(jobs)) + for _, job := range jobs { + if len(job.Errors) < 1 { + continue + } + lastError := job.Errors[len(job.Errors)-1] + if lastError.At.Equal(now) && lastError.Attempt == abandonedActiveJobs[job.ID].executor.JobRow.Attempt && lastError.Error == producerJobAbandonedError { + abandonedJobs = append(abandonedJobs, job) + } + } + + p.Logger.WarnContext(ctx, p.Name+": Abandoned running jobs", slog.Int("num_jobs", len(abandonedJobs)), slog.String("queue", p.config.Queue)) + + if p.config.JobUpdateCallback != nil && len(abandonedJobs) > 0 { + updates := make([]jobcompleter.CompleterJobUpdated, 0, len(abandonedJobs)) + for _, job := range abandonedJobs { + stats := &jobstats.JobStatistics{CompleteDuration: completeDuration} + jobBefore := abandonedActiveJobs[job.ID].executor.JobRow + if jobBefore.AttemptedAt != nil { + stats.QueueWaitDuration = jobBefore.AttemptedAt.Sub(jobBefore.ScheduledAt) + stats.RunDuration = now.Sub(*jobBefore.AttemptedAt) + } + updates = append(updates, jobcompleter.CompleterJobUpdated{Job: job, JobStats: stats, Reason: riverdriver.JobSetStateReasonFailed}) + } + p.config.JobUpdateCallback(ctx, updates) + } + } + } + + numActiveJobSlots := 0 + for id, activeJob := range abandonedActiveJobs { + delete(p.activeJobs, id) + if activeJob.executor.TryCloseSlot() { + numActiveJobSlots++ + } + if p.state != nil { + p.state.JobFinish(activeJob.executor.JobRow) + } + } + p.numJobsActive.Add(-int32(numActiveJobSlots)) +} + +func (p *producer) removeActiveJob(result *producerJobResult) { + // Ignore stale results from executors abandoned out of active tracking. + activeJob := p.activeJobs[result.job.ID] + if activeJob == nil || activeJob.executor != result.executor { + return + } + + delete(p.activeJobs, result.job.ID) + if result.executor == nil || result.executor.TryCloseSlot() { p.numJobsActive.Add(-1) } p.numJobsRan.Add(1) - p.state.JobFinish(job) + if p.state != nil { + p.state.JobFinish(result.job) + } } func (p *producer) handleWorkerStuck(ctx context.Context, executor *jobexecutor.JobExecutor, job *rivertype.JobRow) { @@ -781,11 +991,11 @@ func (p *producer) handleWorkerUnstuck() { } func (p *producer) maybeCancelJob(ctx context.Context, id int64) { - executor, ok := p.activeJobs[id] + activeJob, ok := p.activeJobs[id] if !ok { return } - executor.Cancel(ctx) + activeJob.executor.Cancel(ctx) } func (p *producer) metricEmitHooksFromLookup() []rivertype.HookMetricEmit { @@ -937,7 +1147,7 @@ func (p *producer) startNewExecutors(workCtx context.Context, jobs []*rivertype. Stuck func(ctx context.Context, jobRow *rivertype.JobRow) Unstuck func() }{ - JobDone: p.handleWorkerDone, + JobDone: func(jobRow *rivertype.JobRow) { p.handleWorkerDone(executor, jobRow) }, Stuck: func(ctx context.Context, jobRow *rivertype.JobRow) { p.handleWorkerStuck(ctx, executor, jobRow) }, Unstuck: p.handleWorkerUnstuck, }, @@ -959,8 +1169,11 @@ func (p *producer) maxJobsToFetch() int { return p.config.MaxWorkers - int(p.numJobsActive.Load()) } -func (p *producer) handleWorkerDone(job *rivertype.JobRow) { - p.jobResultCh <- job +func (p *producer) handleWorkerDone(executor *jobexecutor.JobExecutor, job *rivertype.JobRow) { + p.jobResultCh <- &producerJobResult{ + executor: executor, + job: job, + } } func (p *producer) pollForSettingChanges(ctx context.Context, wg *sync.WaitGroup, lastPaused bool, lastMetadata []byte) { diff --git a/producer_test.go b/producer_test.go index f5d17a07..d9750c24 100644 --- a/producer_test.go +++ b/producer_test.go @@ -3,8 +3,10 @@ package river import ( "context" "encoding/json" + "errors" "fmt" "slices" + "sync" "sync/atomic" "testing" "time" @@ -12,6 +14,8 @@ import ( "github.com/stretchr/testify/require" "github.com/riverqueue/river/internal/jobcompleter" + "github.com/riverqueue/river/internal/jobexecutor" + "github.com/riverqueue/river/internal/jobstats" "github.com/riverqueue/river/internal/notifier" "github.com/riverqueue/river/internal/pluginlookup" "github.com/riverqueue/river/internal/rivercommon" @@ -54,6 +58,51 @@ func (p *beforeJobGetAvailablePilot) JobGetAvailable( return p.Pilot.JobGetAvailable(ctx, exec, state, params) } +type blockingJobCompleter struct { + jobcompleter.JobCompleter + + releaseCh chan struct{} + startedCh chan struct{} +} + +func (c *blockingJobCompleter) JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + close(c.startedCh) + <-c.releaseCh + return nil +} + +type jobSetStateIfRunningManyErrorPilot struct { + riverpilot.Pilot + + err error +} + +func (p *jobSetStateIfRunningManyErrorPilot) JobSetStateIfRunningMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) { + return nil, p.err +} + +type jobSetStateIfRunningManyRecordingPilot struct { + riverpilot.Pilot + + paramsCh chan *riverdriver.JobSetStateIfRunningManyParams +} + +func (p *jobSetStateIfRunningManyRecordingPilot) JobSetStateIfRunningMany(ctx context.Context, exec riverdriver.Executor, params *riverdriver.JobSetStateIfRunningManyParams) ([]*rivertype.JobRow, error) { + p.paramsCh <- params + return p.Pilot.JobSetStateIfRunningMany(ctx, exec, params) +} + +type recordingJobCompleter struct { + jobcompleter.JobCompleter + + paramsCh chan *riverdriver.JobSetStateIfRunningParams +} + +func (c *recordingJobCompleter) JobSetStateIfRunning(ctx context.Context, stats *jobstats.JobStatistics, params *riverdriver.JobSetStateIfRunningParams) error { + c.paramsCh <- params + return nil +} + func TestProducer_MetricEmitHook(t *testing.T) { t.Parallel() @@ -187,6 +236,210 @@ func TestProducer_MetricEmitHook(t *testing.T) { }) } +func TestProducer_AbandonActiveJobs(t *testing.T) { + t.Parallel() + + ctx := context.Background() + + type testBundle struct { + exec riverdriver.Executor + jobUpdates chan []jobcompleter.CompleterJobUpdated + producer *producer + schema string + } + + setup := func(t *testing.T, pilot riverpilot.Pilot) *testBundle { + t.Helper() + + var ( + archetype = riversharedtest.BaseServiceArchetype(t) + driver = riverpgxv5.New(riversharedtest.DBPool(ctx, t)) + exec = driver.GetExecutor() + schema = riverdbtest.TestSchema(ctx, t, driver, nil) + ) + if pilot == nil { + pilot = &riverpilot.StandardPilot{} + } + + jobUpdates := make(chan []jobcompleter.CompleterJobUpdated, 1) + completer := jobcompleter.NewInlineCompleter(archetype, schema, exec, pilot, make(chan []jobcompleter.CompleterJobUpdated, 10)) + producer := newProducer(archetype, exec, pilot, &producerConfig{ + ClientID: testClientID, + Completer: completer, + ErrorHandler: newTestErrorHandler(), + FetchCooldown: FetchCooldownDefault, + FetchPollInterval: FetchPollIntervalDefault, + JobTimeout: JobTimeoutDefault, + JobUpdateCallback: func(ctx context.Context, updates []jobcompleter.CompleterJobUpdated) { jobUpdates <- updates }, + MaxWorkers: 10, + PluginLookupByJob: pluginlookup.NewJobPluginLookup(nil), + PluginLookupGlobal: pluginlookup.NewPluginLookup(nil), + Queue: rivercommon.QueueDefault, + QueuePollInterval: queuePollIntervalDefault, + QueueReportInterval: queueReportIntervalDefault, + RetryPolicy: &DefaultClientRetryPolicy{}, + SchedulerInterval: riverinternaltest.SchedulerShortInterval, + Schema: schema, + StaleProducerRetentionPeriod: time.Minute, + Workers: NewWorkers(), + }) + + return &testBundle{exec: exec, jobUpdates: jobUpdates, producer: producer, schema: schema} + } + + t.Run("AbandonDoesNotWinAfterWorkerFinalizationStarts", func(t *testing.T) { + t.Parallel() + + finalizationReleaseCh := make(chan struct{}) + releaseFinalization := sync.OnceFunc(func() { close(finalizationReleaseCh) }) + t.Cleanup(releaseFinalization) + finalizationStartedCh := make(chan struct{}) + executor := &jobexecutor.JobExecutor{ + Completer: &blockingJobCompleter{ + releaseCh: finalizationReleaseCh, + startedCh: finalizationStartedCh, + }, + } + activeJob := newProducerActiveJob(executor) + + finalizationResultCh := make(chan error, 1) + go func() { + finalizationResultCh <- executor.Completer.JobSetStateIfRunning(ctx, &jobstats.JobStatistics{}, riverdriver.JobSetStateCompleted(1, time.Now(), nil)) + }() + + riversharedtest.WaitOrTimeout(t, finalizationStartedCh) + require.False(t, activeJob.tryAbandon()) + + releaseFinalization() + require.NoError(t, riversharedtest.WaitOrTimeout(t, finalizationResultCh)) + require.False(t, activeJob.isAbandoned()) + }) + + t.Run("AbandonWinsBeforeWorkerFinalizationStarts", func(t *testing.T) { + t.Parallel() + + paramsCh := make(chan *riverdriver.JobSetStateIfRunningParams, 1) + executor := &jobexecutor.JobExecutor{ + Completer: &recordingJobCompleter{paramsCh: paramsCh}, + } + activeJob := newProducerActiveJob(executor) + + require.True(t, activeJob.tryAbandon()) + require.NoError(t, executor.Completer.JobSetStateIfRunning(ctx, &jobstats.JobStatistics{}, riverdriver.JobSetStateCompleted(1, time.Now(), nil))) + require.Empty(t, paramsCh) + require.False(t, activeJob.tryAbandon()) + require.False(t, executor.ShouldReportResultFunc()) + }) + + t.Run("DatabaseErrorAbandonsExecutor", func(t *testing.T) { + t.Parallel() + + pilotErr := errors.New("database error") + bundle := setup(t, &jobSetStateIfRunningManyErrorPilot{Pilot: &riverpilot.StandardPilot{}, err: pilotErr}) + + runningState := rivertype.JobStateRunning + job := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{ + Attempt: ptrutil.Ptr(1), + MaxAttempts: ptrutil.Ptr(3), + Schema: bundle.schema, + State: &runningState, + }) + executor := &jobexecutor.JobExecutor{JobRow: job} + bundle.producer.addActiveJob(job.ID, executor) + activeJob := bundle.producer.activeJobs[job.ID] + + bundle.producer.abandon() + bundle.producer.executorShutdownLoop(ctx) + + require.Empty(t, bundle.producer.activeJobs) + require.Zero(t, bundle.producer.numJobsActive.Load()) + require.Empty(t, bundle.jobUpdates) + require.True(t, activeJob.isAbandoned()) + require.False(t, activeJob.tryAbandon()) + + jobAfter, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: job.ID, Schema: bundle.schema}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateRunning, jobAfter.State) + require.Empty(t, jobAfter.Errors) + }) + + t.Run("UpdatesJobsInSingleBatchAndEmitsEvents", func(t *testing.T) { + t.Parallel() + + recordingPilot := &jobSetStateIfRunningManyRecordingPilot{ + Pilot: &riverpilot.StandardPilot{}, + paramsCh: make(chan *riverdriver.JobSetStateIfRunningManyParams, 10), + } + bundle := setup(t, recordingPilot) + + availableState := rivertype.JobStateAvailable + runningState := rivertype.JobStateRunning + alreadyAvailableJob := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{ + Attempt: ptrutil.Ptr(1), + MaxAttempts: ptrutil.Ptr(3), + Schema: bundle.schema, + State: &availableState, + }) + retryableJob := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{ + Attempt: ptrutil.Ptr(1), + MaxAttempts: ptrutil.Ptr(3), + Schema: bundle.schema, + State: &runningState, + }) + discardedJob := testfactory.Job(ctx, t, bundle.exec, &testfactory.JobOpts{ + Attempt: ptrutil.Ptr(3), + MaxAttempts: ptrutil.Ptr(3), + Schema: bundle.schema, + State: &runningState, + }) + + bundle.producer.addActiveJob(alreadyAvailableJob.ID, &jobexecutor.JobExecutor{JobRow: alreadyAvailableJob}) + bundle.producer.addActiveJob(retryableJob.ID, &jobexecutor.JobExecutor{JobRow: retryableJob}) + bundle.producer.addActiveJob(discardedJob.ID, &jobexecutor.JobExecutor{JobRow: discardedJob}) + + bundle.producer.abandon() + bundle.producer.executorShutdownLoop(ctx) + + require.Empty(t, bundle.producer.activeJobs) + require.Zero(t, bundle.producer.numJobsActive.Load()) + batchParams := riversharedtest.WaitOrTimeout(t, recordingPilot.paramsCh) + require.ElementsMatch(t, []int64{alreadyAvailableJob.ID, retryableJob.ID, discardedJob.ID}, batchParams.ID) + require.Empty(t, recordingPilot.paramsCh) + + updates := riversharedtest.WaitOrTimeout(t, bundle.jobUpdates) + require.Len(t, updates, 2) + require.ElementsMatch(t, []int64{retryableJob.ID, discardedJob.ID}, []int64{updates[0].Job.ID, updates[1].Job.ID}) + for _, update := range updates { + require.NotNil(t, update.JobStats) + require.Equal(t, riverdriver.JobSetStateReasonFailed, update.Reason) + } + require.Empty(t, bundle.jobUpdates) + + alreadyAvailableJobAfter, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: alreadyAvailableJob.ID, Schema: bundle.schema}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateAvailable, alreadyAvailableJobAfter.State) + require.Empty(t, alreadyAvailableJobAfter.Errors) + + retryableJobAfter, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: retryableJob.ID, Schema: bundle.schema}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateAvailable, retryableJobAfter.State) + require.Len(t, retryableJobAfter.Errors, 1) + require.Equal(t, producerJobAbandonedError, retryableJobAfter.Errors[0].Error) + require.Equal(t, retryableJob.Attempt, retryableJobAfter.Errors[0].Attempt) + require.Empty(t, retryableJobAfter.Errors[0].Trace) + require.Nil(t, retryableJobAfter.FinalizedAt) + + discardedJobAfter, err := bundle.exec.JobGetByID(ctx, &riverdriver.JobGetByIDParams{ID: discardedJob.ID, Schema: bundle.schema}) + require.NoError(t, err) + require.Equal(t, rivertype.JobStateDiscarded, discardedJobAfter.State) + require.Len(t, discardedJobAfter.Errors, 1) + require.Equal(t, producerJobAbandonedError, discardedJobAfter.Errors[0].Error) + require.Equal(t, discardedJob.Attempt, discardedJobAfter.Errors[0].Attempt) + require.Empty(t, discardedJobAfter.Errors[0].Trace) + require.NotNil(t, discardedJobAfter.FinalizedAt) + }) +} + func TestProducer_PollOnly(t *testing.T) { t.Parallel()