Skip to content

Commit

Permalink
Merge pull request #96 from willdot/handle-context-timeout
Browse files Browse the repository at this point in the history
Allow last error to be returned with context error
  • Loading branch information
JaSei authored Aug 7, 2023
2 parents fdadb7c + b94b74c commit 8d16616
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 11 deletions.
24 changes: 24 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
46 changes: 35 additions & 11 deletions options.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
}
6 changes: 6 additions & 0 deletions retry.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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()
}
}
Expand Down
39 changes: 39 additions & 0 deletions retry_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down

0 comments on commit 8d16616

Please sign in to comment.