feat(storage): add retry config to the Client and HmacKey operations (#5193)

Author: BrennaEpp

storage/bucket.go

@@ -55,18 +55,22 @@ type BucketHandle struct {
 // found at:
 //   https://cloud.google.com/storage/docs/bucket-naming
 func (c *Client) Bucket(name string) *BucketHandle {
+	retry := c.retry.clone()
 	return &BucketHandle{
 		c:    c,
 		name: name,
 		acl: ACLHandle{
 			c:      c,
 			bucket: name,
+			retry:  retry,
 		},
 		defaultObjectACL: ACLHandle{
 			c:         c,
 			bucket:    name,
 			isDefault: true,
+			retry:     retry,
 		},
+		retry: retry,
 	}
 }
 
@@ -1431,9 +1435,19 @@ func (b *BucketHandle) Objects(ctx context.Context, q *Query) *ObjectIterator {
 // on the new handle will use the customized retry configuration.
 // Retry options set on a object handle will take precedence over options set on
 // the bucket handle.
+// These retry options will merge with the client's retry configuration (if set)
+// for the returned handle. Options passed into this method will take precedence
+// over retry options on the client. Note that you must explicitly pass in each
+// option you want to override.
 func (b *BucketHandle) Retryer(opts ...RetryOption) *BucketHandle {
 	b2 := *b
-	retry := &retryConfig{}
+	var retry *retryConfig
+	if b.retry != nil {
+		// merge the options with the existing retry
+		retry = b.retry
+	} else {
+		retry = &retryConfig{}
+	}
 	for _, opt := range opts {
 		opt.apply(retry)
 	}

storage/emulator_test.sh

@@ -58,7 +58,6 @@ trap cleanup EXIT
 # TODO: move to passing once fixed
 FAILING=(   "buckets.setIamPolicy"
             "objects.insert"
-            "hmacKey.update"
         )
 # TODO: remove regex once all tests are passing
 # Unfortunately, there is no simple way to skip specific tests (see https://github.com/golang/go/issues/41583)
@@ -85,6 +84,7 @@ PASSING=(   "buckets.list"
             "hmacKey.list"
             "hmacKey.create"
             "hmacKey.delete"
+            "hmacKey.update"
             "notifications.list"
             "notifications.create"
             "notifications.get"

storage/hmac.go

@@ -89,8 +89,8 @@ type HMACKey struct {
 type HMACKeyHandle struct {
 	projectID string
 	accessID  string
-
-	raw *raw.ProjectsHmacKeysService
+	retry     *retryConfig
+	raw       *raw.ProjectsHmacKeysService
 }
 
 // HMACKeyHandle creates a handle that will be used for HMACKey operations.
@@ -100,6 +100,7 @@ func (c *Client) HMACKeyHandle(projectID, accessID string) *HMACKeyHandle {
 	return &HMACKeyHandle{
 		projectID: projectID,
 		accessID:  accessID,
+		retry:     c.retry,
 		raw:       raw.NewProjectsHmacKeysService(c.raw),
 	}
 }
@@ -126,10 +127,10 @@ func (hkh *HMACKeyHandle) Get(ctx context.Context, opts ...HMACKeyOption) (*HMAC
 
 	var metadata *raw.HmacKeyMetadata
 	var err error
-	err = runWithRetry(ctx, func() error {
+	err = run(ctx, func() error {
 		metadata, err = call.Context(ctx).Do()
 		return err
-	})
+	}, hkh.retry, true)
 	if err != nil {
 		return nil, err
 	}
@@ -156,9 +157,9 @@ func (hkh *HMACKeyHandle) Delete(ctx context.Context, opts ...HMACKeyOption) err
 	}
 	setClientHeader(delCall.Header())
 
-	return runWithRetry(ctx, func() error {
+	return run(ctx, func() error {
 		return delCall.Context(ctx).Do()
-	})
+	}, hkh.retry, true)
 }
 
 func pbHmacKeyToHMACKey(pb *raw.HmacKey, updatedTimeCanBeNil bool) (*HMACKey, error) {
@@ -214,8 +215,13 @@ func (c *Client) CreateHMACKey(ctx context.Context, projectID, serviceAccountEma
 
 	setClientHeader(call.Header())
 
-	hkPb, err := call.Context(ctx).Do()
-	if err != nil {
+	var hkPb *raw.HmacKey
+
+	if err := run(ctx, func() error {
+		h, err := call.Context(ctx).Do()
+		hkPb = h
+		return err
+	}, c.retry, false); err != nil {
 		return nil, err
 	}
 
@@ -257,10 +263,11 @@ func (h *HMACKeyHandle) Update(ctx context.Context, au HMACKeyAttrsToUpdate, opt
 
 	var metadata *raw.HmacKeyMetadata
 	var err error
-	err = runWithRetry(ctx, func() error {
+	isIdempotent := len(au.Etag) > 0
+	err = run(ctx, func() error {
 		metadata, err = call.Context(ctx).Do()
 		return err
-	})
+	}, h.retry, isIdempotent)
 
 	if err != nil {
 		return nil, err
@@ -285,6 +292,7 @@ type HMACKeysIterator struct {
 	nextFunc  func() error
 	index     int
 	desc      hmacKeyDesc
+	retry     *retryConfig
 }
 
 // ListHMACKeys returns an iterator for listing HMACKeys.
@@ -297,6 +305,7 @@ func (c *Client) ListHMACKeys(ctx context.Context, projectID string, opts ...HMA
 		ctx:       ctx,
 		raw:       raw.NewProjectsHmacKeysService(c.raw),
 		projectID: projectID,
+		retry:     c.retry,
 	}
 
 	for _, opt := range opts {
@@ -361,10 +370,10 @@ func (it *HMACKeysIterator) fetch(pageSize int, pageToken string) (token string,
 
 	ctx := it.ctx
 	var resp *raw.HmacKeysMetadata
-	err = runWithRetry(it.ctx, func() error {
+	err = run(it.ctx, func() error {
 		resp, err = call.Context(ctx).Do()
 		return err
-	})
+	}, it.retry, true)
 	if err != nil {
 		return "", err
 	}

storage/retry_conformance_test.go

@@ -252,12 +252,14 @@ var methods = map[string][]retryFunc{
 	"storage.hmacKey.update": {
 		func(ctx context.Context, c *Client, fs *resources, preconditions bool) error {
 			key := c.HMACKeyHandle(projectID, fs.hmacKey.AccessID)
+			uattrs := HMACKeyAttrsToUpdate{State: "INACTIVE"}
 
-			_, err := key.Update(ctx, HMACKeyAttrsToUpdate{State: "INACTIVE"})
-			if err != nil {
-				return err
+			if preconditions {
+				uattrs.Etag = fs.hmacKey.Etag
 			}
-			return fmt.Errorf("Etag preconditions not supported")
+
+			_, err := key.Update(ctx, uattrs)
+			return err
 		},
 	},
 	"storage.objects.compose": {

storage/storage.go

@@ -109,6 +109,7 @@ type Client struct {
 	readHost string
 	// May be nil.
 	creds *google.Credentials
+	retry *retryConfig
 
 	// gc is an optional gRPC-based, GAPIC client.
 	//
@@ -1792,7 +1793,8 @@ func setConditionField(call reflect.Value, name string, value interface{}) bool
 // on the new handle will use the customized retry configuration.
 // These retry options will merge with the bucket's retryer (if set) for the
 // returned handle. Options passed into this method will take precedence over
-// options on the bucket's retryer.
+// retry options on the bucket and client. Note that you must explicitly pass in
+// each option you want to override.
 func (o *ObjectHandle) Retryer(opts ...RetryOption) *ObjectHandle {
 	o2 := *o
 	var retry *retryConfig
@@ -1810,6 +1812,27 @@ func (o *ObjectHandle) Retryer(opts ...RetryOption) *ObjectHandle {
 	return &o2
 }
 
+// SetRetry configures the client with custom retry behavior as specified by the
+// options that are passed to it. All operations using this client will use the
+// customized retry configuration.
+// This should be called once before using the client for network operations, as
+// there could be indeterminate behaviour with operations in progress.
+// Retry options set on a bucket or object handle will take precedence over
+// these options.
+func (c *Client) SetRetry(opts ...RetryOption) {
+	var retry *retryConfig
+	if c.retry != nil {
+		// merge the options with the existing retry
+		retry = c.retry
+	} else {
+		retry = &retryConfig{}
+	}
+	for _, opt := range opts {
+		opt.apply(retry)
+	}
+	c.retry = retry
+}
+
 // RetryOption allows users to configure non-default retry behavior for API
 // calls made to GCS.
 type RetryOption interface {
@@ -1971,10 +1994,10 @@ func (c *Client) ServiceAccount(ctx context.Context, projectID string) (string,
 	r := c.raw.Projects.ServiceAccount.Get(projectID)
 	var res *raw.ServiceAccount
 	var err error
-	err = runWithRetry(ctx, func() error {
+	err = run(ctx, func() error {
 		res, err = r.Context(ctx).Do()
 		return err
-	})
+	}, c.retry, true)
 	if err != nil {
 		return "", err
 	}