feat(storage): Add generation in append_object_spec in resume operation (#15395)

Author: v-pratap

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

@@ -390,7 +390,7 @@ AsyncConnectionImpl::AppendableObjectUploadImpl(AppendableUploadParams p) {
             false);
         return MakeWriterConnectionResumed(std::move(fa), std::move(impl),
                                            std::move(request), std::move(hash),
-                                           *current);
+                                           rpc->first_response, *current);
       });
 }
 

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

@@ -64,12 +64,14 @@ class AsyncWriterConnectionResumedState
       std::unique_ptr<storage_experimental::AsyncWriterConnection> impl,
       google::storage::v2::BidiWriteObjectRequest initial_request,
       std::shared_ptr<storage::internal::HashFunction> hash_function,
+      google::storage::v2::BidiWriteObjectResponse const& first_response,
       Options const& options, std::size_t buffer_size_lwm,
       std::size_t buffer_size_hwm)
       : factory_(std::move(factory)),
         impl_(std::move(impl)),
         initial_request_(std::move(initial_request)),
         hash_function_(std::move(hash_function)),
+        first_response_(std::move(first_response)),
         buffer_size_lwm_(buffer_size_lwm),
         buffer_size_hwm_(buffer_size_hwm) {
     finalized_future_ = finalized_.get_future();
@@ -330,6 +332,7 @@ class AsyncWriterConnectionResumedState
     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());
+    append_object_spec.set_generation(first_response_.resource().generation());
     ApplyWriteRedirectErrors(append_object_spec, std::move(proto_status));
 
     // Capture the finalization state *before* starting the async resume.
@@ -531,6 +534,8 @@ class AsyncWriterConnectionResumedState
 
   google::cloud::internal::ImmutableOptions options_;
 
+  google::storage::v2::BidiWriteObjectResponse first_response_;
+
   // Request a server-side flush if the buffer goes over this threshold.
   std::size_t const buffer_size_lwm_;
 
@@ -636,10 +641,11 @@ class AsyncWriterConnectionResumed
       std::unique_ptr<storage_experimental::AsyncWriterConnection> impl,
       google::storage::v2::BidiWriteObjectRequest initial_request,
       std::shared_ptr<storage::internal::HashFunction> hash_function,
+      google::storage::v2::BidiWriteObjectResponse const& first_response,
       Options const& options)
       : state_(std::make_shared<AsyncWriterConnectionResumedState>(
             std::move(factory), std::move(impl), std::move(initial_request),
-            std::move(hash_function), options,
+            std::move(hash_function), first_response, options,
             options.get<storage_experimental::BufferedUploadLwmOption>(),
             options.get<storage_experimental::BufferedUploadHwmOption>())) {}
 
@@ -683,10 +689,11 @@ MakeWriterConnectionResumed(
     std::unique_ptr<storage_experimental::AsyncWriterConnection> impl,
     google::storage::v2::BidiWriteObjectRequest initial_request,
     std::shared_ptr<storage::internal::HashFunction> hash_function,
+    google::storage::v2::BidiWriteObjectResponse const& first_response,
     Options const& options) {
   return absl::make_unique<AsyncWriterConnectionResumed>(
       std::move(factory), std::move(impl), std::move(initial_request),
-      std::move(hash_function), std::move(options));
+      std::move(hash_function), std::move(first_response), std::move(options));
 }
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

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

@@ -42,6 +42,7 @@ MakeWriterConnectionResumed(
     std::unique_ptr<storage_experimental::AsyncWriterConnection> impl,
     google::storage::v2::BidiWriteObjectRequest initial_request,
     std::shared_ptr<storage::internal::HashFunction> hash_function,
+    google::storage::v2::BidiWriteObjectResponse const& first_response,
     Options const& options);
 
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END

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

@@ -35,6 +35,7 @@ using ::google::cloud::testing_util::AsyncSequencer;
 using ::google::cloud::testing_util::IsOkAndHolds;
 using ::google::cloud::testing_util::IsProtoEqual;
 using ::google::cloud::testing_util::StatusIs;
