fix(storage): Handle request transformation correctly after BidiWriteObject redirects

This commit addresses an issue where BidiWriteObject operations, particularly in cross-region redirect scenarios, would fail with a NOT_FOUND error. The root cause was that bucket information was not propagated in the request headers after the client received a redirect error from the service.

Author: kalragauri

google/cloud/storage/internal/async/connection_impl.cc

@@ -319,9 +319,11 @@ AsyncConnectionImpl::AppendableObjectUploadImpl(AppendableUploadParams p) {
   auto factory = WriteResultFactory(
       [stub = stub_, cq = cq_, retry = std::move(retry),
        // NOLINTNEXTLINE(bugprone-lambda-function-name)
-       backoff = std::move(backoff), current, function_name = __func__](
+       backoff = std::move(backoff), current, function_name = __func__,
+       // Use shared_ptr to propagate RoutingHeaderOptions across retries.
+       current_routing_options = std::make_shared<RoutingHeaderOptions>()](
           google::storage::v2::BidiWriteObjectRequest req) {
-        auto call = [stub, request = std::move(req)](
+        auto call = [stub, request = std::move(req), current_routing_options](
                         CompletionQueue& cq,
                         std::shared_ptr<grpc::ClientContext> context,
                         google::cloud::internal::ImmutableOptions options,
@@ -336,9 +338,11 @@ AsyncConnectionImpl::AppendableObjectUploadImpl(AppendableUploadParams p) {
 
           // Apply the routing header
           if (request.has_write_object_spec())
-            ApplyRoutingHeaders(*context, request.write_object_spec());
+            ApplyRoutingHeaders(*context, request.write_object_spec(),
+                                *current_routing_options);
           else
-            ApplyRoutingHeaders(*context, request.append_object_spec());
+            ApplyRoutingHeaders(*context, request.append_object_spec(),
+                                *current_routing_options);
 
           auto rpc = stub->AsyncBidiWriteObject(cq, std::move(context),
                                                 std::move(options));
@@ -347,18 +351,23 @@ AsyncConnectionImpl::AppendableObjectUploadImpl(AppendableUploadParams p) {
               std::move(rpc));
           request.set_state_lookup(true);
           auto open = std::make_shared<WriteObject>(std::move(rpc), request);
-          return open->Call().then([open, &request](auto f) mutable {
-            open.reset();
-            auto response = f.get();
-            if (!response) {
-              google::rpc::Status grpc_status =
-                  ExtractGrpcStatus(response.status());
-              EnsureFirstMessageAppendObjectSpec(request, grpc_status);
-              ApplyWriteRedirectErrors(*request.mutable_append_object_spec(),
-                                       grpc_status);
-            }
-            return response;
-          });
+          return open->Call().then(
+              [open, &request, current_routing_options](auto f) mutable {
+                open.reset();
+                auto response = f.get();
+                if (!response) {
+                  google::rpc::Status grpc_status =
+                      ExtractGrpcStatus(response.status());
+                  // Handle redirect and get info for updating routing options.
+                  BidiWriteRedirectInfo redirect_info =
+                      HandleBidiWriteRedirect(request, grpc_status);
+
+                  // Update RoutingHeaderOptions for the next attempt.
+                  current_routing_options->routing_token =
+                      redirect_info.routing_token;
+                }
+                return response;
+              });
         };
 
         return google::cloud::internal::AsyncRetryLoop(
@@ -371,7 +380,7 @@ AsyncConnectionImpl::AppendableObjectUploadImpl(AppendableUploadParams p) {
   return pending.then(
       [current, request = std::move(p.request), persisted_size,
        hash = std::move(hash_function), fa = std::move(factory)](auto f) mutable
-      -> StatusOr<std::unique_ptr<storage::AsyncWriterConnection>> {
+          -> StatusOr<std::unique_ptr<storage::AsyncWriterConnection>> {
         auto rpc = f.get();
         if (!rpc) return std::move(rpc).status();
         std::unique_ptr<AsyncWriterConnectionImpl> impl;
@@ -427,7 +436,7 @@ AsyncConnectionImpl::StartBufferedUpload(UploadParams p) {
   return StartUnbufferedUpload(std::move(p))
       .then([current = std::move(current),
              async_write_object = std::move(async_write_object)](auto f) mutable
-            -> StatusOr<std::unique_ptr<storage::AsyncWriterConnection>> {
+                -> StatusOr<std::unique_ptr<storage::AsyncWriterConnection>> {
         auto w = f.get();
         if (!w) return std::move(w).status();
         auto factory = [upload_id = (*w)->UploadId(),
@@ -461,14 +470,15 @@ AsyncConnectionImpl::ResumeBufferedUpload(ResumeUploadParams p) {
   };
 
   auto f = make_unbuffered();
-  return f.then([current = std::move(current),
-                 make_unbuffered = std::move(make_unbuffered)](auto f) mutable
-                -> StatusOr<std::unique_ptr<storage::AsyncWriterConnection>> {
-    auto w = f.get();
-    if (!w) return std::move(w).status();
-    return MakeWriterConnectionBuffered(std::move(make_unbuffered),
-                                        *std::move(w), *current);
-  });
+  return f.then(
+      [current = std::move(current),
+       make_unbuffered = std::move(make_unbuffered)](auto f) mutable
+          -> StatusOr<std::unique_ptr<storage::AsyncWriterConnection>> {
+        auto w = f.get();
+        if (!w) return std::move(w).status();
+        return MakeWriterConnectionBuffered(std::move(make_unbuffered),
+                                            *std::move(w), *current);
+      });
 }
 
 future<StatusOr<google::storage::v2::Object>>