Merge pull request avast#96 from willdot/handle-context-timeout

Allow last error to be returned with context error

Author: JaSei

README.md

@@ -401,6 +401,30 @@ example of augmenting time.After with a print statement
     	   retry.WithTimer(&MyTimer{})
     )
 
+#### func  WrapContextErrorWithLastError
+
+```go
+func WrapContextErrorWithLastError(wrapContextErrorWithLastError bool) Option
+```
+WrapContextErrorWithLastError allows the context error to be returned wrapped
+with the last error that the retried function returned. This is only applicable
+when Attempts is set to 0 to retry indefinitly and when using a context to
+cancel / timeout
+
+default is false
+
+    ctx, cancel := context.WithCancel(context.Background())
+    defer cancel()
+
+    retry.Do(
+    	func() error {
+    		...
+    	},
+    	retry.Context(ctx),
+    	retry.Attempts(0),
+    	retry.WrapContextErrorWithLastError(true),
+    )
+
 #### type RetryIfFunc
 
 ```go

options.go

@@ -23,17 +23,18 @@ type Timer interface {
 }
 
 type Config struct {
-	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
+	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
+	wrapContextErrorWithLastError bool
 
 	maxBackOffN uint
 }
@@ -248,3 +249,26 @@ func WithTimer(t Timer) Option {
 		c.timer = t
 	}
 }
+
+// WrapContextErrorWithLastError allows the context error to be returned wrapped with the last error that the
+// retried function returned. This is only applicable when Attempts is set to 0 to retry indefinitly and when
+// using a context to cancel / timeout
+//
+// default is false
+//
+//	ctx, cancel := context.WithCancel(context.Background())
+//	defer cancel()
+//
+//	retry.Do(
+//		func() error {
+//			...
+//		},
+//		retry.Context(ctx),
+//		retry.Attempts(0),
+//		retry.WrapContextErrorWithLastError(true),
+//	)
+func WrapContextErrorWithLastError(wrapContextErrorWithLastError bool) Option {
+	return func(c *Config) {
+		c.wrapContextErrorWithLastError = wrapContextErrorWithLastError
+	}
+}

retry.go

@@ -136,6 +136,7 @@ func DoWithData[T any](retryableFunc RetryableFuncWithData[T], opts ...Option) (
 	}
 
 	// Setting attempts to 0 means we'll retry until we succeed
+	var lastErr error
 	if config.attempts == 0 {
 		for {
 			t, err := retryableFunc()
@@ -151,11 +152,16 @@ func DoWithData[T any](retryableFunc RetryableFuncWithData[T], opts ...Option) (
 				return emptyT, err
 			}
 
+			lastErr = err
+
 			n++
 			config.onRetry(n, err)
 			select {
 			case <-config.timer.After(delay(config, n, err)):
 			case <-config.context.Done():
+				if config.wrapContextErrorWithLastError {
+					return emptyT, Error{config.context.Err(), lastErr}
+				}
 				return emptyT, config.context.Err()
 			}
 		}

retry_test.go

@@ -464,6 +464,45 @@ func TestContext(t *testing.T) {
 			assert.Equal(t, 2, retrySum, "called at most once")
 		}()
 	})
+
+	t.Run("cancelled on retry infinte attempts - wraps context error with last retried function error", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		retrySum := 0
+		err := Do(
+			func() error { return fooErr{str: fmt.Sprintf("error %d", retrySum+1)} },
+			OnRetry(func(n uint, err error) {
+				retrySum += 1
+				if retrySum == 2 {
+					cancel()
+				}
+			}),
+			Context(ctx),
+			Attempts(0),
+			WrapContextErrorWithLastError(true),
+		)
+		assert.ErrorIs(t, err, context.Canceled)
+		assert.ErrorIs(t, err, fooErr{str: "error 2"})
+	})
+
+	t.Run("timed out on retry infinte attempts - wraps context error with last retried function error", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*500)
+		defer cancel()
+
+		retrySum := 0
+		err := Do(
+			func() error { return fooErr{str: fmt.Sprintf("error %d", retrySum+1)} },
+			OnRetry(func(n uint, err error) {
+				retrySum += 1
+			}),
+			Context(ctx),
+			Attempts(0),
+			WrapContextErrorWithLastError(true),
+		)
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+		assert.ErrorIs(t, err, fooErr{str: "error 2"})
+	})
 }
 
 type testTimer struct {