diff --git a/internal/common/util/retry.go b/internal/common/util/retry.go index c688614e63e..b566b735b43 100644 --- a/internal/common/util/retry.go +++ b/internal/common/util/retry.go @@ -1,9 +1,17 @@ package util import ( + "errors" + "fmt" + "github.com/armadaproject/armada/internal/common/armadacontext" ) +// ErrNonRetryable signals that an error should not be retried, even if remaining +// attempts are available. Sinks can wrap errors with fmt.Errorf("%w: %w", ErrNonRetryable, err) +// to short-circuit RetryUntilSuccessOrExhausted straight to onExhausted. +var ErrNonRetryable = errors.New("non-retryable error") + func RetryUntilSuccess(ctx *armadacontext.Context, performAction func() error, onError func(error)) { for { select { @@ -19,3 +27,46 @@ func RetryUntilSuccess(ctx *armadacontext.Context, performAction func() error, o } } } + +// RetryUntilSuccessOrExhausted behaves like RetryUntilSuccess but gives up after +// maxAttempts consecutive failures, calling onExhausted with the last error +// instead of continuing. Returns true on eventual success, false otherwise. +// onExhausted is NOT called if ctx is cancelled first (shutdown case) - callers +// must distinguish "gave up due to shutdown" from "exhausted attempts" via ctx.Err(). +// If performAction returns an error wrapping ErrNonRetryable, remaining +// attempts are skipped and onExhausted is called immediately with that error. +// If maxAttempts is non-positive, performAction is never called and onExhausted +// is called with a placeholder error describing this. +func RetryUntilSuccessOrExhausted( + ctx *armadacontext.Context, + maxAttempts int, + performAction func() error, + onError func(attempt int, err error), + onExhausted func(lastErr error), +) bool { + lastErr := fmt.Errorf("no attempts were made: maxAttempts was %d", maxAttempts) +attempts: + for attempt := 1; attempt <= maxAttempts; attempt++ { + select { + case <-ctx.Done(): + return false + default: + err := performAction() + if err == nil { + return true + } + lastErr = err + if errors.Is(err, ErrNonRetryable) { + break attempts + } + onError(attempt, err) + } + } + select { + case <-ctx.Done(): + return false + default: + onExhausted(lastErr) + return false + } +} diff --git a/internal/common/util/retry_test.go b/internal/common/util/retry_test.go index 2ad6ea4b300..052262d4b7d 100644 --- a/internal/common/util/retry_test.go +++ b/internal/common/util/retry_test.go @@ -1,6 +1,7 @@ package util import ( + "errors" "fmt" "testing" "time" @@ -84,3 +85,195 @@ func TestSucceedsAfterFailures(t *testing.T) { assert.Equal(t, 5, errorCount) } + +func TestRetryUntilSuccessOrExhausted_SucceedsFirstTry(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 1*time.Second) + defer cancel() + + errorCount := 0 + exhaustedCount := 0 + + ok := RetryUntilSuccessOrExhausted( + ctx, + 5, + func() error { + return nil + }, + func(attempt int, err error) { errorCount++ }, + func(lastErr error) { exhaustedCount++ }, + ) + + assert.True(t, ok) + assert.Equal(t, 0, errorCount) + assert.Equal(t, 0, exhaustedCount) +} + +func TestRetryUntilSuccessOrExhausted_SucceedsWithinBudget(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 1*time.Second) + defer cancel() + + ch := make(chan error, 3) + ch <- fmt.Errorf("dummy error 1") + ch <- fmt.Errorf("dummy error 2") + ch <- nil + + errorCount := 0 + exhaustedCount := 0 + + ok := RetryUntilSuccessOrExhausted( + ctx, + 5, + func() error { + return <-ch + }, + func(attempt int, err error) { errorCount++ }, + func(lastErr error) { exhaustedCount++ }, + ) + + assert.True(t, ok) + assert.Equal(t, 2, errorCount) + assert.Equal(t, 0, exhaustedCount) +} + +func TestRetryUntilSuccessOrExhausted_ExhaustsBudget(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 1*time.Second) + defer cancel() + + dummyErr := fmt.Errorf("dummy error") + errorCount := 0 + var exhaustedErr error + + ok := RetryUntilSuccessOrExhausted( + ctx, + 5, + func() error { + return dummyErr + }, + func(attempt int, err error) { errorCount++ }, + func(lastErr error) { exhaustedErr = lastErr }, + ) + + assert.False(t, ok) + assert.Equal(t, 5, errorCount) + assert.Equal(t, dummyErr, exhaustedErr) + + select { + case <-ctx.Done(): + t.Fatalf("Function exhausted budget but context was also cancelled unexpectedly.") + default: + break + } +} + +func TestRetryUntilSuccessOrExhausted_NonRetryableShortCircuits(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 1*time.Second) + defer cancel() + + dummyErr := fmt.Errorf("dummy error") + nonRetryableErr := fmt.Errorf("%w: %w", ErrNonRetryable, dummyErr) + errorCount := 0 + var exhaustedErr error + + ok := RetryUntilSuccessOrExhausted( + ctx, + 5, + func() error { + return nonRetryableErr + }, + func(attempt int, err error) { errorCount++ }, + func(lastErr error) { exhaustedErr = lastErr }, + ) + + assert.False(t, ok) + assert.Equal(t, 0, errorCount) + assert.Equal(t, nonRetryableErr, exhaustedErr) +} + +func TestRetryUntilSuccessOrExhausted_NonRetryableAfterSomeRetries(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 1*time.Second) + defer cancel() + + retryableErr := fmt.Errorf("retryable error") + nonRetryableErr := fmt.Errorf("%w: %w", ErrNonRetryable, fmt.Errorf("non-retryable error")) + + ch := make(chan error, 3) + ch <- retryableErr + ch <- retryableErr + ch <- nonRetryableErr + + errorCount := 0 + var exhaustedErr error + + ok := RetryUntilSuccessOrExhausted( + ctx, + 5, + func() error { + return <-ch + }, + func(attempt int, err error) { errorCount++ }, + func(lastErr error) { exhaustedErr = lastErr }, + ) + + assert.False(t, ok) + assert.Equal(t, 2, errorCount) + assert.Equal(t, nonRetryableErr, exhaustedErr) +} + +func TestIsNonRetryable(t *testing.T) { + assert.False(t, errors.Is(nil, ErrNonRetryable)) + assert.False(t, errors.Is(fmt.Errorf("plain error"), ErrNonRetryable)) + assert.True(t, errors.Is(fmt.Errorf("%w: %w", ErrNonRetryable, fmt.Errorf("wrapped error")), ErrNonRetryable)) + assert.True(t, errors.Is(fmt.Errorf("outer: %w", fmt.Errorf("%w: %w", ErrNonRetryable, fmt.Errorf("inner"))), ErrNonRetryable)) +} + +func TestRetryUntilSuccessOrExhausted_NonPositiveMaxAttempts(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 1*time.Second) + defer cancel() + + performed := false + var exhaustedErr error + + ok := RetryUntilSuccessOrExhausted( + ctx, + 0, + func() error { + performed = true + return nil + }, + func(attempt int, err error) {}, + func(lastErr error) { exhaustedErr = lastErr }, + ) + + assert.False(t, ok) + assert.False(t, performed) + assert.Error(t, exhaustedErr) +} + +func TestRetryUntilSuccessOrExhausted_CancelledMidRetry(t *testing.T) { + ctx, cancel := armadacontext.WithTimeout(armadacontext.Background(), 50*time.Millisecond) + defer cancel() + + dummyErr := fmt.Errorf("dummy error") + exhaustedCount := 0 + + ok := RetryUntilSuccessOrExhausted( + ctx, + 1000000, + func() error { + time.Sleep(10 * time.Millisecond) + return dummyErr + }, + func(attempt int, err error) {}, + func(lastErr error) { exhaustedCount++ }, + ) + + assert.False(t, ok) + assert.Equal(t, 0, exhaustedCount) + + select { + case <-ctx.Done(): + break + default: + t.Fatalf("Expected context to be done.") + } +}