From 89c863252decbc91176e14632a3285cbd9b0ad5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=A1s=20Pernas=20Maradei?= Date: Tue, 1 Sep 2026 15:42:42 +0200 Subject: [PATCH] retry: add cancelable retriers that cut their waits short on shutdown The retrier used one context for three jobs: the deadline for the operation, the abort signal for the wait between attempts, and the loop's stop condition. Callers that need the operation to run to completion had to pass a context that is never cancelled, which also made every wait uninterruptible. Shutting down then meant sitting out the remaining backoff interval, once per in-flight operation. NewCancelableConstantRetrier and NewCancelableExponentialRetrier take a second context that governs the waits only. It never reaches the operation, so an attempt that has begun always finishes, and the operation still runs at least once. On stop, Retry returns the last attempt's value plus an error wrapping stop.Err(), letting a caller tell a shutdown from a genuine failure. The existing constructors leave the stop context nil and are unchanged: same uninterruptible initial sleep, same WithContextTimeout semantics. Co-Authored-By: Claude Opus 5 (1M context) --- retry.go | 63 ++++++++++++++++-- retry_test.go | 175 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 234 insertions(+), 4 deletions(-) diff --git a/retry.go b/retry.go index e45eb5a..5586a10 100644 --- a/retry.go +++ b/retry.go @@ -2,6 +2,7 @@ package utils import ( "context" + "fmt" "log/slog" "time" @@ -35,24 +36,50 @@ func newRetrier[T any](config retrybuilder[T], b backoff.BackOff) *retrier[T] { } func (c *retrier[T]) Retry(op RetriableOp[T]) (result T, err error) { - time.Sleep(c.config.wait) - ctx, cancel := c.config.context() defer cancel() + + stop := c.waitOrStop(ctx, c.config.wait) + operation := func() (err error) { result, err = op(ctx) return err } - err = backoff.RetryNotify(operation, - backoff.WithContext(c.backoff, ctx), + err = backoff.RetryNotify( + operation, + backoff.WithContext(c.backoff, stop), func(err error, d time.Duration) { fields := append(c.config.fields, "error", err, slog.String("wait", d.String())) c.config.logger.Log(ctx, c.config.lvl, "retrying operation", fields...) }, ) + // backoff always attempts once, so result holds a real value. Report the stop so the caller can + // tell a shutdown from a genuine failure. + if err != nil && c.config.stop != nil && c.config.stop.Err() != nil { + return result, fmt.Errorf("retry stopped: %w", c.config.stop.Err()) + } return result, err } +// waitOrStop waits for d and returns the context governing the waits between attempts. A plain +// retrier sleeps through d and is governed by ctx, which is where WithContextTimeout takes effect; +// a cancelable one cuts d short once its stop context is done, and is governed by that. +func (c *retrier[T]) waitOrStop(ctx context.Context, d time.Duration) context.Context { + if c.config.stop == nil { + time.Sleep(d) + return ctx + } + if d > 0 { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-c.config.stop.Done(): + case <-t.C: + } + } + return c.config.stop +} + func (c *retrier[T]) Reset() { c.backoff.Reset() } @@ -90,6 +117,7 @@ type RetryBuilder[T any] interface { type retrybuilder[T any] struct { ctx context.Context + stop context.Context maxRetries uint64 interval time.Duration logger *slog.Logger @@ -158,6 +186,18 @@ type constant[T any] struct { // NewConstantRetrier builder of a ConstantRetrier which is a wrapper of backoff.ConstantBackOff func NewConstantRetrier[T any](ctx context.Context, logger *slog.Logger) RetryBuilder[T] { + return newConstant[T](ctx, logger) +} + +// NewCancelableConstantRetrier builds a ConstantRetrier whose waits end as soon as stop is done. +// See NewCancelableExponentialRetrier for the contract. +func NewCancelableConstantRetrier[T any](stop, ctx context.Context, logger *slog.Logger) RetryBuilder[T] { + c := newConstant[T](ctx, logger) + c.stop = stop + return c +} + +func newConstant[T any](ctx context.Context, logger *slog.Logger) *constant[T] { c := &constant[T]{} c.retrybuilder = newRetryBuilder(ctx, c, logger) return c @@ -189,6 +229,21 @@ type ExponentialRetryBuilder[T any] interface { // NewExponentialRetrier builder of an ExponentialRetrier which is a wrapper of backoff.ExponentialBackOff func NewExponentialRetrier[T any](ctx context.Context, logger *slog.Logger) ExponentialRetryBuilder[T] { + return newExponential[T](ctx, logger) +} + +// NewCancelableExponentialRetrier builds an ExponentialRetrier whose waits end as soon as stop is +// done, so a long backoff interval never holds up a shutdown. stop never reaches the operation, so +// an attempt that has begun always finishes and the operation runs at least once. On stop, Retry +// returns the last attempt's value and an error wrapping stop.Err(); a successful attempt still +// wins. Does not compose with WithContextTimeout. +func NewCancelableExponentialRetrier[T any](stop, ctx context.Context, logger *slog.Logger) ExponentialRetryBuilder[T] { + c := newExponential[T](ctx, logger) + c.stop = stop + return c +} + +func newExponential[T any](ctx context.Context, logger *slog.Logger) *exponential[T] { c := &exponential[T]{ multiplier: backoff.DefaultMultiplier, maxInterval: backoff.DefaultMaxInterval, diff --git a/retry_test.go b/retry_test.go index c1ac314..f1b0363 100644 --- a/retry_test.go +++ b/retry_test.go @@ -2,6 +2,7 @@ package utils import ( "context" + "errors" "fmt" "log/slog" "math" @@ -225,6 +226,180 @@ func TestExponentialRetrierWithMaxInterval(t *testing.T) { assertInDelta(t, 250*time.Millisecond, intervals[3], tolerance, "fourth interval should be capped at ~250ms") } +func TestCancelableRetrierStopsWhileWaiting(t *testing.T) { + ctx, log := newRetryFixture() + stop, cancel := context.WithCancel(context.Background()) + defer cancel() + + var attempts int + retrier := NewCancelableConstantRetrier[int](stop, ctx, log). + WithLogging(slog.LevelInfo). + WithInterval(time.Hour). // the test hangs if the wait is not interrupted + WithMaxRetries(0). + Build() + + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + start := time.Now() + ret, err := retrier.Retry(func(context.Context) (int, error) { + attempts++ + return 7, fmt.Errorf("boom") + }) + + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("Retry took %v, want it to end as soon as the stop context is cancelled", elapsed) + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want it to wrap context.Canceled", err) + } + if ret != 7 { + t.Fatalf("ret = %d, want the last attempt's value 7", ret) + } + if attempts != 1 { + t.Fatalf("attempts = %d, want 1", attempts) + } +} + +func TestCancelableRetrierAlwaysAttemptsOnce(t *testing.T) { + ctx, log := newRetryFixture() + stop, cancel := context.WithCancel(context.Background()) + cancel() + + var attempts int + retrier := NewCancelableConstantRetrier[int](stop, ctx, log). + WithInitialWait(time.Hour). // the initial wait is interruptible too + WithInterval(time.Hour). + Build() + + start := time.Now() + ret, err := retrier.Retry(func(context.Context) (int, error) { + attempts++ + return 3, fmt.Errorf("boom") + }) + + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("Retry took %v, want the initial wait to be interrupted", elapsed) + } + if attempts != 1 { + t.Fatalf("attempts = %d, want the operation to always run once", attempts) + } + if ret != 3 { + t.Fatalf("ret = %d, want the last attempt's value 3", ret) + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want it to wrap context.Canceled", err) + } +} + +func TestCancelableRetrierDoesNotCancelTheOperation(t *testing.T) { + ctx, log := newRetryFixture() + stop, cancel := context.WithCancel(context.Background()) + cancel() + + retrier := NewCancelableConstantRetrier[int](stop, ctx, log).Build() + + var opCtxErr error + ret, err := retrier.Retry(func(opCtx context.Context) (int, error) { + opCtxErr = opCtx.Err() + return 1, nil + }) + if opCtxErr != nil { + t.Fatalf("operation context error = %v, want the stop context not to reach the operation", opCtxErr) + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if ret != 1 { + t.Fatalf("ret = %d, want a successful attempt to win over the stop context", ret) + } +} + +func TestCancelableExponentialRetrierStopsWhileWaiting(t *testing.T) { + ctx, log := newRetryFixture() + stop, cancel := context.WithCancel(context.Background()) + defer cancel() + + var attempts int + retrier := NewCancelableExponentialRetrier[int](stop, ctx, log). + WithMultiplier(2). + WithMaxInterval(time.Hour). + WithMaxElapsedTime(0). + WithInterval(time.Hour). + WithMaxRetries(0). + Build() + + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + start := time.Now() + ret, err := retrier.Retry(func(context.Context) (int, error) { + attempts++ + return 9, fmt.Errorf("boom") + }) + + if elapsed := time.Since(start); elapsed > 5*time.Second { + t.Fatalf("Retry took %v, want it to end as soon as the stop context is cancelled", elapsed) + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want it to wrap context.Canceled", err) + } + if ret != 9 { + t.Fatalf("ret = %d, want the last attempt's value 9", ret) + } + if attempts != 1 { + t.Fatalf("attempts = %d, want 1", attempts) + } +} + +func TestCancelableRetrierReportsTheOperationErrorWhenNotStopped(t *testing.T) { + ctx, log := newRetryFixture() + stop, cancel := context.WithCancel(context.Background()) + defer cancel() + + retrier := NewCancelableConstantRetrier[int](stop, ctx, log). + WithInterval(time.Millisecond). + WithMaxRetries(2). + Build() + + ret, err := retrier.Retry(func(context.Context) (int, error) { + return 0, fmt.Errorf("boom") + }) + if errors.Is(err, context.Canceled) { + t.Fatalf("err = %v, want the operation error, not a stop", err) + } + if err == nil { + t.Fatal("expected error, got nil") + } + if ret != 0 { + t.Fatalf("ret = %d, want 0", ret) + } +} + +func TestPlainRetrierKeepsAnUninterruptibleInitialWait(t *testing.T) { + log := slog.New(slog.NewTextHandler(os.Stdout, nil)) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + retrier := NewConstantRetrier[int](ctx, log). + WithInitialWait(100 * time.Millisecond). + Build() + + start := time.Now() + if _, err := retrier.Retry(func(context.Context) (int, error) { + return 1, nil + }); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if elapsed := time.Since(start); elapsed < 100*time.Millisecond { + t.Fatalf("initial wait took %v, want the plain retrier to keep sleeping through it", elapsed) + } +} + func assertInDelta(t *testing.T, want, got, delta time.Duration, msg string) { t.Helper() if d := time.Duration(math.Abs(float64(got - want))); d > delta {