fix(storage): always use latest write handle for fast resume in gRPC BidiWriteObject (#15795)

Author: shubham-up-47

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

@@ -178,6 +178,10 @@ AsyncWriterConnectionImpl::MakeRequest() {
     request.clear_append_object_spec();
   }
   request.set_write_offset(offset_);
+  if (latest_write_handle_.has_value()) {
+    *request.mutable_append_object_spec()->mutable_write_handle() =
+        *latest_write_handle_;
+  }
   return request;
 }
 
@@ -242,6 +246,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

@@ -14,6 +14,7 @@
 
 #include "google/cloud/storage/internal/async/writer_connection_impl.h"
 #include "google/cloud/mocks/mock_async_streaming_read_write_rpc.h"
+#include "google/cloud/storage/internal/async/write_object.h"
 #include "google/cloud/storage/internal/crc32c.h"
 #include "google/cloud/storage/internal/grpc/ctype_cord_workaround.h"
 #include "google/cloud/storage/internal/hash_function_impl.h"
@@ -698,6 +699,150 @@ TEST(AsyncWriterConnectionTest, FinalizeAppendableNoChecksum) {
   next.first.set_value(true);
 }
 
+TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterQuery) {
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_unique<MockStream>();
+  std::vector<std::string> seen_handles;
+
+  EXPECT_CALL(*mock, Write)
+      .Times(3)
+      .WillRepeatedly([&](Request const& req, grpc::WriteOptions) {
+        EXPECT_TRUE(req.has_append_object_spec());
+        EXPECT_TRUE(req.append_object_spec().has_write_handle());
+        seen_handles.push_back(
+            req.append_object_spec().write_handle().handle());
+        return sequencer.PushBack("Write");
+      });
+
+  int read_call_count = 0;
+  EXPECT_CALL(*mock, Read).Times(2).WillRepeatedly([&]() {
+    Response resp;
+    if (read_call_count == 0) {
+      resp.mutable_write_handle()->set_handle("handle1");
+      resp.set_persisted_size(42);
+    } else {
+      resp.mutable_write_handle()->set_handle("handle2");
+      resp.set_persisted_size(43);
+    }
+    ++read_call_count;
+    return make_ready_future(absl::make_optional(std::move(resp)));
+  });
+
+  auto hash = std::make_shared<MockHashFunction>();
+  EXPECT_CALL(*mock, Cancel).Times(1);
+  EXPECT_CALL(*mock, Finish).WillOnce([] {
+    return make_ready_future(Status{});
+  });
+  EXPECT_CALL(*hash, Update(_, An<absl::Cord const&>(), _)).Times(3);
+
+  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);
+
+  // First Query sets handle1.
+  EXPECT_THAT(tested->Query().get(), IsOkAndHolds(42));
+
+  // First Write uses handle1.
+  auto result1 = tested->Write(WritePayload("payload1"));
+  auto next1 = sequencer.PopFrontWithName();
+  ASSERT_THAT(next1.second, "Write");
+  next1.first.set_value(true);
+  EXPECT_STATUS_OK(result1.get());
+
+  // Second Query sets handle2.
+  EXPECT_THAT(tested->Query().get(), IsOkAndHolds(43));
+
+  // Second Write uses handle2.
+  auto result2 = tested->Write(WritePayload("payload2"));
+  auto next2 = sequencer.PopFrontWithName();
+  ASSERT_THAT(next2.second, "Write");
+  next2.first.set_value(true);
+  EXPECT_STATUS_OK(result2.get());
+
+  // Third Write also uses handle2.
+  auto result3 = tested->Write(WritePayload("payload3"));
+  auto next3 = sequencer.PopFrontWithName();
+  ASSERT_THAT(next3.second, "Write");
+  next3.first.set_value(true);
+  EXPECT_STATUS_OK(result3.get());
+
+  ASSERT_EQ(seen_handles.size(), 3);
+  EXPECT_EQ(seen_handles[0], "handle1");
+  EXPECT_EQ(seen_handles[1], "handle2");
+  EXPECT_EQ(seen_handles[2], "handle2");
+}
+
+TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterResume) {
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_unique<MockStream>();
+  std::vector<std::string> seen_handles;
+
+  EXPECT_CALL(*mock, Write)
+      .Times(2)
+      .WillRepeatedly([&](Request const& req, grpc::WriteOptions) {
+        EXPECT_TRUE(req.has_append_object_spec());
+        EXPECT_TRUE(req.append_object_spec().has_write_handle());
+        seen_handles.push_back(
+            req.append_object_spec().write_handle().handle());
+        return sequencer.PushBack("Write");
+      });
+
+  EXPECT_CALL(*mock, Read)
+      .WillOnce([&]() {
+        Response resp;
+        resp.mutable_write_handle()->set_handle("handle1");
+        resp.set_persisted_size(42);
+        return make_ready_future(absl::make_optional(std::move(resp)));
+      })
+      .WillOnce([&]() {
+        Response resp;
+        resp.mutable_write_handle()->set_handle("handle2");
+        resp.set_persisted_size(43);
+        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>();
+  EXPECT_CALL(*hash, Update(_, An<absl::Cord const&>(), _)).Times(2);
+
+  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);
+
+  // First Query sets handle1.
+  EXPECT_THAT(tested->Query().get(), IsOkAndHolds(42));
+
+  // First Write uses handle1 but fails.
+  auto result1 = tested->Write(WritePayload("payload1"));
+  auto next1 = sequencer.PopFrontWithName();
+  ASSERT_THAT(next1.second, "Write");
+  next1.first.set_value(false);
+
+  // Simulate resume by calling Query again which returns handle2.
+  EXPECT_THAT(tested->Query().get(), IsOkAndHolds(43));
+
+  // Second Write should use handle2.
+  auto result2 = tested->Write(WritePayload("payload2"));
+  auto next2 = sequencer.PopFrontWithName();
+  ASSERT_THAT(next2.second, "Write");
+  next2.first.set_value(true);
+  EXPECT_STATUS_OK(result2.get());
+
+  ASSERT_EQ(seen_handles.size(), 2);
+  EXPECT_EQ(seen_handles[0], "handle1");
+  EXPECT_EQ(seen_handles[1], "handle2");
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal
 }  // namespace cloud

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

@@ -82,6 +82,13 @@ class AsyncWriterConnectionResumedState
     } else {
       buffer_offset_ = absl::get<std::int64_t>(state);
     }
