Merge pull request #38 from maximebeckman/main

feat: implement WithInfiniteRetry() option

retrier/retrier.go

@@ -11,11 +11,12 @@ import (
 // Retrier implements the "retriable" resiliency pattern, abstracting out the process of retrying a failed action
 // a certain number of times with an optional back-off between each retry.
 type Retrier struct {
-	backoff []time.Duration
-	class   Classifier
-	jitter  float64
-	rand    *rand.Rand
-	randMu  sync.Mutex
+	backoff       []time.Duration
+	infiniteRetry bool
+	class         Classifier
+	jitter        float64
+	rand          *rand.Rand
+	randMu        sync.Mutex
 }
 
 // New constructs a Retrier with the given backoff pattern and classifier. The length of the backoff pattern
@@ -34,9 +35,17 @@ func New(backoff []time.Duration, class Classifier) *Retrier {
 	}
 }
 
+// WithInfiniteRetry set the retrier to loop infinitely on the last backoff duration. Using this option,
+// the program will not exit until the retried function has been executed successfully.
+// WARNING : This may run indefinitely.
+func (r *Retrier) WithInfiniteRetry() *Retrier {
+	r.infiniteRetry = true
+	return r
+}
+
 // Run executes the given work function by executing RunCtx without context.Context.
 func (r *Retrier) Run(work func() error) error {
-	return r.RunCtx(context.Background(), func(ctx context.Context) error {
+	return r.RunFn(context.Background(), func(c context.Context, r int) error {
 		// never use ctx
 		return work()
 	})
@@ -48,15 +57,27 @@ func (r *Retrier) Run(work func() error) error {
 // before retrying. If the total number of retries is exceeded then the return value of the work function
 // is returned to the caller regardless.
 func (r *Retrier) RunCtx(ctx context.Context, work func(ctx context.Context) error) error {
+	return r.RunFn(ctx, func(c context.Context, r int) error {
+		return work(c)
+	})
+}
+
+// RunFn executes the given work function, then classifies its return value based on the classifier used
+// to construct the Retrier. If the result is Succeed or Fail, the return value of the work function is
+// returned to the caller. If the result is Retry, then Run sleeps according to the backoff policy
+// before retrying. If the total number of retries is exceeded then the return value of the work function
+// is returned to the caller regardless. The work function takes 2 args, the context and
+// the number of attempted retries.
+func (r *Retrier) RunFn(ctx context.Context, work func(ctx context.Context, retries int) error) error {
 	retries := 0
 	for {
-		ret := work(ctx)
+		ret := work(ctx, retries)
 
 		switch r.class.Classify(ret) {
 		case Succeed, Fail:
 			return ret
 		case Retry:
-			if retries >= len(r.backoff) {
+			if !r.infiniteRetry && retries >= len(r.backoff) {
 				return ret
 			}
 
@@ -84,6 +105,9 @@ func (r *Retrier) sleep(ctx context.Context, timer *time.Timer) error {
 }
 
 func (r *Retrier) calcSleep(i int) time.Duration {
+	if i >= len(r.backoff) {
+		i = len(r.backoff) - 1
+	}
 	// lock unsafe rand prng
 	r.randMu.Lock()
 	defer r.randMu.Unlock()

retrier/retrier_test.go

@@ -85,6 +85,74 @@ func TestRetrierCtx(t *testing.T) {
 	}
 }
 
+func TestRetrierCtxError(t *testing.T) {
+	ctx := context.Background()
+	r := New([]time.Duration{0, 10 * time.Millisecond}, nil)
+	errExpected := []error{errFoo, errFoo, errBar, errBaz}
+	retries := 0
+	err := r.RunCtx(ctx, func(ctx context.Context) error {
+		if retries >= len(errExpected) {
+			return nil
+		}
+		err := errExpected[retries]
+		retries++
+		return err
+	})
+	if err != errBar {
+		t.Error(err)
+	}
+}
+
+func TestRetrierRunFnError(t *testing.T) {
+	ctx := context.Background()
+	r := New([]time.Duration{0, 10 * time.Millisecond}, nil)
+	errExpected := []error{errFoo, errFoo, errBar, errBaz}
+
+	err := r.RunFn(ctx, func(ctx context.Context, retries int) error {
+		if retries >= len(errExpected) {
+			return nil
+		}
+		return errExpected[retries]
+	})
+	if err != errBar {
+		t.Error(err)
+	}
+}
+
+func TestRetrierCtxWithInfinite(t *testing.T) {
+	ctx := context.Background()
+	r := New([]time.Duration{0, 10 * time.Millisecond}, nil).WithInfiniteRetry()
+	errExpected := []error{errFoo, errFoo, errFoo, errBar, errBaz}
+	retries := 0
+	err := r.RunCtx(ctx, func(ctx context.Context) error {
+		if retries >= len(errExpected) {
+			return nil
+		}
+		err := errExpected[retries]
+		retries++
+		return err
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+func TestRetrierRunFnWithInfinite(t *testing.T) {
+	ctx := context.Background()
+	r := New([]time.Duration{0, 10 * time.Millisecond}, nil).WithInfiniteRetry()
+	errExpected := []error{errFoo, errFoo, errFoo, errBar, errBaz}
+
+	err := r.RunFn(ctx, func(ctx context.Context, retries int) error {
+		if retries >= len(errExpected) {
+			return nil
+		}
+		return errExpected[retries]
+	})
+	if err != nil {
+		t.Error(err)
+	}
+}
+
 func TestRetrierNone(t *testing.T) {
 	r := New(nil, nil)