feat: retry with context (#63)

## Description

Retry with context

Author: AustinAbro321

helpers/misc.go

@@ -4,6 +4,8 @@
 package helpers
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"maps"
 	"math"
@@ -18,26 +20,64 @@ func BoolPtr(b bool) *bool {
 	return &b
 }
 
-// Retry will retry a function until it succeeds or the timeout is reached. timeout == 2^attempt * delay.
-func Retry(fn func() error, retries int, delay time.Duration, logger func(format string, args ...any)) error {
+// RetryWithContext retries a function until it succeeds, the timeout is reached, or the context is done.
+// The delay between attempts increases exponentially as (2^(attempt-1)) * delay.
+// For example, with a delay of one second and three attempts, the timing would be:
+// - First attempt: immediate
+// - Second attempt: after one second
+// - Third attempt: after two seconds
+func RetryWithContext(ctx context.Context, fn func() error, attempts int, delay time.Duration, logger func(format string, args ...any)) error {
+	if attempts < 1 {
+		return errors.New("invalid number of attempts, must be at least 1")
+	}
 	var err error
-	for r := 0; r < retries; r++ {
-		err = fn()
-		if err == nil {
-			break
-		}
+	for r := 0; r < attempts; r++ {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			err = fn()
+			if err == nil {
+				return nil
+			}
+
+			logger("Attempt (%d/%d) failed with: %s", r+1, attempts, err.Error())
+
+			// No reason to wait when we aren't going to retry again
+			if r+1 == attempts {
+				return err
+			}
 
-		pow := math.Pow(2, float64(r))
-		backoff := delay * time.Duration(pow)
+			pow := math.Pow(2, float64(r))
+			backoff := delay * time.Duration(pow)
 
-		logger("Retrying (%d/%d) in %s: %s", r+1, retries, backoff, err.Error())
+			logger("Retrying in %s", backoff)
 
-		time.Sleep(backoff)
+			timer := time.NewTimer(backoff)
+			select {
+			case <-timer.C:
+			case <-ctx.Done():
+				if !timer.Stop() {
+					<-timer.C
+				}
+				return ctx.Err()
+			}
+		}
 	}
 
 	return err
 }
 
+// Retry retries a function until it succeeds, the timeout is reached, or the context is done.
+// The delay between attempts increases exponentially as (2^(attempt-1)) * delay.
+// For example, with a delay of one second and three attempts, the timing would be:
+// - First attempt: immediate
+// - Second attempt: after one second
+// - Third attempt: after two seconds
+func Retry(fn func() error, attempts int, delay time.Duration, logger func(format string, args ...any)) error {
+	return RetryWithContext(context.Background(), fn, attempts, delay, logger)
+}
+
 // TransformMapKeys takes a map and transforms its keys using the provided function.
 func TransformMapKeys[T any](m map[string]T, transform func(string) string) (r map[string]T) {
 	r = map[string]T{}

helpers/misc_test.go

@@ -4,10 +4,12 @@
 package helpers
 
 import (
+	"context"
 	"errors"
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -46,33 +48,109 @@ func (suite *TestMiscSuite) SetupSuite() {
 	}
 }
 
-func (suite *TestMiscSuite) TestRetry() {
-	var count int
-	countFn := func() error {
-		count++
-		if count < 4 {
-			return errors.New("count exceeded")
+func TestRetry(t *testing.T) {
+	t.Run("RetriesWhenThereAreFailures", func(t *testing.T) {
+		count := 0
+		logCount := 0
+		returnedErr := errors.New("always fail")
+		countFn := func() error {
+			count++
+			return returnedErr
 		}
-		return nil
-	}
-	var logCount int
-	loggerFn := func(_ string, _ ...any) {
-		logCount++
-	}
+		loggerFn := func(_ string, _ ...any) {
+			logCount++
+		}
+
+		err := Retry(countFn, 3, 0, loggerFn)
+		require.ErrorIs(t, err, returnedErr)
+		require.Equal(t, 3, count)
+		require.Equal(t, 5, logCount)
+	})
+
+	t.Run("ContextCancellationBeforeStart", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		count := 0
+		fn := func() error {
+			count++
+			return errors.New("Never here since context got cancelled")
+		}
+		logger := func(_ string, _ ...any) {}
+
+		waitThatsNotCalled := 1000000 * time.Minute
+		err := RetryWithContext(ctx, fn, 5, waitThatsNotCalled, logger)
+		require.Equal(t, 0, count)
+		require.ErrorIs(t, err, context.Canceled)
+	})
+
+	t.Run("ContextCancellationDuringExecution", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+
+		count := 0
+		fn := func() error {
+			count++
+			if count < 2 {
+				return errors.New("fail")
+			}
+			cancel()
+			return errors.New("don't care about this error since we've cancelled and there is still another retry")
+		}
+
+		logger := func(_ string, _ ...any) {}
+
+		err := RetryWithContext(ctx, fn, 3, 0, logger)
+		require.Equal(t, 2, count)
+		require.ErrorIs(t, err, context.Canceled)
+	})
+
+	t.Run("NoErr", func(t *testing.T) {
+		count := 0
+		fn := func() error {
+			count++
+			return nil
+		}
+
+		logger := func(_ string, _ ...any) {}
+
+		err := RetryWithContext(context.TODO(), fn, 3, 0, logger)
+		require.ErrorIs(t, err, nil)
+		require.Equal(t, 1, count)
+	})
+
+	t.Run("InvalidAttempts", func(t *testing.T) {
+		count := 0
+		fn := func() error {
+			count++
+			return nil
+		}
+
+		logger := func(_ string, _ ...any) {}
+
+		err := RetryWithContext(context.TODO(), fn, 0, 0, logger)
+		require.Error(t, err)
+		require.Equal(t, 0, count)
+	})
+
+	t.Run("ContextCancellationDeadline", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Duration(2*time.Second))
+		defer cancel()
+
+		count := 0
+		fn := func() error {
+			count++
+			return errors.New("Always fail")
+		}
+
+		logger := func(_ string, _ ...any) {}
+
+		err := RetryWithContext(ctx, fn, 3, 1*time.Second, logger)
+		// fn should be called twice, it will wait one second after the first attempt
+		// and tries to wait two seconds after the second attempt
+		// but the context will cancel before the third attempt is called
+		require.Equal(t, 2, count)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+	})
 
-	count = 0
-	logCount = 0
-	err := Retry(countFn, 3, 0, loggerFn)
-	suite.Error(err)
-	suite.Equal(3, count)
-	suite.Equal(3, logCount)
-
-	count = 0
-	logCount = 0
-	err = Retry(countFn, 4, 0, loggerFn)
-	suite.NoError(err)
-	suite.Equal(4, count)
-	suite.Equal(3, logCount)
 }
 
 func (suite *TestMiscSuite) TestTransformMapKeys() {