feat(storage): allow retry ErrorFunc configs (#5166)

Allows users to supply a custom ErrorFunc which will be used
to determine whether the error is retryable.

storage/invoke.go

@@ -45,9 +45,13 @@ func run(ctx context.Context, call func() error, retry *retryConfig, isIdempoten
 		bo.Initial = retry.backoff.Initial
 		bo.Max = retry.backoff.Max
 	}
+	var errorFunc func(err error) bool = shouldRetry
+	if retry.shouldRetry != nil {
+		errorFunc = retry.shouldRetry
+	}
 	return internal.Retry(ctx, bo, func() (stop bool, err error) {
 		err = call()
-		return !shouldRetry(err), err
+		return !errorFunc(err), err
 	})
 }
 

storage/storage.go

@@ -1862,9 +1862,39 @@ func (ws *withPolicy) apply(config *retryConfig) {
 	config.policy = ws.policy
 }
 
+// WithErrorFunc allows users to pass a custom function to the retryer. Errors
+// will be retried if and only if `shouldRetry(err)` returns true.
+// By default, the following errors are retried (see invoke.go for the default
+// shouldRetry function):
+//
+// - HTTP responses with codes 429, 502, 503, and 504.
+//
+// - Transient network errors such as connection reset and io.ErrUnexpectedEOF.
+//
+// - Errors which are considered transient using the Temporary() interface.
+//
+// - Wrapped versions of these errors.
+//
+// This option can be used to retry on a different set of errors than the
+// default.
+func WithErrorFunc(shouldRetry func(err error) bool) RetryOption {
+	return &withErrorFunc{
+		shouldRetry: shouldRetry,
+	}
+}
+
+type withErrorFunc struct {
+	shouldRetry func(err error) bool
+}
+
+func (wef *withErrorFunc) apply(config *retryConfig) {
+	config.shouldRetry = wef.shouldRetry
+}
+
 type retryConfig struct {
-	backoff *gax.Backoff
-	policy  RetryPolicy
+	backoff     *gax.Backoff
+	policy      RetryPolicy
+	shouldRetry func(err error) bool
 }
 
 // composeSourceObj wraps a *raw.ComposeRequestSourceObjects, but adds the methods

storage/storage_test.go

@@ -810,15 +810,17 @@ func TestRetryer(t *testing.T) {
 						Max:        30 * time.Second,
 						Multiplier: 3,
 					}),
-					WithPolicy(RetryAlways))
+					WithPolicy(RetryAlways),
+					WithErrorFunc(func(err error) bool { return false }))
 			},
 			want: &retryConfig{
 				backoff: &gax.Backoff{
 					Initial:    2 * time.Second,
 					Max:        30 * time.Second,
 					Multiplier: 3,
 				},
-				policy: RetryAlways,
+				policy:      RetryAlways,
+				shouldRetry: func(err error) bool { return false },
 			},
 		},
 		{
@@ -843,11 +845,30 @@ func TestRetryer(t *testing.T) {
 				policy: RetryNever,
 			},
 		},
+		{
+			name: "set ErrorFunc only",
+			call: func(o *ObjectHandle) *ObjectHandle {
+				return o.Retryer(
+					WithErrorFunc(func(err error) bool { return false }))
+			},
+			want: &retryConfig{
+				shouldRetry: func(err error) bool { return false },
+			},
+		},
 	}
 	for _, tc := range testCases {
 		t.Run(tc.name, func(s *testing.T) {
 			o := tc.call(&ObjectHandle{})
-			if diff := cmp.Diff(o.retry, tc.want, cmp.AllowUnexported(retryConfig{}, gax.Backoff{})); diff != "" {
+			if diff := cmp.Diff(
+				o.retry,
+				tc.want,
+				cmp.AllowUnexported(retryConfig{}, gax.Backoff{}),
+				// ErrorFunc cannot be compared directly, but we check if both are
+				// either nil or non-nil.
+				cmp.Comparer(func(a, b func(err error) bool) bool {
+					return (a == nil && b == nil) || (a != nil && b != nil)
+				}),
+			); diff != "" {
 				s.Fatalf("retry not configured correctly: %v", diff)
 			}
 		})