From 76b464dc0ca7c12830376c37f7c4300fe343c6ab Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Fri, 11 Sep 2026 07:34:45 +0000 Subject: [PATCH 1/6] fix(storage): ensure thread-safe stream resumption and callback hygiene in ObjectDescriptorImpl --- .../internal/async/multi_stream_manager.h | 7 + .../async/multi_stream_manager_test.cc | 15 + .../internal/async/object_descriptor_impl.cc | 141 ++++- .../internal/async/object_descriptor_impl.h | 16 +- .../async/object_descriptor_impl_test.cc | 491 ++++++++++++++++++ 5 files changed, 638 insertions(+), 32 deletions(-) diff --git a/google/cloud/storage/internal/async/multi_stream_manager.h b/google/cloud/storage/internal/async/multi_stream_manager.h index bf26914127b88..aead9d8af0598 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager.h +++ b/google/cloud/storage/internal/async/multi_stream_manager.h @@ -159,6 +159,13 @@ class MultiStreamManager { return false; } + bool Contains(StreamIterator target) const { + for (auto it = streams_.begin(); it != streams_.end(); ++it) { + if (it == target) return true; + } + return false; + } + bool Empty() const { return streams_.empty(); } ConstStreamIterator End() const { return streams_.end(); } std::size_t Size() const { return streams_.size(); } diff --git a/google/cloud/storage/internal/async/multi_stream_manager_test.cc b/google/cloud/storage/internal/async/multi_stream_manager_test.cc index b17fa4d32173a..f890b1a1e7d65 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager_test.cc +++ b/google/cloud/storage/internal/async/multi_stream_manager_test.cc @@ -233,6 +233,21 @@ TEST(MultiStreamManagerTest, EmptyAndSizeTransitions) { EXPECT_EQ(mgr.Size(), 1U); } +TEST(MultiStreamManagerTest, ContainsTracksStreamMembership) { + auto mgr = MultiStreamManagerTest::MakeManager(); + auto it1 = mgr.GetFirstStream(); + EXPECT_TRUE(mgr.Contains(it1)); + + auto s2 = std::make_shared(); + auto it2 = mgr.AddStream(s2); + EXPECT_TRUE(mgr.Contains(it1)); + EXPECT_TRUE(mgr.Contains(it2)); + + mgr.RemoveStreamAndNotifyRanges(it1, Status()); + EXPECT_FALSE(mgr.Contains(it1)); + EXPECT_TRUE(mgr.Contains(it2)); +} + GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal } // namespace cloud diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.cc b/google/cloud/storage/internal/async/object_descriptor_impl.cc index 194f288c66744..dcb7c5093724a 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl.cc @@ -77,10 +77,14 @@ ObjectDescriptorImpl::ObjectDescriptorImpl( options_(std::move(options)), has_initial_read_ranges_(options_.has()), transport_ok_(std::move(transport_ok)) { + auto initial_read_stream = std::make_shared( + std::move(stream), resume_policy_prototype_->clone()); + // Notify the resume policy that the initial stream was established + // successfully. + initial_read_stream->resume_policy->OnStartSuccess(); stream_manager_ = std::make_unique( []() -> std::shared_ptr { return nullptr; }, // NOLINT - std::make_shared(std::move(stream), - resume_policy_prototype_->clone())); + std::move(initial_read_stream)); // Initialize the pacing limit from options if configured. if (options_.has()) { max_prewarmed_buffer_size_ = options_.get(); @@ -122,8 +126,9 @@ void ObjectDescriptorImpl::Start( std::unique_lock lk(mu_); auto it = stream_manager_->GetFirstStream(); if (it == stream_manager_->End()) return; + auto current_stream = it->stream->stream; lk.unlock(); - OnRead(it, std::move(first_response)); + OnRead(it, current_stream, std::move(first_response)); // Acquire lock and queue the background stream if multi-stream optimization // is enabled. if (options_.get()) { @@ -204,14 +209,16 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { auto read_stream = std::make_shared(std::move(stream_result->stream), self->resume_policy_prototype_->clone()); + read_stream->resume_policy->OnStartSuccess(); auto new_it = self->stream_manager_->AddStream(std::move(read_stream)); // Now that we consumed pending_stream_, queue the next one immediately. self->AssurePendingStreamQueued(lk); + auto new_stream = new_it->stream->stream; lk.unlock(); - self->OnRead(new_it, std::move(stream_result->first_response)); + self->OnRead(new_it, new_stream, std::move(stream_result->first_response)); }); } @@ -400,7 +407,7 @@ ObjectDescriptorImpl::CreateHashValidator(bool is_full_read) const { void ObjectDescriptorImpl::Flush(std::unique_lock lk, StreamIterator it) { - if (it->stream->write_pending || + if (it->stream->resuming || it->stream->write_pending || it->stream->next_request.read_ranges().empty()) { return; } @@ -414,14 +421,22 @@ void ObjectDescriptorImpl::Flush(std::unique_lock lk, auto current_stream = it->stream->stream; lk.unlock(); current_stream->Write(std::move(request)) - .then([w = WeakFromThis(), it](auto f) { - if (auto self = w.lock()) self->OnWrite(it, f.get()); + .then([w = WeakFromThis(), it, current_stream](auto f) { + if (auto self = w.lock()) self->OnWrite(it, current_stream, f.get()); }); } -void ObjectDescriptorImpl::OnWrite(StreamIterator it, bool ok) { +void ObjectDescriptorImpl::OnWrite(StreamIterator it, + std::shared_ptr const& stream, + bool ok) { std::unique_lock lk(mu_); - if (!ok) return DoFinish(std::move(lk), it); + // Discard callbacks from stale or removed streams (e.g. if the stream was + // replaced during reconnection or removed after an error). + if (!stream_manager_->Contains(it) || !it->stream || + it->stream->stream != stream) { + return; + } + if (!ok) return DoFinish(std::move(lk), it, stream); it->stream->write_pending = false; Flush(std::move(lk), it); } @@ -436,18 +451,24 @@ void ObjectDescriptorImpl::DoRead(std::unique_lock lk, // end of the block. auto current_stream = it->stream->stream; lk.unlock(); - current_stream->Read().then([w = WeakFromThis(), it](auto f) { - if (auto self = w.lock()) self->OnRead(it, f.get()); + current_stream->Read().then([w = WeakFromThis(), it, current_stream](auto f) { + if (auto self = w.lock()) self->OnRead(it, current_stream, f.get()); }); } void ObjectDescriptorImpl::OnRead( - StreamIterator it, + StreamIterator it, std::shared_ptr const& stream, std::optional response) { std::unique_lock lk(mu_); + // Discard callbacks from stale or removed streams (e.g. if the stream was + // replaced during reconnection or removed after an error). + if (!stream_manager_->Contains(it) || !it->stream || + it->stream->stream != stream) { + return; + } it->stream->read_pending = false; - if (!response) return DoFinish(std::move(lk), it); + if (!response) return DoFinish(std::move(lk), it, stream); if (response->has_metadata()) { metadata_ = std::move(*response->mutable_metadata()); } @@ -527,12 +548,19 @@ void ObjectDescriptorImpl::OnRead( } } lk.lock(); + if (!stream_manager_->Contains(it) || !it->stream) return; stream_manager_->CleanupDoneRanges(it); DoRead(std::move(lk), it); } void ObjectDescriptorImpl::DoFinish(std::unique_lock lk, - StreamIterator it) { + StreamIterator it, + std::shared_ptr const& stream) { + // Discard finish requests if the stream was already replaced or removed. + if (!stream_manager_->Contains(it) || !it->stream || + it->stream->stream != stream) { + return; + } it->stream->read_pending = false; // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the @@ -541,16 +569,35 @@ void ObjectDescriptorImpl::DoFinish(std::unique_lock lk, lk.unlock(); auto pending = current_stream->Finish(); if (!pending.valid()) return; - pending.then([w = WeakFromThis(), it](auto f) { - if (auto self = w.lock()) self->OnFinish(it, f.get()); + pending.then([w = WeakFromThis(), it, current_stream](auto f) { + if (auto self = w.lock()) self->OnFinish(it, current_stream, f.get()); }); } -void ObjectDescriptorImpl::OnFinish(StreamIterator it, Status const& status) { +void ObjectDescriptorImpl::OnFinish(StreamIterator it, + std::shared_ptr const& stream, + Status const& status) { + { + std::unique_lock lk(mu_); + // Discard callbacks if cancelled or from stale/removed streams. + if (cancelled_ || !stream_manager_->Contains(it) || !it->stream || + it->stream->stream != stream) { + return; + } + } auto proto_status = ExtractGrpcStatus(status); - if (IsResumable(it, status, proto_status)) return Resume(it, proto_status); + if (IsResumable(it, status, proto_status)) { + return Resume(it, proto_status); + } std::unique_lock lk(mu_); + // Re-verify stream identity under lock because IsResumable() releases and + // re-acquires the mutex while notifying range callbacks, during which time + // another thread or callback could have modified or replaced the stream. + if (cancelled_ || !stream_manager_->Contains(it) || !it->stream || + it->stream->stream != stream) { + return; + } stream_manager_->RemoveStreamAndNotifyRanges(it, status); // Since a stream died, we might want to ensure a replacement is queued. AssurePendingStreamQueued(lk); @@ -559,6 +606,12 @@ void ObjectDescriptorImpl::OnFinish(StreamIterator it, Status const& status) { void ObjectDescriptorImpl::Resume(StreamIterator it, google::rpc::Status const& proto_status) { std::unique_lock lk(mu_); + if (cancelled_ || !stream_manager_->Contains(it) || !it->stream) return; + // Set resuming flag to true to prevent any concurrent Flush() from writing + // to the dying stream while we establish a new one. + it->stream->resuming = true; + it->stream->next_request.Clear(); + auto current_stream = it->stream->stream; // This call needs to happen inside the lock, as it may modify // `read_object_spec_`. ApplyRedirectErrors(read_object_spec_, proto_status); @@ -570,31 +623,58 @@ void ObjectDescriptorImpl::Resume(StreamIterator it, *request.add_read_ranges() = *std::move(range); } lk.unlock(); - make_stream_(std::move(request)).then([w = WeakFromThis(), it](auto f) { - if (auto self = w.lock()) self->OnResume(it, f.get()); - }); + make_stream_(std::move(request)) + .then([w = WeakFromThis(), it, current_stream](auto f) { + if (auto self = w.lock()) self->OnResume(it, current_stream, f.get()); + }); } -void ObjectDescriptorImpl::OnResume(StreamIterator it, - StatusOr result) { - if (!result) return OnFinish(it, std::move(result).status()); +void ObjectDescriptorImpl::OnResume( + StreamIterator it, std::shared_ptr const& old_stream, + StatusOr result) { + { + std::unique_lock lk(mu_); + if (cancelled_) { + if (result && result->stream) result->stream->Cancel(); + return; + } + } + if (!result) return OnFinish(it, old_stream, std::move(result).status()); std::unique_lock lk(mu_); - if (cancelled_) return; + // Discard resume responses if cancelled or if the stream entry was removed or + // already replaced. + if (cancelled_ || !stream_manager_->Contains(it) || !it->stream || + it->stream->stream != old_stream) { + if (result->stream) result->stream->Cancel(); + return; + } + + // Preserve any Read() range requests that arrived concurrently while the + // reconnection was in flight. + auto queued_request = std::move(it->stream->next_request); + // Replace the old stream with the new stream and reset policy/state. it->stream = std::make_shared(std::move(result->stream), resume_policy_prototype_->clone()); + it->stream->resume_policy->OnStartSuccess(); it->stream->write_pending = false; it->stream->read_pending = false; + it->stream->resuming = false; + it->stream->next_request = std::move(queued_request); + auto new_stream = it->stream->stream; // TODO(#15105) - this should be done without release the lock. + // Flush any queued range requests onto the newly active stream. Flush(std::move(lk), it); - OnRead(it, std::move(result->first_response)); + // Process the first response received during stream establishment. + OnRead(it, new_stream, std::move(result->first_response)); } bool ObjectDescriptorImpl::IsResumable( StreamIterator it, Status const& status, google::rpc::Status const& proto_status) { std::unique_lock lk(mu_); + if (cancelled_ || !stream_manager_->Contains(it) || !it->stream) return false; for (auto const& any : proto_status.details()) { auto error = google::storage::v2::BidiReadObjectError{}; if (!any.UnpackTo(&error)) continue; @@ -614,10 +694,17 @@ bool ObjectDescriptorImpl::IsResumable( if (l != copy.end()) l->second->OnFinish(p.second); } lk.lock(); + if (cancelled_ || !stream_manager_->Contains(it) || !it->stream) { + return true; + } stream_manager_->CleanupDoneRanges(it); return true; } - return it->stream->resume_policy->OnFinish(status) == + auto effective_status = status; + if (status.code() == StatusCode::kCancelled) { + effective_status = Status(StatusCode::kUnavailable, status.message()); + } + return it->stream->resume_policy->OnFinish(effective_status) == storage::ResumePolicy::kContinue; } diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.h b/google/cloud/storage/internal/async/object_descriptor_impl.h index d0cfa98487c05..55a121694a75b 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.h +++ b/google/cloud/storage/internal/async/object_descriptor_impl.h @@ -53,6 +53,7 @@ struct ReadStream : public storage_internal::StreamBase { google::storage::v2::BidiReadObjectRequest next_request; bool write_pending = false; bool read_pending = false; + bool resuming = false; }; class ObjectDescriptorImpl @@ -101,15 +102,20 @@ class ObjectDescriptorImpl void AssurePendingStreamQueued(std::unique_lock const&); void Flush(std::unique_lock lk, StreamIterator it); - void OnWrite(StreamIterator it, bool ok); + void OnWrite(StreamIterator it, std::shared_ptr const& stream, + bool ok); void DoRead(std::unique_lock lk, StreamIterator it); void OnRead( - StreamIterator it, + StreamIterator it, std::shared_ptr const& stream, std::optional response); - void DoFinish(std::unique_lock lk, StreamIterator it); - void OnFinish(StreamIterator it, Status const& status); + void DoFinish(std::unique_lock lk, StreamIterator it, + std::shared_ptr const& stream); + void OnFinish(StreamIterator it, std::shared_ptr const& stream, + Status const& status); void Resume(StreamIterator it, google::rpc::Status const& proto_status); - void OnResume(StreamIterator it, StatusOr result); + void OnResume(StreamIterator it, + std::shared_ptr const& old_stream, + StatusOr result); bool IsResumable(StreamIterator it, Status const& status, google::rpc::Status const& proto_status); diff --git a/google/cloud/storage/internal/async/object_descriptor_impl_test.cc b/google/cloud/storage/internal/async/object_descriptor_impl_test.cc index d8ac5abd97024..7c98b3c4d2665 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl_test.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl_test.cc @@ -24,6 +24,7 @@ #include "google/cloud/storage/internal/async/default_options.h" #include "google/cloud/storage/options.h" #include "google/cloud/storage/testing/canonical_errors.h" +#include "google/cloud/storage/testing/mock_resume_policy.h" #include "google/cloud/storage/testing/mock_storage_stub.h" #include "google/cloud/testing_util/async_sequencer.h" #include "google/cloud/testing_util/is_proto_equal.h" @@ -41,6 +42,7 @@ namespace storage_internal { GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN namespace { +using ::google::cloud::storage::testing::MockResumePolicy; using ::google::cloud::storage::testing::canonical_errors::PermanentError; using ::google::cloud::storage::testing::canonical_errors::TransientError; using ::google::cloud::testing_util::AsyncSequencer; @@ -3144,6 +3146,495 @@ TEST(ObjectDescriptorImpl, DuplicateInitialRangesDeduplication) { next.first.set_value(true); } +/// @test Verify that when a stream fails and triggers resumption, concurrent +/// Read() requests queued while reconnecting are not dropped and are flushed to +/// the newly connected stream. +TEST(ObjectDescriptorImpl, ResumeRacesWithConcurrentRead) { + AsyncSequencer sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillRepeatedly([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[1]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto f) { + if (f.get()) return TransientError(); + return PermanentError(); + }); + }); + + auto stream2 = std::make_unique(); + EXPECT_CALL(*stream2, Read) + .WillOnce([&sequencer]() { + return sequencer.PushBack("Read[2]").then( + [](auto) { return std::optional{}; }); + }) + .WillRepeatedly( + []() { return make_ready_future(std::optional{}); }); + EXPECT_CALL(*stream2, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[2]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream2, Finish).WillRepeatedly([]() { + return make_ready_future(Status{}); + }); + + MockFactory factory; + EXPECT_CALL(factory, Call).WillRepeatedly([](Request const&) { + return make_ready_future(StatusOr(TransientError())); + }); + EXPECT_CALL( + factory, + Call(ResultOf([](Request const& r) { return !r.read_ranges().empty(); }, + true))) + .WillOnce([&stream2](Request const&) { + auto response = Response{}; + return make_ready_future(make_status_or( + OpenStreamResult{std::make_shared(std::move(stream2)), + std::move(response)})); + }); + + auto tested = MakeTested(storage::LimitedErrorCountResumePolicy(3)(), + factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1))); + + tested->Start(Response{}); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_EQ(read1.second, "Read[1]"); + + auto s1 = tested->Read({1000, 100}); + ASSERT_THAT(s1, NotNull()); + auto write1 = sequencer.PopFrontWithName(); + EXPECT_EQ(write1.second, "Write[1]"); + write1.first.set_value(true); + + // Trigger stream 1 failure and Resume + read1.first.set_value(false); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Finish[1]"); + next.first.set_value(true); + + // Concurrently issue Read for range 2 while Resume is connecting stream 2 + auto s2 = tested->Read({2000, 100}); + ASSERT_THAT(s2, NotNull()); + + // Stream 2 should be read and receive the flushed Write[2] with range 2 + auto read2 = sequencer.PopFrontWithName(); + EXPECT_EQ(read2.second, "Read[2]"); + + auto write2 = sequencer.PopFrontWithName(); + EXPECT_EQ(write2.second, "Write[2]"); + write2.first.set_value(true); + + tested.reset(); + read2.first.set_value(false); +} + +/// @test Verify that callbacks from discarded/stale stream instances are +/// ignored and do not affect the newly active stream. +TEST(ObjectDescriptorImpl, IgnoresCallbacksFromStaleStreams) { + AsyncSequencer sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillRepeatedly([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[1]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto f) { + if (f.get()) return TransientError(); + return PermanentError(); + }); + }); + + auto stream2 = std::make_unique(); + EXPECT_CALL(*stream2, Read) + .WillOnce([&sequencer]() { + return sequencer.PushBack("Read[2]").then( + [](auto) { return std::optional{}; }); + }) + .WillRepeatedly( + []() { return make_ready_future(std::optional{}); }); + EXPECT_CALL(*stream2, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[2]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream2, Finish).WillRepeatedly([]() { + return make_ready_future(Status{}); + }); + + MockFactory factory; + EXPECT_CALL(factory, Call).WillRepeatedly([](Request const&) { + return make_ready_future(StatusOr(TransientError())); + }); + EXPECT_CALL( + factory, + Call(ResultOf([](Request const& r) { return !r.read_ranges().empty(); }, + true))) + .WillOnce([&stream2](Request const&) { + auto response = Response{}; + return make_ready_future(make_status_or( + OpenStreamResult{std::make_shared(std::move(stream2)), + std::move(response)})); + }); + + auto tested = MakeTested(storage::LimitedErrorCountResumePolicy(3)(), + factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1))); + + tested->Start(Response{}); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_EQ(read1.second, "Read[1]"); + + auto s1 = tested->Read({1000, 100}); + ASSERT_THAT(s1, NotNull()); + auto write1 = sequencer.PopFrontWithName(); + EXPECT_EQ(write1.second, "Write[1]"); + + // Stream 1 fails on read and resumes stream 2 + read1.first.set_value(false); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Finish[1]"); + next.first.set_value(true); + + // Stream 2 is now active + auto read2 = sequencer.PopFrontWithName(); + EXPECT_EQ(read2.second, "Read[2]"); + + // A stale Write[1] callback from stream 1 arrives late with false (error). + // It should NOT call Finish[2] on stream 2! + write1.first.set_value(false); + + EXPECT_TRUE(sequencer.empty()); + + tested.reset(); + read2.first.set_value(false); +} + +/// @test Verify that when a stream is cancelled due to a stall watchdog timeout +/// (StatusCode::kCancelled), it is treated as resumable by the resume policy +/// and automatically resumed. +TEST(ObjectDescriptorImpl, ResumeOnStallTimeout) { + AsyncSequencer sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillRepeatedly([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[1]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto f) { + if (f.get()) return Status(StatusCode::kCancelled, "Stream stalled"); + return PermanentError(); + }); + }); + + auto stream2 = std::make_unique(); + EXPECT_CALL(*stream2, Read) + .WillOnce([&sequencer]() { + return sequencer.PushBack("Read[2]").then( + [](auto) { return std::optional{}; }); + }) + .WillRepeatedly( + []() { return make_ready_future(std::optional{}); }); + EXPECT_CALL(*stream2, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[2]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream2, Finish).WillRepeatedly([]() { + return make_ready_future(Status{}); + }); + + MockFactory factory; + EXPECT_CALL(factory, Call).WillRepeatedly([](Request const&) { + return make_ready_future(StatusOr(TransientError())); + }); + EXPECT_CALL( + factory, + Call(ResultOf([](Request const& r) { return !r.read_ranges().empty(); }, + true))) + .WillOnce([&stream2](Request const&) { + auto response = Response{}; + return make_ready_future(make_status_or( + OpenStreamResult{std::make_shared(std::move(stream2)), + std::move(response)})); + }); + + auto tested = MakeTested(storage::LimitedErrorCountResumePolicy(3)(), + factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1))); + + tested->Start(Response{}); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_EQ(read1.second, "Read[1]"); + + auto s1 = tested->Read({1000, 100}); + ASSERT_THAT(s1, NotNull()); + auto write1 = sequencer.PopFrontWithName(); + EXPECT_EQ(write1.second, "Write[1]"); + write1.first.set_value(true); + + // Watchdog cancels stream 1 on read stall + read1.first.set_value(false); + + auto next = sequencer.PopFrontWithName(); + EXPECT_EQ(next.second, "Finish[1]"); + next.first.set_value(true); + + // If resumed, stream 2 is connected and Read[2] is invoked + auto read2 = sequencer.PopFrontWithName(); + EXPECT_EQ(read2.second, "Read[2]"); + + tested.reset(); + read2.first.set_value(false); +} + +/// @test Verify that when a stream starts successfully, the resume policy is +/// notified via OnStartSuccess(). +TEST(ObjectDescriptorImpl, NotifiesResumePolicyOnStartSuccess) { + auto mock_policy = std::make_unique(); + EXPECT_CALL(*mock_policy, clone).WillOnce([]() { + auto p = std::make_unique(); + EXPECT_CALL(*p, OnStartSuccess).Times(::testing::AtLeast(1)); + return p; + }); + + MockFactory factory; + EXPECT_CALL(factory, Call).WillRepeatedly([](Request const&) { + return make_ready_future(StatusOr(TransientError())); + }); + auto stream = std::make_unique(); + EXPECT_CALL(*stream, Read).WillRepeatedly([]() { + return make_ready_future(std::optional{}); + }); + EXPECT_CALL(*stream, Write) + .WillRepeatedly([](Request const&, grpc::WriteOptions) { + return make_ready_future(true); + }); + EXPECT_CALL(*stream, Finish).WillRepeatedly([]() { + return make_ready_future(Status{}); + }); + + auto tested = MakeTested(std::move(mock_policy), factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream))); + + tested->Start(Response{}); +} + +/// @test Verify that when the user explicitly cancels ObjectDescriptorImpl, +/// streams are cancelled and no resume attempts are made. +TEST(ObjectDescriptorImpl, UserCancelStopsStreamWithoutResume) { + AsyncSequencer sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Cancel).Times(::testing::AtLeast(1)); + EXPECT_CALL(*stream1, Write) + .Times(AtMost(1)) + .WillRepeatedly([](Request const&, grpc::WriteOptions) { + return make_ready_future(true); + }); + EXPECT_CALL(*stream1, Read).WillOnce([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto) { + return Status(StatusCode::kCancelled, "Cancelled by user"); + }); + }); + + MockFactory factory; + // Factory should NEVER be called on user cancellation. + EXPECT_CALL(factory, Call).Times(0); + + Options options; + options.set(false); + auto tested = std::make_shared( + storage::LimitedErrorCountResumePolicy(2)(), factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1)), options); + + Response start_resp; + tested->Start(std::move(start_resp)); + + auto r1 = tested->Read({100, 50}); + ASSERT_THAT(r1, NotNull()); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_EQ(read1.second, "Read[1]"); + + // User explicitly cancels. + tested->Cancel(); + EXPECT_FALSE(tested->IsOpen()); + + read1.first.set_value(true); + + auto finish1 = sequencer.PopFrontWithName(); + EXPECT_EQ(finish1.second, "Finish[1]"); + finish1.first.set_value(true); + + // No resume factory calls happen. + EXPECT_TRUE(sequencer.empty()); + EXPECT_FALSE(tested->IsOpen()); +} + +/// @test Verify that range requests arriving during multiple consecutive +/// transient reconnection retries are preserved across all attempts and +/// flushed to the new stream once reconnected. +TEST(ObjectDescriptorImpl, ConcurrentReadPreservedAcrossMultipleResumeRetries) { + AsyncSequencer sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillRepeatedly([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[1]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto f) { + if (f.get()) return TransientError(); + return PermanentError(); + }); + }); + + auto stream2 = std::make_unique(); + EXPECT_CALL(*stream2, Read) + .WillOnce([&sequencer]() { + return sequencer.PushBack("Read[2]").then( + [](auto) { return std::optional{}; }); + }) + .WillRepeatedly( + []() { return make_ready_future(std::optional{}); }); + EXPECT_CALL(*stream2, Write) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[2]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream2, Finish).WillRepeatedly([]() { + return make_ready_future(Status{}); + }); + + MockFactory factory; + int retry_count = 0; + EXPECT_CALL(factory, Call).WillRepeatedly([&](Request const& r) { + ++retry_count; + if (retry_count == 1) { + return sequencer.PushBack("Factory[1]").then([](auto f) { + if (f.get()) { + return StatusOr(TransientError()); + } + return StatusOr(PermanentError()); + }); + } + // Second attempt must contain both range 1 and range 2 + EXPECT_EQ(r.read_ranges_size(), 2); + return sequencer.PushBack("Factory[2]").then([&stream2](auto f) { + if (f.get()) { + auto response = Response{}; + return make_status_or( + OpenStreamResult{std::make_shared(std::move(stream2)), + std::move(response)}); + } + return StatusOr(PermanentError()); + }); + }); + + auto tested = std::make_shared( + storage::LimitedErrorCountResumePolicy(3)(), factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1)), + Options{}.set(false)); + + tested->Start(Response{}); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_EQ(read1.second, "Read[1]"); + + auto s1 = tested->Read({1000, 100}); + ASSERT_THAT(s1, NotNull()); + auto write1 = sequencer.PopFrontWithName(); + EXPECT_EQ(write1.second, "Write[1]"); + write1.first.set_value(true); + + // Trigger stream 1 failure + read1.first.set_value(false); + + auto finish1 = sequencer.PopFrontWithName(); + EXPECT_EQ(finish1.second, "Finish[1]"); + finish1.first.set_value(true); + + // Factory attempt 1 is invoked + auto factory1 = sequencer.PopFrontWithName(); + EXPECT_EQ(factory1.second, "Factory[1]"); + + // While reconnection attempt 1 is in-flight, add range 2 concurrently + auto s2 = tested->Read({2000, 100}); + ASSERT_THAT(s2, NotNull()); + + // Factory attempt 1 fails with TransientError + factory1.first.set_value(true); + + // Factory attempt 2 is invoked with both range 1 and range 2 + auto factory2 = sequencer.PopFrontWithName(); + EXPECT_EQ(factory2.second, "Factory[2]"); + + // While factory attempt 2 is in-flight, add range 3 concurrently + auto s3 = tested->Read({3000, 100}); + ASSERT_THAT(s3, NotNull()); + + // Factory attempt 2 succeeds with stream2 + factory2.first.set_value(true); + + // Range 3 (queued during factory attempt 2) is flushed to stream 2 first + auto write2 = sequencer.PopFrontWithName(); + EXPECT_EQ(write2.second, "Write[2]"); + write2.first.set_value(true); + + // Stream 2 is read + auto read2 = sequencer.PopFrontWithName(); + EXPECT_EQ(read2.second, "Read[2]"); + + EXPECT_EQ(retry_count, 2); + + tested->Cancel(); + read2.first.set_value(false); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal From c7c7788eaadf6b1071a1fb123150e5ea5331881f Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Fri, 11 Sep 2026 09:14:18 +0000 Subject: [PATCH 2/6] Address feedback from Gemini code assistant --- .../internal/async/multi_stream_manager.h | 6 +- .../async/multi_stream_manager_test.cc | 13 +- .../internal/async/object_descriptor_impl.cc | 256 ++++++++++-------- .../internal/async/object_descriptor_impl.h | 29 +- 4 files changed, 176 insertions(+), 128 deletions(-) diff --git a/google/cloud/storage/internal/async/multi_stream_manager.h b/google/cloud/storage/internal/async/multi_stream_manager.h index aead9d8af0598..2204de1b0fdd0 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager.h +++ b/google/cloud/storage/internal/async/multi_stream_manager.h @@ -159,11 +159,11 @@ class MultiStreamManager { return false; } - bool Contains(StreamIterator target) const { + StreamIterator Find(std::shared_ptr const& target) { for (auto it = streams_.begin(); it != streams_.end(); ++it) { - if (it == target) return true; + if (it->stream == target) return it; } - return false; + return streams_.end(); } bool Empty() const { return streams_.empty(); } diff --git a/google/cloud/storage/internal/async/multi_stream_manager_test.cc b/google/cloud/storage/internal/async/multi_stream_manager_test.cc index f890b1a1e7d65..4a32a3380bf97 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager_test.cc +++ b/google/cloud/storage/internal/async/multi_stream_manager_test.cc @@ -233,19 +233,20 @@ TEST(MultiStreamManagerTest, EmptyAndSizeTransitions) { EXPECT_EQ(mgr.Size(), 1U); } -TEST(MultiStreamManagerTest, ContainsTracksStreamMembership) { +TEST(MultiStreamManagerTest, FindTracksStreamMembership) { auto mgr = MultiStreamManagerTest::MakeManager(); auto it1 = mgr.GetFirstStream(); - EXPECT_TRUE(mgr.Contains(it1)); + auto s1 = it1->stream; + EXPECT_NE(mgr.Find(s1), mgr.End()); auto s2 = std::make_shared(); auto it2 = mgr.AddStream(s2); - EXPECT_TRUE(mgr.Contains(it1)); - EXPECT_TRUE(mgr.Contains(it2)); + EXPECT_NE(mgr.Find(s1), mgr.End()); + EXPECT_NE(mgr.Find(s2), mgr.End()); mgr.RemoveStreamAndNotifyRanges(it1, Status()); - EXPECT_FALSE(mgr.Contains(it1)); - EXPECT_TRUE(mgr.Contains(it2)); + EXPECT_EQ(mgr.Find(s1), mgr.End()); + EXPECT_NE(mgr.Find(s2), mgr.End()); } GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.cc b/google/cloud/storage/internal/async/object_descriptor_impl.cc index dcb7c5093724a..193f6d32454a3 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl.cc @@ -126,9 +126,10 @@ void ObjectDescriptorImpl::Start( std::unique_lock lk(mu_); auto it = stream_manager_->GetFirstStream(); if (it == stream_manager_->End()) return; - auto current_stream = it->stream->stream; + auto read_stream = it->stream; + auto current_stream = read_stream->stream; lk.unlock(); - OnRead(it, current_stream, std::move(first_response)); + OnRead(read_stream, current_stream, std::move(first_response)); // Acquire lock and queue the background stream if multi-stream optimization // is enabled. if (options_.get()) { @@ -211,14 +212,15 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { self->resume_policy_prototype_->clone()); read_stream->resume_policy->OnStartSuccess(); - auto new_it = self->stream_manager_->AddStream(std::move(read_stream)); + self->stream_manager_->AddStream(read_stream); // Now that we consumed pending_stream_, queue the next one immediately. self->AssurePendingStreamQueued(lk); - auto new_stream = new_it->stream->stream; + auto new_stream = read_stream->stream; lk.unlock(); - self->OnRead(new_it, new_stream, std::move(stream_result->first_response)); + self->OnRead(read_stream, new_stream, + std::move(stream_result->first_response)); }); } @@ -315,13 +317,14 @@ std::unique_ptr ObjectDescriptorImpl::Read( } auto it = stream_manager_->GetLeastBusyStream(); + auto read_stream = it->stream; auto const id = ++read_id_generator_; it->active_ranges.emplace(id, range); - auto& read_range = *it->stream->next_request.add_read_ranges(); + auto& read_range = *read_stream->next_request.add_read_ranges(); read_range.set_read_id(id); read_range.set_read_offset(p.start); read_range.set_read_length(p.length); - Flush(std::move(lk), it); + Flush(std::move(lk), read_stream); if (!internal::TracingEnabled(options_)) { return std::unique_ptr( @@ -405,70 +408,80 @@ ObjectDescriptorImpl::CreateHashValidator(bool is_full_read) const { return hash_validator; } -void ObjectDescriptorImpl::Flush(std::unique_lock lk, - StreamIterator it) { - if (it->stream->resuming || it->stream->write_pending || - it->stream->next_request.read_ranges().empty()) { +void ObjectDescriptorImpl::Flush( + std::unique_lock lk, + std::shared_ptr const& read_stream) { + if (!read_stream || read_stream->resuming || read_stream->write_pending || + read_stream->next_request.read_ranges().empty()) { return; } - it->stream->write_pending = true; + read_stream->write_pending = true; google::storage::v2::BidiReadObjectRequest request; - request.Swap(&it->stream->next_request); + request.Swap(&read_stream->next_request); // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the // end of the block. - auto current_stream = it->stream->stream; + auto current_stream = read_stream->stream; lk.unlock(); current_stream->Write(std::move(request)) - .then([w = WeakFromThis(), it, current_stream](auto f) { - if (auto self = w.lock()) self->OnWrite(it, current_stream, f.get()); + .then([w = WeakFromThis(), read_stream, current_stream](auto f) { + if (auto self = w.lock()) { + self->OnWrite(read_stream, current_stream, f.get()); + } }); } -void ObjectDescriptorImpl::OnWrite(StreamIterator it, - std::shared_ptr const& stream, - bool ok) { +void ObjectDescriptorImpl::OnWrite( + std::shared_ptr const& read_stream, + std::shared_ptr const& stream, bool ok) { std::unique_lock lk(mu_); // Discard callbacks from stale or removed streams (e.g. if the stream was // replaced during reconnection or removed after an error). - if (!stream_manager_->Contains(it) || !it->stream || + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream || it->stream->stream != stream) { return; } - if (!ok) return DoFinish(std::move(lk), it, stream); + if (!ok) return DoFinish(std::move(lk), read_stream, stream); it->stream->write_pending = false; - Flush(std::move(lk), it); + Flush(std::move(lk), read_stream); } -void ObjectDescriptorImpl::DoRead(std::unique_lock lk, - StreamIterator it) { - if (it->stream->read_pending) return; - it->stream->read_pending = true; +void ObjectDescriptorImpl::DoRead( + std::unique_lock lk, + std::shared_ptr const& read_stream) { + if (!read_stream || read_stream->read_pending) return; + read_stream->read_pending = true; // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the // end of the block. - auto current_stream = it->stream->stream; + auto current_stream = read_stream->stream; lk.unlock(); - current_stream->Read().then([w = WeakFromThis(), it, current_stream](auto f) { - if (auto self = w.lock()) self->OnRead(it, current_stream, f.get()); - }); + current_stream->Read().then( + [w = WeakFromThis(), read_stream, current_stream](auto f) { + if (auto self = w.lock()) { + self->OnRead(read_stream, current_stream, f.get()); + } + }); } void ObjectDescriptorImpl::OnRead( - StreamIterator it, std::shared_ptr const& stream, + std::shared_ptr const& read_stream, + std::shared_ptr const& stream, std::optional response) { std::unique_lock lk(mu_); // Discard callbacks from stale or removed streams (e.g. if the stream was // replaced during reconnection or removed after an error). - if (!stream_manager_->Contains(it) || !it->stream || + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream || it->stream->stream != stream) { return; } it->stream->read_pending = false; - if (!response) return DoFinish(std::move(lk), it, stream); + if (!response) return DoFinish(std::move(lk), read_stream, stream); if (response->has_metadata()) { metadata_ = std::move(*response->mutable_metadata()); } @@ -486,35 +499,6 @@ void ObjectDescriptorImpl::OnRead( // Release the lock while notifying the ranges. The notifications may trigger // application code, and that code may callback on this class. lk.unlock(); - auto apply_pacing_and_check_eviction = [this](std::int64_t id, - std::size_t chunk_size, - StreamIterator it) { - auto unclaimed_it = unclaimed_ranges_.find(id); - if (unclaimed_it == unclaimed_ranges_.end()) return false; - - if (total_prewarmed_bytes_buffered_ + chunk_size > - max_prewarmed_buffer_size_) { - // Evict the range if it exceeds the pacing limit. - total_prewarmed_bytes_buffered_ -= unclaimed_it->second.bytes_buffered; - // Cap tombstone set size to prevent unbounded memory growth in long-lived - // descriptors where pre-warmed ranges are evicted but never requested. - if (evicted_ranges_.size() < 1000) { - evicted_ranges_.insert(unclaimed_it->second.cache_it->first); - } - prewarmed_ranges_.erase(unclaimed_it->second.cache_it); - unclaimed_ranges_.erase(unclaimed_it); - - // Erasing from active_ranges ensures we ignore any subsequent GCS chunks - // for this range. - it->active_ranges.erase(id); - return true; - } - - // Track buffered data size for pacing. - unclaimed_it->second.bytes_buffered += chunk_size; - total_prewarmed_bytes_buffered_ += chunk_size; - return false; - }; for (auto& range_data : *response->mutable_object_data_ranges()) { auto id = range_data.read_range().read_id(); @@ -529,9 +513,11 @@ void ObjectDescriptorImpl::OnRead( // Verify the range is still active under the lock. Because `OnRead` // processes chunks in batches, an earlier chunk in the same batch could // breach the pacing limit and evict a subsequent chunk's range. - bool active = it->active_ranges.count(id) != 0; + auto it_curr = stream_manager_->Find(read_stream); + bool active = (it_curr != stream_manager_->End()) && + (it_curr->active_ranges.count(id) != 0); if (active) { - evict = apply_pacing_and_check_eviction(id, chunk_size, it); + evict = ApplyPacingAndCheckEviction(id, chunk_size, it_curr); } lk.unlock(); if (active) { @@ -548,16 +534,49 @@ void ObjectDescriptorImpl::OnRead( } } lk.lock(); - if (!stream_manager_->Contains(it) || !it->stream) return; - stream_manager_->CleanupDoneRanges(it); - DoRead(std::move(lk), it); + auto it_final = stream_manager_->Find(read_stream); + if (it_final == stream_manager_->End() || !it_final->stream) return; + stream_manager_->CleanupDoneRanges(it_final); + DoRead(std::move(lk), read_stream); } -void ObjectDescriptorImpl::DoFinish(std::unique_lock lk, - StreamIterator it, - std::shared_ptr const& stream) { +bool ObjectDescriptorImpl::ApplyPacingAndCheckEviction(std::int64_t id, + std::size_t chunk_size, + StreamIterator it) { + auto unclaimed_it = unclaimed_ranges_.find(id); + if (unclaimed_it == unclaimed_ranges_.end()) return false; + + if (total_prewarmed_bytes_buffered_ + chunk_size > + max_prewarmed_buffer_size_) { + // Evict the range if it exceeds the pacing limit. + total_prewarmed_bytes_buffered_ -= unclaimed_it->second.bytes_buffered; + // Cap tombstone set size to prevent unbounded memory growth in long-lived + // descriptors where pre-warmed ranges are evicted but never requested. + if (evicted_ranges_.size() < 1000) { + evicted_ranges_.insert(unclaimed_it->second.cache_it->first); + } + prewarmed_ranges_.erase(unclaimed_it->second.cache_it); + unclaimed_ranges_.erase(unclaimed_it); + + // Erasing from active_ranges ensures we ignore any subsequent GCS chunks + // for this range. + it->active_ranges.erase(id); + return true; + } + + // Track buffered data size for pacing. + unclaimed_it->second.bytes_buffered += chunk_size; + total_prewarmed_bytes_buffered_ += chunk_size; + return false; +} + +void ObjectDescriptorImpl::DoFinish( + std::unique_lock lk, + std::shared_ptr const& read_stream, + std::shared_ptr const& stream) { // Discard finish requests if the stream was already replaced or removed. - if (!stream_manager_->Contains(it) || !it->stream || + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream || it->stream->stream != stream) { return; } @@ -569,32 +588,38 @@ void ObjectDescriptorImpl::DoFinish(std::unique_lock lk, lk.unlock(); auto pending = current_stream->Finish(); if (!pending.valid()) return; - pending.then([w = WeakFromThis(), it, current_stream](auto f) { - if (auto self = w.lock()) self->OnFinish(it, current_stream, f.get()); + pending.then([w = WeakFromThis(), read_stream, current_stream](auto f) { + if (auto self = w.lock()) { + self->OnFinish(read_stream, current_stream, f.get()); + } }); } -void ObjectDescriptorImpl::OnFinish(StreamIterator it, - std::shared_ptr const& stream, - Status const& status) { +void ObjectDescriptorImpl::OnFinish( + std::shared_ptr const& read_stream, + std::shared_ptr const& stream, Status const& status) { { std::unique_lock lk(mu_); // Discard callbacks if cancelled or from stale/removed streams. - if (cancelled_ || !stream_manager_->Contains(it) || !it->stream || + if (cancelled_) return; + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream || it->stream->stream != stream) { return; } } auto proto_status = ExtractGrpcStatus(status); - if (IsResumable(it, status, proto_status)) { - return Resume(it, proto_status); + if (IsResumable(read_stream, status, proto_status)) { + return Resume(read_stream, proto_status); } std::unique_lock lk(mu_); // Re-verify stream identity under lock because IsResumable() releases and // re-acquires the mutex while notifying range callbacks, during which time // another thread or callback could have modified or replaced the stream. - if (cancelled_ || !stream_manager_->Contains(it) || !it->stream || + if (cancelled_) return; + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream || it->stream->stream != stream) { return; } @@ -603,10 +628,13 @@ void ObjectDescriptorImpl::OnFinish(StreamIterator it, AssurePendingStreamQueued(lk); } -void ObjectDescriptorImpl::Resume(StreamIterator it, - google::rpc::Status const& proto_status) { +void ObjectDescriptorImpl::Resume( + std::shared_ptr const& read_stream, + google::rpc::Status const& proto_status) { std::unique_lock lk(mu_); - if (cancelled_ || !stream_manager_->Contains(it) || !it->stream) return; + if (cancelled_) return; + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream) return; // Set resuming flag to true to prevent any concurrent Flush() from writing // to the dying stream while we establish a new one. it->stream->resuming = true; @@ -617,20 +645,23 @@ void ObjectDescriptorImpl::Resume(StreamIterator it, ApplyRedirectErrors(read_object_spec_, proto_status); auto request = google::storage::v2::BidiReadObjectRequest{}; *request.mutable_read_object_spec() = read_object_spec_; - for (auto const& kv : it->active_ranges) { - auto range = kv.second->RangeForResume(kv.first); + for (auto const& [read_id, range_ptr] : it->active_ranges) { + auto range = range_ptr->RangeForResume(read_id); if (!range) continue; *request.add_read_ranges() = *std::move(range); } lk.unlock(); make_stream_(std::move(request)) - .then([w = WeakFromThis(), it, current_stream](auto f) { - if (auto self = w.lock()) self->OnResume(it, current_stream, f.get()); + .then([w = WeakFromThis(), read_stream, current_stream](auto f) { + if (auto self = w.lock()) { + self->OnResume(read_stream, current_stream, f.get()); + } }); } void ObjectDescriptorImpl::OnResume( - StreamIterator it, std::shared_ptr const& old_stream, + std::shared_ptr const& old_read_stream, + std::shared_ptr const& old_stream, StatusOr result) { { std::unique_lock lk(mu_); @@ -639,11 +670,14 @@ void ObjectDescriptorImpl::OnResume( return; } } - if (!result) return OnFinish(it, old_stream, std::move(result).status()); + if (!result) { + return OnFinish(old_read_stream, old_stream, std::move(result).status()); + } std::unique_lock lk(mu_); // Discard resume responses if cancelled or if the stream entry was removed or // already replaced. - if (cancelled_ || !stream_manager_->Contains(it) || !it->stream || + auto it = stream_manager_->Find(old_read_stream); + if (cancelled_ || it == stream_manager_->End() || !it->stream || it->stream->stream != old_stream) { if (result->stream) result->stream->Cancel(); return; @@ -654,27 +688,31 @@ void ObjectDescriptorImpl::OnResume( auto queued_request = std::move(it->stream->next_request); // Replace the old stream with the new stream and reset policy/state. - it->stream = std::make_shared(std::move(result->stream), - resume_policy_prototype_->clone()); - it->stream->resume_policy->OnStartSuccess(); - it->stream->write_pending = false; - it->stream->read_pending = false; - it->stream->resuming = false; - it->stream->next_request = std::move(queued_request); + auto new_read_stream = std::make_shared( + std::move(result->stream), resume_policy_prototype_->clone()); + new_read_stream->resume_policy->OnStartSuccess(); + new_read_stream->write_pending = false; + new_read_stream->read_pending = false; + new_read_stream->resuming = false; + new_read_stream->next_request = std::move(queued_request); + + it->stream = new_read_stream; + auto new_stream = new_read_stream->stream; - auto new_stream = it->stream->stream; // TODO(#15105) - this should be done without release the lock. // Flush any queued range requests onto the newly active stream. - Flush(std::move(lk), it); + Flush(std::move(lk), new_read_stream); // Process the first response received during stream establishment. - OnRead(it, new_stream, std::move(result->first_response)); + OnRead(new_read_stream, new_stream, std::move(result->first_response)); } bool ObjectDescriptorImpl::IsResumable( - StreamIterator it, Status const& status, + std::shared_ptr const& read_stream, Status const& status, google::rpc::Status const& proto_status) { std::unique_lock lk(mu_); - if (cancelled_ || !stream_manager_->Contains(it) || !it->stream) return false; + if (cancelled_) return false; + auto it = stream_manager_->Find(read_stream); + if (it == stream_manager_->End() || !it->stream) return false; for (auto const& any : proto_status.details()) { auto error = google::storage::v2::BidiReadObjectError{}; if (!any.UnpackTo(&error)) continue; @@ -689,15 +727,15 @@ bool ObjectDescriptorImpl::IsResumable( auto copy = it->active_ranges; lk.unlock(); - for (auto const& p : notify) { - auto l = copy.find(p.first); - if (l != copy.end()) l->second->OnFinish(p.second); + for (auto const& [read_id, range_status] : notify) { + auto l = copy.find(read_id); + if (l != copy.end()) l->second->OnFinish(range_status); } lk.lock(); - if (cancelled_ || !stream_manager_->Contains(it) || !it->stream) { - return true; - } - stream_manager_->CleanupDoneRanges(it); + if (cancelled_) return true; + auto it_curr = stream_manager_->Find(read_stream); + if (it_curr == stream_manager_->End() || !it_curr->stream) return true; + stream_manager_->CleanupDoneRanges(it_curr); return true; } auto effective_status = status; diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.h b/google/cloud/storage/internal/async/object_descriptor_impl.h index 55a121694a75b..238764d8ecdc7 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.h +++ b/google/cloud/storage/internal/async/object_descriptor_impl.h @@ -101,23 +101,32 @@ class ObjectDescriptorImpl // invoked while holding `mu_`. void AssurePendingStreamQueued(std::unique_lock const&); - void Flush(std::unique_lock lk, StreamIterator it); - void OnWrite(StreamIterator it, std::shared_ptr const& stream, - bool ok); - void DoRead(std::unique_lock lk, StreamIterator it); + void Flush(std::unique_lock lk, + std::shared_ptr const& read_stream); + void OnWrite(std::shared_ptr const& read_stream, + std::shared_ptr const& stream, bool ok); + void DoRead(std::unique_lock lk, + std::shared_ptr const& read_stream); void OnRead( - StreamIterator it, std::shared_ptr const& stream, + std::shared_ptr const& read_stream, + std::shared_ptr const& stream, std::optional response); - void DoFinish(std::unique_lock lk, StreamIterator it, + void DoFinish(std::unique_lock lk, + std::shared_ptr const& read_stream, std::shared_ptr const& stream); - void OnFinish(StreamIterator it, std::shared_ptr const& stream, + void OnFinish(std::shared_ptr const& read_stream, + std::shared_ptr const& stream, Status const& status); - void Resume(StreamIterator it, google::rpc::Status const& proto_status); - void OnResume(StreamIterator it, + void Resume(std::shared_ptr const& read_stream, + google::rpc::Status const& proto_status); + void OnResume(std::shared_ptr const& old_read_stream, std::shared_ptr const& old_stream, StatusOr result); - bool IsResumable(StreamIterator it, Status const& status, + bool IsResumable(std::shared_ptr const& read_stream, + Status const& status, google::rpc::Status const& proto_status); + bool ApplyPacingAndCheckEviction(std::int64_t id, std::size_t chunk_size, + StreamIterator it); std::shared_ptr CreateHashFunction( bool is_full_read) const; From 620951f0cd223a35a975565078af41df5376e431 Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Fri, 11 Sep 2026 09:30:44 +0000 Subject: [PATCH 3/6] Fix failing checks --- .../async/multi_stream_manager_test.cc | 2 +- .../internal/async/object_descriptor_impl.cc | 59 ++++++++++--------- 2 files changed, 32 insertions(+), 29 deletions(-) diff --git a/google/cloud/storage/internal/async/multi_stream_manager_test.cc b/google/cloud/storage/internal/async/multi_stream_manager_test.cc index 4a32a3380bf97..755cd23d05c18 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager_test.cc +++ b/google/cloud/storage/internal/async/multi_stream_manager_test.cc @@ -240,7 +240,7 @@ TEST(MultiStreamManagerTest, FindTracksStreamMembership) { EXPECT_NE(mgr.Find(s1), mgr.End()); auto s2 = std::make_shared(); - auto it2 = mgr.AddStream(s2); + mgr.AddStream(s2); EXPECT_NE(mgr.Find(s1), mgr.End()); EXPECT_NE(mgr.Find(s2), mgr.End()); diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.cc b/google/cloud/storage/internal/async/object_descriptor_impl.cc index 193f6d32454a3..67672d5d15a55 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl.cc @@ -126,8 +126,8 @@ void ObjectDescriptorImpl::Start( std::unique_lock lk(mu_); auto it = stream_manager_->GetFirstStream(); if (it == stream_manager_->End()) return; - auto read_stream = it->stream; - auto current_stream = read_stream->stream; + std::shared_ptr read_stream = it->stream; + std::shared_ptr current_stream = read_stream->stream; lk.unlock(); OnRead(read_stream, current_stream, std::move(first_response)); // Acquire lock and queue the background stream if multi-stream optimization @@ -197,7 +197,7 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { auto self = w.lock(); if (!self) return; - auto stream_result = f.get(); + StatusOr stream_result = f.get(); if (!stream_result) { // Stream creation failed. // The next call to AssurePendingStreamQueued will retry creation. @@ -207,7 +207,7 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { std::unique_lock lk(self->mu_); if (self->cancelled_) return; - auto read_stream = + std::shared_ptr read_stream = std::make_shared(std::move(stream_result->stream), self->resume_policy_prototype_->clone()); read_stream->resume_policy->OnStartSuccess(); @@ -217,7 +217,7 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { // Now that we consumed pending_stream_, queue the next one immediately. self->AssurePendingStreamQueued(lk); - auto new_stream = read_stream->stream; + std::shared_ptr new_stream = read_stream->stream; lk.unlock(); self->OnRead(read_stream, new_stream, std::move(stream_result->first_response)); @@ -317,10 +317,11 @@ std::unique_ptr ObjectDescriptorImpl::Read( } auto it = stream_manager_->GetLeastBusyStream(); - auto read_stream = it->stream; - auto const id = ++read_id_generator_; + std::shared_ptr read_stream = it->stream; + std::int64_t const id = ++read_id_generator_; it->active_ranges.emplace(id, range); - auto& read_range = *read_stream->next_request.add_read_ranges(); + google::storage::v2::ReadRange& read_range = + *read_stream->next_request.add_read_ranges(); read_range.set_read_id(id); read_range.set_read_offset(p.start); read_range.set_read_length(p.length); @@ -422,7 +423,7 @@ void ObjectDescriptorImpl::Flush( // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the // end of the block. - auto current_stream = read_stream->stream; + std::shared_ptr current_stream = read_stream->stream; lk.unlock(); current_stream->Write(std::move(request)) .then([w = WeakFromThis(), read_stream, current_stream](auto f) { @@ -457,7 +458,7 @@ void ObjectDescriptorImpl::DoRead( // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the // end of the block. - auto current_stream = read_stream->stream; + std::shared_ptr current_stream = read_stream->stream; lk.unlock(); current_stream->Read().then( [w = WeakFromThis(), read_stream, current_stream](auto f) { @@ -501,12 +502,12 @@ void ObjectDescriptorImpl::OnRead( lk.unlock(); for (auto& range_data : *response->mutable_object_data_ranges()) { - auto id = range_data.read_range().read_id(); + std::int64_t id = range_data.read_range().read_id(); auto const l = copy.find(id); if (l == copy.end()) continue; auto range = l->second; - auto chunk_size = range_data.checksummed_data().content().size(); + std::size_t chunk_size = range_data.checksummed_data().content().size(); bool evict = false; lk.lock(); @@ -584,9 +585,9 @@ void ObjectDescriptorImpl::DoFinish( // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the // end of the block. - auto current_stream = it->stream->stream; + std::shared_ptr current_stream = it->stream->stream; lk.unlock(); - auto pending = current_stream->Finish(); + future pending = current_stream->Finish(); if (!pending.valid()) return; pending.then([w = WeakFromThis(), read_stream, current_stream](auto f) { if (auto self = w.lock()) { @@ -608,7 +609,7 @@ void ObjectDescriptorImpl::OnFinish( return; } } - auto proto_status = ExtractGrpcStatus(status); + google::rpc::Status proto_status = ExtractGrpcStatus(status); if (IsResumable(read_stream, status, proto_status)) { return Resume(read_stream, proto_status); @@ -639,14 +640,15 @@ void ObjectDescriptorImpl::Resume( // to the dying stream while we establish a new one. it->stream->resuming = true; it->stream->next_request.Clear(); - auto current_stream = it->stream->stream; + std::shared_ptr current_stream = it->stream->stream; // This call needs to happen inside the lock, as it may modify // `read_object_spec_`. ApplyRedirectErrors(read_object_spec_, proto_status); - auto request = google::storage::v2::BidiReadObjectRequest{}; + google::storage::v2::BidiReadObjectRequest request; *request.mutable_read_object_spec() = read_object_spec_; - for (auto const& [read_id, range_ptr] : it->active_ranges) { - auto range = range_ptr->RangeForResume(read_id); + for (auto const& kv : it->active_ranges) { + std::optional range = + kv.second->RangeForResume(kv.first); if (!range) continue; *request.add_read_ranges() = *std::move(range); } @@ -685,10 +687,11 @@ void ObjectDescriptorImpl::OnResume( // Preserve any Read() range requests that arrived concurrently while the // reconnection was in flight. - auto queued_request = std::move(it->stream->next_request); + google::storage::v2::BidiReadObjectRequest queued_request = + std::move(it->stream->next_request); // Replace the old stream with the new stream and reset policy/state. - auto new_read_stream = std::make_shared( + std::shared_ptr new_read_stream = std::make_shared( std::move(result->stream), resume_policy_prototype_->clone()); new_read_stream->resume_policy->OnStartSuccess(); new_read_stream->write_pending = false; @@ -697,7 +700,7 @@ void ObjectDescriptorImpl::OnResume( new_read_stream->next_request = std::move(queued_request); it->stream = new_read_stream; - auto new_stream = new_read_stream->stream; + std::shared_ptr new_stream = new_read_stream->stream; // TODO(#15105) - this should be done without release the lock. // Flush any queued range requests onto the newly active stream. @@ -713,8 +716,8 @@ bool ObjectDescriptorImpl::IsResumable( if (cancelled_) return false; auto it = stream_manager_->Find(read_stream); if (it == stream_manager_->End() || !it->stream) return false; - for (auto const& any : proto_status.details()) { - auto error = google::storage::v2::BidiReadObjectError{}; + for (google::protobuf::Any const& any : proto_status.details()) { + google::storage::v2::BidiReadObjectError error; if (!any.UnpackTo(&error)) continue; std::vector> notify; @@ -727,9 +730,9 @@ bool ObjectDescriptorImpl::IsResumable( auto copy = it->active_ranges; lk.unlock(); - for (auto const& [read_id, range_status] : notify) { - auto l = copy.find(read_id); - if (l != copy.end()) l->second->OnFinish(range_status); + for (auto const& p : notify) { + auto l = copy.find(p.first); + if (l != copy.end()) l->second->OnFinish(p.second); } lk.lock(); if (cancelled_) return true; @@ -738,7 +741,7 @@ bool ObjectDescriptorImpl::IsResumable( stream_manager_->CleanupDoneRanges(it_curr); return true; } - auto effective_status = status; + Status effective_status = status; if (status.code() == StatusCode::kCancelled) { effective_status = Status(StatusCode::kUnavailable, status.message()); } From 08e9fd2d3fd6325a6bc222e4e9dcf16583496f30 Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Fri, 11 Sep 2026 10:06:59 +0000 Subject: [PATCH 4/6] Revert unrelated changes --- .../internal/async/object_descriptor_impl.cc | 22 +++++++++---------- 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.cc b/google/cloud/storage/internal/async/object_descriptor_impl.cc index 67672d5d15a55..0a9f17f242b85 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl.cc @@ -197,7 +197,7 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { auto self = w.lock(); if (!self) return; - StatusOr stream_result = f.get(); + auto stream_result = f.get(); if (!stream_result) { // Stream creation failed. // The next call to AssurePendingStreamQueued will retry creation. @@ -320,8 +320,7 @@ std::unique_ptr ObjectDescriptorImpl::Read( std::shared_ptr read_stream = it->stream; std::int64_t const id = ++read_id_generator_; it->active_ranges.emplace(id, range); - google::storage::v2::ReadRange& read_range = - *read_stream->next_request.add_read_ranges(); + auto& read_range = *read_stream->next_request.add_read_ranges(); read_range.set_read_id(id); read_range.set_read_offset(p.start); read_range.set_read_length(p.length); @@ -585,9 +584,9 @@ void ObjectDescriptorImpl::DoFinish( // Assign CurrentStream to a temporary variable to prevent // lifetime extension which can cause the lock to be held until the // end of the block. - std::shared_ptr current_stream = it->stream->stream; + auto current_stream = it->stream->stream; lk.unlock(); - future pending = current_stream->Finish(); + auto pending = current_stream->Finish(); if (!pending.valid()) return; pending.then([w = WeakFromThis(), read_stream, current_stream](auto f) { if (auto self = w.lock()) { @@ -609,7 +608,7 @@ void ObjectDescriptorImpl::OnFinish( return; } } - google::rpc::Status proto_status = ExtractGrpcStatus(status); + auto proto_status = ExtractGrpcStatus(status); if (IsResumable(read_stream, status, proto_status)) { return Resume(read_stream, proto_status); @@ -640,15 +639,14 @@ void ObjectDescriptorImpl::Resume( // to the dying stream while we establish a new one. it->stream->resuming = true; it->stream->next_request.Clear(); - std::shared_ptr current_stream = it->stream->stream; + auto current_stream = it->stream->stream; // This call needs to happen inside the lock, as it may modify // `read_object_spec_`. ApplyRedirectErrors(read_object_spec_, proto_status); - google::storage::v2::BidiReadObjectRequest request; + auto request = google::storage::v2::BidiReadObjectRequest{}; *request.mutable_read_object_spec() = read_object_spec_; for (auto const& kv : it->active_ranges) { - std::optional range = - kv.second->RangeForResume(kv.first); + auto range = kv.second->RangeForResume(kv.first); if (!range) continue; *request.add_read_ranges() = *std::move(range); } @@ -716,8 +714,8 @@ bool ObjectDescriptorImpl::IsResumable( if (cancelled_) return false; auto it = stream_manager_->Find(read_stream); if (it == stream_manager_->End() || !it->stream) return false; - for (google::protobuf::Any const& any : proto_status.details()) { - google::storage::v2::BidiReadObjectError error; + for (auto const& any : proto_status.details()) { + auto error = google::storage::v2::BidiReadObjectError{}; if (!any.UnpackTo(&error)) continue; std::vector> notify; From df32470ad8b9e5b78cb5b497d185726ebcb84bbd Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Mon, 14 Sep 2026 05:40:27 +0000 Subject: [PATCH 5/6] Handle additional corner cases --- .../internal/async/multi_stream_manager.h | 32 ++- .../async/multi_stream_manager_test.cc | 32 +++ .../internal/async/object_descriptor_impl.cc | 34 ++- .../async/object_descriptor_impl_test.cc | 264 ++++++++++++++++++ 4 files changed, 341 insertions(+), 21 deletions(-) diff --git a/google/cloud/storage/internal/async/multi_stream_manager.h b/google/cloud/storage/internal/async/multi_stream_manager.h index 2204de1b0fdd0..809c9af8c71cd 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager.h +++ b/google/cloud/storage/internal/async/multi_stream_manager.h @@ -19,6 +19,7 @@ #include "google/cloud/version.h" #include #include +#include #include #include #include @@ -89,18 +90,21 @@ class MultiStreamManager { return streams_.begin(); } - StreamIterator GetLeastBusyStream() { + // Returns an iterator to the stream with the fewest active ranges matching + // the given predicate. Returns End() if no stream satisfies the predicate. + // Strict less-than ensures stability by preferring earlier (older) streams if + // tied. + template + StreamIterator GetLeastBusyStream(Pred pred) { if (streams_.empty()) return streams_.end(); - auto least_busy_stream_it = streams_.begin(); - // Track min_ranges to avoid calling .size() repeatedly if possible, - // though for std::unordered_map .size() is O(1). - std::size_t min_ranges = least_busy_stream_it->active_ranges.size(); - if (min_ranges == 0) return least_busy_stream_it; - - // Start checking from the second element - for (auto it = std::next(streams_.begin()); it != streams_.end(); ++it) { - // Strict less-than ensures stability (preferring older streams if tied) - auto size = it->active_ranges.size(); + StreamIterator least_busy_stream_it = streams_.end(); + // Track min_ranges to avoid calling .size() repeatedly. + std::size_t min_ranges = (std::numeric_limits::max)(); + + for (StreamIterator it = streams_.begin(); it != streams_.end(); ++it) { + if (!pred(*it)) continue; + std::size_t const size = it->active_ranges.size(); + // Strict less-than ensures stability (preferring older streams if tied). if (size < min_ranges) { least_busy_stream_it = it; min_ranges = size; @@ -110,6 +114,12 @@ class MultiStreamManager { return least_busy_stream_it; } + // Overload of `GetLeastBusyStream` without predicate that selects the stream + // with the fewest active ranges across all managed streams. + StreamIterator GetLeastBusyStream() { + return GetLeastBusyStream([](Stream const&) { return true; }); + } + StreamIterator AddStream(std::shared_ptr stream) { streams_.emplace_front(Stream{std::move(stream), {}}); return streams_.begin(); diff --git a/google/cloud/storage/internal/async/multi_stream_manager_test.cc b/google/cloud/storage/internal/async/multi_stream_manager_test.cc index 755cd23d05c18..805f8ee142e26 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager_test.cc +++ b/google/cloud/storage/internal/async/multi_stream_manager_test.cc @@ -101,6 +101,38 @@ TEST(MultiStreamManagerTest, GetLeastBusyPrefersFewestActiveRanges) { EXPECT_EQ(it_least->active_ranges.size(), 1U); } +/// @test Verifies that GetLeastBusyStream with a predicate filters candidates +/// based on the provided predicate, correctly returning End() if no stream +/// matches or the least busy stream among those that satisfy the predicate. +TEST(MultiStreamManagerTest, GetLeastBusyStreamWithPredicate) { + auto mgr = MultiStreamManagerTest::MakeManager(); + mgr.GetFirstStream()->stream->write_pending = true; + + auto s1 = std::make_shared(); + auto s2 = std::make_shared(); + mgr.AddStream(s1); + Manager::StreamIterator it2 = mgr.AddStream(s2); + + // s1 has 0 ranges, but write_pending = true. + s1->write_pending = true; + // s2 has 1 range, write_pending = false. + s2->write_pending = false; + it2->active_ranges.emplace(1, std::make_shared()); + + // Predicate filtering out write_pending streams selects s2 even though it has + // more ranges than s1. + Manager::StreamIterator it_pred = + mgr.GetLeastBusyStream([](Manager::Stream const& s) { + return s.stream != nullptr && !s.stream->write_pending; + }); + EXPECT_THAT(it_pred, ::testing::Eq(it2)); + + // If predicate matches no stream, returns End(). + Manager::StreamIterator it_none = + mgr.GetLeastBusyStream([](Manager::Stream const&) { return false; }); + EXPECT_THAT(it_none, ::testing::Eq(mgr.End())); +} + TEST(MultiStreamManagerTest, CleanupDoneRangesRemovesFinished) { auto mgr = MultiStreamManagerTest::MakeManager(); auto it = mgr.GetFirstStream(); diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.cc b/google/cloud/storage/internal/async/object_descriptor_impl.cc index 0a9f17f242b85..bce9cf7de7a2e 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl.cc @@ -177,12 +177,14 @@ void ObjectDescriptorImpl::MakeSubsequentStream() { } std::unique_lock lk(mu_); - // Reuse an idle stream if possible. + // Reuse an idle stream if possible. A stream undergoing reconnection + // (resuming == true) must not be treated as idle, as it cannot accept new + // read ranges immediately. if (stream_manager_->ReuseIdleStreamToFront( [](StreamManager::Stream const& s) { auto const* rs = s.stream.get(); return rs != nullptr && s.active_ranges.empty() && - !rs->write_pending; + !rs->write_pending && !rs->resuming; })) { return; } @@ -316,7 +318,18 @@ std::unique_ptr ObjectDescriptorImpl::Read( CacheStatusToString(cache_status)); } - auto it = stream_manager_->GetLeastBusyStream(); + // Prioritize selecting a healthy stream that is not undergoing reconnection. + // If all streams are currently reconnecting, fall back to the least busy + // resuming stream so that the range is queued in next_request and dispatched + // upon reconnection completion in OnResume(). + StreamManager::StreamIterator it = + stream_manager_->GetLeastBusyStream([](StreamManager::Stream const& s) { + auto const* rs = s.stream.get(); + return rs != nullptr && !rs->resuming; + }); + if (it == stream_manager_->End()) { + it = stream_manager_->GetLeastBusyStream(); + } std::shared_ptr read_stream = it->stream; std::int64_t const id = ++read_id_generator_; it->active_ranges.emplace(id, range); @@ -688,9 +701,11 @@ void ObjectDescriptorImpl::OnResume( google::storage::v2::BidiReadObjectRequest queued_request = std::move(it->stream->next_request); - // Replace the old stream with the new stream and reset policy/state. + // Replace the old stream with the new stream and preserve the existing + // resume policy so failure budgets and error counts are maintained across + // reconnects. std::shared_ptr new_read_stream = std::make_shared( - std::move(result->stream), resume_policy_prototype_->clone()); + std::move(result->stream), std::move(it->stream->resume_policy)); new_read_stream->resume_policy->OnStartSuccess(); new_read_stream->write_pending = false; new_read_stream->read_pending = false; @@ -739,11 +754,10 @@ bool ObjectDescriptorImpl::IsResumable( stream_manager_->CleanupDoneRanges(it_curr); return true; } - Status effective_status = status; - if (status.code() == StatusCode::kCancelled) { - effective_status = Status(StatusCode::kUnavailable, status.message()); - } - return it->stream->resume_policy->OnFinish(effective_status) == + // Pass the original status directly to the resume policy without rewriting + // status codes. This allows custom resume policies (e.g., detecting stall + // cancellations) to observe the exact failure cause. + return it->stream->resume_policy->OnFinish(status) == storage::ResumePolicy::kContinue; } diff --git a/google/cloud/storage/internal/async/object_descriptor_impl_test.cc b/google/cloud/storage/internal/async/object_descriptor_impl_test.cc index 7c98b3c4d2665..3d39bfaac9ea2 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl_test.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl_test.cc @@ -51,8 +51,12 @@ using ::google::cloud::testing_util::IsProtoEqual; using ::google::cloud::testing_util::StatusIs; using ::google::protobuf::TextFormat; using ::testing::_; +using ::testing::AnyNumber; using ::testing::AtMost; using ::testing::ElementsAre; +using ::testing::Eq; +using ::testing::IsFalse; +using ::testing::IsTrue; using ::testing::NotNull; using ::testing::Optional; using ::testing::ResultOf; @@ -3635,6 +3639,266 @@ TEST(ObjectDescriptorImpl, ConcurrentReadPreservedAcrossMultipleResumeRetries) { read2.first.set_value(false); } +/// @test Verify that a stream waiting for a reconnection is not offered as an +/// "idle" stream, and that new reads are not parked on it. +TEST(ObjectDescriptorImpl, ResumingStreamIsNotReusedAsIdleStream) { + AsyncSequencer sequencer; + AsyncSequencer factory_sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillOnce([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto) { + return TransientError(); + }); + }); + EXPECT_CALL(*stream1, Cancel).Times(AtMost(1)); + + auto stream2 = std::make_unique(); + EXPECT_CALL(*stream2, Read).WillRepeatedly([&sequencer]() { + return sequencer.PushBack("Read[2]").then( + [](auto) { return std::optional{}; }); + }); + // The read issued while stream 1 is reconnecting must be written to the + // healthy stream. + EXPECT_CALL(*stream2, Write) + .WillOnce([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[2]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream2, Finish).Times(AtMost(1)).WillRepeatedly([&sequencer]() { + return sequencer.PushBack("Finish[2]").then([](auto) { return Status{}; }); + }); + EXPECT_CALL(*stream2, Cancel).Times(AtMost(1)); + + // The first factory call is the proactive background stream queued by + // `Start()`; it produces stream 2. Any later call (including the + // reconnection attempt for stream 1) fails, and the reconnection is left + // in flight for most of the test. + int factory_calls = 0; + MockFactory factory; + EXPECT_CALL(factory, Call).WillRepeatedly([&](Request const&) { + int const call = ++factory_calls; + char const* name = call == 1 ? "Factory[background]" : "Factory[resume]"; + return factory_sequencer.PushBack(name).then( + [&stream2, call](auto f) -> StatusOr { + if (call != 1 || !f.get()) return TransientError(); + return make_status_or(OpenStreamResult{ + std::make_shared(std::move(stream2)), Response{}}); + }); + }); + + auto tested = MakeTested(storage::LimitedErrorCountResumePolicy(3)(), + factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1))); + + tested->Start(Response{}); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_THAT(read1.second, Eq("Read[1]")); + auto background = factory_sequencer.PopFrontWithName(); + EXPECT_THAT(background.second, Eq("Factory[background]")); + + // Stream 1 fails while it has no active ranges, e.g. the stall watchdog + // cancelled an idle stream. The descriptor starts a reconnection which does + // not complete yet, leaving the entry in the `resuming` state. + read1.first.set_value(false); + auto finish1 = sequencer.PopFrontWithName(); + EXPECT_THAT(finish1.second, Eq("Finish[1]")); + finish1.first.set_value(true); + auto reconnect = factory_sequencer.PopFrontWithName(); + EXPECT_THAT(reconnect.second, Eq("Factory[resume]")); + + // The application asks for an additional stream. The only entry is waiting + // to reconnect, so the descriptor must consume the background stream + // instead of reporting that it reused an idle one. + tested->MakeSubsequentStream(); + background.first.set_value(true); + EXPECT_THAT(tested->StreamSize(), Eq(std::size_t{2})); + + // The new read must go to the healthy stream. If it is queued on the + // reconnecting stream `Flush()` is skipped and nothing reaches the wire. + std::unique_ptr reader = + tested->Read({0, 100}); + EXPECT_THAT(reader, NotNull()); + EXPECT_THAT(sequencer.empty(), IsFalse()); + + tested.reset(); + reconnect.first.set_value(false); + while (!sequencer.empty()) sequencer.PopFront().set_value(false); + while (!factory_sequencer.empty()) { + factory_sequencer.PopFront().set_value(false); + } +} + +/// @test Verify the resume policy is told the actual status of the stream. +TEST(ObjectDescriptorImpl, ResumePolicyObservesCancelledStatus) { + AsyncSequencer sequencer; + std::optional policy_status; + auto prototype = std::make_unique(); + EXPECT_CALL(*prototype, clone).WillRepeatedly([&policy_status]() { + auto policy = std::make_unique(); + EXPECT_CALL(*policy, OnStartSuccess).Times(AnyNumber()); + EXPECT_CALL(*policy, OnFinish) + .WillRepeatedly([&policy_status](Status const& status) { + policy_status = status; + return storage::ResumePolicy::kStop; + }); + return policy; + }); + + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillOnce([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto) { + return Status(StatusCode::kCancelled, "Stream stalled"); + }); + }); + EXPECT_CALL(*stream1, Cancel).Times(AtMost(1)); + + MockFactory factory; + EXPECT_CALL(factory, Call).WillRepeatedly([](Request const&) { + return make_ready_future(StatusOr(PermanentError())); + }); + + Options options; + options.set(false); + auto tested = std::make_shared( + std::move(prototype), factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1)), options); + + tested->Start(Response{}); + + auto read1 = sequencer.PopFrontWithName(); + EXPECT_THAT(read1.second, Eq("Read[1]")); + read1.first.set_value(false); + + auto finish1 = sequencer.PopFrontWithName(); + EXPECT_THAT(finish1.second, Eq("Finish[1]")); + finish1.first.set_value(true); + + EXPECT_THAT(policy_status.has_value(), IsTrue()); + if (policy_status.has_value()) { + EXPECT_THAT(*policy_status, + StatusIs(StatusCode::kCancelled, "Stream stalled")); + } + + tested.reset(); + while (!sequencer.empty()) sequencer.PopFront().set_value(false); +} + +/// @test Verify the resume budget is consumed across successful reconnects. +TEST(ObjectDescriptorImpl, ResumeBudgetIsNotResetByASuccessfulReconnect) { + AsyncSequencer sequencer; + auto stream1 = std::make_unique(); + EXPECT_CALL(*stream1, Read).WillOnce([&sequencer]() { + return sequencer.PushBack("Read[1]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream1, Write) + .WillOnce([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[1]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream1, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[1]").then([](auto) { + return TransientError(); + }); + }); + EXPECT_CALL(*stream1, Cancel).Times(AtMost(1)); + + auto stream2 = std::make_unique(); + EXPECT_CALL(*stream2, Read).WillOnce([&sequencer]() { + return sequencer.PushBack("Read[2]").then( + [](auto) { return std::optional{}; }); + }); + EXPECT_CALL(*stream2, Write) + .Times(AtMost(1)) + .WillRepeatedly([&sequencer](Request const&, grpc::WriteOptions) { + return sequencer.PushBack("Write[2]").then([](auto f) { + return f.get(); + }); + }); + EXPECT_CALL(*stream2, Finish).WillOnce([&sequencer]() { + return sequencer.PushBack("Finish[2]").then([](auto) { + return TransientError(); + }); + }); + EXPECT_CALL(*stream2, Cancel).Times(AtMost(1)); + + // Reconnection attempts carry the ranges to resume; the proactive background + // streams do not. Only the former are counted. + int resume_attempts = 0; + MockFactory factory; + EXPECT_CALL(factory, Call) + .WillRepeatedly( + [&](Request const& request) -> future> { + if (request.read_ranges().empty()) { + return make_ready_future( + StatusOr(TransientError())); + } + if (++resume_attempts != 1) { + return make_ready_future( + StatusOr(TransientError())); + } + return make_ready_future(make_status_or(OpenStreamResult{ + std::make_shared(std::move(stream2)), Response{}})); + }); + + Options options; + options.set(false); + // A budget of exactly one resume for the lifetime of this descriptor. + auto tested = std::make_shared( + storage::LimitedErrorCountResumePolicy(1)(), factory.AsStdFunction(), + google::storage::v2::BidiReadObjectSpec{}, + std::make_shared(std::move(stream1)), options); + + tested->Start(Response{}); + + // `Start()` leaves a read outstanding on stream 1. + auto read1 = sequencer.PopFrontWithName(); + EXPECT_THAT(read1.second, Eq("Read[1]")); + + std::unique_ptr reader = + tested->Read({0, 100}); + EXPECT_THAT(reader, NotNull()); + auto write1 = sequencer.PopFrontWithName(); + EXPECT_THAT(write1.second, Eq("Write[1]")); + write1.first.set_value(true); + + // First failure: consumes the single resume allowed by the policy. + read1.first.set_value(false); + auto finish1 = sequencer.PopFrontWithName(); + EXPECT_THAT(finish1.second, Eq("Finish[1]")); + finish1.first.set_value(true); + EXPECT_THAT(resume_attempts, Eq(1)); + + // Second failure, on the replacement stream. The budget is exhausted, so the + // descriptor must give up instead of reconnecting again. + auto read2 = sequencer.PopFrontWithName(); + EXPECT_THAT(read2.second, Eq("Read[2]")); + read2.first.set_value(false); + + auto finish2 = sequencer.PopFrontWithName(); + EXPECT_THAT(finish2.second, Eq("Finish[2]")); + finish2.first.set_value(true); + + EXPECT_THAT(resume_attempts, Eq(1)); + EXPECT_THAT(tested->IsOpen(), IsFalse()); + + tested.reset(); + while (!sequencer.empty()) sequencer.PopFront().set_value(false); +} + } // namespace GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END } // namespace storage_internal From 4dcc5902aba4d2fe6f6713c68572c7dbcb39ad64 Mon Sep 17 00:00:00 2001 From: Gauri Kalra Date: Mon, 14 Sep 2026 08:18:53 +0000 Subject: [PATCH 6/6] Fix failing checks --- .../storage/internal/async/multi_stream_manager.h | 4 ++-- .../internal/async/multi_stream_manager_test.cc | 11 +++++------ .../storage/internal/async/object_descriptor_impl.cc | 7 ++++--- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/google/cloud/storage/internal/async/multi_stream_manager.h b/google/cloud/storage/internal/async/multi_stream_manager.h index 809c9af8c71cd..876f228d5ada3 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager.h +++ b/google/cloud/storage/internal/async/multi_stream_manager.h @@ -97,11 +97,11 @@ class MultiStreamManager { template StreamIterator GetLeastBusyStream(Pred pred) { if (streams_.empty()) return streams_.end(); - StreamIterator least_busy_stream_it = streams_.end(); + auto least_busy_stream_it = streams_.end(); // Track min_ranges to avoid calling .size() repeatedly. std::size_t min_ranges = (std::numeric_limits::max)(); - for (StreamIterator it = streams_.begin(); it != streams_.end(); ++it) { + for (auto it = streams_.begin(); it != streams_.end(); ++it) { if (!pred(*it)) continue; std::size_t const size = it->active_ranges.size(); // Strict less-than ensures stability (preferring older streams if tied). diff --git a/google/cloud/storage/internal/async/multi_stream_manager_test.cc b/google/cloud/storage/internal/async/multi_stream_manager_test.cc index 805f8ee142e26..d3dd69c9ef3da 100644 --- a/google/cloud/storage/internal/async/multi_stream_manager_test.cc +++ b/google/cloud/storage/internal/async/multi_stream_manager_test.cc @@ -111,7 +111,7 @@ TEST(MultiStreamManagerTest, GetLeastBusyStreamWithPredicate) { auto s1 = std::make_shared(); auto s2 = std::make_shared(); mgr.AddStream(s1); - Manager::StreamIterator it2 = mgr.AddStream(s2); + auto it2 = mgr.AddStream(s2); // s1 has 0 ranges, but write_pending = true. s1->write_pending = true; @@ -121,14 +121,13 @@ TEST(MultiStreamManagerTest, GetLeastBusyStreamWithPredicate) { // Predicate filtering out write_pending streams selects s2 even though it has // more ranges than s1. - Manager::StreamIterator it_pred = - mgr.GetLeastBusyStream([](Manager::Stream const& s) { - return s.stream != nullptr && !s.stream->write_pending; - }); + auto it_pred = mgr.GetLeastBusyStream([](Manager::Stream const& s) { + return s.stream != nullptr && !s.stream->write_pending; + }); EXPECT_THAT(it_pred, ::testing::Eq(it2)); // If predicate matches no stream, returns End(). - Manager::StreamIterator it_none = + auto it_none = mgr.GetLeastBusyStream([](Manager::Stream const&) { return false; }); EXPECT_THAT(it_none, ::testing::Eq(mgr.End())); } diff --git a/google/cloud/storage/internal/async/object_descriptor_impl.cc b/google/cloud/storage/internal/async/object_descriptor_impl.cc index bce9cf7de7a2e..4e06e316e24e0 100644 --- a/google/cloud/storage/internal/async/object_descriptor_impl.cc +++ b/google/cloud/storage/internal/async/object_descriptor_impl.cc @@ -322,7 +322,7 @@ std::unique_ptr ObjectDescriptorImpl::Read( // If all streams are currently reconnecting, fall back to the least busy // resuming stream so that the range is queued in next_request and dispatched // upon reconnection completion in OnResume(). - StreamManager::StreamIterator it = + auto it = stream_manager_->GetLeastBusyStream([](StreamManager::Stream const& s) { auto const* rs = s.stream.get(); return rs != nullptr && !rs->resuming; @@ -755,8 +755,9 @@ bool ObjectDescriptorImpl::IsResumable( return true; } // Pass the original status directly to the resume policy without rewriting - // status codes. This allows custom resume policies (e.g., detecting stall - // cancellations) to observe the exact failure cause. + // status codes (such as StatusCode::kCancelled). This allows custom resume + // policies (e.g., detecting stall cancellations) to observe the exact failure + // cause. return it->stream->resume_policy->OnFinish(status) == storage::ResumePolicy::kContinue; }