From c3d4fe00614c33a8e820b4b5f8d1aa1e0f1426f0 Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Tue, 29 Sep 2026 09:47:00 +0000 Subject: [PATCH 1/3] fix(storage): mark async writer idle before invoking flush callbacks --- .../async/writer_connection_buffered.cc | 38 +-- .../async/writer_connection_buffered_test.cc | 234 +++++++++++++++++ .../async/writer_connection_resumed.cc | 39 +-- .../async/writer_connection_resumed_test.cc | 246 ++++++++++++++++++ 4 files changed, 527 insertions(+), 30 deletions(-) diff --git a/google/cloud/storage/internal/async/writer_connection_buffered.cc b/google/cloud/storage/internal/async/writer_connection_buffered.cc index db062aa37f326..45d0d2edf6efd 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered.cc @@ -383,28 +383,28 @@ class AsyncWriterConnectionBufferedState write_offset_ -= static_cast(n); } } - // If the buffer is small enough, collect all the handlers to notify them. - auto const handlers = ClearHandlersIfEmpty(lk); if (is_resume) { // We are resuming. The pending flush promises (if any) should not be // satisfied yet, because we haven't actually flushed the data on the new // connection. The `WriteLoop` will trigger a flush (potentially empty) // if `flush_` is still true, which will satisfy the promises when it - // completes. However, we still need to notify any handlers waiting for - // the buffer to shrink, and we need to restart the write loop. + // completes. + auto const handlers = ClearHandlersIfEmpty(lk); + // Mark the writer idle under the lock so any operation chained from a + // handler below sees an idle writer and is dispatched immediately, and + // `writing_` is never modified without holding `mu_`. resuming_ = false; - lk.unlock(); + writing_ = false; + lk.unlock(); // Release lock before notifying. + // The notifications are deferred until the lock is released, as they + // might call back and try to acquire the lock. for (auto const& h : handlers) h->Execute(Status{}); - WriteLoop(std::unique_lock(mu_)); + // Re-acquire the lock to restart the write loop. This is a no-op if a + // handler above already restarted it. + StartWriting(std::unique_lock(mu_)); return; } - // SetFlushed will release the lock before returning. SetFlushed(std::move(lk), Status{}, persisted_size); - // Re-acquire the lock to re-enter the write loop. - WriteLoop(std::unique_lock(mu_)); - // The notifications are deferred until the lock is released, as they might - // call back and try to acquire the lock. - for (auto const& h : handlers) h->Execute(Status{}); } void WriteStep(std::unique_lock lk, absl::Cord payload) { @@ -575,11 +575,19 @@ class AsyncWriterConnectionBufferedState if (pending_flush_promises_.empty()) { flush_ = false; } - lk.unlock(); // Unlock only once before notifying - // Notify handlers and the specific flush promises *after* releasing the - // lock. + // Mark the writer idle under the lock so any operation chained from a + // callback below sees an idle writer and is dispatched immediately, and + // `writing_` is never modified without holding `mu_`. + writing_ = false; + lk.unlock(); // Release lock before notifying. + // Notify handlers and satisfied flush promises before restarting the + // write loop so callbacks cannot be overtaken by a queued flush that + // completes inline. for (auto& h : handlers) h->Execute(Status{}); for (auto& f : flushes_to_complete) f.set_value(result); + // Re-acquire the lock to resume writing any remaining buffered data. + // This is a no-op if a callback above already restarted the write loop. + StartWriting(std::unique_lock(mu_)); } void SetError(std::unique_lock lk, Status const& status) { diff --git a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc index 6321ad92150f6..1b26fd4d88cfa 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc @@ -21,6 +21,13 @@ #include "google/cloud/testing_util/is_proto_equal.h" #include "google/cloud/testing_util/status_matchers.h" #include +#include +#include +#include +#include +#include +#include +#include namespace google { namespace cloud { @@ -36,6 +43,7 @@ using ::google::cloud::testing_util::IsOkAndHolds; using ::google::cloud::testing_util::IsProtoEqual; using ::google::cloud::testing_util::StatusIs; using ::testing::AtLeast; +using ::testing::ElementsAre; using ::testing::Eq; using ::testing::InSequence; using ::testing::IsSupersetOf; @@ -1931,6 +1939,232 @@ TEST(WriteConnectionBuffered, CustomRetryPolicyOption) { EXPECT_THAT(write2.get(), StatusIs(StatusCode::kInvalidArgument)); } +/// Test case for an operation issued from inside a `Flush()` continuation. +struct ChainedOpCase { + std::string name; + std::function(storage::AsyncWriterConnection&)> start_op; + std::string expected_call; +}; + +class WriterConnectionBufferedFlushCallbackTest + : public ::testing::TestWithParam {}; + +/// @test Verify that an operation chained from a `Flush()` continuation is +/// dispatched to the underlying connection before the continuation returns. +TEST_P(WriterConnectionBufferedFlushCallbackTest, + ChainedOpDispatchedBeforeCallbackReturns) { + AsyncSequencer sequencer; + auto mock = std::make_unique(); + auto mock_persisted_size = std::make_shared(0); + int call_count = 0; + + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState).WillRepeatedly([mock_persisted_size] { + return MakePersistedState(*mock_persisted_size); + }); + EXPECT_CALL(*mock, Write).WillRepeatedly([&](auto const&) { + ++call_count; + return sequencer.PushBack("Write").then([](auto) { return Status{}; }); + }); + EXPECT_CALL(*mock, Flush).WillRepeatedly([&](storage::WritePayload const& p) { + ++call_count; + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush").then( + [mock_persisted_size, size](auto f) { + if (!f.get()) return TransientError(); + *mock_persisted_size += size; + return Status{}; + }); + }); + EXPECT_CALL(*mock, Close).WillRepeatedly([&](auto const&) { + ++call_count; + return sequencer.PushBack("Close").then([](auto) { return Status{}; }); + }); + EXPECT_CALL(*mock, Finalize).WillRepeatedly([&](auto const&) { + ++call_count; + return sequencer.PushBack("Finalize").then([](auto) { + return make_status_or(TestObject()); + }); + }); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto flush = connection->Flush(TestPayload(1024)); + ASSERT_EQ(call_count, 1); + + // Chain the next operation from `flush`'s continuation. Because `flush` is + // already satisfied when this callback runs, the writer must already be idle + // so the chained operation dispatches synchronously (`call_count == 2`) + // before the callback returns. + future chained_op; + auto callback_done = flush.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained_op = GetParam().start_op(*connection); + EXPECT_EQ(call_count, 2); + }); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, GetParam().expected_call); + next.first.set_value(true); + EXPECT_STATUS_OK(chained_op.get()); +} + +INSTANTIATE_TEST_SUITE_P( + WriteConnectionBuffered, WriterConnectionBufferedFlushCallbackTest, + ::testing::Values(ChainedOpCase{"Flush", + [](storage::AsyncWriterConnection& c) { + return c.Flush(TestPayload(1024)); + }, + "Flush"}, + ChainedOpCase{"EmptyFlush", + [](storage::AsyncWriterConnection& c) { + return c.Flush(storage::WritePayload{}); + }, + "Flush"}, + ChainedOpCase{"Write", + [](storage::AsyncWriterConnection& c) { + return c.Write(TestPayload(1024)); + }, + "Write"}, + ChainedOpCase{"Close", + [](storage::AsyncWriterConnection& c) { + return c.Close(storage::WritePayload{}); + }, + "Close"}, + ChainedOpCase{"Finalize", + [](storage::AsyncWriterConnection& c) { + return c.Finalize(storage::WritePayload{}) + .then([](auto f) { + return f.get().status(); + }); + }, + "Finalize"}), + [](::testing::TestParamInfo const& info) { + return info.param.name; + }); + +/// @test Verify that a `Flush()` continuation handing the next `Flush()` to +/// another thread does not race with `OnQuery()` on `writing_`. +TEST(WriteConnectionBuffered, FlushCallbackCrossThreadFlush) { + AsyncSequencer sequencer; + auto mock = std::make_unique(); + auto mock_persisted_size = std::make_shared(0); + + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState).WillRepeatedly([mock_persisted_size] { + return MakePersistedState(*mock_persisted_size); + }); + EXPECT_CALL(*mock, Flush).WillRepeatedly([&](storage::WritePayload const& p) { + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush").then( + [mock_persisted_size, size](auto f) { + if (!f.get()) return TransientError(); + *mock_persisted_size += size; + return Status{}; + }); + }); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto f1 = connection->Flush(TestPayload(1024)); + auto next = sequencer.PopFrontWithName(); + ASSERT_EQ(next.second, "Flush"); + + std::thread worker; + future f2; + auto callback_done = f1.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + promise started; + worker = std::thread([&] { + started.set_value(); + f2 = connection->Flush(TestPayload(1024)); + }); + started.get_future().wait(); + // Give the worker time to enter the writer before this continuation + // returns and `SetFlushed()` restarts the write loop. Joining here would + // add a happens-before edge and hide any race from TSAN. The sleep only + // makes the interleaving likely: the assertions below hold for any + // interleaving, so a slow scheduler reduces race coverage but cannot + // make this test fail. + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + }); + next.first.set_value(true); + worker.join(); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + EXPECT_STATUS_OK(f2.get()); +} + +/// @test Verify that `Flush()` futures complete in the order they were issued, +/// even if the underlying connection completes the next flush inline. +TEST(WriteConnectionBuffered, FlushFuturesCompleteInOrderWithInlineCompletion) { + AsyncSequencer sequencer; + auto mock = std::make_unique(); + auto mock_persisted_size = std::make_shared(0); + + EXPECT_CALL(*mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(*mock, PersistedState).WillRepeatedly([mock_persisted_size] { + return MakePersistedState(*mock_persisted_size); + }); + EXPECT_CALL(*mock, Flush) + .WillOnce([&](storage::WritePayload const& p) { + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush1").then( + [mock_persisted_size, size](auto) { + *mock_persisted_size += size; + return Status{}; + }); + }) + // Return an already-satisfied future so the second flush's continuation + // runs inline when `FlushStep()` dispatches it. + .WillOnce([mock_persisted_size](storage::WritePayload const& p) { + *mock_persisted_size += static_cast(p.size()); + return make_ready_future(Status{}); + }); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto f1 = connection->Flush(TestPayload(1024)); + auto f2 = connection->Flush(TestPayload(1024)); + + std::vector completion_order; + auto done1 = f1.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + completion_order.push_back(1); + }); + auto done2 = f2.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + completion_order.push_back(2); + }); + + // Complete the first flush; the write loop then dispatches the second flush, + // which completes inline. `f1` must still complete before `f2`. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush1"); + next.first.set_value(true); + + ASSERT_TRUE(done1.is_ready()); + ASSERT_TRUE(done2.is_ready()); + EXPECT_THAT(completion_order, ElementsAre(1, 2)); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal diff --git a/google/cloud/storage/internal/async/writer_connection_resumed.cc b/google/cloud/storage/internal/async/writer_connection_resumed.cc index 8e77e9ecdbb60..fe0e5a7e995f8 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed.cc @@ -404,28 +404,29 @@ class AsyncWriterConnectionResumedState write_offset_ -= static_cast(n); } } - // If the buffer is small enough, collect all the handlers to notify them. - auto const handlers = ClearHandlersIfEmpty(lk); if (is_resume) { + // We are resuming. The pending flush promises (if any) should not be + // satisfied yet, because we haven't actually flushed the data on the new + // connection. The `WriteLoop` will trigger a flush (potentially empty) + // if `flush_` is still true, which will satisfy the promises when it + // completes. + auto const handlers = ClearHandlersIfEmpty(lk); + // Mark the writer idle under the lock so any operation chained from a + // handler below sees an idle writer and is dispatched immediately, and + // `state_` is never modified without holding `mu_`. state_ = State::kIdle; - StartWriting(std::move(lk)); + lk.unlock(); // Release lock before notifying. // The notifications are deferred until the lock is released, as they // might call back and try to acquire the lock. for (auto const& h : handlers) { h->Execute(Status{}); } + // Re-acquire the lock to restart the write loop. This is a no-op if a + // handler above already restarted it. + StartWriting(std::unique_lock(mu_)); return; } - // SetFlushed will release the lock before returning. SetFlushed(std::move(lk), Status{}, persisted_size); - // Re-acquire the lock to resume writing now that flush_ has been updated. - state_ = State::kIdle; - StartWriting(std::unique_lock(mu_)); - // The notifications are deferred until the lock is released, as they might - // call back and try to acquire the lock. - for (auto const& h : handlers) { - h->Execute(Status{}); - } } void WriteStep(std::unique_lock lk, absl::Cord payload) { @@ -625,11 +626,19 @@ class AsyncWriterConnectionResumedState if (pending_flush_promises_.empty()) { flush_ = false; } - lk.unlock(); // Unlock only once before notifying - // Notify handlers and the specific flush promises *after* releasing the - // lock. + // Mark the writer idle under the lock so any operation chained from a + // callback below sees an idle writer and is dispatched immediately, and + // `state_` is never modified without holding `mu_`. + state_ = State::kIdle; + lk.unlock(); // Release lock before notifying. + // Notify handlers and satisfied flush promises before restarting the + // write loop so callbacks cannot be overtaken by a queued flush that + // completes inline. for (auto& h : handlers) h->Execute(Status{}); for (auto& f : flushes_to_complete) f.set_value(result); + // Re-acquire the lock to resume writing any remaining buffered data. + // This is a no-op if a callback above already restarted the write loop. + StartWriting(std::unique_lock(mu_)); } void SetError(std::unique_lock lk, Status const& status) { diff --git a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc index a4f7b9bf7f090..309145022276b 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc @@ -25,7 +25,12 @@ #include "google/storage/v2/storage.pb.h" #include #include +#include +#include +#include +#include #include +#include namespace google { namespace cloud { @@ -41,6 +46,7 @@ using ::google::cloud::testing_util::IsOkAndHolds; using ::google::cloud::testing_util::IsProtoEqual; using ::google::cloud::testing_util::StatusIs; using ::testing::_; +using ::testing::ElementsAre; using ::testing::Eq; using ::testing::ResultOf; using ::testing::Return; @@ -1312,6 +1318,246 @@ TEST(WriteConnectionResumed, CustomRetryPolicyOption) { EXPECT_THAT(write.get(), StatusIs(StatusCode::kInvalidArgument)); } +/// Test case for an operation issued from inside a `Flush()` continuation. +struct ChainedOpCase { + std::string name; + std::function(storage::AsyncWriterConnection&)> start_op; + std::string expected_call; +}; + +class WriterConnectionResumedFlushCallbackTest + : public ::testing::TestWithParam {}; + +/// @test Verify that an operation chained from a `Flush()` continuation is +/// dispatched to the underlying connection before the continuation returns. +TEST_P(WriterConnectionResumedFlushCallbackTest, + ChainedOpDispatchedBeforeCallbackReturns) { + AsyncSequencer sequencer; + auto mock = std::make_unique(); + auto mock_persisted_size = std::make_shared(0); + int call_count = 0; + + EXPECT_CALL(*mock, PersistedState).WillRepeatedly([mock_persisted_size] { + return MakePersistedState(*mock_persisted_size); + }); + EXPECT_CALL(*mock, WriteHandle).WillRepeatedly(Return(std::nullopt)); + EXPECT_CALL(*mock, Write).WillRepeatedly([&](auto const&) { + ++call_count; + return sequencer.PushBack("Write").then([](auto) { return Status{}; }); + }); + EXPECT_CALL(*mock, Flush).WillRepeatedly([&](storage::WritePayload const& p) { + ++call_count; + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush").then( + [mock_persisted_size, size](auto f) { + if (!f.get()) return TransientError(); + *mock_persisted_size += size; + return Status{}; + }); + }); + EXPECT_CALL(*mock, Close).WillRepeatedly([&](auto const&) { + ++call_count; + return sequencer.PushBack("Close").then([](auto) { return Status{}; }); + }); + EXPECT_CALL(*mock, Finalize).WillRepeatedly([&](auto const&) { + ++call_count; + return sequencer.PushBack("Finalize").then([](auto) { + return make_status_or(TestObject()); + }); + }); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto initial_request = google::storage::v2::BidiWriteObjectRequest{}; + auto first_response = google::storage::v2::BidiWriteObjectResponse{}; + // Use non-zero watermarks so a small `Write()` is dispatched as a `Write` + // rather than a `Flush` (the default low-water mark of 0 always flushes). + auto options = Options{} + .set(16 * 1024) + .set(32 * 1024); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, + first_response, options); + + auto flush = connection->Flush(TestPayload(1024)); + ASSERT_EQ(call_count, 1); + + // Chain the next operation from `flush`'s continuation. Because `flush` is + // already satisfied when this callback runs, the writer must already be idle + // so the chained operation dispatches synchronously (`call_count == 2`) + // before the callback returns. + future chained_op; + auto callback_done = flush.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained_op = GetParam().start_op(*connection); + EXPECT_EQ(call_count, 2); + }); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, GetParam().expected_call); + next.first.set_value(true); + EXPECT_STATUS_OK(chained_op.get()); +} + +INSTANTIATE_TEST_SUITE_P( + WriteConnectionResumed, WriterConnectionResumedFlushCallbackTest, + ::testing::Values(ChainedOpCase{"Flush", + [](storage::AsyncWriterConnection& c) { + return c.Flush(TestPayload(1024)); + }, + "Flush"}, + ChainedOpCase{"EmptyFlush", + [](storage::AsyncWriterConnection& c) { + return c.Flush(storage::WritePayload{}); + }, + "Flush"}, + ChainedOpCase{"Write", + [](storage::AsyncWriterConnection& c) { + return c.Write(TestPayload(1024)); + }, + "Write"}, + ChainedOpCase{"Close", + [](storage::AsyncWriterConnection& c) { + return c.Close(storage::WritePayload{}); + }, + "Close"}, + ChainedOpCase{"Finalize", + [](storage::AsyncWriterConnection& c) { + return c.Finalize(storage::WritePayload{}) + .then([](auto f) { + return f.get().status(); + }); + }, + "Finalize"}), + [](::testing::TestParamInfo const& info) { + return info.param.name; + }); + +/// @test Verify that a `Flush()` continuation handing the next `Flush()` to +/// another thread does not race with `OnQuery()` on `state_`. +TEST(WriteConnectionResumed, FlushCallbackCrossThreadFlush) { + AsyncSequencer sequencer; + auto mock = std::make_unique(); + auto mock_persisted_size = std::make_shared(0); + + EXPECT_CALL(*mock, PersistedState).WillRepeatedly([mock_persisted_size] { + return MakePersistedState(*mock_persisted_size); + }); + EXPECT_CALL(*mock, WriteHandle).WillRepeatedly(Return(std::nullopt)); + EXPECT_CALL(*mock, Flush).WillRepeatedly([&](storage::WritePayload const& p) { + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush").then( + [mock_persisted_size, size](auto f) { + if (!f.get()) return TransientError(); + *mock_persisted_size += size; + return Status{}; + }); + }); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto initial_request = google::storage::v2::BidiWriteObjectRequest{}; + auto first_response = google::storage::v2::BidiWriteObjectResponse{}; + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, + first_response, Options{}); + + auto f1 = connection->Flush(TestPayload(1024)); + auto next = sequencer.PopFrontWithName(); + ASSERT_EQ(next.second, "Flush"); + + std::thread worker; + future f2; + auto callback_done = f1.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + promise started; + worker = std::thread([&] { + started.set_value(); + f2 = connection->Flush(TestPayload(1024)); + }); + started.get_future().wait(); + // Give the worker time to enter the writer before this continuation + // returns and `SetFlushed()` restarts the write loop. Joining here would + // add a happens-before edge and hide any race from TSAN. The sleep only + // makes the interleaving likely: the assertions below hold for any + // interleaving, so a slow scheduler reduces race coverage but cannot + // make this test fail. + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + }); + next.first.set_value(true); + worker.join(); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + EXPECT_STATUS_OK(f2.get()); +} + +/// @test Verify that `Flush()` futures complete in the order they were issued, +/// even if the underlying connection completes the next flush inline. +TEST(WriteConnectionResumed, FlushFuturesCompleteInOrderWithInlineCompletion) { + AsyncSequencer sequencer; + auto mock = std::make_unique(); + auto mock_persisted_size = std::make_shared(0); + + EXPECT_CALL(*mock, PersistedState).WillRepeatedly([mock_persisted_size] { + return MakePersistedState(*mock_persisted_size); + }); + EXPECT_CALL(*mock, WriteHandle).WillRepeatedly(Return(std::nullopt)); + EXPECT_CALL(*mock, Flush) + .WillOnce([&](storage::WritePayload const& p) { + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush1").then( + [mock_persisted_size, size](auto) { + *mock_persisted_size += size; + return Status{}; + }); + }) + // Return an already-satisfied future so the second flush's continuation + // runs inline when `FlushStep()` dispatches it. + .WillOnce([mock_persisted_size](storage::WritePayload const& p) { + *mock_persisted_size += static_cast(p.size()); + return make_ready_future(Status{}); + }); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto initial_request = google::storage::v2::BidiWriteObjectRequest{}; + auto first_response = google::storage::v2::BidiWriteObjectResponse{}; + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, + first_response, Options{}); + + auto f1 = connection->Flush(TestPayload(1024)); + auto f2 = connection->Flush(TestPayload(1024)); + + std::vector completion_order; + auto done1 = f1.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + completion_order.push_back(1); + }); + auto done2 = f2.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + completion_order.push_back(2); + }); + + // Complete the first flush; the write loop then dispatches the second flush, + // which completes inline. `f1` must still complete before `f2`. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush1"); + next.first.set_value(true); + + ASSERT_TRUE(done1.is_ready()); + ASSERT_TRUE(done2.is_ready()); + EXPECT_THAT(completion_order, ElementsAre(1, 2)); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal From 41f2fb7ddb2eaaf80e319e6c586608a1e65498e1 Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Tue, 29 Sep 2026 11:12:38 +0000 Subject: [PATCH 2/3] Address feedback from Gemini code assistant --- .../internal/async/writer_connection_buffered_test.cc | 5 ++++- .../storage/internal/async/writer_connection_resumed_test.cc | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc index 1b26fd4d88cfa..cdca1f10bca2b 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc @@ -2100,8 +2100,11 @@ TEST(WriteConnectionBuffered, FlushCallbackCrossThreadFlush) { std::this_thread::sleep_for(std::chrono::milliseconds(10)); }); next.first.set_value(true); - worker.join(); + // Completing the sequencer promise runs the whole chain inline, including + // the `f1` continuation, so `worker` is assigned before `set_value()` + // returns. ASSERT_TRUE(callback_done.is_ready()); + worker.join(); next = sequencer.PopFrontWithName(); EXPECT_EQ(next.second, "Flush"); diff --git a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc index 309145022276b..c95730b0093e7 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc @@ -1490,8 +1490,11 @@ TEST(WriteConnectionResumed, FlushCallbackCrossThreadFlush) { std::this_thread::sleep_for(std::chrono::milliseconds(10)); }); next.first.set_value(true); - worker.join(); + // Completing the sequencer promise runs the whole chain inline, including + // the `f1` continuation, so `worker` is assigned before `set_value()` + // returns. ASSERT_TRUE(callback_done.is_ready()); + worker.join(); next = sequencer.PopFrontWithName(); EXPECT_EQ(next.second, "Flush"); From c0f601c0b52629345ae3293a4f7eee311a1d6171 Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Mon, 5 Oct 2026 06:38:16 +0000 Subject: [PATCH 3/3] Address feedback about small writes and test coverage --- .../async/writer_connection_buffered.cc | 40 +- .../async/writer_connection_buffered_test.cc | 331 ++++++++++++++ .../async/writer_connection_resumed.cc | 40 +- .../async/writer_connection_resumed_test.cc | 410 +++++++++++++++++- 4 files changed, 804 insertions(+), 17 deletions(-) diff --git a/google/cloud/storage/internal/async/writer_connection_buffered.cc b/google/cloud/storage/internal/async/writer_connection_buffered.cc index 45d0d2edf6efd..52f679a2b51f4 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered.cc @@ -217,7 +217,11 @@ class AsyncWriterConnectionBufferedState bool flush = false) { if (!resume_status_.ok()) return make_ready_future(resume_status_); auto const buffer_size = resend_buffer_.size(); - flush_ = (buffer_size >= buffer_size_lwm_) || flush; + // `NeedsFlush()` stays true while `Flush()` calls are pending, so a small + // `Write()`, possibly chained from a callback, cannot turn a queued flush + // into a plain write. The explicit low-water mark check preserves the + // existing behavior for an empty `Write()` when the low-water mark is 0. + flush_ = flush || NeedsFlush(lk) || buffer_size >= buffer_size_lwm_; auto result = make_ready_future(Status{}); if (buffer_size >= buffer_size_hwm_) { auto p = promise(); @@ -345,11 +349,29 @@ class AsyncWriterConnectionBufferedState auto ClearHandlersIfEmpty(std::unique_lock const& /* lk */) { decltype(flush_handlers_) tmp; - if (resend_buffer_.size() >= buffer_size_lwm_) return tmp; + // Release the waiters once the buffer is empty or below the low-water + // mark. The emptiness check matters when the low-water mark is 0. + if (!resend_buffer_.empty() && resend_buffer_.size() >= buffer_size_lwm_) { + return tmp; + } flush_handlers_.swap(tmp); return tmp; } + /** + * Returns true if the write loop must send data with `Flush()`. + * + * That is the case while `Flush()` calls are pending, or while the buffer is + * non-empty and at or above the low-water mark. An empty buffer only needs a + * flush if a `Flush()` is pending, which avoids an endless loop of empty + * flushes when the low-water mark is 0. + */ + bool NeedsFlush(std::unique_lock const& /* lk */) const { + return !pending_flush_promises_.empty() || + (!resend_buffer_.empty() && + resend_buffer_.size() >= buffer_size_lwm_); + } + void OnQuery(std::unique_lock lk, std::int64_t persisted_size, bool is_resume = false) { if (persisted_size < buffer_offset_) { @@ -395,7 +417,7 @@ class AsyncWriterConnectionBufferedState // `writing_` is never modified without holding `mu_`. resuming_ = false; writing_ = false; - lk.unlock(); // Release lock before notifying. + lk.unlock(); // Release lock before notifying // The notifications are deferred until the lock is released, as they // might call back and try to acquire the lock. for (auto const& h : handlers) h->Execute(Status{}); @@ -564,7 +586,9 @@ class AsyncWriterConnectionBufferedState std::int64_t persisted_size) { if (!result.ok()) return SetError(std::move(lk), std::move(result)); // Do NOT reset finalize_ or finalizing_ here. - auto handlers = ClearHandlers(lk); + // Only release the high-water mark waiters once the buffer drops below the + // low-water mark (or is empty). + auto handlers = ClearHandlersIfEmpty(lk); std::vector> flushes_to_complete; while (!pending_flush_promises_.empty() && pending_flush_promises_.front().target_offset <= persisted_size) { @@ -572,14 +596,14 @@ class AsyncWriterConnectionBufferedState std::move(pending_flush_promises_.front().p)); pending_flush_promises_.pop_front(); } - if (pending_flush_promises_.empty()) { - flush_ = false; - } + // Keep flushing while `NeedsFlush()` holds, so any high-water mark waiters + // not released above are released by a later flush. + flush_ = NeedsFlush(lk); // Mark the writer idle under the lock so any operation chained from a // callback below sees an idle writer and is dispatched immediately, and // `writing_` is never modified without holding `mu_`. writing_ = false; - lk.unlock(); // Release lock before notifying. + lk.unlock(); // Release lock before notifying // Notify handlers and satisfied flush promises before restarting the // write loop so callbacks cannot be overtaken by a queued flush that // completes inline. diff --git a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc index cdca1f10bca2b..fe98e0a736891 100644 --- a/google/cloud/storage/internal/async/writer_connection_buffered_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_buffered_test.cc @@ -78,6 +78,39 @@ storage::WritePayload TestPayload(std::size_t n) { return storage::WritePayload(std::string(n, 'A')); } +/// Configures `mock` so every `Write()` and `Flush()` is recorded in +/// `sequencer` (as `"Write" + tag` / `"Flush" + tag`) and counted in `*calls`. +/// A successful `Flush()` advances `*persisted` by the payload size. +void RecordWritesAndFlushes(MockAsyncWriterConnection& mock, + AsyncSequencer& sequencer, + std::shared_ptr const& persisted, + std::shared_ptr const& calls, + std::string const& tag) { + EXPECT_CALL(mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(mock, PersistedState).WillRepeatedly([persisted] { + return MakePersistedState(*persisted); + }); + EXPECT_CALL(mock, Write) + .WillRepeatedly([&sequencer, calls, tag](storage::WritePayload const&) { + ++*calls; + return sequencer.PushBack("Write" + tag).then([](auto) { + return Status{}; + }); + }); + EXPECT_CALL(mock, Flush) + .WillRepeatedly( + [&sequencer, persisted, calls, tag](storage::WritePayload const& p) { + ++*calls; + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush" + tag) + .then([persisted, size](auto f) { + if (!f.get()) return TransientError(); + *persisted += size; + return Status{}; + }); + }); +} + TEST(WriteConnectionBuffered, FinalizeEmpty) { AsyncSequencer sequencer; auto mock = std::make_unique(); @@ -2168,6 +2201,304 @@ TEST(WriteConnectionBuffered, FlushFuturesCompleteInOrderWithInlineCompletion) { EXPECT_THAT(completion_order, ElementsAre(1, 2)); } +/// @test Verify that a small `Write()` issued while a second `Flush()` is +/// queued does not cancel that pending flush. +TEST(WriteConnectionBuffered, SmallWriteDoesNotCancelQueuedFlush) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls, ""); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto f1 = connection->Flush(TestPayload(1024)); + auto f2 = connection->Flush(TestPayload(1024)); + // Below the low-water mark: on its own this `Write()` would not need a flush. + auto w = connection->Write(TestPayload(1024)); + ASSERT_TRUE(w.is_ready()); + EXPECT_STATUS_OK(w.get()); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(f1.is_ready()); + EXPECT_STATUS_OK(f1.get()); + ASSERT_FALSE(f2.is_ready()); + + // `f2` is still pending, so the remaining data must be flushed, not written. + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(f2.is_ready()); + EXPECT_STATUS_OK(f2.get()); +} + +/// @test Verify that a small `Write()` chained from a `Flush()` continuation +/// does not cancel a second `Flush()` that is still pending. +TEST(WriteConnectionBuffered, ChainedSmallWriteDoesNotCancelQueuedFlush) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls, ""); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto f1 = connection->Flush(TestPayload(1024)); + auto f2 = connection->Flush(TestPayload(1024)); + future chained; + auto callback_done = f1.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + }); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(callback_done.is_ready()); + ASSERT_FALSE(f2.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(f2.is_ready()); + EXPECT_STATUS_OK(f2.get()); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that `Write()` calls blocked at the high-water mark stay +/// blocked, and the buffer keeps being flushed, until the buffer drops below +/// the low-water mark. +TEST(WriteConnectionBuffered, HwmWaitersStayBlockedUntilBelowLwm) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls, ""); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto w1 = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w1.is_ready()); + auto w2 = connection->Write(TestPayload(20 * 1024)); + ASSERT_FALSE(w2.is_ready()); + + // Persisting the first 32 KiB leaves 20 KiB buffered, which is still above + // the 16 KiB low-water mark. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + EXPECT_FALSE(w1.is_ready()); + EXPECT_FALSE(w2.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(w1.is_ready()); + EXPECT_STATUS_OK(w1.get()); + ASSERT_TRUE(w2.is_ready()); + EXPECT_STATUS_OK(w2.get()); +} + +/// @test Verify that an operation chained from a `Write()` blocked at the +/// high-water mark is dispatched before the continuation returns. +TEST(WriteConnectionBuffered, HwmWriteCallbackChainedOpDispatchedBeforeReturn) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls, ""); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), TestOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + future chained; + auto callback_done = w.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + EXPECT_EQ(*calls, 2); + }); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Write"); + next.first.set_value(true); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that an operation chained from a `Write()` blocked at the +/// high-water mark is dispatched before the continuation returns when the +/// buffer shrinks as part of a `Resume()`. +TEST(WriteConnectionBuffered, HwmWriteCallbackChainedOpDispatchedAfterResume) { + AsyncSequencer sequencer; + auto persisted1 = std::make_shared(0); + auto calls1 = std::make_shared(0); + auto mock1 = std::make_unique(); + RecordWritesAndFlushes(*mock1, sequencer, persisted1, calls1, "1"); + // The resumed connection reports that all 32 KiB were persisted. + auto persisted2 = std::make_shared(32 * 1024); + auto calls2 = std::make_shared(0); + auto mock2 = std::make_unique(); + RecordWritesAndFlushes(*mock2, sequencer, persisted2, calls2, "2"); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).WillOnce([&] { + return make_ready_future(make_status_or( + std::unique_ptr(std::move(mock2)))); + }); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock1), TestOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + future chained; + auto callback_done = w.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + EXPECT_EQ(*calls2, 1); + }); + + // Fail the flush to trigger a `Resume()`. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush1"); + next.first.set_value(false); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Write2"); + next.first.set_value(true); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that a small `Write()` chained from a high-water mark +/// continuation during a `Resume()` does not cancel a pending `Flush()`. +TEST(WriteConnectionBuffered, ChainedSmallWriteAfterResumeKeepsPendingFlush) { + AsyncSequencer sequencer; + auto persisted1 = std::make_shared(0); + auto calls1 = std::make_shared(0); + auto mock1 = std::make_unique(); + RecordWritesAndFlushes(*mock1, sequencer, persisted1, calls1, "1"); + auto persisted2 = std::make_shared(32 * 1024); + auto calls2 = std::make_shared(0); + auto mock2 = std::make_unique(); + RecordWritesAndFlushes(*mock2, sequencer, persisted2, calls2, "2"); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).WillOnce([&] { + return make_ready_future(make_status_or( + std::unique_ptr(std::move(mock2)))); + }); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock1), TestOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + auto flush = connection->Flush(TestPayload(1024)); + future chained; + auto callback_done = w.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + }); + + // Fail the flush. The resumed connection reports 32 KiB persisted, leaving + // 1 KiB (below the low-water mark) buffered for the pending `Flush()`. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush1"); + next.first.set_value(false); + ASSERT_TRUE(callback_done.is_ready()); + ASSERT_FALSE(flush.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush2"); + next.first.set_value(true); + ASSERT_TRUE(flush.is_ready()); + EXPECT_STATUS_OK(flush.get()); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that a `Resume()` does not release `Write()` calls blocked at +/// the high-water mark while the buffer is still above the low-water mark. +TEST(WriteConnectionBuffered, HwmWaitersStayBlockedAfterResumeUntilBelowLwm) { + AsyncSequencer sequencer; + auto persisted1 = std::make_shared(0); + auto calls1 = std::make_shared(0); + auto mock1 = std::make_unique(); + RecordWritesAndFlushes(*mock1, sequencer, persisted1, calls1, "1"); + // The resumed connection reports that nothing was persisted. + auto persisted2 = std::make_shared(0); + auto calls2 = std::make_shared(0); + auto mock2 = std::make_unique(); + RecordWritesAndFlushes(*mock2, sequencer, persisted2, calls2, "2"); + + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).WillOnce([&] { + return make_ready_future(make_status_or( + std::unique_ptr(std::move(mock2)))); + }); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock1), TestOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + + // Fail the flush to trigger a `Resume()`. All 32 KiB are still buffered, + // which is above the 16 KiB low-water mark, so `w` must stay blocked. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush1"); + next.first.set_value(false); + EXPECT_FALSE(w.is_ready()); + + // Once the resumed connection persists the data, `w` is released. + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush2"); + next.first.set_value(true); + ASSERT_TRUE(w.is_ready()); + EXPECT_STATUS_OK(w.get()); +} + +/// @test Verify that, with a zero low-water mark, a `Write()` blocked at the +/// high-water mark is released once the buffer drains, and that the writer +/// does not keep issuing empty flushes afterwards. +TEST(WriteConnectionBuffered, HwmWaiterReleasedWhenBufferDrainsWithZeroLwm) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls, ""); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionBuffered( + mock_factory.AsStdFunction(), std::move(mock), + Options{} + .set(0) + .set(1024)); + + auto w = connection->Write(TestPayload(1024)); + ASSERT_FALSE(w.is_ready()); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(w.is_ready()); + EXPECT_STATUS_OK(w.get()); + // The buffer is empty and no `Flush()` is pending, so no further calls. + EXPECT_EQ(*calls, 1); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal diff --git a/google/cloud/storage/internal/async/writer_connection_resumed.cc b/google/cloud/storage/internal/async/writer_connection_resumed.cc index fe0e5a7e995f8..81f8264dc22bf 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed.cc @@ -224,7 +224,11 @@ class AsyncWriterConnectionResumedState bool flush = false) { if (!resume_status_.ok()) return make_ready_future(resume_status_); auto const buffer_size = resend_buffer_.size(); - flush_ = (buffer_size >= buffer_size_lwm_) || flush; + // `NeedsFlush()` stays true while `Flush()` calls are pending, so a small + // `Write()`, possibly chained from a callback, cannot turn a queued flush + // into a plain write. The explicit low-water mark check preserves the + // existing behavior for an empty `Write()` when the low-water mark is 0. + flush_ = flush || NeedsFlush(lk) || buffer_size >= buffer_size_lwm_; auto result = make_ready_future(Status{}); if (buffer_size >= buffer_size_hwm_) { auto p = promise(); @@ -363,11 +367,29 @@ class AsyncWriterConnectionResumedState auto ClearHandlersIfEmpty(std::unique_lock const& /* lk */) { decltype(flush_handlers_) tmp; - if (resend_buffer_.size() >= buffer_size_lwm_) return tmp; + // Release the waiters once the buffer is empty or below the low-water + // mark. The emptiness check matters when the low-water mark is 0. + if (!resend_buffer_.empty() && resend_buffer_.size() >= buffer_size_lwm_) { + return tmp; + } flush_handlers_.swap(tmp); return tmp; } + /** + * Returns true if the write loop must send data with `Flush()`. + * + * That is the case while `Flush()` calls are pending, or while the buffer is + * non-empty and at or above the low-water mark. An empty buffer only needs a + * flush if a `Flush()` is pending, which avoids an endless loop of empty + * flushes when the low-water mark is 0. + */ + bool NeedsFlush(std::unique_lock const& /* lk */) const { + return !pending_flush_promises_.empty() || + (!resend_buffer_.empty() && + resend_buffer_.size() >= buffer_size_lwm_); + } + void OnQuery(std::unique_lock lk, std::int64_t persisted_size, bool is_resume = false) { auto handle = impl_->WriteHandle(); @@ -415,7 +437,7 @@ class AsyncWriterConnectionResumedState // handler below sees an idle writer and is dispatched immediately, and // `state_` is never modified without holding `mu_`. state_ = State::kIdle; - lk.unlock(); // Release lock before notifying. + lk.unlock(); // Release lock before notifying // The notifications are deferred until the lock is released, as they // might call back and try to acquire the lock. for (auto const& h : handlers) { @@ -615,7 +637,9 @@ class AsyncWriterConnectionResumedState std::int64_t persisted_size) { if (!result.ok()) return SetError(std::move(lk), std::move(result)); // Do NOT reset finalize_ or finalizing_ here. - auto handlers = ClearHandlers(lk); + // Only release the high-water mark waiters once the buffer drops below the + // low-water mark (or is empty). + auto handlers = ClearHandlersIfEmpty(lk); std::vector> flushes_to_complete; while (!pending_flush_promises_.empty() && pending_flush_promises_.front().target_offset <= persisted_size) { @@ -623,14 +647,14 @@ class AsyncWriterConnectionResumedState std::move(pending_flush_promises_.front().p)); pending_flush_promises_.pop_front(); } - if (pending_flush_promises_.empty()) { - flush_ = false; - } + // Keep flushing while `NeedsFlush()` holds, so any high-water mark waiters + // not released above are released by a later flush. + flush_ = NeedsFlush(lk); // Mark the writer idle under the lock so any operation chained from a // callback below sees an idle writer and is dispatched immediately, and // `state_` is never modified without holding `mu_`. state_ = State::kIdle; - lk.unlock(); // Release lock before notifying. + lk.unlock(); // Release lock before notifying // Notify handlers and satisfied flush promises before restarting the // write loop so callbacks cannot be overtaken by a queued flush that // completes inline. diff --git a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc index c95730b0093e7..e8e7cfd595dcb 100644 --- a/google/cloud/storage/internal/async/writer_connection_resumed_test.cc +++ b/google/cloud/storage/internal/async/writer_connection_resumed_test.cc @@ -46,6 +46,7 @@ using ::google::cloud::testing_util::IsOkAndHolds; using ::google::cloud::testing_util::IsProtoEqual; using ::google::cloud::testing_util::StatusIs; using ::testing::_; +using ::testing::An; using ::testing::ElementsAre; using ::testing::Eq; using ::testing::ResultOf; @@ -59,6 +60,10 @@ using MockFactory = using MockStreamingRpc = ::testing::MockFunction()>; +using MockStream = ::google::cloud::mocks::MockAsyncStreamingReadWriteRpc< + google::storage::v2::BidiWriteObjectRequest, + google::storage::v2::BidiWriteObjectResponse>; + std::variant MakePersistedState( std::int64_t persisted_size) { return persisted_size; @@ -68,6 +73,13 @@ storage::WritePayload TestPayload(std::size_t n) { return storage::WritePayload(std::string(n, 'A')); } +/// Options with non-zero watermarks, so small writes are not always flushed. +Options WatermarkOptions() { + return Options{} + .set(16 * 1024) + .set(32 * 1024); +} + auto TestObject() { auto object = google::storage::v2::Object{}; object.set_bucket("projects/_/buckets/test-bucket"); @@ -75,6 +87,76 @@ auto TestObject() { return object; } +/// Configures `mock` so every `Write()` and `Flush()` is recorded in +/// `sequencer` (as `"Write"` / `"Flush"`) and counted in `*calls`. A +/// successful `Flush()` advances `*persisted` by the payload size. +void RecordWritesAndFlushes(MockAsyncWriterConnection& mock, + AsyncSequencer& sequencer, + std::shared_ptr const& persisted, + std::shared_ptr const& calls) { + EXPECT_CALL(mock, UploadId).WillRepeatedly(Return("test-upload-id")); + EXPECT_CALL(mock, WriteHandle).WillRepeatedly(Return(std::nullopt)); + EXPECT_CALL(mock, PersistedState).WillRepeatedly([persisted] { + return MakePersistedState(*persisted); + }); + EXPECT_CALL(mock, Write) + .WillRepeatedly([&sequencer, calls](storage::WritePayload const&) { + ++*calls; + return sequencer.PushBack("Write").then([](auto) { return Status{}; }); + }); + EXPECT_CALL(mock, Flush) + .WillRepeatedly( + [&sequencer, persisted, calls](storage::WritePayload const& p) { + ++*calls; + auto const size = static_cast(p.size()); + return sequencer.PushBack("Flush").then([persisted, size](auto f) { + if (!f.get()) return TransientError(); + *persisted += size; + return Status{}; + }); + }); +} + +/// Configures the stream used after `Resume()`. Each stream `Write()` is +/// recorded as `"StreamFlush"` or `"StreamWrite"` and counted in `*calls`. A +/// flush persists all bytes up to the end of its payload. +void RecordStreamWrites(MockStream& stream, AsyncSequencer& sequencer, + std::shared_ptr const& persisted, + std::shared_ptr const& calls) { + EXPECT_CALL(stream, Write) + .WillRepeatedly([&sequencer, persisted, calls]( + google::storage::v2::BidiWriteObjectRequest const& r, + grpc::WriteOptions) { + ++*calls; + bool const flush = r.flush(); + std::int64_t const end = + r.write_offset() + + static_cast(GetContent(r.checksummed_data()).size()); + return sequencer.PushBack(flush ? "StreamFlush" : "StreamWrite") + .then([persisted, flush, end](auto) { + if (flush) *persisted = end; + return true; + }); + }); + EXPECT_CALL(stream, Read).WillRepeatedly([persisted] { + google::storage::v2::BidiWriteObjectResponse response; + response.set_persisted_size(*persisted); + return make_ready_future(std::make_optional(response)); + }); + EXPECT_CALL(stream, Finish).WillRepeatedly([] { + return make_ready_future(Status{}); + }); + EXPECT_CALL(stream, Cancel).WillRepeatedly(Return()); +} + +std::shared_ptr MakeMockHash() { + auto hash = std::make_shared(); + EXPECT_CALL(*hash, Update(An(), An(), + An())) + .WillRepeatedly(Return(Status())); + return hash; +} + TEST(WriteConnectionResumed, FinalizeEmpty) { AsyncSequencer sequencer; auto mock = std::make_unique(); @@ -495,9 +577,12 @@ TEST(WriteConnectionResumed, NoConcurrentWritesWhenFlushAndWriteRace) { MockFactory mock_factory; EXPECT_CALL(mock_factory, Call).Times(0); + // Use non-zero watermarks so the 1 KiB `Write()` below is not blocked as a + // high-water mark waiter. With `Options{}` it would only complete once the + // buffer drains, which this mock (persisted size fixed at 0) never reports. auto connection = MakeWriterConnectionResumed( mock_factory.AsStdFunction(), std::move(mock), initial_request, nullptr, - first_response, Options{}); + first_response, WatermarkOptions()); // Start a flush which will call impl->Flush() and block. auto flush_future = connection->Flush({}); @@ -1561,6 +1646,329 @@ TEST(WriteConnectionResumed, FlushFuturesCompleteInOrderWithInlineCompletion) { EXPECT_THAT(completion_order, ElementsAre(1, 2)); } +/// @test Verify that a small `Write()` issued while a second `Flush()` is +/// queued does not cancel that pending flush. +TEST(WriteConnectionResumed, SmallWriteDoesNotCancelQueuedFlush) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, nullptr, + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto f1 = connection->Flush(TestPayload(1024)); + auto f2 = connection->Flush(TestPayload(1024)); + // Below the low-water mark: on its own this `Write()` would not need a flush. + auto w = connection->Write(TestPayload(1024)); + ASSERT_TRUE(w.is_ready()); + EXPECT_STATUS_OK(w.get()); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(f1.is_ready()); + EXPECT_STATUS_OK(f1.get()); + ASSERT_FALSE(f2.is_ready()); + + // `f2` is still pending, so the remaining data must be flushed, not written. + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(f2.is_ready()); + EXPECT_STATUS_OK(f2.get()); +} + +/// @test Verify that a small `Write()` chained from a `Flush()` continuation +/// does not cancel a second `Flush()` that is still pending. +TEST(WriteConnectionResumed, ChainedSmallWriteDoesNotCancelQueuedFlush) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, nullptr, + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto f1 = connection->Flush(TestPayload(1024)); + auto f2 = connection->Flush(TestPayload(1024)); + future chained; + auto callback_done = f1.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + }); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(callback_done.is_ready()); + ASSERT_FALSE(f2.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(f2.is_ready()); + EXPECT_STATUS_OK(f2.get()); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that `Write()` calls blocked at the high-water mark stay +/// blocked, and the buffer keeps being flushed, until the buffer drops below +/// the low-water mark. +TEST(WriteConnectionResumed, HwmWaitersStayBlockedUntilBelowLwm) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, nullptr, + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto w1 = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w1.is_ready()); + auto w2 = connection->Write(TestPayload(20 * 1024)); + ASSERT_FALSE(w2.is_ready()); + + // Persisting the first 32 KiB leaves 20 KiB buffered, which is still above + // the 16 KiB low-water mark. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + EXPECT_FALSE(w1.is_ready()); + EXPECT_FALSE(w2.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(w1.is_ready()); + EXPECT_STATUS_OK(w1.get()); + ASSERT_TRUE(w2.is_ready()); + EXPECT_STATUS_OK(w2.get()); +} + +/// @test Verify that an operation chained from a `Write()` blocked at the +/// high-water mark is dispatched before the continuation returns. +TEST(WriteConnectionResumed, HwmWriteCallbackChainedOpDispatchedBeforeReturn) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, nullptr, + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + future chained; + auto callback_done = w.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + EXPECT_EQ(*calls, 2); + }); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Write"); + next.first.set_value(true); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that an operation chained from a `Write()` blocked at the +/// high-water mark is dispatched before the continuation returns when the +/// buffer shrinks as part of a `Resume()`. +TEST(WriteConnectionResumed, HwmWriteCallbackChainedOpDispatchedAfterResume) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + + // The resumed stream reports that all 32 KiB were persisted. + auto stream_persisted = std::make_shared(32 * 1024); + auto stream_calls = std::make_shared(0); + auto stream = std::make_unique(); + RecordStreamWrites(*stream, sequencer, stream_persisted, stream_calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).WillOnce([&](auto const&) { + WriteObject::WriteResult result; + result.stream = std::move(stream); + result.first_response.set_persisted_size(32 * 1024); + return make_ready_future( + StatusOr(std::move(result))); + }); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, MakeMockHash(), + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + future chained; + auto callback_done = w.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + EXPECT_EQ(*stream_calls, 1); + }); + + // Fail the flush to trigger a `Resume()`. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(false); + ASSERT_TRUE(callback_done.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "StreamWrite"); + next.first.set_value(true); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that a small `Write()` chained from a high-water mark +/// continuation during a `Resume()` does not cancel a pending `Flush()`. +TEST(WriteConnectionResumed, ChainedSmallWriteAfterResumeKeepsPendingFlush) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + + auto stream_persisted = std::make_shared(32 * 1024); + auto stream_calls = std::make_shared(0); + auto stream = std::make_unique(); + RecordStreamWrites(*stream, sequencer, stream_persisted, stream_calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).WillOnce([&](auto const&) { + WriteObject::WriteResult result; + result.stream = std::move(stream); + result.first_response.set_persisted_size(32 * 1024); + return make_ready_future( + StatusOr(std::move(result))); + }); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, MakeMockHash(), + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + auto flush = connection->Flush(TestPayload(1024)); + future chained; + auto callback_done = w.then([&](future f) { + EXPECT_STATUS_OK(f.get()); + chained = connection->Write(TestPayload(1024)); + }); + + // Fail the flush. The resumed stream reports 32 KiB persisted, leaving + // 1 KiB (below the low-water mark) buffered for the pending `Flush()`. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(false); + ASSERT_TRUE(callback_done.is_ready()); + ASSERT_FALSE(flush.is_ready()); + + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "StreamFlush"); + next.first.set_value(true); + ASSERT_TRUE(flush.is_ready()); + EXPECT_STATUS_OK(flush.get()); + EXPECT_STATUS_OK(chained.get()); +} + +/// @test Verify that a `Resume()` does not release `Write()` calls blocked at +/// the high-water mark while the buffer is still above the low-water mark. +TEST(WriteConnectionResumed, HwmWaitersStayBlockedAfterResumeUntilBelowLwm) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + + // The resumed stream reports that nothing was persisted. + auto stream_persisted = std::make_shared(0); + auto stream_calls = std::make_shared(0); + auto stream = std::make_unique(); + RecordStreamWrites(*stream, sequencer, stream_persisted, stream_calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).WillOnce([&](auto const&) { + WriteObject::WriteResult result; + result.stream = std::move(stream); + result.first_response.set_persisted_size(0); + return make_ready_future( + StatusOr(std::move(result))); + }); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, MakeMockHash(), + google::storage::v2::BidiWriteObjectResponse{}, WatermarkOptions()); + + auto w = connection->Write(TestPayload(32 * 1024)); + ASSERT_FALSE(w.is_ready()); + + // Fail the flush to trigger a `Resume()`. All 32 KiB are still buffered, + // which is above the 16 KiB low-water mark, so `w` must stay blocked. + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(false); + EXPECT_FALSE(w.is_ready()); + + // Once the resumed stream persists the data, `w` is released. + next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "StreamFlush"); + next.first.set_value(true); + ASSERT_TRUE(w.is_ready()); + EXPECT_STATUS_OK(w.get()); +} + +/// @test Verify that, with a zero low-water mark, a `Write()` blocked at the +/// high-water mark is released once the buffer drains, and that the writer +/// does not keep issuing empty flushes afterwards. +TEST(WriteConnectionResumed, HwmWaiterReleasedWhenBufferDrainsWithZeroLwm) { + AsyncSequencer sequencer; + auto persisted = std::make_shared(0); + auto calls = std::make_shared(0); + auto mock = std::make_unique(); + RecordWritesAndFlushes(*mock, sequencer, persisted, calls); + MockFactory mock_factory; + EXPECT_CALL(mock_factory, Call).Times(0); + auto connection = MakeWriterConnectionResumed( + mock_factory.AsStdFunction(), std::move(mock), + google::storage::v2::BidiWriteObjectRequest{}, nullptr, + google::storage::v2::BidiWriteObjectResponse{}, + Options{} + .set(0) + .set(1024)); + + auto w = connection->Write(TestPayload(1024)); + ASSERT_FALSE(w.is_ready()); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Flush"); + next.first.set_value(true); + ASSERT_TRUE(w.is_ready()); + EXPECT_STATUS_OK(w.get()); + // The buffer is empty and no `Flush()` is pending, so no further calls. + EXPECT_EQ(*calls, 1); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal