From cd91bae5bb006250eee5df6c626e93f19a6bf4c2 Mon Sep 17 00:00:00 2001 From: Brandur Date: Thu, 23 Jul 2026 22:32:19 -0500 Subject: [PATCH] Use context helpers more broadly in the codebase This one follows up #1277 just to apply the same pattern more broadly. It's not strictly needed, but it should result in some improved error messages, and give us some material for a blog post. I told Codex not to bother applying it to any places where we'd have to deviate too much from the normal pattern or do anything weird, so included are just the spots where `WithTimeout`/`WithTimeoutV` applies cleanly. --- internal/jobcompleter/job_completer.go | 6 +- internal/leadership/elector.go | 98 +++++++++---------- internal/maintenance/job_cleaner.go | 9 +- internal/maintenance/job_rescuer.go | 20 ++-- internal/maintenance/job_scheduler.go | 9 +- internal/maintenance/queue_cleaner.go | 9 +- .../sqlite_notification_cleaner.go | 28 +++--- .../sqlite_notification_cleaner_test.go | 11 +++ internal/notifier/notifier.go | 52 +++++----- producer.go | 81 +++++++-------- rivershared/util/dbutil/db_util.go | 6 +- 11 files changed, 158 insertions(+), 171 deletions(-) diff --git a/internal/jobcompleter/job_completer.go b/internal/jobcompleter/job_completer.go index 57b827cb..5e6c7c74 100644 --- a/internal/jobcompleter/job_completer.go +++ b/internal/jobcompleter/job_completer.go @@ -16,6 +16,7 @@ import ( "github.com/riverqueue/river/rivershared/riverpilot" "github.com/riverqueue/river/rivershared/startstop" "github.com/riverqueue/river/rivershared/util/serviceutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivertype" ) @@ -683,10 +684,7 @@ func withRetries[T any](logCtx context.Context, baseService *baseservice.BaseSer for attempt := 1; attempt <= numRetries; attempt++ { // I've found that we want at least ten seconds for a large batch, // although it usually doesn't need that long. - ctx, cancel := context.WithTimeout(uncancelledCtx, rivercommon.HotOperationTimeout) - defer cancel() - - retVal, err := retryFunc(ctx) + retVal, err := timeoututil.WithTimeoutV(uncancelledCtx, rivercommon.HotOperationTimeout, baseService.Name+".withRetries", retryFunc) if err != nil { // A cancelled context or a closed pool will never succeed. if isNonRetryableCompleterError(err) { diff --git a/internal/leadership/elector.go b/internal/leadership/elector.go index c9008582..c79391a7 100644 --- a/internal/leadership/elector.go +++ b/internal/leadership/elector.go @@ -20,6 +20,7 @@ import ( "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivertype" ) @@ -609,25 +610,24 @@ func (e *Elector) attemptResign(ctx context.Context, attempt int, term leadershi // Wait one second longer each time we try to resign: timeout := time.Duration(attempt) * time.Second - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - resigned, err := e.exec.LeaderResign(ctx, &riverdriver.LeaderResignParams{ - ElectedAt: term.electedAt, - LeaderID: term.clientID, - LeadershipTopic: string(notifier.NotificationTopicLeadership), - Schema: e.config.Schema, - }) - if err != nil { - return err - } + return timeoututil.WithTimeout(ctx, timeout, e.Name+".attemptResign", func(ctx context.Context) error { + resigned, err := e.exec.LeaderResign(ctx, &riverdriver.LeaderResignParams{ + ElectedAt: term.electedAt, + LeaderID: term.clientID, + LeadershipTopic: string(notifier.NotificationTopicLeadership), + Schema: e.config.Schema, + }) + if err != nil { + return err + } - if resigned { - e.Logger.DebugContext(ctx, e.Name+": Resigned leadership successfully", "client_id", e.config.ClientID) - e.testSignals.ResignedLeadership.Signal(struct{}{}) - } + if resigned { + e.Logger.DebugContext(ctx, e.Name+": Resigned leadership successfully", "client_id", e.config.ClientID) + e.testSignals.ResignedLeadership.Signal(struct{}{}) + } - return nil + return nil + }) } // Produces a common set of key/value pairs for logging when an error occurs. @@ -767,44 +767,42 @@ func attemptElect(ctx context.Context, exec riverdriver.Executor, params *riverd } func attemptElectWithTimeout(ctx context.Context, exec riverdriver.Executor, params *riverdriver.LeaderElectParams, timeout time.Duration) (*riverdriver.Leader, error) { - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - execTx, err := exec.Begin(ctx) - if err != nil { - var additionalDetail string - if errors.Is(err, context.DeadlineExceeded) { - additionalDetail = " (a common cause of this is a database pool that's at its connection limit; you may need to increase maximum connections)" - } + return timeoututil.WithTimeoutV(ctx, timeout, "leadership.attemptElect", func(ctx context.Context) (*riverdriver.Leader, error) { + execTx, err := exec.Begin(ctx) + if err != nil { + var additionalDetail string + if errors.Is(err, context.DeadlineExceeded) { + additionalDetail = " (a common cause of this is a database pool that's at its connection limit; you may need to increase maximum connections)" + } - return nil, fmt.Errorf("error beginning transaction: %w%s", err, additionalDetail) - } - defer dbutil.RollbackWithoutCancel(ctx, execTx) + return nil, fmt.Errorf("error beginning transaction: %w%s", err, additionalDetail) + } + defer dbutil.RollbackWithoutCancel(ctx, execTx) - if _, err := execTx.LeaderDeleteExpired(ctx, &riverdriver.LeaderDeleteExpiredParams{ - Now: params.Now, - Schema: params.Schema, - }); err != nil { - return nil, err - } + if _, err := execTx.LeaderDeleteExpired(ctx, &riverdriver.LeaderDeleteExpiredParams{ + Now: params.Now, + Schema: params.Schema, + }); err != nil { + return nil, err + } - leader, err := execTx.LeaderAttemptElect(ctx, params) - if err != nil && !errors.Is(err, rivertype.ErrNotFound) { - return nil, err - } - if err := execTx.Commit(ctx); err != nil { - return nil, fmt.Errorf("error committing transaction: %w", err) - } - if err != nil { - return nil, err - } + leader, err := execTx.LeaderAttemptElect(ctx, params) + if err != nil && !errors.Is(err, rivertype.ErrNotFound) { + return nil, err + } + if err := execTx.Commit(ctx); err != nil { + return nil, fmt.Errorf("error committing transaction: %w", err) + } + if err != nil { + return nil, err + } - return leader, nil + return leader, nil + }) } func attemptReelectWithTimeout(ctx context.Context, exec riverdriver.Executor, params *riverdriver.LeaderReelectParams, timeout time.Duration) (*riverdriver.Leader, error) { - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - return exec.LeaderAttemptReelect(ctx, params) + return timeoututil.WithTimeoutV(ctx, timeout, "leadership.attemptReelect", func(ctx context.Context) (*riverdriver.Leader, error) { + return exec.LeaderAttemptReelect(ctx, params) + }) } diff --git a/internal/maintenance/job_cleaner.go b/internal/maintenance/job_cleaner.go index f3216561..e70fe300 100644 --- a/internal/maintenance/job_cleaner.go +++ b/internal/maintenance/job_cleaner.go @@ -17,6 +17,7 @@ import ( "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivershared/util/timeutil" ) @@ -183,8 +184,7 @@ func (s *JobCleaner) runOnce(ctx context.Context) (*jobCleanerRunOnceResult, err res := &jobCleanerRunOnceResult{} for { - // Wrapped in a function so that defers run as expected. - numDeleted, err := func() (int, error) { + numDeleted, err := timeoututil.WithTimeoutV(ctx, s.Config.Timeout, s.Name+".runOnce", func(ctx context.Context) (int, error) { // In the special case that all retentions are indefinite, don't // bother issuing the query at all as an optimization. if s.Config.CompletedJobRetentionPeriod == -1 && @@ -193,9 +193,6 @@ func (s *JobCleaner) runOnce(ctx context.Context) (*jobCleanerRunOnceResult, err return 0, nil } - ctx, cancelFunc := context.WithTimeout(ctx, s.Config.Timeout) - defer cancelFunc() - numDeleted, err := s.exec.JobDeleteBefore(ctx, &riverdriver.JobDeleteBeforeParams{ CancelledDoDelete: s.Config.CancelledJobRetentionPeriod != -1, CancelledFinalizedAtHorizon: time.Now().Add(-s.Config.CancelledJobRetentionPeriod), @@ -214,7 +211,7 @@ func (s *JobCleaner) runOnce(ctx context.Context) (*jobCleanerRunOnceResult, err s.reducedBatchSizeBreaker.ResetIfNotOpen() return numDeleted, nil - }() + }) if err != nil { if errors.Is(err, context.DeadlineExceeded) { s.reducedBatchSizeBreaker.Trip() diff --git a/internal/maintenance/job_rescuer.go b/internal/maintenance/job_rescuer.go index 595e4e6d..586b7893 100644 --- a/internal/maintenance/job_rescuer.go +++ b/internal/maintenance/job_rescuer.go @@ -21,6 +21,7 @@ import ( "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivershared/util/timeutil" "github.com/riverqueue/river/rivertype" ) @@ -294,9 +295,6 @@ func (s *JobRescuer) runOnce(ctx context.Context) (*rescuerRunOnceResult, error) } func (s *JobRescuer) getStuckJobs(ctx context.Context, afterID int64, batchSize int, stuckHorizon time.Time) ([]*rivertype.JobRow, error) { - ctx, cancelFunc := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault) - defer cancelFunc() - params := &riverdriver.JobGetStuckParams{ AfterID: afterID, Max: batchSize, @@ -304,14 +302,16 @@ func (s *JobRescuer) getStuckJobs(ctx context.Context, afterID int64, batchSize StuckHorizon: stuckHorizon, } - if pilot, ok := s.Config.Pilot.(riverpilot.PilotJobRescuer); ok { - return pilot.JobGetStuck(ctx, s.exec, params) - } + return timeoututil.WithTimeoutV(ctx, riversharedmaintenance.TimeoutDefault, s.Name+".getStuckJobs", func(ctx context.Context) ([]*rivertype.JobRow, error) { + if pilot, ok := s.Config.Pilot.(riverpilot.PilotJobRescuer); ok { + return pilot.JobGetStuck(ctx, s.exec, params) + } - // Compatibility fallback for Pilot implementations from before - // PilotJobRescuer. Once Pilot embeds PilotJobRescuer, replace the assertion - // above and this fallback with a direct call to s.Config.Pilot.JobGetStuck. - return s.exec.JobGetStuck(ctx, params) + // Compatibility fallback for Pilot implementations from before + // PilotJobRescuer. Once Pilot embeds PilotJobRescuer, replace the assertion + // above and this fallback with a direct call to s.Config.Pilot.JobGetStuck. + return s.exec.JobGetStuck(ctx, params) + }) } // jobRetryDecision is a signal from makeRetryDecision as to what to do with a diff --git a/internal/maintenance/job_scheduler.go b/internal/maintenance/job_scheduler.go index f593ae35..a2d289a2 100644 --- a/internal/maintenance/job_scheduler.go +++ b/internal/maintenance/job_scheduler.go @@ -18,6 +18,7 @@ import ( "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivershared/util/timeutil" ) @@ -163,11 +164,7 @@ func (s *JobScheduler) runOnce(ctx context.Context) (*schedulerRunOnceResult, er res := &schedulerRunOnceResult{} for { - // Wrapped in a function so that defers run as expected. - numScheduled, err := func() (int, error) { - ctx, cancelFunc := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault) - defer cancelFunc() - + numScheduled, err := timeoututil.WithTimeoutV(ctx, riversharedmaintenance.TimeoutDefault, s.Name+".runOnce", func(ctx context.Context) (int, error) { execTx, err := s.exec.Begin(ctx) if err != nil { return 0, fmt.Errorf("error starting transaction: %w", err) @@ -212,7 +209,7 @@ func (s *JobScheduler) runOnce(ctx context.Context) (*schedulerRunOnceResult, er } return len(scheduledJobResults), execTx.Commit(ctx) - }() + }) if err != nil { if errors.Is(err, context.DeadlineExceeded) { s.reducedBatchSizeBreaker.Trip() diff --git a/internal/maintenance/queue_cleaner.go b/internal/maintenance/queue_cleaner.go index f64e97db..e417558b 100644 --- a/internal/maintenance/queue_cleaner.go +++ b/internal/maintenance/queue_cleaner.go @@ -18,6 +18,7 @@ import ( "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivershared/util/timeutil" ) @@ -156,11 +157,7 @@ func (s *QueueCleaner) runOnce(ctx context.Context) (*queueCleanerRunOnceResult, res := &queueCleanerRunOnceResult{QueuesDeleted: make([]string, 0, 10)} for { - // Wrapped in a function so that defers run as expected. - queuesDeleted, err := func() ([]string, error) { - ctx, cancelFunc := context.WithTimeout(ctx, riversharedmaintenance.TimeoutDefault) - defer cancelFunc() - + queuesDeleted, err := timeoututil.WithTimeoutV(ctx, riversharedmaintenance.TimeoutDefault, s.Name+".runOnce", func(ctx context.Context) ([]string, error) { queuesDeleted, err := s.exec.QueueDeleteExpired(ctx, &riverdriver.QueueDeleteExpiredParams{ Max: s.batchSize(), Schema: s.Config.Schema, @@ -173,7 +170,7 @@ func (s *QueueCleaner) runOnce(ctx context.Context) (*queueCleanerRunOnceResult, s.reducedBatchSizeBreaker.ResetIfNotOpen() return queuesDeleted, nil - }() + }) if err != nil { if errors.Is(err, context.DeadlineExceeded) { s.reducedBatchSizeBreaker.Trip() diff --git a/internal/maintenance/sqlite_notification_cleaner.go b/internal/maintenance/sqlite_notification_cleaner.go index 2fd07f7c..de9e9526 100644 --- a/internal/maintenance/sqlite_notification_cleaner.go +++ b/internal/maintenance/sqlite_notification_cleaner.go @@ -13,6 +13,7 @@ import ( "github.com/riverqueue/river/rivershared/startstop" "github.com/riverqueue/river/rivershared/testsignal" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivershared/util/timeutil" ) @@ -133,20 +134,19 @@ type sqliteNotificationCleanerRunOnceResult struct { } func (s *SQLiteNotificationCleaner) runOnce(ctx context.Context) (*sqliteNotificationCleanerRunOnceResult, error) { - ctx, cancelFunc := context.WithTimeout(ctx, s.Config.Timeout) - defer cancelFunc() - - numDeleted, err := s.exec.NotificationDeleteBefore(ctx, &riverdriver.NotificationDeleteBeforeParams{ - CreatedAtHorizon: time.Now().Add(-s.Config.RetentionPeriod), - Schema: s.Config.Schema, - }) - if err != nil { - return nil, err - } + return timeoututil.WithTimeoutV(ctx, s.Config.Timeout, s.Name+".runOnce", func(ctx context.Context) (*sqliteNotificationCleanerRunOnceResult, error) { + numDeleted, err := s.exec.NotificationDeleteBefore(ctx, &riverdriver.NotificationDeleteBeforeParams{ + CreatedAtHorizon: time.Now().Add(-s.Config.RetentionPeriod), + Schema: s.Config.Schema, + }) + if err != nil { + return nil, err + } - s.TestSignals.DeletedBatch.Signal(struct{}{}) + s.TestSignals.DeletedBatch.Signal(struct{}{}) - return &sqliteNotificationCleanerRunOnceResult{ - NumNotificationsDeleted: numDeleted, - }, nil + return &sqliteNotificationCleanerRunOnceResult{ + NumNotificationsDeleted: numDeleted, + }, nil + }) } diff --git a/internal/maintenance/sqlite_notification_cleaner_test.go b/internal/maintenance/sqlite_notification_cleaner_test.go index 348dacce..47238850 100644 --- a/internal/maintenance/sqlite_notification_cleaner_test.go +++ b/internal/maintenance/sqlite_notification_cleaner_test.go @@ -102,4 +102,15 @@ func TestSQLiteNotificationCleaner(t *testing.T) { startstoptest.Stress(ctx, t, cleaner) }) + + t.Run("TimeoutErrorIncludesOperation", func(t *testing.T) { + t.Parallel() + + cleaner, _ := setup(t) + cleaner.Config.Timeout = time.Nanosecond + + _, err := cleaner.runOnce(ctx) + require.ErrorContains(t, err, cleaner.Name+".runOnce timed out after 1ns") + require.ErrorIs(t, err, context.DeadlineExceeded) + }) } diff --git a/internal/notifier/notifier.go b/internal/notifier/notifier.go index bf12ff29..8b48fbdd 100644 --- a/internal/notifier/notifier.go +++ b/internal/notifier/notifier.go @@ -19,6 +19,7 @@ import ( "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/sliceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" ) type NotificationTopic string @@ -309,20 +310,19 @@ func (n *Notifier) listenerConnect(ctx context.Context, skipLock bool) error { return nil } - ctx, cancel := context.WithTimeout(ctx, listenerTimeout) - defer cancel() + return timeoututil.WithTimeout(ctx, listenerTimeout, n.Name+".listenerConnect", func(ctx context.Context) error { + n.Logger.DebugContext(ctx, n.Name+": Listener connecting") + if err := n.listener.Connect(ctx); err != nil { + if !errors.Is(err, context.Canceled) { + n.Logger.ErrorContext(ctx, n.Name+": Error connecting listener", "err", err) + } - n.Logger.DebugContext(ctx, n.Name+": Listener connecting") - if err := n.listener.Connect(ctx); err != nil { - if !errors.Is(err, context.Canceled) { - n.Logger.ErrorContext(ctx, n.Name+": Error connecting listener", "err", err) + return err } - return err - } - - n.isConnected = true - return nil + n.isConnected = true + return nil + }) } // Listens on a topic with an appropriate logging statement. Should be preferred @@ -331,15 +331,14 @@ func (n *Notifier) listenerConnect(ctx context.Context, skipLock bool) error { // Not protected by mutex because it doesn't modify any notifier state and the // underlying listener has a mutex around its operations. func (n *Notifier) listenerListen(ctx context.Context, topic NotificationTopic) error { - ctx, cancel := context.WithTimeout(ctx, listenerTimeout) - defer cancel() - - n.Logger.DebugContext(ctx, n.Name+": Listening on topic", "topic", topic) - if err := n.listener.Listen(ctx, string(topic)); err != nil { - return fmt.Errorf("error listening on topic %q: %w", topic, err) - } + return timeoututil.WithTimeout(ctx, listenerTimeout, n.Name+".listenerListen", func(ctx context.Context) error { + n.Logger.DebugContext(ctx, n.Name+": Listening on topic", "topic", topic) + if err := n.listener.Listen(ctx, string(topic)); err != nil { + return fmt.Errorf("error listening on topic %q: %w", topic, err) + } - return nil + return nil + }) } // Unlistens on a topic with an appropriate logging statement. Should be @@ -348,15 +347,14 @@ func (n *Notifier) listenerListen(ctx context.Context, topic NotificationTopic) // Not protected by mutex because it doesn't modify any notifier state and the // underlying listener has a mutex around its operations. func (n *Notifier) listenerUnlisten(ctx context.Context, topic NotificationTopic) error { - ctx, cancel := context.WithTimeout(ctx, listenerTimeout) - defer cancel() - - n.Logger.DebugContext(ctx, n.Name+": Unlistening on topic", "topic", topic) - if err := n.listener.Unlisten(ctx, string(topic)); err != nil { - return fmt.Errorf("error unlistening on topic %q: %w", topic, err) - } + return timeoututil.WithTimeout(ctx, listenerTimeout, n.Name+".listenerUnlisten", func(ctx context.Context) error { + n.Logger.DebugContext(ctx, n.Name+": Unlistening on topic", "topic", topic) + if err := n.listener.Unlisten(ctx, string(topic)); err != nil { + return fmt.Errorf("error unlistening on topic %q: %w", topic, err) + } - return nil + return nil + }) } // Enters a single blocking wait for notifications on the underlying listener. diff --git a/producer.go b/producer.go index 6c24baa3..11a6efee 100644 --- a/producer.go +++ b/producer.go @@ -28,6 +28,7 @@ import ( "github.com/riverqueue/river/rivershared/util/randutil" "github.com/riverqueue/river/rivershared/util/serviceutil" "github.com/riverqueue/river/rivershared/util/testutil" + "github.com/riverqueue/river/rivershared/util/timeoututil" "github.com/riverqueue/river/rivershared/util/timeutil" "github.com/riverqueue/river/rivertype" ) @@ -298,10 +299,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { return errors.Is(err, startstop.ErrStop) || strings.HasSuffix(err.Error(), "conn closed") || fetchCtx.Err() != nil } - fetchedQueue, err := func() (*rivertype.Queue, error) { - ctx, cancel := context.WithTimeout(fetchCtx, 10*time.Second) - defer cancel() - + fetchedQueue, err := timeoututil.WithTimeoutV(fetchCtx, 10*time.Second, p.Name+".StartWorkContext", func(ctx context.Context) (*rivertype.Queue, error) { p.Logger.DebugContext(ctx, p.Name+": Fetching initial queue settings", slog.String("queue", p.config.Queue)) return p.exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ Metadata: []byte("{}"), @@ -309,7 +307,7 @@ func (p *producer) StartWorkContext(fetchCtx, workCtx context.Context) error { Now: p.Time.NowOrNil(), Schema: p.config.Schema, }) - }() + }) if err != nil { stopped() if isExpectedShutdownError(err) { @@ -696,23 +694,22 @@ func (p *producer) finalizeShutdown(ctx context.Context) { ) attemptShutdown := func(timeout time.Duration) error { - ctx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - if err := p.pilot.ProducerShutdown(ctx, p.exec, &riverpilot.ProducerShutdownParams{ - ProducerID: p.id.Load(), - Queue: p.config.Queue, - Schema: p.config.Schema, - }); err != nil { - // Don't retry on these errors: - // - context.Canceled: parent context is canceled, so retrying with a new timeout won't help - // - ErrClosedPool: the database connection pool is closed, so retrying won't succeed - if errors.Is(err, context.Canceled) || errors.Is(err, riverdriver.ErrClosedPool) { - return nil + return timeoututil.WithTimeout(ctx, timeout, p.Name+".finalizeShutdown", func(ctx context.Context) error { + if err := p.pilot.ProducerShutdown(ctx, p.exec, &riverpilot.ProducerShutdownParams{ + ProducerID: p.id.Load(), + Queue: p.config.Queue, + Schema: p.config.Schema, + }); err != nil { + // Don't retry on these errors: + // - context.Canceled: parent context is canceled, so retrying with a new timeout won't help + // - ErrClosedPool: the database connection pool is closed, so retrying won't succeed + if errors.Is(err, context.Canceled) || errors.Is(err, riverdriver.ErrClosedPool) { + return nil + } + return err } - return err - } - return nil + return nil + }) } // Progressive retry with increasing timeouts: @@ -973,15 +970,12 @@ func (p *producer) pollForSettingChanges(ctx context.Context, wg *sync.WaitGroup case <-ctx.Done(): return case <-ticker.C: - updatedQueue, err := func() (*rivertype.Queue, error) { - ctx, cancel := context.WithTimeout(ctx, 10*time.Second) - defer cancel() - + updatedQueue, err := timeoututil.WithTimeoutV(ctx, 10*time.Second, p.Name+".pollForSettingChanges", func(ctx context.Context) (*rivertype.Queue, error) { return p.exec.QueueGet(ctx, &riverdriver.QueueGetParams{ Name: p.config.Queue, Schema: p.config.Schema, }) - }() + }) if err != nil { // Don't log if this is part of a standard shutdown. if !errors.Is(context.Cause(ctx), startstop.ErrStop) { @@ -1060,15 +1054,14 @@ func (p *producer) reportProducerStatusLoop(ctx context.Context, wg *sync.WaitGr } func (p *producer) reportProducerStatusOnce(ctx context.Context) { - ctx, cancel := context.WithTimeout(ctx, 10*time.Second) - defer cancel() - - p.Logger.DebugContext(ctx, p.Name+": Reporting producer status", slog.Int64("id", p.id.Load()), slog.String("queue", p.config.Queue)) - err := p.pilot.ProducerKeepAlive(ctx, p.exec, &riverdriver.ProducerKeepAliveParams{ - ID: p.id.Load(), - QueueName: p.config.Queue, - Schema: p.config.Schema, - StaleUpdatedAtHorizon: p.Time.Now().Add(-p.config.StaleProducerRetentionPeriod), + err := timeoututil.WithTimeout(ctx, 10*time.Second, p.Name+".reportProducerStatusOnce", func(ctx context.Context) error { + p.Logger.DebugContext(ctx, p.Name+": Reporting producer status", slog.Int64("id", p.id.Load()), slog.String("queue", p.config.Queue)) + return p.pilot.ProducerKeepAlive(ctx, p.exec, &riverdriver.ProducerKeepAliveParams{ + ID: p.id.Load(), + QueueName: p.config.Queue, + Schema: p.config.Schema, + StaleUpdatedAtHorizon: p.Time.Now().Add(-p.config.StaleProducerRetentionPeriod), + }) }) if err != nil && errors.Is(context.Cause(ctx), startstop.ErrStop) { return @@ -1101,15 +1094,15 @@ func (p *producer) reportQueueStatusLoop(ctx context.Context, wg *sync.WaitGroup } func (p *producer) reportQueueStatusOnce(ctx context.Context) { - ctx, cancel := context.WithTimeout(ctx, 10*time.Second) - defer cancel() - - p.Logger.DebugContext(ctx, p.Name+": Reporting queue status", slog.String("queue", p.config.Queue)) - _, err := p.exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ - Metadata: []byte("{}"), - Name: p.config.Queue, - Now: p.Time.NowOrNil(), - Schema: p.config.Schema, + err := timeoututil.WithTimeout(ctx, 10*time.Second, p.Name+".reportQueueStatusOnce", func(ctx context.Context) error { + p.Logger.DebugContext(ctx, p.Name+": Reporting queue status", slog.String("queue", p.config.Queue)) + _, err := p.exec.QueueCreateOrSetUpdatedAt(ctx, &riverdriver.QueueCreateOrSetUpdatedAtParams{ + Metadata: []byte("{}"), + Name: p.config.Queue, + Now: p.Time.NowOrNil(), + Schema: p.config.Schema, + }) + return err }) if err != nil && errors.Is(context.Cause(ctx), startstop.ErrStop) { return diff --git a/rivershared/util/dbutil/db_util.go b/rivershared/util/dbutil/db_util.go index 830ada96..9853d868 100644 --- a/rivershared/util/dbutil/db_util.go +++ b/rivershared/util/dbutil/db_util.go @@ -7,6 +7,7 @@ import ( "time" "github.com/riverqueue/river/riverdriver" + "github.com/riverqueue/river/rivershared/util/timeoututil" ) // SafeIdentifier returns a safely quoted identifier (e.g. a table or schem @@ -27,9 +28,6 @@ func SafeIdentifier(ident string) string { func RollbackWithoutCancel[TExec riverdriver.ExecutorTx](ctx context.Context, execTx TExec) error { ctxWithoutCancel := context.WithoutCancel(ctx) - ctx, cancel := context.WithTimeout(ctxWithoutCancel, 5*time.Second) - defer cancel() - // It might not be the worst idea to log an unexpected error on rollback // here instead of returning it. I had this in place initially, but there's // a number of common errors that need to be ignored like "conn closed", @@ -40,7 +38,7 @@ func RollbackWithoutCancel[TExec riverdriver.ExecutorTx](ctx context.Context, ex // function like `ShouldIgnoreRollbackError` that'd need a lot of plumbing // and it becomes questionable as to whether it's all worth it as Rollback // producing a non-standard error would be quite unusual. - return execTx.Rollback(ctx) + return timeoututil.WithTimeout(ctxWithoutCancel, 5*time.Second, "dbutil.RollbackWithoutCancel", execTx.Rollback) } // WithTx starts and commits a transaction on a driver executor around