feat(storage): add retry config to BucketHandle

storage/bucket.go

@@ -44,6 +44,7 @@ type BucketHandle struct {
 	defaultObjectACL ACLHandle
 	conds            *BucketConditions
 	userProject      string // project for Requester Pays buckets
+	retry            *retryConfig
 }
 
 // Bucket returns a BucketHandle, which provides operations on the named bucket.
@@ -210,12 +211,20 @@ func (b *BucketHandle) Update(ctx context.Context, uattrs BucketAttrsToUpdate) (
 	if uattrs.PredefinedDefaultObjectACL != "" {
 		req.PredefinedDefaultObjectAcl(uattrs.PredefinedDefaultObjectACL)
 	}
-	// TODO(jba): retry iff metagen is set?
-	rb, err := req.Context(ctx).Do()
-	if err != nil {
+
+	isIdempotent := b.conds != nil && b.conds.MetagenerationMatch != 0
+
+	var rawBucket *raw.Bucket
+	call := func() error {
+		rb, err := req.Context(ctx).Do()
+		rawBucket = rb
+		return err
+	}
+
+	if err := run(ctx, call, b.retry, isIdempotent); err != nil {
 		return nil, err
 	}
-	return newBucket(rb)
+	return newBucket(rawBucket)
 }
 
 func (b *BucketHandle) newPatchCall(uattrs *BucketAttrsToUpdate) (*raw.BucketsPatchCall, error) {
@@ -1367,6 +1376,19 @@ func (b *BucketHandle) Objects(ctx context.Context, q *Query) *ObjectIterator {
 	return it
 }
 
+// Retryer returns a bucket handle that is configured with custom retry
+// behavior as specified by the options that are passed to it. All operations
+// on the new handle will use the customized retry configuration.
+func (b *BucketHandle) Retryer(opts ...RetryOption) *BucketHandle {
+	b2 := *b
+	retry := &retryConfig{}
+	for _, opt := range opts {
+		opt.apply(retry)
+	}
+	b2.retry = retry
+	return &b2
+}
+
 // An ObjectIterator is an iterator over ObjectAttrs.
 //
 // Note: This iterator is not safe for concurrent operations without explicit synchronization.

storage/bucket_test.go

@@ -23,6 +23,7 @@ import (
 
 	"cloud.google.com/go/internal/testutil"
 	"github.com/google/go-cmp/cmp"
+	gax "github.com/googleapis/gax-go/v2"
 	"google.golang.org/api/googleapi"
 	raw "google.golang.org/api/storage/v1"
 )
@@ -717,3 +718,90 @@ func TestNewBucket(t *testing.T) {
 		t.Errorf("got=-, want=+:\n%s", diff)
 	}
 }
+
+func TestBucketRetryer(t *testing.T) {
+	testCases := []struct {
+		name string
+		call func(b *BucketHandle) *BucketHandle
+		want *retryConfig
+	}{
+		{
+			name: "all defaults",
+			call: func(b *BucketHandle) *BucketHandle {
+				return b.Retryer()
+			},
+			want: &retryConfig{},
+		},
+		{
+			name: "set all options",
+			call: func(b *BucketHandle) *BucketHandle {
+				return b.Retryer(
+					WithBackoff(gax.Backoff{
+						Initial:    2 * time.Second,
+						Max:        30 * time.Second,
+						Multiplier: 3,
+					}),
+					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,
+				shouldRetry: func(err error) bool { return false },
+			},
+		},
+		{
+			name: "set some backoff options",
+			call: func(b *BucketHandle) *BucketHandle {
+				return b.Retryer(
+					WithBackoff(gax.Backoff{
+						Multiplier: 3,
+					}))
+			},
+			want: &retryConfig{
+				backoff: &gax.Backoff{
+					Multiplier: 3,
+				}},
+		},
+		{
+			name: "set policy only",
+			call: func(b *BucketHandle) *BucketHandle {
+				return b.Retryer(WithPolicy(RetryNever))
+			},
+			want: &retryConfig{
+				policy: RetryNever,
+			},
+		},
+		{
+			name: "set ErrorFunc only",
+			call: func(b *BucketHandle) *BucketHandle {
+				return b.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) {
+			b := tc.call(&BucketHandle{})
+			if diff := cmp.Diff(
+				b.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)
+			}
+		})
+	}
+}