Dealing with a canceled Context

retry.go

@@ -49,6 +49,13 @@ func (e *retryableError) Error() string {
 // Do wraps a function with a backoff to retry. The provided context is the same
 // context passed to the RetryFunc.
 func Do(ctx context.Context, b Backoff, f RetryFunc) error {
+	// If ctx is already canceled, then return immediately without calling f
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
+	}
+
 	for {
 		err := f(ctx)
 		if err == nil {
@@ -66,6 +73,13 @@ func Do(ctx context.Context, b Backoff, f RetryFunc) error {
 			return rerr.Unwrap()
 		}
 
+		// ctx.Done() has priority, so we test it alone first
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
 		select {
 		case <-ctx.Done():
 			return ctx.Err()

retry_test.go

@@ -2,6 +2,7 @@ package retry_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -173,3 +174,34 @@ func ExampleDo_customRetry() {
 		// handle error
 	}
 }
+
+func TestCancel(t *testing.T) {
+	for i := 0; i < 100000; i++ {
+		ctx, cancel := context.WithCancel(context.Background())
+
+		calls := 0
+		rf := func(ctx context.Context) error {
+			calls++
+			// Never succeed.
+			// Always return a RetryableError
+			return retry.RetryableError(errors.New("nope"))
+		}
+
+		const delay time.Duration = time.Millisecond
+		b := retry.NewConstant(delay)
+
+		const maxRetries = 5
+		b = retry.WithMaxRetries(maxRetries, b)
+
+		const jitter time.Duration = 5 * time.Millisecond
+		b = retry.WithJitter(jitter, b)
+
+		// Here we cancel the Context *before* the call to Do
+		cancel()
+		retry.Do(ctx, b, rf)
+
+		if calls > 1 {
+			t.Errorf("rf was called %d times instead of 0 or 1", calls)
+		}
+	}
+}