Always Use Latest Write Handle for Fast Resume in gRPC BidiWriteObject

Author: shubham-up-47

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

@@ -177,6 +177,12 @@ AsyncWriterConnectionImpl::MakeRequest() {
     request.clear_append_object_spec();
   }
   request.set_write_offset(offset_);
+  if (latest_write_handle_.has_value()) {
+    if (request.has_append_object_spec()) {
+      *request.mutable_append_object_spec()->mutable_write_handle() =
+          *latest_write_handle_;
+    }
+  }
   return request;
 }
 
@@ -241,6 +247,9 @@ future<StatusOr<std::int64_t>> AsyncWriterConnectionImpl::OnQuery(
           return StatusOr<std::int64_t>(std::move(result));
         });
   }
+  if (response->has_write_handle()) {
+    latest_write_handle_ = response->write_handle();
+  }
   if (response->has_persisted_size()) {
     persisted_state_ = response->persisted_size();
     return make_ready_future(make_status_or(response->persisted_size()));

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

@@ -105,6 +105,9 @@ class AsyncWriterConnectionImpl
   // that `on_finish_` is called when `Finish()` is completed.
   google::cloud::promise<void> on_finish_;
   google::cloud::future<void> finished_;
+
+  // Track the latest write handle seen in responses.
+  absl::optional<google::storage::v2::BidiWriteHandle> latest_write_handle_;
 };
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

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

@@ -648,6 +648,44 @@ TEST(AsyncWriterConnectionTest, UnexpectedQueryFailsWithoutError) {
   EXPECT_THAT(query.get(), StatusIs(StatusCode::kInternal));
 }
 
+TEST(AsyncWriterConnectionImpl, WriteHandleIsPropagatedAfterQuery) {
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_unique<MockStream>();
+  EXPECT_CALL(*mock, Write)
+      .WillOnce([&](Request const& req, grpc::WriteOptions) {
+        EXPECT_TRUE(req.has_append_object_spec());
+        EXPECT_TRUE(req.append_object_spec().has_write_handle());
+        EXPECT_EQ(req.append_object_spec().write_handle().handle(),
+                  "test-handle");
+        return sequencer.PushBack("Write");
+      });
+  EXPECT_CALL(*mock, Read).WillOnce([&]() {
+    Response resp;
+    resp.mutable_write_handle()->set_handle("test-handle");
+    resp.set_persisted_size(42);
+    return make_ready_future(absl::make_optional(std::move(resp)));
+  });
+  EXPECT_CALL(*mock, Cancel).Times(1);
+  EXPECT_CALL(*mock, Finish).WillOnce([] {
+    return make_ready_future(Status{});
+  });
+
+  auto hash = std::make_shared<MockHashFunction>();
+  google::storage::v2::BidiWriteObjectRequest req;
+  req.mutable_append_object_spec()->set_bucket("bucket");
+  req.mutable_append_object_spec()->set_object("object");
+
+  auto tested = std::make_unique<AsyncWriterConnectionImpl>(
+      TestOptions(), req, std::move(mock), hash, 0);
+
+  EXPECT_THAT(tested->Query().get(), IsOkAndHolds(42));
+  auto result = tested->Write(WritePayload("payload"));
+  auto next = sequencer.PopFrontWithName();
+  ASSERT_THAT(next.second, "Write");
+  next.first.set_value(true);
+  EXPECT_STATUS_OK(result.get());
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal
 }  // namespace cloud

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

@@ -363,6 +363,11 @@ class AsyncWriterConnectionResumedState
   void OnResume(Status const& original_status, bool was_finalizing,
                 StatusOr<WriteObject::WriteResult> res) {
     std::unique_lock<std::mutex> lk(mu_);
+    // Update write_handle from any resume response that contains it.
+    if (res && res.value().first_response.has_write_handle()) {
+      *first_response_.mutable_write_handle() =
+          res.value().first_response.write_handle();
+    }
 
     if (was_finalizing) {
       // If resuming due to a finalization error, we *must* complete the