adding test for multiple responses having latest handle

Author: shubham-up-47

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

@@ -177,7 +177,7 @@ AsyncWriterConnectionImpl::MakeRequest() {
     request.clear_append_object_spec();
   }
   request.set_write_offset(offset_);
-  if (latest_write_handle_.has_value() && request.has_append_object_spec()) {
+  if (latest_write_handle_.has_value()) {
     *request.mutable_append_object_spec()->mutable_write_handle() =
         *latest_write_handle_;
   }

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"
@@ -648,55 +649,148 @@ TEST(AsyncWriterConnectionTest, UnexpectedQueryFailsWithoutError) {
   EXPECT_THAT(query.get(), StatusIs(StatusCode::kInternal));
 }
 
-TEST(AsyncWriterConnectionImpl, WriteHandleIsPropagatedAfterQuery) {
+TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterQuery) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockStream>();
-  int write_call_count = 0;
+  std::vector<std::string> seen_handles;
+
   EXPECT_CALL(*mock, Write)
-      .Times(2)
+      .Times(3)
       .WillRepeatedly([&](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");
-        ++write_call_count;
+        seen_handles.push_back(
+            req.append_object_spec().write_handle().handle());
         return sequencer.PushBack("Write");
       });
-  EXPECT_CALL(*mock, Read).WillOnce([&]() {
+
+  int read_call_count = 0;
+  EXPECT_CALL(*mock, Read).Times(2).WillRepeatedly([&]() {
     Response resp;
-    resp.mutable_write_handle()->set_handle("test-handle");
-    resp.set_persisted_size(42);
+    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);
 
-  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);
 
+  // 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 write should also include write_handle.
+  // 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());
 
-  EXPECT_EQ(write_call_count, 2);
+  ASSERT_EQ(seen_handles.size(), 2);
+  EXPECT_EQ(seen_handles[0], "handle1");
+  EXPECT_EQ(seen_handles[1], "handle2");
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

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));
@@ -365,8 +371,7 @@ class AsyncWriterConnectionResumedState
     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()) {
-      *first_response_.mutable_write_handle() =
-          res->first_response.write_handle();
+      latest_write_handle_ = res->first_response.write_handle();
     }
 
     if (was_finalizing) {
@@ -607,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,67 @@ 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