Merge pull request #34 from avast/context_support

add context support

README.md

@@ -58,9 +58,8 @@ complicated interface.
 * [rafaeljesus/retry-go](https://github.com/rafaeljesus/retry-go) - looks good,
 slightly similar as this package, don't have 'simple' `Retry` method
 
-* [matryer/try](https://github.com/matryer/try) - popular package (probably second most popular), nonintuitive interface (for me)
-
-* [sethvargo/go-retry](https://github.com/sethvargo/go-retry) - popular package (probably third most popular), from my point of view, haven't a simple API as this module
+* [matryer/try](https://github.com/matryer/try) - very popular package,
+nonintuitive interface (for me)
 
 
 ### BREAKING CHANGES
@@ -92,6 +91,7 @@ var (
 	DefaultRetryIf       = IsRecoverable
 	DefaultDelayType     = CombineDelay(BackOffDelay, RandomDelay)
 	DefaultLastErrorOnly = false
+	DefaultContext       = context.Background()
 )
 ```
 
@@ -208,6 +208,24 @@ func Attempts(attempts uint) Option
 ```
 Attempts set count of retry default is 10
 
+#### func  Context
+
+```go
+func Context(ctx context.Context) Option
+```
+Context allow to set context of retry default are Background context
+
+    example of immediately cancellation (maybe it isn't the best example, but it describes behavior enough; I hope)
+    ctx, cancel := context.WithCancel(context.Background())
+    cancel()
+
+    retry.Do(
+    	func() error {
+    		...
+    	},
+    	retry.Context(ctx),
+    )
+
 #### func  Delay
 
 ```go

VERSION

@@ -1 +1 @@
-2.6.1
+2.7.0

options.go

@@ -1,6 +1,7 @@
 package retry
 
 import (
+	"context"
 	"math"
 	"math/rand"
 	"time"
@@ -24,6 +25,7 @@ type Config struct {
 	retryIf       RetryIfFunc
 	delayType     DelayTypeFunc
 	lastErrorOnly bool
+	context       context.Context
 
 	maxBackOffN uint
 }
@@ -170,3 +172,22 @@ func RetryIf(retryIf RetryIfFunc) Option {
 		c.retryIf = retryIf
 	}
 }
+
+// Context allow to set context of retry
+// default are Background context
+//
+//	example of immediately cancellation (maybe it isn't the best example, but it describes behavior enough; I hope)
+//	ctx, cancel := context.WithCancel(context.Background())
+//	cancel()
+//
+//	retry.Do(
+//		func() error {
+//			...
+//		},
+//		retry.Context(ctx),
+//	)
+func Context(ctx context.Context) Option {
+	return func(c *Config) {
+		c.context = ctx
+	}
+}

retry.go

@@ -65,6 +65,7 @@ BREAKING CHANGES
 package retry
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -81,6 +82,7 @@ var (
 	DefaultRetryIf       = IsRecoverable
 	DefaultDelayType     = CombineDelay(BackOffDelay, RandomDelay)
 	DefaultLastErrorOnly = false
+	DefaultContext       = context.Background()
 )
 
 func Do(retryableFunc RetryableFunc, opts ...Option) error {
@@ -95,13 +97,18 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 		retryIf:       DefaultRetryIf,
 		delayType:     DefaultDelayType,
 		lastErrorOnly: DefaultLastErrorOnly,
+		context:       DefaultContext,
 	}
 
 	//apply opts
 	for _, opt := range opts {
 		opt(config)
 	}
 
+	if err := config.context.Err(); err != nil {
+		return err
+	}
+
 	var errorLog Error
 	if !config.lastErrorOnly {
 		errorLog = make(Error, config.attempts)
@@ -131,7 +138,13 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 			if config.maxDelay > 0 && delayTime > config.maxDelay {
 				delayTime = config.maxDelay
 			}
-			time.Sleep(delayTime)
+
+			select {
+			case <-time.After(delayTime):
+			case <-config.context.Done():
+				return config.context.Err()
+			}
+
 		} else {
 			return nil
 		}

retry_test.go

@@ -1,12 +1,12 @@
 package retry
 
 import (
+	"context"
 	"errors"
+	"fmt"
 	"testing"
 	"time"
 
-	"fmt"
-
 	"github.com/stretchr/testify/assert"
 )
 
@@ -260,3 +260,40 @@ func TestCombineDelay(t *testing.T) {
 		)
 	}
 }
+
+func TestContext(t *testing.T) {
+	t.Run("cancel before", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		retrySum := 0
+		start := time.Now()
+		err := Do(
+			func() error { return errors.New("test") },
+			OnRetry(func(n uint, err error) { retrySum += 1 }),
+			Context(ctx),
+		)
+		dur := time.Since(start)
+		assert.Error(t, err)
+		assert.True(t, dur < DefaultDelay, "immediately cancellation")
+		assert.Equal(t, 0, retrySum, "called at most once")
+	})
+
+	t.Run("cancel in retry progress", func(t *testing.T) {
+		ctx, cancel := context.WithCancel(context.Background())
+
+		retrySum := 0
+		err := Do(
+			func() error { return errors.New("test") },
+			OnRetry(func(n uint, err error) {
+				retrySum += 1
+				if retrySum > 1 {
+					cancel()
+				}
+			}),
+			Context(ctx),
+		)
+		assert.Error(t, err)
+		assert.Equal(t, 2, retrySum, "called at most once")
+	})
+}