impl(ACv2): Add retry with routing_token for rpc Start (#91)

Author: bajajneha27

google/cloud/storage/google_cloud_cpp_storage_grpc.bzl

@@ -59,6 +59,7 @@ google_cloud_cpp_storage_grpc_hdrs = [
     "internal/async/rewriter_connection_impl.h",
     "internal/async/rewriter_connection_tracing.h",
     "internal/async/token_impl.h",
+    "internal/async/write_object.h",
     "internal/async/write_payload_fwd.h",
     "internal/async/write_payload_impl.h",
     "internal/async/writer_connection_buffered.h",
@@ -132,6 +133,7 @@ google_cloud_cpp_storage_grpc_srcs = [
     "internal/async/rewriter_connection_impl.cc",
     "internal/async/rewriter_connection_tracing.cc",
     "internal/async/token_impl.cc",
+    "internal/async/write_object.cc",
     "internal/async/writer_connection_buffered.cc",
     "internal/async/writer_connection_finalized.cc",
     "internal/async/writer_connection_impl.cc",

google/cloud/storage/google_cloud_cpp_storage_grpc.cmake

@@ -141,6 +141,8 @@ add_library(
     internal/async/rewriter_connection_tracing.h
     internal/async/token_impl.cc
     internal/async/token_impl.h
+    internal/async/write_object.cc
+    internal/async/write_object.h
     internal/async/write_payload_fwd.h
     internal/async/write_payload_impl.h
     internal/async/writer_connection_buffered.cc

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

@@ -28,6 +28,7 @@
 #include "google/cloud/storage/internal/async/reader_connection_impl.h"
 #include "google/cloud/storage/internal/async/reader_connection_resume.h"
 #include "google/cloud/storage/internal/async/rewriter_connection_impl.h"
+#include "google/cloud/storage/internal/async/write_object.h"
 #include "google/cloud/storage/internal/async/write_payload_impl.h"
 #include "google/cloud/storage/internal/async/writer_connection_buffered.h"
 #include "google/cloud/storage/internal/async/writer_connection_finalized.h"
@@ -292,44 +293,46 @@ AsyncConnectionImpl::ReadObjectRange(ReadObjectParams p) {
 future<StatusOr<std::unique_ptr<storage_experimental::AsyncWriterConnection>>>
 AsyncConnectionImpl::StartAppendableObjectUpload(AppendableUploadParams p) {
   auto current = internal::MakeImmutableOptions(std::move(p.options));
+  auto request = p.request;
   std::int64_t persisted_size = 0;
   auto hash_function = CreateHashFunction(*current);
-  using StreamingRpc = AsyncWriterConnectionImpl::StreamingRpc;
+  auto retry = retry_policy(*current);
+  auto backoff = backoff_policy(*current);
   struct RequestPlaceholder {};
 
-  auto call = [stub = stub_](CompletionQueue& cq,
-                             std::shared_ptr<grpc::ClientContext> context,
-                             google::cloud::internal::ImmutableOptions options,
-                             RequestPlaceholder const&)
-      -> future<StatusOr<std::unique_ptr<StreamingRpc>>> {
+  auto call = [stub = stub_, request = std::move(request)](
+                  CompletionQueue& cq,
+                  std::shared_ptr<grpc::ClientContext> context,
+                  google::cloud::internal::ImmutableOptions options,
+                  RequestPlaceholder const&) mutable
+      -> future<StatusOr<WriteObject::WriteResult>> {
     auto rpc =
         stub->AsyncBidiWriteObject(cq, std::move(context), std::move(options));
-    auto start = rpc->Start();
-    return start.then([rpc = std::move(rpc)](auto f) mutable {
-      if (f.get()) return make_ready_future(make_status_or(std::move(rpc)));
-      auto pending = rpc->Finish();
-      return pending.then([rpc = std::move(rpc)](auto f) mutable {
-        rpc.reset();
-        auto status = f.get();
-        return StatusOr<std::unique_ptr<StreamingRpc>>(std::move(status));
-      });
+    request.set_state_lookup(true);
+    auto open = std::make_shared<WriteObject>(std::move(rpc), request);
+    return open->Call().then([open, &request](auto f) mutable {
+      auto response = f.get();
+      if (!response) {
+        EnsureFirstMessageAppendObjectSpec(request);
+        ApplyWriteRedirectErrors(*request.mutable_append_object_spec(),
+                                 ExtractGrpcStatus(response.status()));
+      }
+      return response;
     });
   };
 
-  auto transform = [current, request = std::move(p.request), persisted_size,
+  auto transform = [current, request, persisted_size,
                     hash = std::move(hash_function)](auto f) mutable
       -> StatusOr<
           std::unique_ptr<storage_experimental::AsyncWriterConnection>> {
     auto rpc = f.get();
     if (!rpc) return std::move(rpc).status();
     return std::unique_ptr<storage_experimental::AsyncWriterConnection>(
         std::make_unique<AsyncWriterConnectionImpl>(
-            current, std::move(request), *std::move(rpc), std::move(hash),
-            persisted_size));
+            current, std::move(request), std::move(rpc->stream),
+            std::move(hash), persisted_size, false));
   };
 
-  auto retry = retry_policy(*current);
-  auto backoff = backoff_policy(*current);
   return google::cloud::internal::AsyncRetryLoop(
              std::move(retry), std::move(backoff), Idempotency::kIdempotent,
              cq_, std::move(call), std::move(current), RequestPlaceholder{},
@@ -716,7 +719,7 @@ AsyncConnectionImpl::UnbufferedUploadImpl(
     return std::unique_ptr<storage_experimental::AsyncWriterConnection>(
         std::make_unique<AsyncWriterConnectionImpl>(
             current, std::move(request), *std::move(rpc), std::move(hash),
-            persisted_size));
+            persisted_size, true));
   };
 
   auto retry = retry_policy(*current);

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

@@ -20,6 +20,16 @@ namespace cloud {
 namespace storage_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
+void EnsureFirstMessageAppendObjectSpec(
+    google::storage::v2::BidiWriteObjectRequest& request) {
+  if (request.has_write_object_spec()) {
+    auto spec = request.write_object_spec();
+    auto& append_object_spec = *request.mutable_append_object_spec();
+    append_object_spec.set_bucket(spec.resource().bucket());
+    append_object_spec.set_object(spec.resource().name());
+  }
+}
+
 google::rpc::Status ExtractGrpcStatus(Status const& status) {
   auto proto_status = google::rpc::Status{};
   auto payload = google::cloud::internal::GetPayload(
@@ -38,6 +48,16 @@ void ApplyRedirectErrors(google::storage::v2::BidiReadObjectSpec& spec,
   }
 }
 
+void ApplyWriteRedirectErrors(google::storage::v2::AppendObjectSpec& spec,
+                              google::rpc::Status const& rpc_status) {
+  for (auto const& any : rpc_status.details()) {
+    auto error = google::storage::v2::BidiWriteObjectRedirectedError{};
+    if (!any.UnpackTo(&error)) continue;
+    *spec.mutable_write_handle() = std::move(*error.mutable_write_handle());
+    *spec.mutable_routing_token() = std::move(*error.mutable_routing_token());
+  }
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal
 }  // namespace cloud

google/cloud/storage/internal/async/handle_redirect_error.h

@@ -25,11 +25,17 @@ namespace cloud {
 namespace storage_internal {
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
 
+void EnsureFirstMessageAppendObjectSpec(
+    google::storage::v2::BidiWriteObjectRequest& request);
+
 google::rpc::Status ExtractGrpcStatus(Status const& status);
 
 void ApplyRedirectErrors(google::storage::v2::BidiReadObjectSpec& spec,
                          google::rpc::Status const& rpc_status);
 
+void ApplyWriteRedirectErrors(google::storage::v2::AppendObjectSpec& spec,
+                              google::rpc::Status const& rpc_status);
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal
 }  // namespace cloud

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

@@ -0,0 +1,80 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "google/cloud/storage/internal/async/write_object.h"
+#include "google/cloud/internal/absl_str_cat_quiet.h"
+#include "google/cloud/internal/make_status.h"
+#include <utility>
+
+namespace google {
+namespace cloud {
+namespace storage_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+WriteObject::WriteObject(std::unique_ptr<StreamingRpc> rpc,
+                         google::storage::v2::BidiWriteObjectRequest request)
+    : rpc_(std::move(rpc)), initial_request_(std::move(request)) {}
+
+future<StatusOr<WriteObject::WriteResult>> WriteObject::Call() {
+  auto future = promise_.get_future();
+  rpc_->Start().then([w = WeakFromThis()](auto f) {
+    if (auto self = w.lock()) self->OnStart(f.get());
+  });
+  return future;
+}
+
+std::weak_ptr<WriteObject> WriteObject::WeakFromThis() {
+  return shared_from_this();
+}
+
+void WriteObject::OnStart(bool ok) {
+  if (!ok) return DoFinish();
+  rpc_->Write(initial_request_, grpc::WriteOptions{})
+      .then([w = WeakFromThis()](auto f) {
+        if (auto self = w.lock()) self->OnWrite(f.get());
+      });
+}
+
+void WriteObject::OnWrite(bool ok) {
+  if (!ok) return DoFinish();
+  rpc_->Read().then([w = WeakFromThis()](auto f) {
+    if (auto self = w.lock()) self->OnRead(f.get());
+  });
+}
+
+void WriteObject::OnRead(
+    absl::optional<google::storage::v2::BidiWriteObjectResponse> response) {
+  if (!response) return DoFinish();
+  promise_.set_value(WriteResult{std::move(rpc_), std::move(*response)});
+}
+
+void WriteObject::DoFinish() {
+  rpc_->Finish().then([w = WeakFromThis()](auto f) {
+    if (auto self = w.lock()) self->OnFinish(f.get());
+  });
+}
+
+void WriteObject::OnFinish(Status status) {
+  if (!status.ok()) return promise_.set_value(std::move(status));
+  // This should not happen, it indicates an EOF on the stream, but we
+  // did not ask to close it.
+  promise_.set_value(google::cloud::internal::InternalError(
+      "could not open stream, but the stream closed successfully",
+      GCP_ERROR_INFO()));
+}
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace storage_internal
+}  // namespace cloud
+}  // namespace google

google/cloud/storage/internal/async/write_object.h

@@ -0,0 +1,102 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_STORAGE_INTERNAL_ASYNC_WRITE_OBJECT_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_STORAGE_INTERNAL_ASYNC_WRITE_OBJECT_H
+
+#include "google/cloud/storage/internal/storage_stub.h"
+#include "google/cloud/completion_queue.h"
+#include "google/cloud/future.h"
+#include "google/cloud/options.h"
+#include "google/cloud/status.h"
+#include "google/cloud/status_or.h"
+#include "google/cloud/version.h"
+#include <google/storage/v2/storage.pb.h>
+#include <grpcpp/grpcpp.h>
+#include <memory>
+#include <string>
+
+namespace google {
+namespace cloud {
+namespace storage_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+/**
+ * Performs a single attempt to open a bidi-streaming write RPC.
+ *
+ * Before we can use a bidi-streaming write RPC we must call `Start()`, and
+ * then call `Read()` to check the RPC start was successful.
+ *
+ * Using C++20 coroutines we would write this as:
+ *
+ * @code
+ * using StreamingRpc = google::cloud::AsyncStreamingReadWriteRpc<
+ *     google::storage::v2::BidiWriteObjectRequest,
+ *     google::storage::v2::BidiWriteObjectResponse>;
+ *
+ *
+ * future<StatusOr<google::storage::v2::BidiWriteObjectResponse>> Call(
+ *     StreamingRpc rpc,
+ *     google::storage::v2::BidiWriteObjectRequest request) {
+ *   auto start = co_await rpc->Start();
+ *   if (!start) co_return co_await rpc->Finish();
+ *   auto read = co_await rpc->Read();
+ *   if (!read) co_return co_await rpc->Finish();
+ *   co_return std::move(*read);
+ * }
+ * @endcode
+ *
+ * As usual, all `co_await` calls become a callback. And all `co_return` calls
+ * must set the value in an explicit `google::cloud::promise<>` object.
+ */
+class WriteObject : public std::enable_shared_from_this<WriteObject> {
+ public:
+  using StreamingRpc = google::cloud::AsyncStreamingReadWriteRpc<
+      google::storage::v2::BidiWriteObjectRequest,
+      google::storage::v2::BidiWriteObjectResponse>;
+
+  using ReturnType = google::storage::v2::BidiWriteObjectResponse;
+  /// Create a coroutine to create an open a bidi streaming write RPC.
+  WriteObject(std::unique_ptr<StreamingRpc> rpc,
+              google::storage::v2::BidiWriteObjectRequest request);
+
+  struct WriteResult {
+    std::unique_ptr<StreamingRpc> stream;
+    google::storage::v2::BidiWriteObjectResponse first_response;
+  };
+
+  /// Start the coroutine.
+  future<StatusOr<WriteResult>> Call();
+
+ private:
+  std::weak_ptr<WriteObject> WeakFromThis();
+
+  void OnStart(bool ok);
+  void OnWrite(bool ok);
+  void OnRead(
+      absl::optional<google::storage::v2::BidiWriteObjectResponse> response);
+  void DoFinish();
+  void OnFinish(Status status);
+
+  std::unique_ptr<StreamingRpc> rpc_;
+  promise<StatusOr<WriteResult>> promise_;
+  google::storage::v2::BidiWriteObjectRequest initial_request_;
+};
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace storage_internal
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_STORAGE_INTERNAL_ASYNC_WRITE_OBJECT_H