diff --git a/google/cloud/storage/internal/async/writer_connection_impl.cc b/google/cloud/storage/internal/async/writer_connection_impl.cc index fe62edbafea0d..5e0ee2e2376a3 100644 --- a/google/cloud/storage/internal/async/writer_connection_impl.cc +++ b/google/cloud/storage/internal/async/writer_connection_impl.cc @@ -196,7 +196,7 @@ future AsyncWriterConnectionImpl::Flush(storage::WritePayload payload) { return coro->Start().then([coro, size, this](auto f) mutable { coro.reset(); // breaks the cycle between the completion queue and coro - return OnPartialUpload(size, f.get()); + return OnFlush(size, f.get()); }); } @@ -215,10 +215,15 @@ future AsyncWriterConnectionImpl::Close(storage::WritePayload payload) { future> AsyncWriterConnectionImpl::Query() { std::unique_lock lk(mu_); - auto impl = impl_; - lk.unlock(); - return impl->Read().then( - [this](auto f) { return OnQuery(std::move(f).get()); }); + if (auto* size = absl::get_if(&persisted_state_)) { + return make_ready_future(make_status_or(*size)); + } + if (auto* meta = + absl::get_if(&persisted_state_)) { + return make_ready_future( + make_status_or(static_cast(meta->size()))); + } + return make_ready_future(make_status_or(static_cast(0))); } RpcMetadata AsyncWriterConnectionImpl::GetRequestMetadata() { @@ -252,11 +257,13 @@ future AsyncWriterConnectionImpl::OnPartialUpload( return Finish().then( HandleFinishAfterError("Expected Finish() error after non-ok Write()")); } + std::unique_lock lk(mu_); offset_ += upload_size; + lk.unlock(); return make_ready_future(Status{}); } -future AsyncWriterConnectionImpl::OnClose(std::size_t upload_size, +future AsyncWriterConnectionImpl::OnFlush(std::size_t upload_size, StatusOr success) { if (!success) { return Finish().then(HandleFinishAfterError(std::move(success).status())); @@ -265,8 +272,57 @@ future AsyncWriterConnectionImpl::OnClose(std::size_t upload_size, return Finish().then( HandleFinishAfterError("Expected Finish() error after non-ok Write()")); } + std::unique_lock lk(mu_); offset_ += upload_size; + auto expected_offset = offset_; + auto impl = impl_; + lk.unlock(); + + struct ConsumeLoop { + AsyncWriterConnectionImpl* self; + std::shared_ptr impl; + std::int64_t expected_offset; + + future operator()( + future> + f) { + auto response = std::move(f).get(); + if (!response.has_value()) { + return self->Finish().then(HandleFinishAfterError( + "Expected error in Finish() after non-ok Read()")); + } + std::unique_lock lk(self->mu_); + if (response->has_write_handle()) { + self->latest_write_handle_ = response->write_handle(); + } + if (response->has_persisted_size()) { + self->persisted_state_ = response->persisted_size(); + if (response->persisted_size() >= expected_offset) { + return make_ready_future(Status{}); + } + } + if (response->has_resource()) { + self->persisted_state_ = response->resource(); + return make_ready_future(Status{}); + } + lk.unlock(); + return impl->Read().then(std::move(*this)); + } + }; + return impl->Read().then(ConsumeLoop{this, impl, expected_offset}); +} + +future AsyncWriterConnectionImpl::OnClose(std::size_t upload_size, + StatusOr success) { + if (!success) { + return Finish().then(HandleFinishAfterError(std::move(success).status())); + } + if (!*success) { + return Finish().then( + HandleFinishAfterError("Expected Finish() error after non-ok Write()")); + } std::unique_lock lk(mu_); + offset_ += upload_size; auto impl = impl_; lk.unlock(); @@ -304,9 +360,9 @@ AsyncWriterConnectionImpl::OnFinalUpload(std::size_t upload_size, "Expected error in Finish() after non-ok Write()")) .then(transform); } + std::unique_lock lk(mu_); offset_ += upload_size; - std::unique_lock lk(mu_); auto impl = impl_; lk.unlock(); @@ -352,41 +408,6 @@ AsyncWriterConnectionImpl::OnFinalUpload(std::size_t upload_size, }); } -future> AsyncWriterConnectionImpl::OnQuery( - std::optional response) { - if (!response.has_value()) { - return Finish() - .then(HandleFinishAfterError( - "Expected error in Finish() after non-ok Read()")) - .then([this](auto g) { - auto result = g.get(); - google::rpc::Status grpc_status = ExtractGrpcStatus(result); - HandleBidiWriteRedirect(request_, grpc_status); - return StatusOr(std::move(result)); - }); - } - - std::shared_ptr impl; - { - std::unique_lock lk(mu_); - 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())); - } - if (response->has_resource()) { - persisted_state_ = response->resource(); - return make_ready_future(make_status_or(response->resource().size())); - } - impl = impl_; - } - - return impl->Read().then( - [this](auto f) { return OnQuery(std::move(f).get()); }); -} - future AsyncWriterConnectionImpl::Finish() { std::unique_lock lk(mu_); if (std::exchange(finish_called_, true)) { diff --git a/google/cloud/storage/internal/async/writer_connection_impl.h b/google/cloud/storage/internal/async/writer_connection_impl.h index 97bbbd854f6e9..f88ad99f56311 100644 --- a/google/cloud/storage/internal/async/writer_connection_impl.h +++ b/google/cloud/storage/internal/async/writer_connection_impl.h @@ -85,11 +85,10 @@ class AsyncWriterConnectionImpl : public storage::AsyncWriterConnection { future OnPartialUpload(std::size_t upload_size, StatusOr success); + future OnFlush(std::size_t upload_size, StatusOr success); future OnClose(std::size_t upload_size, StatusOr success); future> OnFinalUpload( std::size_t upload_size, StatusOr success); - future> OnQuery( - std::optional response); future Finish(); google::cloud::internal::ImmutableOptions options_; diff --git a/google/cloud/storage/internal/async/writer_connection_impl_test.cc b/google/cloud/storage/internal/async/writer_connection_impl_test.cc index 15ce7b757a95b..7e16808831278 100644 --- a/google/cloud/storage/internal/async/writer_connection_impl_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_impl_test.cc @@ -543,12 +543,12 @@ TEST(AsyncWriterConnectionTest, FlushEmpty) { auto next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Write"); next.first.set_value(true); - EXPECT_THAT(flush.get(), IsOk()); - - auto query = tested->Query(); next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Read"); next.first.set_value(true); + EXPECT_THAT(flush.get(), IsOk()); + + auto query = tested->Query(); EXPECT_THAT(query.get(), IsOkAndHolds(16384)); tested = {}; @@ -615,10 +615,13 @@ TEST(AsyncWriterConnectionTest, UnexpectedFlushFailsWithoutError) { EXPECT_THAT(response.get(), StatusIs(StatusCode::kInternal)); } -TEST(AsyncWriterConnectionTest, QueryFails) { +TEST(AsyncWriterConnectionTest, FlushFailsOnRead) { AsyncSequencer sequencer; auto mock = std::make_unique(); EXPECT_CALL(*mock, Cancel).Times(1); + EXPECT_CALL(*mock, Write).WillOnce([&](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write"); + }); EXPECT_CALL(*mock, Read).WillOnce([&]() { return sequencer.PushBack("Read").then( [](auto) { return std::optional(); }); @@ -630,25 +633,32 @@ TEST(AsyncWriterConnectionTest, QueryFails) { }); }); auto hash = std::make_shared(); - EXPECT_CALL(*hash, Update(_, An(), _)).Times(0); - EXPECT_CALL(*hash, Finish).Times(0); + EXPECT_CALL(*hash, Update(_, An(), _)).Times(1); auto tested = std::make_unique( TestOptions(), MakeRequest(), std::move(mock), hash, 1024); - auto query = tested->Query(); + auto flush = tested->Flush(WritePayload{std::string("fake-payload-data")}); auto next = sequencer.PopFrontWithName(); + ASSERT_THAT(next.second, "Write"); + next.first.set_value(true); + + next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Read"); next.first.set_value(false); // Detect error from Read() + next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Finish"); next.first.set_value(false); // Return error from Finish() - EXPECT_THAT(query.get(), StatusIs(PermanentError().code())); + EXPECT_THAT(flush.get(), StatusIs(PermanentError().code())); } -TEST(AsyncWriterConnectionTest, UnexpectedQueryFailsWithoutError) { +TEST(AsyncWriterConnectionTest, UnexpectedFlushFailsOnReadWithoutError) { AsyncSequencer sequencer; auto mock = std::make_unique(); EXPECT_CALL(*mock, Cancel).Times(1); + EXPECT_CALL(*mock, Write).WillOnce([&](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write"); + }); EXPECT_CALL(*mock, Read).WillOnce([&]() { return sequencer.PushBack("Read").then( [](auto) { return std::optional(); }); @@ -660,25 +670,32 @@ TEST(AsyncWriterConnectionTest, UnexpectedQueryFailsWithoutError) { }); }); auto hash = std::make_shared(); - EXPECT_CALL(*hash, Update(_, An(), _)).Times(0); - EXPECT_CALL(*hash, Finish).Times(0); + EXPECT_CALL(*hash, Update(_, An(), _)).Times(1); auto tested = std::make_unique( TestOptions(), MakeRequest(), std::move(mock), hash, 1024); - auto query = tested->Query(); + auto flush = tested->Flush(WritePayload{std::string("fake-payload-data")}); auto next = sequencer.PopFrontWithName(); + ASSERT_THAT(next.second, "Write"); + next.first.set_value(true); + + next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Read"); next.first.set_value(false); // Detect error from Read() + next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Finish"); next.first.set_value(true); // Return success from Finish() - EXPECT_THAT(query.get(), StatusIs(StatusCode::kInternal)); + EXPECT_THAT(flush.get(), StatusIs(StatusCode::kInternal)); } -TEST(AsyncWriterConnectionTest, QueryFailsWithRedirect) { +TEST(AsyncWriterConnectionTest, FlushFailsWithRedirect) { AsyncSequencer sequencer; auto mock = std::make_unique(); EXPECT_CALL(*mock, Cancel).Times(1); + EXPECT_CALL(*mock, Write).WillOnce([&](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write"); + }); EXPECT_CALL(*mock, Read).WillOnce([&]() { return sequencer.PushBack("Read").then( [](auto) { return std::optional(); }); @@ -705,17 +722,23 @@ TEST(AsyncWriterConnectionTest, QueryFailsWithRedirect) { }); }); auto hash = std::make_shared(); + EXPECT_CALL(*hash, Update(_, An(), _)).Times(1); auto tested = std::make_unique( TestOptions(), MakeRequest(), std::move(mock), hash, 1024); - auto query = tested->Query(); + auto flush = tested->Flush(WritePayload{std::string("fake-payload-data")}); auto next = sequencer.PopFrontWithName(); + ASSERT_THAT(next.second, "Write"); + next.first.set_value(true); + + next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Read"); next.first.set_value(false); // Detect error from Read() + next = sequencer.PopFrontWithName(); ASSERT_THAT(next.second, "Finish"); next.first.set_value(false); // Return error from Finish() - EXPECT_THAT(query.get(), StatusIs(StatusCode::kAborted)); + EXPECT_THAT(flush.get(), StatusIs(StatusCode::kAborted)); } TEST(AsyncWriterConnectionTest, FinalizeAppendableNoChecksum) { @@ -1097,56 +1120,6 @@ TEST(AsyncWriterConnectionTest, ResumeWithHandle) { ASSERT_EQ(seen_handles.size(), 1); EXPECT_EQ(seen_handles[0], "test-handle"); } -TEST(AsyncWriterConnectionTest, QueryUpdatesHandle) { - AsyncSequencer sequencer; - auto mock = std::make_unique(); - std::vector seen_handles; - - EXPECT_CALL(*mock, Write) - .Times(1) - .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("queried-handle"); - resp.set_persisted_size(42); - return make_ready_future(std::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(); - EXPECT_CALL(*hash, Update(_, An(), _)).Times(1); - - 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( - TestOptions(), req, std::move(mock), hash, 0); - - // Query should update the internal handle. - EXPECT_THAT(tested->Query().get(), IsOkAndHolds(42)); - - // 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"); -} TEST(AsyncWriterConnectionTest, CloseEmpty) { AsyncSequencer sequencer; @@ -1311,16 +1284,20 @@ TEST(AsyncWriterConnectionTest, CloseMultipleResponsesBeforeEOF) { tested = {}; } -TEST(AsyncWriterConnectionTest, QueryConsumesIntermediateMessagesBeforeSize) { +TEST(AsyncWriterConnectionTest, FlushConsumesIntermediateMessagesBeforeSize) { AsyncSequencer sequencer; auto mock = std::make_unique(); EXPECT_CALL(*mock, Cancel).Times(1); - // We simulate a scenario where Flush(state_lookup=true) causes GCS to send + EXPECT_CALL(*mock, Write).WillOnce([&](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write"); + }); + + // We simulate a scenario where Flush() causes GCS to send // multiple responses. The first is an intermediate message (no size). // The second is the actual persisted_size response. EXPECT_CALL(*mock, Read) - .WillOnce([&] { + .WillOnce([&]() { return sequencer.PushBack("Read1").then([](auto f) { f.get(); auto r = Response{}; @@ -1328,7 +1305,7 @@ TEST(AsyncWriterConnectionTest, QueryConsumesIntermediateMessagesBeforeSize) { return std::make_optional(r); }); }) - .WillOnce([&] { + .WillOnce([&]() { return sequencer.PushBack("Read2").then([](auto f) { f.get(); auto r = Response{}; @@ -1344,6 +1321,9 @@ TEST(AsyncWriterConnectionTest, QueryConsumesIntermediateMessagesBeforeSize) { }); }); + auto hash = std::make_shared(); + EXPECT_CALL(*hash, Update(_, An(), _)).Times(1); + Request request; request.mutable_write_object_spec()->mutable_resource()->set_bucket( "test-bucket"); @@ -1351,14 +1331,18 @@ TEST(AsyncWriterConnectionTest, QueryConsumesIntermediateMessagesBeforeSize) { "test-object"); auto tested = std::make_unique( - TestOptions(), request, std::move(mock), /*hash_function=*/nullptr, - /*persisted_size=*/1024); + TestOptions(), request, std::move(mock), hash, + /*persisted_size=*/0); + + auto flush = tested->Flush(WritePayload{std::string("fake-payload")}); - auto query_future = tested->Query(); + auto next = sequencer.PopFrontWithName(); + ASSERT_EQ(next.second, "Write"); + next.first.set_value(true); // The client requests the first read. We satisfy it with the intermediate // message. - auto next = sequencer.PopFrontWithName(); + next = sequencer.PopFrontWithName(); ASSERT_EQ(next.second, "Read1"); next.first.set_value(true); @@ -1367,9 +1351,9 @@ TEST(AsyncWriterConnectionTest, QueryConsumesIntermediateMessagesBeforeSize) { ASSERT_EQ(next.second, "Read2"); next.first.set_value(true); - // The query should finally resolve with the correct size (1024) instead of 0. - auto result = query_future.get(); - EXPECT_THAT(result, IsOkAndHolds(1024)); + // The flush should finally resolve successfully. + auto result = flush.get(); + EXPECT_STATUS_OK(result); tested.reset(); next = sequencer.PopFrontWithName();