feat(storage): add maxAttempts RetryOption

storage/bucket_test.go

@@ -1071,6 +1071,7 @@ func TestBucketRetryer(t *testing.T) {
 						Multiplier: 3,
 					}),
 					WithPolicy(RetryAlways),
+					WithMaxAttempts(5),
 					WithErrorFunc(func(err error) bool { return false }))
 			},
 			want: &retryConfig{
@@ -1080,6 +1081,7 @@ func TestBucketRetryer(t *testing.T) {
 					Multiplier: 3,
 				},
 				policy:      RetryAlways,
+				maxAttempts: expectedAttempts(5),
 				shouldRetry: func(err error) bool { return false },
 			},
 		},
@@ -1105,6 +1107,15 @@ func TestBucketRetryer(t *testing.T) {
 				policy: RetryNever,
 			},
 		},
+		{
+			name: "set max retry attempts only",
+			call: func(b *BucketHandle) *BucketHandle {
+				return b.Retryer(WithMaxAttempts(5))
+			},
+			want: &retryConfig{
+				maxAttempts: expectedAttempts(5),
+			},
+		},
 		{
 			name: "set ErrorFunc only",
 			call: func(b *BucketHandle) *BucketHandle {

storage/invoke.go

@@ -70,6 +70,9 @@ func run(ctx context.Context, call func(ctx context.Context) error, retry *retry
 	return internal.Retry(ctx, bo, func() (stop bool, err error) {
 		ctxWithHeaders := setInvocationHeaders(ctx, invocationID, attempts)
 		err = call(ctxWithHeaders)
+		if retry.maxAttempts != nil && attempts >= *retry.maxAttempts {
+			return true, err
+		}
 		attempts++
 		return !errorFunc(err), err
 	})

storage/invoke_test.go

@@ -105,7 +105,6 @@ func TestInvoke(t *testing.T) {
 			expectFinalErr:    false,
 		},
 		{
-
 			desc:              "non-idempotent retriable error retried when policy is RetryAlways",
 			count:             2,
 			initialErr:        &googleapi.Error{Code: 500},
@@ -132,7 +131,6 @@ func TestInvoke(t *testing.T) {
 			retry:             &retryConfig{policy: RetryAlways},
 			expectFinalErr:    false,
 		},
-
 		{
 			desc:              "non-retriable error retried with custom fn",
 			count:             2,
@@ -173,6 +171,65 @@ func TestInvoke(t *testing.T) {
 			},
 			expectFinalErr: false,
 		},
+		{
+			desc:              "non-idempotent retriable error retried when policy is RetryAlways till maxAttempts",
+			count:             4,
+			initialErr:        &googleapi.Error{Code: 500},
+			finalErr:          nil,
+			isIdempotentValue: false,
+			retry:             &retryConfig{policy: RetryAlways, maxAttempts: expectedAttempts(2)},
+			expectFinalErr:    false,
+		},
+		{
+			desc:              "non-idempotent retriable error not retried when policy is RetryNever with maxAttempts set",
+			count:             4,
+			initialErr:        &googleapi.Error{Code: 500},
+			finalErr:          nil,
+			isIdempotentValue: false,
+			retry:             &retryConfig{policy: RetryNever, maxAttempts: expectedAttempts(2)},
+			expectFinalErr:    false,
+		},
+		{
+			desc:              "non-retriable error retried with custom fn till maxAttempts",
+			count:             4,
+			initialErr:        io.ErrNoProgress,
+			finalErr:          nil,
+			isIdempotentValue: true,
+			retry: &retryConfig{
+				shouldRetry: func(err error) bool {
+					return err == io.ErrNoProgress
+				},
+				maxAttempts: expectedAttempts(2),
+			},
+			expectFinalErr: false,
+		},
+		{
+			desc:              "non-idempotent retriable error retried when policy is RetryAlways till maxAttempts where count equals to maxAttempts-1",
+			count:             3,
+			initialErr:        &googleapi.Error{Code: 500},
+			finalErr:          nil,
+			isIdempotentValue: false,
+			retry:             &retryConfig{policy: RetryAlways, maxAttempts: expectedAttempts(4)},
+			expectFinalErr:    true,
+		},
+		{
+			desc:              "non-idempotent retriable error retried when policy is RetryAlways till maxAttempts where count equals to maxAttempts",
+			count:             4,
+			initialErr:        &googleapi.Error{Code: 500},
+			finalErr:          nil,
+			isIdempotentValue: true,
+			retry:             &retryConfig{policy: RetryAlways, maxAttempts: expectedAttempts(4)},
+			expectFinalErr:    false,
+		},
+		{
+			desc:              "non-idempotent retriable error not retried when policy is RetryAlways with maxAttempts equals to zero",
+			count:             4,
+			initialErr:        &googleapi.Error{Code: 500},
+			finalErr:          nil,
+			isIdempotentValue: true,
+			retry:             &retryConfig{maxAttempts: expectedAttempts(0), policy: RetryAlways},
+			expectFinalErr:    false,
+		},
 	} {
 		t.Run(test.desc, func(s *testing.T) {
 			counter := 0
@@ -203,6 +260,10 @@ func TestInvoke(t *testing.T) {
 			if !test.expectFinalErr {
 				wantAttempts = 1
 			}
+			if test.retry != nil && test.retry.maxAttempts != nil && *test.retry.maxAttempts != 0 && test.retry.policy != RetryNever {
+				wantAttempts = *test.retry.maxAttempts
+			}
+
 			wantClientHeader := strings.ReplaceAll(initialClientHeader, "gccl-attempt-count/1", fmt.Sprintf("gccl-attempt-count/%v", wantAttempts))
 			if gotClientHeader != wantClientHeader {
 				t.Errorf("case %q, retry header:\ngot %v\nwant %v", test.desc, gotClientHeader, wantClientHeader)

storage/storage.go

@@ -2076,6 +2076,26 @@ func (wb *withBackoff) apply(config *retryConfig) {
 	config.backoff = &wb.backoff
 }
 
+// WithMaxAttempts configures the maximum number of times an API call can be made
+// in the case of retryable errors.
+// For example, if you set WithMaxAttempts(5), the operation will be attempted up to 5
+// times total (initial call plus 4 retries).
+// Without this setting, operations will continue retrying indefinitely
+// until either the context is canceled or a deadline is reached.
+func WithMaxAttempts(maxAttempts int) RetryOption {
+	return &withMaxAttempts{
+		maxAttempts: maxAttempts,
+	}
+}
+
+type withMaxAttempts struct {
+	maxAttempts int
+}
+
+func (wb *withMaxAttempts) apply(config *retryConfig) {
+	config.maxAttempts = &wb.maxAttempts
+}
+
 // RetryPolicy describes the available policies for which operations should be
 // retried. The default is `RetryIdempotent`.
 type RetryPolicy int
@@ -2148,6 +2168,7 @@ type retryConfig struct {
 	backoff     *gax.Backoff
 	policy      RetryPolicy
 	shouldRetry func(err error) bool
+	maxAttempts *int
 }
 
 func (r *retryConfig) clone() *retryConfig {
@@ -2168,6 +2189,7 @@ func (r *retryConfig) clone() *retryConfig {
 		backoff:     bo,
 		policy:      r.policy,
 		shouldRetry: r.shouldRetry,
+		maxAttempts: r.maxAttempts,
 	}
 }
 

storage/storage_test.go

@@ -953,6 +953,10 @@ func TestConditionErrors(t *testing.T) {
 	}
 }
 
+func expectedAttempts(value int) *int {
+	return &value
+}
+
 // Test that ObjectHandle.Retryer correctly configures the retry configuration
 // in the ObjectHandle.
 func TestObjectRetryer(t *testing.T) {
@@ -977,6 +981,7 @@ func TestObjectRetryer(t *testing.T) {
 						Max:        30 * time.Second,
 						Multiplier: 3,
 					}),
+					WithMaxAttempts(5),
 					WithPolicy(RetryAlways),
 					WithErrorFunc(func(err error) bool { return false }))
 			},
@@ -986,6 +991,7 @@ func TestObjectRetryer(t *testing.T) {
 					Max:        30 * time.Second,
 					Multiplier: 3,
 				},
+				maxAttempts: expectedAttempts(5),
 				policy:      RetryAlways,
 				shouldRetry: func(err error) bool { return false },
 			},
@@ -1012,6 +1018,15 @@ func TestObjectRetryer(t *testing.T) {
 				policy: RetryNever,
 			},
 		},
+		{
+			name: "set max retry attempts only",
+			call: func(o *ObjectHandle) *ObjectHandle {
+				return o.Retryer(WithMaxAttempts(11))
+			},
+			want: &retryConfig{
+				maxAttempts: expectedAttempts(11),
+			},
+		},
 		{
 			name: "set ErrorFunc only",
 			call: func(o *ObjectHandle) *ObjectHandle {
@@ -1063,6 +1078,7 @@ func TestClientSetRetry(t *testing.T) {
 					Max:        30 * time.Second,
 					Multiplier: 3,
 				}),
+				WithMaxAttempts(5),
 				WithPolicy(RetryAlways),
 				WithErrorFunc(func(err error) bool { return false }),
 			},
@@ -1072,6 +1088,7 @@ func TestClientSetRetry(t *testing.T) {
 					Max:        30 * time.Second,
 					Multiplier: 3,
 				},
+				maxAttempts: expectedAttempts(5),
 				policy:      RetryAlways,
 				shouldRetry: func(err error) bool { return false },
 			},
@@ -1097,6 +1114,15 @@ func TestClientSetRetry(t *testing.T) {
 				policy: RetryNever,
 			},
 		},
+		{
+			name: "set max retry attempts only",
+			clientOptions: []RetryOption{
+				WithMaxAttempts(7),
+			},
+			want: &retryConfig{
+				maxAttempts: expectedAttempts(7),
+			},
+		},
 		{
 			name: "set ErrorFunc only",
 			clientOptions: []RetryOption{
@@ -1150,10 +1176,12 @@ func TestRetryer(t *testing.T) {
 			name: "object retryer configures retry",
 			objectOptions: []RetryOption{
 				WithPolicy(RetryAlways),
+				WithMaxAttempts(5),
 				WithErrorFunc(ShouldRetry),
 			},
 			want: &retryConfig{
 				shouldRetry: ShouldRetry,
+				maxAttempts: expectedAttempts(5),
 				policy:      RetryAlways,
 			},
 		},
@@ -1166,6 +1194,7 @@ func TestRetryer(t *testing.T) {
 					Multiplier: 6,
 				}),
 				WithPolicy(RetryAlways),
+				WithMaxAttempts(11),
 				WithErrorFunc(ShouldRetry),
 			},
 			want: &retryConfig{
@@ -1175,6 +1204,7 @@ func TestRetryer(t *testing.T) {
 					Multiplier: 6,
 				},
 				shouldRetry: ShouldRetry,
+				maxAttempts: expectedAttempts(11),
 				policy:      RetryAlways,
 			},
 		},
