fix context behaviour

retry.go

@@ -105,6 +105,10 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 		opt(config)
 	}
 
+	if err := config.context.Err(); err != nil {
+		return err
+	}
+
 	var errorLog Error
 	if !config.lastErrorOnly {
 		errorLog = make(Error, config.attempts)
@@ -138,6 +142,7 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 			select {
 			case <-time.After(delayTime):
 			case <-config.context.Done():
+				return config.context.Err()
 			}
 
 		} else {

retry_test.go

@@ -262,14 +262,38 @@ func TestCombineDelay(t *testing.T) {
 }
 
 func TestContext(t *testing.T) {
-	start := time.Now()
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel()
-	err := Do(
-		func() error { return errors.New("test") },
-		Context(ctx),
-	)
-	dur := time.Since(start)
-	assert.Error(t, err)
-	assert.True(t, dur < time.Millisecond, "immediately cancellation")
+	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")
+	})
 }