+using ::testing::_;
 using ::testing::Return;
 using ::testing::VariantWith;
 
@@ -65,6 +66,7 @@ TEST(WriteConnectionResumed, FinalizeEmpty) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockAsyncWriterConnection>();
   auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
+  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
 
   EXPECT_CALL(*mock, UploadId).WillOnce(Return("test-upload-id"));
   EXPECT_CALL(*mock, PersistedState)
@@ -79,9 +81,9 @@ TEST(WriteConnectionResumed, FinalizeEmpty) {
   MockFactory mock_factory;
   EXPECT_CALL(mock_factory, Call).Times(0);
 
-  auto connection =
-      MakeWriterConnectionResumed(mock_factory.AsStdFunction(), std::move(mock),
-                                  initial_request, nullptr, Options{});
+  auto connection = MakeWriterConnectionResumed(
+      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
+      first_response, Options{});
   EXPECT_EQ(connection->UploadId(), "test-upload-id");
   EXPECT_THAT(connection->PersistedState(), VariantWith<std::int64_t>(0));
 
@@ -95,6 +97,7 @@ TEST(WriteConnectionResumed, FinalizeEmpty) {
 TEST(WriteConnectionResumed, FinalizedOnConstruction) {
   AsyncSequencer<bool> sequencer;
   auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
+  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
   auto mock = std::make_unique<MockAsyncWriterConnection>();
   EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
   EXPECT_CALL(*mock, PersistedState).WillRepeatedly(Return(TestObject()));
@@ -103,9 +106,9 @@ TEST(WriteConnectionResumed, FinalizedOnConstruction) {
   MockFactory mock_factory;
   EXPECT_CALL(mock_factory, Call).Times(0);
 
-  auto connection =
-      MakeWriterConnectionResumed(mock_factory.AsStdFunction(), std::move(mock),
-                                  initial_request, nullptr, Options{});
+  auto connection = MakeWriterConnectionResumed(
+      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
+      first_response, Options{});
   EXPECT_EQ(connection->UploadId(), "test-upload-id");
   EXPECT_THAT(
       connection->PersistedState(),
@@ -119,6 +122,7 @@ TEST(WriteConnectionResumed, FinalizedOnConstruction) {
 TEST(WriteConnectionResumed, Cancel) {
   AsyncSequencer<bool> sequencer;
   auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
+  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
   auto mock = std::make_unique<MockAsyncWriterConnection>();
   EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id"));
   EXPECT_CALL(*mock, PersistedState)
@@ -132,9 +136,9 @@ TEST(WriteConnectionResumed, Cancel) {
   MockFactory mock_factory;
   EXPECT_CALL(mock_factory, Call).Times(0);
 
-  auto connection =
-      MakeWriterConnectionResumed(mock_factory.AsStdFunction(), std::move(mock),
-                                  initial_request, nullptr, Options{});
+  auto connection = MakeWriterConnectionResumed(
+      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
+      first_response, Options{});
 
   auto write = connection->Write(TestPayload(64 * 1024));
   ASSERT_FALSE(write.is_ready());
@@ -152,6 +156,7 @@ TEST(WriteConnectionResumed, Cancel) {
 TEST(WriterConnectionResumed, FlushEmpty) {
   AsyncSequencer<bool> sequencer;
   auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
+  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
 
   auto mock = std::make_unique<MockAsyncWriterConnection>();
   EXPECT_CALL(*mock, PersistedState)
@@ -174,9 +179,9 @@ TEST(WriterConnectionResumed, FlushEmpty) {
   MockFactory mock_factory;
   EXPECT_CALL(mock_factory, Call).Times(0);
 
-  auto connection =
-      MakeWriterConnectionResumed(mock_factory.AsStdFunction(), std::move(mock),
-                                  initial_request, nullptr, Options{});
+  auto connection = MakeWriterConnectionResumed(
+      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
+      first_response, Options{});
   EXPECT_THAT(connection->PersistedState(), VariantWith<std::int64_t>(0));
 
   auto flush = connection->Flush({});
@@ -195,6 +200,7 @@ TEST(WriteConnectionResumed, FlushNonEmpty) {
   AsyncSequencer<bool> sequencer;
   auto mock = std::make_unique<MockAsyncWriterConnection>();
   auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
+  auto first_response = google::storage::v2::BidiWriteObjectResponse{};
   auto const payload = TestPayload(1024);
 
   EXPECT_CALL(*mock, PersistedState)
@@ -217,9 +223,9 @@ TEST(WriteConnectionResumed, FlushNonEmpty) {
   MockFactory mock_factory;
   EXPECT_CALL(mock_factory, Call).Times(0);
 
-  auto connection =
-      MakeWriterConnectionResumed(mock_factory.AsStdFunction(), std::move(mock),
-                                  initial_request, nullptr, Options{});
+  auto connection = MakeWriterConnectionResumed(
+      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
+      first_response, Options{});
   EXPECT_THAT(connection->PersistedState(), VariantWith<std::int64_t>(0));
 
   auto write = connection->Write(payload);
@@ -243,6 +249,71 @@ TEST(WriteConnectionResumed, FlushNonEmpty) {
   EXPECT_THAT(write.get(), StatusIs(StatusCode::kOk));
 }
 
+TEST(WriteConnectionResumed, ResumeUsesGenerationFromFirstResponse) {
+  AsyncSequencer<bool> sequencer;
+  auto mock = std::make_unique<MockAsyncWriterConnection>();
+  auto initial_request = google::storage::v2::BidiWriteObjectRequest{};
+  initial_request.mutable_write_object_spec()->mutable_resource()->set_bucket(
+      "projects/_/buckets/test-bucket");
+  initial_request.mutable_write_object_spec()->mutable_resource()->set_name(
+      "test-object");
+
+  google::storage::v2::BidiWriteObjectResponse first_response;
+  first_response.mutable_resource()->set_generation(12345);
+
+  EXPECT_CALL(*mock, PersistedState)
+      .WillRepeatedly(Return(MakePersistedState(0)));
+
+  // Expect Flush to be called because flush_ is true by default.
+  // Make it fail to trigger Resume().
+  EXPECT_CALL(*mock, Flush(_)).WillOnce([&](auto) {
+    return sequencer.PushBack("Flush").then([](auto f) {
+      if (f.get()) return google::cloud::Status{};  // Should not be true
+      return TransientError();
+    });
+  });
+
+  MockFactory mock_factory;
+  google::storage::v2::BidiWriteObjectRequest captured_request;
+  EXPECT_CALL(mock_factory, Call(_))
+      .WillOnce([&](google::storage::v2::BidiWriteObjectRequest request) {
+        captured_request = std::move(request);
+        return sequencer.PushBack("Factory").then([](auto) {
+          return StatusOr<WriteObject::WriteResult>(
+              internal::AbortedError("stop test", GCP_ERROR_INFO()));
+        });
+      });
+
+  auto connection = MakeWriterConnectionResumed(
+      mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr,
+      first_response, Options{});
+
+  // This will call FlushStep -> mock->Flush()
+  auto write = connection->Write(TestPayload(1));
+  ASSERT_FALSE(write.is_ready());
+
+  // Trigger the Flush error in the mock
+  auto next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Flush");
+  next.first.set_value(false);  // This makes the lambda return TransientError
+
+  // The error in OnFlush triggers Resume(), which calls the mock_factory.
+  // Allow the factory callback to proceed.
+  next = sequencer.PopFrontWithName();
+  EXPECT_EQ(next.second, "Factory");
+  next.first.set_value(true);
+
+  // The write future should now be ready, containing the error from the
+  // factory.
+  EXPECT_THAT(write.get(), StatusIs(StatusCode::kAborted));
+
+  EXPECT_TRUE(captured_request.has_append_object_spec());
+  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(),
+            "projects/_/buckets/test-bucket");
+}
+
 }  // namespace
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal