Merge pull request avast#71 from Hrily/hrishi/attempts-for-error

feat: add support for attempts based on error types

Author: JaSei

options.go

@@ -23,16 +23,17 @@ type Timer interface {
 }
 
 type Config struct {
-	attempts      uint
-	delay         time.Duration
-	maxDelay      time.Duration
-	maxJitter     time.Duration
-	onRetry       OnRetryFunc
-	retryIf       RetryIfFunc
-	delayType     DelayTypeFunc
-	lastErrorOnly bool
-	context       context.Context
-	timer         Timer
+	attempts         uint
+	attemptsForError map[error]uint
+	delay            time.Duration
+	maxDelay         time.Duration
+	maxJitter        time.Duration
+	onRetry          OnRetryFunc
+	retryIf          RetryIfFunc
+	delayType        DelayTypeFunc
+	lastErrorOnly    bool
+	context          context.Context
+	timer            Timer
 
 	maxBackOffN uint
 }
@@ -58,6 +59,15 @@ func Attempts(attempts uint) Option {
 	}
 }
 
+// AttemptsForError sets count of retry in case execution results in given `err`
+// Retries for the given `err` are also counted against total retries.
+// The retry will stop if any of given retries is exhausted.
+func AttemptsForError(attempts uint, err error) Option {
+	return func(c *Config) {
+		c.attemptsForError[err] = attempts
+	}
+}
+
 // Delay set delay between retry
 // default is 100ms
 func Delay(delay time.Duration) Option {

retry.go

@@ -117,8 +117,14 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 		errorLog = make(Error, 1)
 	}
 
+	attemptsForError := make(map[error]uint, len(config.attemptsForError))
+	for err, attempts := range config.attemptsForError {
+		attemptsForError[err] = attempts
+	}
+
 	lastErrIndex := n
-	for n < config.attempts {
+	shouldRetry := true
+	for shouldRetry {
 		err := retryableFunc()
 
 		if err != nil {
@@ -130,6 +136,14 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 
 			config.onRetry(n, err)
 
+			for errToCheck, attempts := range attemptsForError {
+				if errors.Is(err, errToCheck) {
+					attempts--
+					attemptsForError[errToCheck] = attempts
+					shouldRetry = shouldRetry && attempts > 0
+				}
+			}
+
 			// if this is last attempt - don't wait
 			if n == config.attempts-1 {
 				break
@@ -150,6 +164,8 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 		}
 
 		n++
+		shouldRetry = shouldRetry && n < config.attempts
+
 		if !config.lastErrorOnly {
 			lastErrIndex = n
 		}
@@ -163,15 +179,16 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 
 func newDefaultRetryConfig() *Config {
 	return &Config{
-		attempts:      uint(10),
-		delay:         100 * time.Millisecond,
-		maxJitter:     100 * time.Millisecond,
-		onRetry:       func(n uint, err error) {},
-		retryIf:       IsRecoverable,
-		delayType:     CombineDelay(BackOffDelay, RandomDelay),
-		lastErrorOnly: false,
-		context:       context.Background(),
-		timer:         &timerImpl{},
+		attempts:         uint(10),
+		attemptsForError: make(map[error]uint),
+		delay:            100 * time.Millisecond,
+		maxJitter:        100 * time.Millisecond,
+		onRetry:          func(n uint, err error) {},
+		retryIf:          IsRecoverable,
+		delayType:        CombineDelay(BackOffDelay, RandomDelay),
+		lastErrorOnly:    false,
+		context:          context.Background(),
+		timer:            &timerImpl{},
 	}
 }
 

retry_test.go

@@ -110,6 +110,22 @@ func TestZeroAttemptsWithoutError(t *testing.T) {
 	assert.Equal(t, count, 1)
 }
 
+func TestAttemptsForError(t *testing.T) {
+	count := uint(0)
+	testErr := os.ErrInvalid
+	attemptsForTestError := uint(3)
+	err := Do(
+		func() error {
+			count++
+			return testErr
+		},
+		AttemptsForError(attemptsForTestError, testErr),
+		Attempts(5),
+	)
+	assert.Error(t, err)
+	assert.Equal(t, attemptsForTestError, count)
+}
+
 func TestDefaultSleep(t *testing.T) {
 	start := time.Now()
 	err := Do(