Handle context cancellations in retry middleware

retry/retry.go

@@ -2,6 +2,7 @@ package retry
 
 import (
 	"context"
+	"github.com/slok/goresilience/errors"
 	"math"
 	"math/rand"
 	"time"
@@ -52,7 +53,7 @@ func NewMiddleware(cfg Config) goresilience.Middleware {
 			var err error
 			metricsRecorder, _ := metrics.RecorderFromContext(ctx)
 
-			// Start the attemps. (it's 1 + the number of retries.)
+			// Start the attempts. (it's 1 + the number of retries.)
 			for i := 0; i <= cfg.Times; i++ {
 				// Only measure the retries.
 				if i != 0 {
@@ -64,6 +65,11 @@ func NewMiddleware(cfg Config) goresilience.Middleware {
 					return nil
 				}
 
+				// If the context is cancelled no retry will succeed
+				if err == errors.ErrContextCanceled {
+					return err
+				}
+
 				// We need to sleep before making a retry.
 				waitDuration := cfg.WaitBase
 

retry/retry_test.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	goresilienceerrors "github.com/slok/goresilience/errors"
+	"math"
 	"testing"
 	"time"
 
@@ -193,3 +195,27 @@ func TestBackoffJitterRetry(t *testing.T) {
 		})
 	}
 }
+
+func TestRetryWithContextCancel(t *testing.T) {
+	t.Run("Context is cancelled pre-execution", func(t *testing.T) {
+		assert := assert.New(t)
+		ctx, cancel := context.WithCancel(context.TODO())
+		cancel()
+		exec := retry.New(retry.Config{Times: math.MaxInt32})
+		err := exec.Run(ctx, func(ctx context.Context) error {
+			return fmt.Errorf("test error")
+		})
+		assert.EqualError(err, goresilienceerrors.ErrContextCanceled.Error())
+	})
+
+	t.Run("Context is cancel after deadline", func(t *testing.T) {
+		assert := assert.New(t)
+		ctx, cancel := context.WithTimeout(context.TODO(), 1*time.Nanosecond)
+		defer cancel()
+		exec := retry.New(retry.Config{Times: math.MaxInt32})
+		err := exec.Run(ctx, func(ctx context.Context) error {
+			return fmt.Errorf("test error")
+		})
+		assert.EqualError(err, goresilienceerrors.ErrContextCanceled.Error())
+	})
+}