Add support for contexts during infinite attempts

retry.go

@@ -93,8 +93,11 @@ func Do(retryableFunc RetryableFunc, opts ...Option) error {
 			n++
 
 			config.onRetry(n, err)
-
-			<-time.After(delay(config, n, err))
+			select {
+			case <-time.After(delay(config, n, err)):
+			case <-config.context.Done():
+				return nil
+			}
 		}
 
 		return nil

retry_test.go

@@ -359,4 +359,43 @@ func TestContext(t *testing.T) {
 
 		assert.Equal(t, 2, retrySum, "called at most once")
 	})
+
+	t.Run("cancel in retry progress - infinite attempts", func(t *testing.T) {
+		testFailedInRetry := make(chan bool)
+		testEnded := make(chan bool)
+
+		go func() {
+			ctx, cancel := context.WithCancel(context.Background())
+
+			retrySum := 0
+			err := Do(
+				func() error { return errors.New("test") },
+				OnRetry(func(n uint, err error) {
+					fmt.Println(n)
+					retrySum += 1
+					if retrySum > 1 {
+						cancel()
+					}
+
+					if retrySum > 2 {
+						testFailedInRetry <- true
+					}
+
+				}),
+				Context(ctx),
+				Attempts(0),
+			)
+
+			assert.NoError(t, err, "infinite attempts should not report error")
+			assert.Error(t, ctx.Err(), "immediately canceled after context cancel called")
+			testEnded <- true
+		}()
+
+		select {
+		case <-testFailedInRetry:
+			t.Error("Test ran longer than expected, cancel did not work")
+		case <-testEnded:
+		}
+
+	})
 }