Add RetryUntilSuccessOrExhausted and NonRetryableError to internal/common/util

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
+	}
+}

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.")
+	}
+}