fix(storage): add backoff to gRPC write retries

storage/bucket_test.go

@@ -1116,11 +1116,11 @@ func TestBucketRetryer(t *testing.T) {
 					WithErrorFunc(func(err error) bool { return false }))
 			},
 			want: &retryConfig{
-				backoff: &gax.Backoff{
+				backoff: gaxBackoffFromStruct(&gax.Backoff{
 					Initial:    2 * time.Second,
 					Max:        30 * time.Second,
 					Multiplier: 3,
-				},
+				}),
 				policy:      RetryAlways,
 				maxAttempts: expectedAttempts(5),
 				shouldRetry: func(err error) bool { return false },
@@ -1135,9 +1135,9 @@ func TestBucketRetryer(t *testing.T) {
 					}))
 			},
 			want: &retryConfig{
-				backoff: &gax.Backoff{
+				backoff: gaxBackoffFromStruct(&gax.Backoff{
 					Multiplier: 3,
-				}},
+				})},
 		},
 		{
 			name: "set policy only",

storage/client_test.go

@@ -1396,7 +1396,7 @@ func TestRetryMaxAttemptsEmulated(t *testing.T) {
 		instructions := map[string][]string{"storage.buckets.get": {"return-503", "return-503", "return-503", "return-503", "return-503"}}
 		testID := createRetryTest(t, client, instructions)
 		ctx = callctx.SetHeaders(ctx, "x-retry-test-id", testID)
-		config := &retryConfig{maxAttempts: expectedAttempts(3), backoff: &gax.Backoff{Initial: 10 * time.Millisecond}}
+		config := &retryConfig{maxAttempts: expectedAttempts(3), backoff: gaxBackoffFromStruct(&gax.Backoff{Initial: 10 * time.Millisecond})}
 		_, err = client.GetBucket(ctx, bucket, nil, idempotent(true), withRetryConfig(config))
 
 		var ae *apierror.APIError
@@ -1421,7 +1421,7 @@ func TestTimeoutErrorEmulated(t *testing.T) {
 		ctx, cancel := context.WithTimeout(ctx, time.Nanosecond)
 		defer cancel()
 		time.Sleep(5 * time.Nanosecond)
-		config := &retryConfig{backoff: &gax.Backoff{Initial: 10 * time.Millisecond}}
+		config := &retryConfig{backoff: gaxBackoffFromStruct(&gax.Backoff{Initial: 10 * time.Millisecond})}
 		_, err := client.GetBucket(ctx, bucket, nil, idempotent(true), withRetryConfig(config))
 
 		// Error may come through as a context.DeadlineExceeded (HTTP) or status.DeadlineExceeded (gRPC)
@@ -1447,7 +1447,7 @@ func TestRetryDeadlineExceedeEmulated(t *testing.T) {
 		instructions := map[string][]string{"storage.buckets.get": {"return-504", "return-504"}}
 		testID := createRetryTest(t, client, instructions)
 		ctx = callctx.SetHeaders(ctx, "x-retry-test-id", testID)
-		config := &retryConfig{maxAttempts: expectedAttempts(4), backoff: &gax.Backoff{Initial: 10 * time.Millisecond}}
+		config := &retryConfig{maxAttempts: expectedAttempts(4), backoff: gaxBackoffFromStruct(&gax.Backoff{Initial: 10 * time.Millisecond})}
 		if _, err := client.GetBucket(ctx, bucket, nil, idempotent(true), withRetryConfig(config)); err != nil {
 			t.Fatalf("GetBucket: got unexpected error %v, want nil", err)
 		}

storage/grpc_client.go

@@ -1184,8 +1184,6 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 		gw.ctx = setUserProjectMetadata(gw.ctx, s.userProject)
 	}
 
-	gw.initializeRetryConfig()
-
 	// This function reads the data sent to the pipe and sends sets of messages
 	// on the gRPC client-stream as the buffer is filled.
 	go func() {
@@ -1226,7 +1224,7 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 				}
 			}
 
-			o, off, err := gw.uploadBuffer(recvd, offset, doneReading)
+			o, off, err := gw.uploadBuffer(recvd, offset, doneReading, newUploadBufferRetryConfig(gw.settings))
 			if err != nil {
 				err = checkCanceled(err)
 				errorf(err)
@@ -2047,12 +2045,6 @@ type gRPCWriter struct {
 
 	// The Resumable Upload ID started by a gRPC-based Writer.
 	upid string
-
-	// Retry tracking vars.
-	backoff      gax.Backoff // backoff for all retries of write calls
-	attempts     int
-	invocationID string
-	lastErr      error
 }
 
 // startResumableUpload initializes a Resumable Upload with gRPC and sets the
@@ -2100,7 +2092,7 @@ func (w *gRPCWriter) queryProgress() (int64, error) {
 // completed.
 //
 // Returns object, persisted size, and any error that is not retriable.
-func (w *gRPCWriter) uploadBuffer(recvd int, start int64, doneReading bool) (*storagepb.Object, int64, error) {
+func (w *gRPCWriter) uploadBuffer(recvd int, start int64, doneReading bool, retryConfig *uploadBufferRetryConfig) (*storagepb.Object, int64, error) {
 	var err error
 	var lastWriteOfEntireObject bool
 
@@ -2147,7 +2139,7 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 		if w.stream == nil {
 			hds := []string{"x-goog-request-params", fmt.Sprintf("bucket=projects/_/buckets/%s", url.QueryEscape(w.bucket))}
 			ctx := gax.InsertMetadataIntoOutgoingContext(w.ctx, hds...)
-			ctx = setInvocationHeaders(ctx, w.invocationID, w.attempts)
+			ctx = setInvocationHeaders(ctx, retryConfig.invocationID, retryConfig.attempts)
 
 			w.stream, err = w.c.raw.BidiWriteObject(ctx)
 			if err != nil {
@@ -2172,7 +2164,6 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 				req.ObjectChecksums = toProtoChecksums(w.sendCRC32C, w.attrs)
 			}
 		}
-		w.stream.Context()
 
 		err = w.stream.Send(req)
 		if err == io.EOF {
@@ -2194,9 +2185,11 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 			// Retriable errors mean we should start over and attempt to
 			// resend the entire buffer via a new stream.
 			// If not retriable, falling through will return the error received.
-			var shouldRetry bool
-			shouldRetry, err = w.shouldRetry(w.ctx, err)
-			if shouldRetry {
+			err = retryConfig.retriable(w.ctx, err)
+
+			if err == nil {
+				retryConfig.doBackOff(w.ctx)
+
 				// TODO: Add test case for failure modes of querying progress.
 				writeOffset, err = w.determineOffset(start)
 				if err != nil {
@@ -2238,13 +2231,17 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 		if !lastWriteOfEntireObject {
 			resp, err := w.stream.Recv()
 
-			// Retriable errors mean we should start over and attempt to
-			// resend the entire buffer via a new stream.
-			// If not retriable, falling through will return the error received
-			// from closing the stream.
-			var shouldRetry bool
-			shouldRetry, err = w.shouldRetry(w.ctx, err)
-			if shouldRetry {
+			if err != nil {
+				// Retriable errors mean we should start over and attempt to
+				// resend the entire buffer via a new stream.
+				// If not retriable, falling through will return the error received
+				// from closing the stream.
+				err = retryConfig.retriable(w.ctx, err)
+				if err != nil {
+					return nil, 0, err
+				}
+
+				retryConfig.doBackOff(w.ctx)
 				writeOffset, err = w.determineOffset(start)
 				if err != nil {
 					return nil, 0, err
@@ -2256,9 +2253,6 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 
 				continue sendBytes
 			}
-			if err != nil {
-				return nil, 0, err
-			}
 
 			if resp.GetPersistedSize() != writeOffset {
 				// Retry if not all bytes were persisted.
@@ -2284,9 +2278,14 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 			var obj *storagepb.Object
 			for obj == nil {
 				resp, err := w.stream.Recv()
-				var shouldRetry bool
-				shouldRetry, err = w.shouldRetry(w.ctx, err)
-				if shouldRetry {
+
+				if err != nil {
+					err = retryConfig.retriable(w.ctx, err)
+					if err != nil {
+						return nil, 0, err
+					}
+					retryConfig.doBackOff(w.ctx)
+
 					writeOffset, err = w.determineOffset(start)
 					if err != nil {
 						return nil, 0, err
@@ -2295,9 +2294,6 @@ sendBytes: // label this loop so that we can use a continue statement from a nes
 					w.stream = nil
 					continue sendBytes
 				}
-				if err != nil {
-					return nil, 0, err
-				}
 
 				obj = resp.GetResource()
 			}
@@ -2383,77 +2379,88 @@ func checkCanceled(err error) error {
 	return err
 }
 
-func (w *gRPCWriter) initializeRetryConfig() {
-	if w.attempts == 0 {
-		w.attempts = 1
+type uploadBufferRetryConfig struct {
+	attempts     int
+	invocationID string
+	config       *retryConfig
+	lastErr      error
+}
 
-		if w.settings.retry == nil {
-			w.settings.retry = defaultRetry
-		}
+func newUploadBufferRetryConfig(settings *settings) *uploadBufferRetryConfig {
+	config := settings.retry
 
-		if w.settings.retry.shouldRetry == nil {
-			w.settings.retry.shouldRetry = ShouldRetry
-		}
+	if config == nil {
+		config = defaultRetry.clone()
+	}
 
-		w.invocationID = uuid.New().String()
+	if config.shouldRetry == nil {
+		config.shouldRetry = ShouldRetry
+	}
 
-		if w.settings.retry.backoff != nil {
-			w.backoff.Multiplier = w.settings.retry.backoff.Multiplier
-			w.backoff.Initial = w.settings.retry.backoff.Initial
-			w.backoff.Max = w.settings.retry.backoff.Max
-		}
+	if config.backoff == nil {
+		config.backoff = &gaxBackoff{}
+	} else {
+		config.backoff.SetMultiplier(settings.retry.backoff.GetMultiplier())
+		config.backoff.SetInitial(settings.retry.backoff.GetInitial())
+		config.backoff.SetMax(settings.retry.backoff.GetMax())
+	}
+
+	return &uploadBufferRetryConfig{
+		attempts:     1,
+		invocationID: uuid.New().String(),
+		config:       config,
 	}
 }
 
-// shouldRetry determines if a retry is necessary and if so waits the appropriate
-// amount of time. It returns true if the error is retryable or the error to be
-// surfaced to the user if not.
-func (w *gRPCWriter) shouldRetry(ctx context.Context, err error) (bool, error) {
+// retriable determines if a retry is necessary and if so returns a nil error;
+// otherwise it returns the error to be surfaced to the user.
+func (retry *uploadBufferRetryConfig) retriable(ctx context.Context, err error) error {
 	if err == nil {
 		// a nil err does not need to be retried
-		return false, nil
+		return nil
 	}
 	if err != context.Canceled && err != context.DeadlineExceeded {
-		w.lastErr = err
+		retry.lastErr = err
 	}
 
-	retryConfig := w.settings.retry
-
-	if retryConfig.policy == RetryNever {
-		return false, err
+	if retry.config.policy == RetryNever {
+		return err
 	}
 
-	if retryConfig.maxAttempts != nil && w.attempts >= *retryConfig.maxAttempts {
-		return false, fmt.Errorf("storage: retry failed after %v attempts; last error: %w", w.attempts, err)
+	if retry.config.maxAttempts != nil && retry.attempts >= *retry.config.maxAttempts {
+		return fmt.Errorf("storage: retry failed after %v attempts; last error: %w", retry.attempts, err)
 	}
 
-	w.attempts++
+	retry.attempts++
 
-	retryable := retryConfig.shouldRetry(err)
 	// Explicitly check context cancellation so that we can distinguish between a
 	// DEADLINE_EXCEEDED error from the server and a user-set context deadline.
 	// Unfortunately gRPC will codes.DeadlineExceeded (which may be retryable if it's
 	// sent by the server) in both cases.
 	ctxErr := ctx.Err()
 	if errors.Is(ctxErr, context.Canceled) || errors.Is(ctxErr, context.DeadlineExceeded) {
-		retryable = false
-
-		if w.lastErr != nil {
-			return false, fmt.Errorf("retry failed with %v; last error: %w", ctxErr, w.lastErr)
+		if retry.lastErr != nil {
+			return fmt.Errorf("retry failed with %v; last error: %w", ctxErr, retry.lastErr)
 		}
-		return false, ctxErr
+		return ctxErr
+	}
+
+	if !retry.config.shouldRetry(err) {
+		return err
 	}
+	return nil
+}
 
-	if retryable {
-		p := w.backoff.Pause()
+// doBackOff pauses for the appropriate amount of time; it should be called after
+// encountering a retriable error.
+func (retry *uploadBufferRetryConfig) doBackOff(ctx context.Context) error {
+	p := retry.config.backoff.Pause()
 
-		if ctxErr := gax.Sleep(ctx, p); ctxErr != nil {
-			if w.lastErr != nil {
-				return false, fmt.Errorf("retry failed with %v; last error: %w", ctxErr, w.lastErr)
-			}
-			return false, ctxErr
+	if ctxErr := gax.Sleep(ctx, p); ctxErr != nil {
+		if retry.lastErr != nil {
+			return fmt.Errorf("retry failed with %v; last error: %w", ctxErr, retry.lastErr)
 		}
+		return ctxErr
 	}
-
-	return retryable, err
+	return nil
 }