fix(storage): append object spec should only be set in the first request

Author: v-pratap

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

@@ -172,16 +172,16 @@ AsyncWriterConnectionImpl::MakeRequest() {
   auto request = request_;
   if (first_request_) {
     first_request_ = false;
+    if (latest_write_handle_.has_value()) {
+      *request.mutable_append_object_spec()->mutable_write_handle() =
+          *latest_write_handle_;
+    }
   } else {
     request.clear_upload_id();
     request.clear_write_object_spec();
     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;
 }
 

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

@@ -699,13 +699,13 @@ TEST(AsyncWriterConnectionTest, FinalizeAppendableNoChecksum) {
   next.first.set_value(true);
 }
 
-TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterQuery) {
+TEST(AsyncWriterConnectionTest, ResumeWithHandle) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockStream>();
   std::vector<std::string> seen_handles;
 
   EXPECT_CALL(*mock, Write)
-      .Times(3)
+      .Times(1)
       .WillRepeatedly([&](Request const& req, grpc::WriteOptions) {
         EXPECT_TRUE(req.has_append_object_spec());
         EXPECT_TRUE(req.append_object_spec().has_write_handle());
@@ -714,74 +714,40 @@ TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterQuery) {
         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);
+
+  auto hash = std::make_shared<MockHashFunction>();
+  EXPECT_CALL(*hash, Update(_, An<absl::Cord const&>(), _)).Times(1);
+  EXPECT_CALL(*hash, Finish).Times(0);
 
   google::storage::v2::BidiWriteObjectRequest req;
   req.mutable_append_object_spec()->set_bucket("bucket");
   req.mutable_append_object_spec()->set_object("object");
+  req.mutable_append_object_spec()->mutable_write_handle()->set_handle(
+      "test-handle");
 
   auto tested = std::make_unique<AsyncWriterConnectionImpl>(
       TestOptions(), req, std::move(mock), hash, 0);
 
-  // First Query sets handle1.
-  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());
 
-  // 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");
+  ASSERT_EQ(seen_handles.size(), 1);
+  EXPECT_EQ(seen_handles[0], "test-handle");
 }
-
-TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterResume) {
+TEST(AsyncWriterConnectionTest, QueryUpdatesHandle) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockStream>();
   std::vector<std::string> seen_handles;
 
   EXPECT_CALL(*mock, Write)
-      .Times(2)
+      .Times(1)
       .WillRepeatedly([&](Request const& req, grpc::WriteOptions) {
         EXPECT_TRUE(req.has_append_object_spec());
         EXPECT_TRUE(req.append_object_spec().has_write_handle());
@@ -790,27 +756,20 @@ TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterResume) {
         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, Read).WillOnce([&]() {
+    Response resp;
+    resp.mutable_write_handle()->set_handle("queried-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>();
-  EXPECT_CALL(*hash, Update(_, An<absl::Cord const&>(), _)).Times(2);
+  EXPECT_CALL(*hash, Update(_, An<absl::Cord const&>(), _)).Times(1);
 
   google::storage::v2::BidiWriteObjectRequest req;
   req.mutable_append_object_spec()->set_bucket("bucket");
@@ -819,28 +778,18 @@ TEST(AsyncWriterConnectionTest, WriteHandleIsUpdatedAfterResume) {
   auto tested = std::make_unique<AsyncWriterConnectionImpl>(
       TestOptions(), req, std::move(mock), hash, 0);
 
-  // First Query sets handle1.
+  // Query should update the internal handle.
   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");
+  // Write should now use the handle from the Query response.
+  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());
+
+  ASSERT_EQ(seen_handles.size(), 1);
+  EXPECT_EQ(seen_handles[0], "queried-handle");
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END