@@ -1187,6 +1217,7 @@ func TestRetryer(t *testing.T) {
 					Multiplier: 6,
 				}),
 				WithPolicy(RetryAlways),
+				WithMaxAttempts(7),
 				WithErrorFunc(ShouldRetry),
 			},
 			want: &retryConfig{
@@ -1196,6 +1227,7 @@ func TestRetryer(t *testing.T) {
 					Multiplier: 6,
 				},
 				shouldRetry: ShouldRetry,
+				maxAttempts: expectedAttempts(7),
 				policy:      RetryAlways,
 			},
 		},
@@ -1206,10 +1238,12 @@ func TestRetryer(t *testing.T) {
 			},
 			objectOptions: []RetryOption{
 				WithPolicy(RetryNever),
+				WithMaxAttempts(5),
 				WithErrorFunc(ShouldRetry),
 			},
 			want: &retryConfig{
 				policy:      RetryNever,
+				maxAttempts: expectedAttempts(5),
 				shouldRetry: ShouldRetry,
 			},
 		},
@@ -1220,10 +1254,12 @@ func TestRetryer(t *testing.T) {
 			},
 			objectOptions: []RetryOption{
 				WithPolicy(RetryNever),
+				WithMaxAttempts(11),
 				WithErrorFunc(ShouldRetry),
 			},
 			want: &retryConfig{
 				policy:      RetryNever,
+				maxAttempts: expectedAttempts(11),
 				shouldRetry: ShouldRetry,
 			},
 		},
@@ -1243,9 +1279,11 @@ func TestRetryer(t *testing.T) {
 					Max:     time.Microsecond,
 				}),
 				WithErrorFunc(ShouldRetry),
+				WithMaxAttempts(5),
 			},
 			want: &retryConfig{
 				policy:      RetryAlways,
+				maxAttempts: expectedAttempts(5),
 				shouldRetry: ShouldRetry,
 				backoff: &gax.Backoff{
 					Initial: time.Nanosecond,
@@ -1280,6 +1318,7 @@ func TestRetryer(t *testing.T) {
 			bucketOptions: []RetryOption{
 				WithPolicy(RetryNever),
 				WithErrorFunc(ShouldRetry),
+				WithMaxAttempts(5),
 			},
 			objectOptions: []RetryOption{
 				WithBackoff(gax.Backoff{
@@ -1289,6 +1328,7 @@ func TestRetryer(t *testing.T) {
 			},
 			want: &retryConfig{
 				policy:      RetryNever,
+				maxAttempts: expectedAttempts(5),
 				shouldRetry: ShouldRetry,
 				backoff: &gax.Backoff{
 					Initial: time.Nanosecond,