fix(storage): Fix retries for redirection errors. (#12093)

cjc25 · tritone · web-flow · commit 3e177e755f5b · 2025-05-01T15:28:39.000-04:00
* fix(storage): Fix retries for redirection errors.

Retrying inline might resending buffers where aren't stable yet. Surface
redirects to the stream-level retry machinery to catch the resend logic.

* fix check against idempotency to not lose the shouldRetry function if !s.idempotent but policy == RetryAlways

---------

Co-authored-by: Chris Cotter &lt;cjcotter@google.com&gt;
diff --git a/storage/grpc_client.go b/storage/grpc_client.go
@@ -1419,11 +1419,7 @@ func (mrd *gRPCBidiReader) activeRange() []mrdRange {
 
 // retryStream cancel's stream and reopen the stream again.
 func (mrd *gRPCBidiReader) retryStream(err error) error {
-	var shouldRetry = ShouldRetry
-	if mrd.settings.retry != nil && mrd.settings.retry.shouldRetry != nil {
-		shouldRetry = mrd.settings.retry.shouldRetry
-	}
-	if shouldRetry(err) {
+	if mrd.settings.retry.runShouldRetry(err) {
 		// This will "close" the existing stream and immediately attempt to
 		// reopen the stream, but will backoff if further attempts are necessary.
 		// When Reopening the stream only failed readID will be added to stream.
@@ -2032,11 +2028,7 @@ func (r *gRPCReader) Close() error {
 func (r *gRPCReader) recv() error {
 	databufs := mem.BufferSlice{}
 	err := r.stream.RecvMsg(&databufs)
-	var shouldRetry = ShouldRetry
-	if r.settings.retry != nil && r.settings.retry.shouldRetry != nil {
-		shouldRetry = r.settings.retry.shouldRetry
-	}
-	if err != nil && shouldRetry(err) {
+	if err != nil && r.settings.retry.runShouldRetry(err) {
 		// This will "close" the existing stream and immediately attempt to
 		// reopen the stream, but will backoff if further attempts are necessary.
 		// Reopening the stream Recvs the first message, so if retrying is
diff --git a/storage/grpc_reader.go b/storage/grpc_reader.go
@@ -424,11 +424,7 @@ func (r *gRPCReadObjectReader) recv() error {
 	databufs := mem.BufferSlice{}
 	err := r.stream.RecvMsg(&databufs)
 
-	var shouldRetry = ShouldRetry
-	if r.settings.retry != nil && r.settings.retry.shouldRetry != nil {
-		shouldRetry = r.settings.retry.shouldRetry
-	}
-	if err != nil && shouldRetry(err) {
+	if err != nil && r.settings.retry.runShouldRetry(err) {
 		// This will "close" the existing stream and immediately attempt to
 		// reopen the stream, but will backoff if further attempts are necessary.
 		// Reopening the stream Recvs the first message, so if retrying is
diff --git a/storage/grpc_writer.go b/storage/grpc_writer.go
@@ -42,6 +42,40 @@ const (
 	maxPerMessageWriteSize int = int(storagepb.ServiceConstants_MAX_WRITE_CHUNK_BYTES)
 )
 
+func withBidiWriteObjectRedirectionErrorRetries(s *settings) (newr *retryConfig) {
+	oldr := s.retry
+	newr = oldr.clone()
+	if newr == nil {
+		newr = &retryConfig{}
+	}
+	if (oldr.policy == RetryIdempotent && !s.idempotent) || oldr.policy == RetryNever {
+		// We still retry redirection errors even when settings indicate not to
+		// retry.
+		//
+		// The protocol requires us to respect redirection errors, so RetryNever has
+		// to ignore them.
+		//
+		// Idempotency is always protected by redirection errors: they either
+		// contain a handle which can be used as idempotency information, or they do
+		// not contain a handle and are "affirmative failures" which indicate that
+		// no server-side action occurred.
+		newr.policy = RetryAlways
+		newr.shouldRetry = func(err error) bool {
+			return errors.Is(err, bidiWriteObjectRedirectionError{})
+		}
+		return newr
+	}
+	// If retry settings allow retries normally, fall back to that behavior.
+	newr.shouldRetry = func(err error) bool {
+		if errors.Is(err, bidiWriteObjectRedirectionError{}) {
+			return true
+		}
+		v := oldr.runShouldRetry(err)
+		return v
+	}
+	return newr
+}
+
 func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storageOption) (*io.PipeWriter, error) {
 	var offset int64
 	errorf := params.setError
@@ -58,6 +92,9 @@ func (c *grpcStorageClient) OpenWriter(params *openWriterParams, opts ...storage
 	if s.retry == nil {
 		s.retry = defaultRetry.clone()
 	}
+	if params.append {
+		s.retry = withBidiWriteObjectRedirectionErrorRetries(s)
+	}
 	s.retry.maxRetryDuration = retryDeadline
 
 	// Set Flush func for use by exported Writer.Flush.
@@ -723,6 +760,10 @@ func (s *gRPCAppendBidiWriteBufferSender) connect(ctx context.Context) (err erro
 		if s.firstMessage.GetAppendObjectSpec().GetGeneration() != 0 {
 			return nil
 		}
+		// Also always ok to reconnect if we've seen a redirect token
+		if s.routingToken != nil {
+			return nil
+		}
 
 		// We can also reconnect if the first message has an if_generation_match or
 		// if_metageneration_match condition. Note that negative conditions like
@@ -805,7 +846,7 @@ func (s *gRPCAppendBidiWriteBufferSender) maybeUpdateFirstMessage(resp *storagep
 type bidiWriteObjectRedirectionError struct{}
 
 func (e bidiWriteObjectRedirectionError) Error() string {
-	return "BidiWriteObjectRedirectedError"
+	return ""
 }
 
 func (s *gRPCAppendBidiWriteBufferSender) handleRedirectionError(e *storagepb.BidiWriteObjectRedirectedError) bool {
@@ -850,10 +891,10 @@ func (s *gRPCAppendBidiWriteBufferSender) receiveMessages(resps chan<- *storagep
 	if st, ok := status.FromError(err); ok && st.Code() == codes.Aborted {
 		for _, d := range st.Details() {
 			if e, ok := d.(*storagepb.BidiWriteObjectRedirectedError); ok {
-				// If we can handle this error, replace it with the sentinel. Otherwise,
-				// report it to the user.
+				// If we can handle this error, wrap it with the sentinel so it gets
+				// retried.
 				if ok := s.handleRedirectionError(e); ok {
-					err = bidiWriteObjectRedirectionError{}
+					err = fmt.Errorf("%w%w", bidiWriteObjectRedirectionError{}, err)
 				}
 			}
 		}
@@ -971,12 +1012,6 @@ func (s *gRPCAppendBidiWriteBufferSender) sendBuffer(ctx context.Context, buf []
 			err = s.recvErr
 		}
 		s.stream = nil
-
-		// Retry transparently on a redirection error
-		if _, ok := err.(bidiWriteObjectRedirectionError); ok {
-			s.forceFirstMessage = true
-			continue
-		}
 		return
 	}
 }
diff --git a/storage/invoke.go b/storage/invoke.go
@@ -53,6 +53,13 @@ var (
 	})
 )
 
+func (r *retryConfig) runShouldRetry(err error) bool {
+	if r == nil || r.shouldRetry == nil {
+		return ShouldRetry(err)
+	}
+	return r.shouldRetry(err)
+}
+
 // run determines whether a retry is necessary based on the config and
 // idempotency information. It then calls the function with or without retries
 // as appropriate, using the configured settings.
@@ -73,10 +80,6 @@ func run(ctx context.Context, call func(ctx context.Context) error, retry *retry
 		bo.Initial = retry.backoff.Initial
 		bo.Max = retry.backoff.Max
 	}
-	var errorFunc func(err error) bool = ShouldRetry
-	if retry.shouldRetry != nil {
-		errorFunc = retry.shouldRetry
-	}
 
 	var quitAfterTimer *time.Timer
 	if retry.maxRetryDuration != 0 {
@@ -103,7 +106,7 @@ func run(ctx context.Context, call func(ctx context.Context) error, retry *retry
 			return true, fmt.Errorf("storage: retry failed after %v attempts; last error: %w", *retry.maxAttempts, lastErr)
 		}
 		attempts++
-		retryable := errorFunc(lastErr)
+		retryable := retry.runShouldRetry(lastErr)
 		// 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
