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