Add a SurfaceWorkErrors() opt to the retrier

retrier/retrier.go

@@ -11,12 +11,13 @@ import (
 // Retrier implements the "retriable" resiliency pattern, abstracting out the process of retrying a failed action
 // a certain number of times with an optional back-off between each retry.
 type Retrier struct {
-	backoff       []time.Duration
-	infiniteRetry bool
-	class         Classifier
-	jitter        float64
-	rand          *rand.Rand
-	randMu        sync.Mutex
+	backoff           []time.Duration
+	infiniteRetry     bool
+	surfaceWorkErrors bool
+	class             Classifier
+	jitter            float64
+	rand              *rand.Rand
+	randMu            sync.Mutex
 }
 
 // New constructs a Retrier with the given backoff pattern and classifier. The length of the backoff pattern
@@ -43,6 +44,13 @@ func (r *Retrier) WithInfiniteRetry() *Retrier {
 	return r
 }
 
+// WithSurfaceWorkErrors configures the retrier to always return the last non-nil error received from work()
+// even if a context timeout/deadline is hit.
+func (r *Retrier) WithSurfaceWorkErrors() *Retrier {
+	r.surfaceWorkErrors = true
+	return r
+}
+
 // Run executes the given work function by executing RunCtx without context.Context.
 func (r *Retrier) Run(work func() error) error {
 	return r.RunFn(context.Background(), func(c context.Context, r int) error {
@@ -83,6 +91,9 @@ func (r *Retrier) RunFn(ctx context.Context, work func(ctx context.Context, retr
 
 			timer := time.NewTimer(r.calcSleep(retries))
 			if err := r.sleep(ctx, timer); err != nil {
+				if r.surfaceWorkErrors && ret != nil {
+					return ret
+				}
 				return err
 			}
 

retrier/retrier_test.go

@@ -153,6 +153,44 @@ func TestRetrierRunFnWithInfinite(t *testing.T) {
 	}
 }
 
+func TestRetrierCtxSurfaceWorkErrors(t *testing.T) {
+	// Will timeout before getting to errBar.
+	ctx, cancel := context.WithTimeout(context.Background(), 7*time.Millisecond)
+	defer cancel()
+	r := New([]time.Duration{0, 10 * time.Millisecond}, nil).WithSurfaceWorkErrors()
+	errExpected := []error{errFoo, errFoo, errBar, errBaz}
+	retries := 0
+	err := r.RunCtx(ctx, func(ctx context.Context) error {
+		if retries >= len(errExpected) {
+			return nil
+		}
+		err := errExpected[retries]
+		retries++
+		return err
+	})
+	if err != errFoo {
+		t.Error(err)
+	}
+}
+
+func TestRetrierRunFnWithSurfaceWorkErrors(t *testing.T) {
+	// Will timeout before getting to errBar.
+	ctx, cancel := context.WithTimeout(context.Background(), 7*time.Millisecond)
+	defer cancel()
+	r := New([]time.Duration{0, 10 * time.Millisecond}, nil).WithSurfaceWorkErrors()
+	errExpected := []error{errFoo, errFoo, errFoo, errBar, errBaz}
+
+	err := r.RunFn(ctx, func(ctx context.Context, retries int) error {
+		if retries >= len(errExpected) {
+			return nil
+		}
+		return errExpected[retries]
+	})
+	if err != errFoo {
+		t.Error(err)
+	}
+}
+
 func TestRetrierNone(t *testing.T) {
 	r := New(nil, nil)