From 0076c454cda9015f33de6864ca630e7ab543edd9 Mon Sep 17 00:00:00 2001 From: Maurice Yap Date: Mon, 3 Aug 2026 14:23:40 +0100 Subject: [PATCH 1/3] Add RetryUntilSuccessOrExhausted and NonRetryableError to internal/common/util Add a bounded-retry helpter which retries `performAction` up to `maxAttempts` times. It calls `onExhausted` with the last error if the budget runs out. It short-circuits immediately if the error is wrapped in the new `NonRetryableError`. This will be used in the ingestion pipeline to decide when to dead-leatter a message after retrying it. Signed-off-by: Maurice Yap --- internal/common/util/retry.go | 67 ++++++++++++ internal/common/util/retry_test.go | 169 +++++++++++++++++++++++++++++ 2 files changed, 236 insertions(+) diff --git a/internal/common/util/retry.go b/internal/common/util/retry.go index c688614e63e..1323622399d 100644 --- a/internal/common/util/retry.go +++ b/internal/common/util/retry.go @@ -1,9 +1,35 @@ package util import ( + "errors" + "github.com/armadaproject/armada/internal/common/armadacontext" ) +// NonRetryableError wraps an error to signal that it should not be retried, even if +// remaining attempts are available. Sinks can wrap errors with this to short-circuit +// RetryUntilSuccessOrExhausted straight to onExhausted. +type NonRetryableError struct { + err error +} + +func NewNonRetryableError(err error) *NonRetryableError { + return &NonRetryableError{err: err} +} + +func (e *NonRetryableError) Error() string { + return e.err.Error() +} + +func (e *NonRetryableError) Unwrap() error { + return e.err +} + +func IsNonRetryable(err error) bool { + var nonRetryable *NonRetryableError + return errors.As(err, &nonRetryable) +} + func RetryUntilSuccess(ctx *armadacontext.Context, performAction func() error, onError func(error)) { for { select { @@ -19,3 +45,44 @@ 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 wrapped with NewNonRetryableError, remaining +// attempts are skipped and onExhausted is called immediately with that error. +func RetryUntilSuccessOrExhausted( + ctx *armadacontext.Context, + maxAttempts int, + performAction func() error, + onError func(attempt int, err error), + onExhausted func(lastErr error), +) bool { + var lastErr error +attempts: + for attempt := 1; attempt <= maxAttempts; attempt++ { + select { + case <-ctx.Done(): + return false + default: + err := performAction() + if err == nil { + return true + } + lastErr = err + if IsNonRetryable(err) { + 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..65bfd20d832 100644 --- a/internal/common/util/retry_test.go +++ b/internal/common/util/retry_test.go @@ -84,3 +84,172 @@ 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 := NewNonRetryableError(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 := NewNonRetryableError(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, IsNonRetryable(nil)) + assert.False(t, IsNonRetryable(fmt.Errorf("plain error"))) + assert.True(t, IsNonRetryable(NewNonRetryableError(fmt.Errorf("wrapped error")))) + assert.True(t, IsNonRetryable(fmt.Errorf("outer: %w", NewNonRetryableError(fmt.Errorf("inner"))))) +} + +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.") + } +} From e8c3d9e542b1e96a54fdcb80fdf97cb51f82bc46 Mon Sep 17 00:00:00 2001 From: Maurice Yap Date: Mon, 3 Aug 2026 15:20:05 +0100 Subject: [PATCH 2/3] comments Signed-off-by: Maurice Yap --- internal/common/util/retry.go | 8 +++++++- internal/common/util/retry_test.go | 31 ++++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/internal/common/util/retry.go b/internal/common/util/retry.go index 1323622399d..76a2d711a9d 100644 --- a/internal/common/util/retry.go +++ b/internal/common/util/retry.go @@ -2,6 +2,7 @@ package util import ( "errors" + "fmt" "github.com/armadaproject/armada/internal/common/armadacontext" ) @@ -14,6 +15,9 @@ type NonRetryableError struct { } func NewNonRetryableError(err error) *NonRetryableError { + if err == nil { + err = errors.New("non-retryable error") + } return &NonRetryableError{err: err} } @@ -53,6 +57,8 @@ func RetryUntilSuccess(ctx *armadacontext.Context, performAction func() error, o // must distinguish "gave up due to shutdown" from "exhausted attempts" via ctx.Err(). // If performAction returns an error wrapped with NewNonRetryableError, 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, @@ -60,7 +66,7 @@ func RetryUntilSuccessOrExhausted( onError func(attempt int, err error), onExhausted func(lastErr error), ) bool { - var lastErr error + lastErr := fmt.Errorf("no attempts were made: maxAttempts was %d", maxAttempts) attempts: for attempt := 1; attempt <= maxAttempts; attempt++ { select { diff --git a/internal/common/util/retry_test.go b/internal/common/util/retry_test.go index 65bfd20d832..159a3d02fd5 100644 --- a/internal/common/util/retry_test.go +++ b/internal/common/util/retry_test.go @@ -225,6 +225,37 @@ func TestIsNonRetryable(t *testing.T) { assert.True(t, IsNonRetryable(fmt.Errorf("outer: %w", NewNonRetryableError(fmt.Errorf("inner"))))) } +func TestNewNonRetryableError_NilCauseDoesNotPanic(t *testing.T) { + err := NewNonRetryableError(nil) + + assert.NotPanics(t, func() { + _ = err.Error() + }) +} + +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() From b5dfb494b5477851991f3a13eeaf8ee6d4df0aff Mon Sep 17 00:00:00 2001 From: Maurice Yap Date: Tue, 4 Aug 2026 16:37:29 +0100 Subject: [PATCH 3/3] ErrNonRetryable Signed-off-by: Maurice Yap --- internal/common/util/retry.go | 34 ++++++------------------------ internal/common/util/retry_test.go | 21 ++++++------------ 2 files changed, 13 insertions(+), 42 deletions(-) diff --git a/internal/common/util/retry.go b/internal/common/util/retry.go index 76a2d711a9d..b566b735b43 100644 --- a/internal/common/util/retry.go +++ b/internal/common/util/retry.go @@ -7,32 +7,10 @@ import ( "github.com/armadaproject/armada/internal/common/armadacontext" ) -// NonRetryableError wraps an error to signal that it should not be retried, even if -// remaining attempts are available. Sinks can wrap errors with this to short-circuit -// RetryUntilSuccessOrExhausted straight to onExhausted. -type NonRetryableError struct { - err error -} - -func NewNonRetryableError(err error) *NonRetryableError { - if err == nil { - err = errors.New("non-retryable error") - } - return &NonRetryableError{err: err} -} - -func (e *NonRetryableError) Error() string { - return e.err.Error() -} - -func (e *NonRetryableError) Unwrap() error { - return e.err -} - -func IsNonRetryable(err error) bool { - var nonRetryable *NonRetryableError - return errors.As(err, &nonRetryable) -} +// 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 { @@ -55,7 +33,7 @@ func RetryUntilSuccess(ctx *armadacontext.Context, performAction func() error, o // 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 wrapped with NewNonRetryableError, remaining +// 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. @@ -78,7 +56,7 @@ attempts: return true } lastErr = err - if IsNonRetryable(err) { + if errors.Is(err, ErrNonRetryable) { break attempts } onError(attempt, err) diff --git a/internal/common/util/retry_test.go b/internal/common/util/retry_test.go index 159a3d02fd5..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" @@ -169,7 +170,7 @@ func TestRetryUntilSuccessOrExhausted_NonRetryableShortCircuits(t *testing.T) { defer cancel() dummyErr := fmt.Errorf("dummy error") - nonRetryableErr := NewNonRetryableError(dummyErr) + nonRetryableErr := fmt.Errorf("%w: %w", ErrNonRetryable, dummyErr) errorCount := 0 var exhaustedErr error @@ -193,7 +194,7 @@ func TestRetryUntilSuccessOrExhausted_NonRetryableAfterSomeRetries(t *testing.T) defer cancel() retryableErr := fmt.Errorf("retryable error") - nonRetryableErr := NewNonRetryableError(fmt.Errorf("non-retryable error")) + nonRetryableErr := fmt.Errorf("%w: %w", ErrNonRetryable, fmt.Errorf("non-retryable error")) ch := make(chan error, 3) ch <- retryableErr @@ -219,18 +220,10 @@ func TestRetryUntilSuccessOrExhausted_NonRetryableAfterSomeRetries(t *testing.T) } func TestIsNonRetryable(t *testing.T) { - assert.False(t, IsNonRetryable(nil)) - assert.False(t, IsNonRetryable(fmt.Errorf("plain error"))) - assert.True(t, IsNonRetryable(NewNonRetryableError(fmt.Errorf("wrapped error")))) - assert.True(t, IsNonRetryable(fmt.Errorf("outer: %w", NewNonRetryableError(fmt.Errorf("inner"))))) -} - -func TestNewNonRetryableError_NilCauseDoesNotPanic(t *testing.T) { - err := NewNonRetryableError(nil) - - assert.NotPanics(t, func() { - _ = err.Error() - }) + 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) {