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 {