diff --git a/storage/bucket_test.go b/storage/bucket_test.go index 6ba32daf1463..d574d19862e1 100644 --- a/storage/bucket_test.go +++ b/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 { diff --git a/storage/invoke.go b/storage/invoke.go index dc79fd88bbc4..1b52eb5d2c65 100644 --- a/storage/invoke.go +++ b/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 }) diff --git a/storage/invoke_test.go b/storage/invoke_test.go index 005cff117109..a89695a86734 100644 --- a/storage/invoke_test.go +++ b/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) diff --git a/storage/storage.go b/storage/storage.go index 78ecbf0e892a..7af32ac6fcd0 100644 --- a/storage/storage.go +++ b/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, } } diff --git a/storage/storage_test.go b/storage/storage_test.go index edceac229c04..e5aba513daee 100644 --- a/storage/storage_test.go +++ b/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,