retry: add cancelable retriers that cut their waits short on shutdown

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,

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 {