+    if (first_response_.has_write_handle()) {
+      latest_write_handle_ = first_response_.write_handle();
+    } else if (initial_request_.has_append_object_spec() &&
+               initial_request_.append_object_spec().has_write_handle()) {
+      latest_write_handle_ =
+          initial_request_.append_object_spec().write_handle();
+    }
   }
 
   void Cancel() {
@@ -336,9 +343,8 @@ class AsyncWriterConnectionResumedState
     }
     // Include write_handle to enable fast resume instead of slow
     // takeover. Without handle, server performs full state validation.
-    if (first_response_.has_write_handle()) {
-      *append_object_spec.mutable_write_handle() =
-          first_response_.write_handle();
+    if (latest_write_handle_) {
+      *append_object_spec.mutable_write_handle() = *latest_write_handle_;
     }
     append_object_spec.set_generation(first_response_.resource().generation());
     ApplyWriteRedirectErrors(append_object_spec, std::move(proto_status));
@@ -363,6 +369,10 @@ 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->first_response.has_write_handle()) {
+      latest_write_handle_ = res->first_response.write_handle();
+    }
 
     if (was_finalizing) {
       // If resuming due to a finalization error, we *must* complete the
@@ -602,6 +612,9 @@ class AsyncWriterConnectionResumedState
 
   // Tracks if the final promise (`finalized_`) has been completed.
   bool finalized_promise_completed_ = false;
+
+  // Track the latest write handle seen in responses.
+  absl::optional<google::storage::v2::BidiWriteHandle> latest_write_handle_;
 };
 
 /**

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

@@ -268,7 +268,7 @@ TEST(WriteConnectionResumed, ResumeUsesWriteObjectSpecFromInitialRequest) {
       "test-object");
 
   google::storage::v2::BidiWriteObjectResponse first_response;
-  first_response.mutable_write_handle();
+  first_response.mutable_write_handle()->set_handle("test-handle");
   first_response.mutable_resource()->set_generation(12345);
 
   EXPECT_CALL(*mock, PersistedState)
@@ -320,6 +320,8 @@ TEST(WriteConnectionResumed, ResumeUsesWriteObjectSpecFromInitialRequest) {
   EXPECT_FALSE(captured_request.has_write_object_spec());
   EXPECT_TRUE(captured_request.has_append_object_spec());
   EXPECT_TRUE(captured_request.append_object_spec().has_write_handle());
+  EXPECT_EQ(captured_request.append_object_spec().write_handle().handle(),
+            "test-handle");
   EXPECT_EQ(captured_request.append_object_spec().generation(), 12345);
   EXPECT_EQ(captured_request.append_object_spec().object(), "test-object");
   EXPECT_EQ(captured_request.append_object_spec().bucket(),
@@ -392,6 +394,71 @@ TEST(WriteConnectionResumed, ResumeUsesAppendObjectSpecFromInitialRequest) {
             "projects/_/buckets/test-bucket");
 }
 
+TEST(WriteConnectionResumed, WriteHandleAssignmentAfterResume) {
+  struct {
+    bool use_write_object_spec;
+    std::string bucket, object, handle;
+    std::int64_t generation;
+  } cases[] = {
+      {false, "projects/_/buckets/test-bucket", "test-object",
+       "expected-handle", 12345},
+      {true, "bucket1", "object1", "handle1", 111},
+      {false, "bucket2", "object2", "handle2", 222},
+  };
+
+  for (auto const& tc : cases) {
+    AsyncSequencer<bool> sequencer;
+    auto mock = std::make_unique<MockAsyncWriterConnection>();
+    google::storage::v2::BidiWriteObjectRequest req;
+    if (tc.use_write_object_spec) {
+      req.mutable_write_object_spec()->mutable_resource()->set_bucket(
+          tc.bucket);
+      req.mutable_write_object_spec()->mutable_resource()->set_name(tc.object);
+    } else {
+      req.mutable_append_object_spec()->set_bucket(tc.bucket);
+      req.mutable_append_object_spec()->set_object(tc.object);
+    }
+    google::storage::v2::BidiWriteObjectResponse resp;
+    resp.mutable_write_handle()->set_handle(tc.handle);
+    resp.mutable_resource()->set_generation(tc.generation);
+
+    EXPECT_CALL(*mock, PersistedState)
+        .WillRepeatedly(Return(MakePersistedState(0)));
+    EXPECT_CALL(*mock, Flush(_)).WillOnce([&](auto) {
+      return sequencer.PushBack("Flush").then([](auto f) {
+        if (f.get()) return google::cloud::Status{};
+        return TransientError();
+      });
+    });
+
+    MockFactory mock_factory;
+    google::storage::v2::BidiWriteObjectRequest captured;
+    EXPECT_CALL(mock_factory, Call(_))
+        .WillOnce([&](google::storage::v2::BidiWriteObjectRequest r) {
+          captured = std::move(r);
+          return sequencer.PushBack("Factory").then([](auto) {
+            return StatusOr<WriteObject::WriteResult>(
+                internal::AbortedError("stop test", GCP_ERROR_INFO()));
+          });
+        });
+
+    auto conn = MakeWriterConnectionResumed(mock_factory.AsStdFunction(),
+                                            std::move(mock), req, nullptr, resp,
+                                            Options{});
+    auto write = conn->Write(TestPayload(1));
+    sequencer.PopFrontWithName().first.set_value(false);
+    sequencer.PopFrontWithName().first.set_value(true);
+
+    EXPECT_THAT(write.get(), StatusIs(StatusCode::kAborted));
+    EXPECT_TRUE(captured.has_append_object_spec());
+    EXPECT_EQ(captured.append_object_spec().bucket(), tc.bucket);
+    EXPECT_EQ(captured.append_object_spec().object(), tc.object);
+    EXPECT_EQ(captured.append_object_spec().generation(), tc.generation);
+    EXPECT_TRUE(captured.append_object_spec().has_write_handle());
+    EXPECT_EQ(captured.append_object_spec().write_handle().handle(), tc.handle);
+  }
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal