fix(storage): ensure thread-safe stream resumption in ObjectDescriptorImpl

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

@@ -19,6 +19,7 @@
 #include "google/cloud/version.h"
 #include <cstdint>
 #include <functional>
+#include <limits>
 #include <list>
 #include <memory>
 #include <unordered_map>
@@ -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 <typename Pred>
+  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();
+    auto least_busy_stream_it = streams_.end();
+    // Track min_ranges to avoid calling .size() repeatedly.
+    std::size_t min_ranges = (std::numeric_limits<std::size_t>::max)();
+
+    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).
       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<StreamT> stream) {
     streams_.emplace_front(Stream{std::move(stream), {}});
     return streams_.begin();
@@ -159,6 +169,13 @@ class MultiStreamManager {
     return false;
   }
 
+  StreamIterator Find(std::shared_ptr<StreamT> const& target) {
+    for (auto it = streams_.begin(); it != streams_.end(); ++it) {
+      if (it->stream == target) return it;
+    }
+    return streams_.end();
+  }
+
   bool Empty() const { return streams_.empty(); }
   ConstStreamIterator End() const { return streams_.end(); }
   std::size_t Size() const { return streams_.size(); }

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

@@ -101,6 +101,37 @@ 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<FakeStream>();
+  auto s2 = std::make_shared<FakeStream>();
+  mgr.AddStream(s1);
+  auto 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<FakeRange>());
+
+  // Predicate filtering out write_pending streams selects s2 even though it has
+  // more ranges than s1.
+  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().
+  auto 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();
@@ -233,6 +264,22 @@ TEST(MultiStreamManagerTest, EmptyAndSizeTransitions) {
   EXPECT_EQ(mgr.Size(), 1U);
 }
 
+TEST(MultiStreamManagerTest, FindTracksStreamMembership) {
+  auto mgr = MultiStreamManagerTest::MakeManager();
+  auto it1 = mgr.GetFirstStream();
+  auto s1 = it1->stream;
+  EXPECT_NE(mgr.Find(s1), mgr.End());
+
+  auto s2 = std::make_shared<FakeStream>();
+  mgr.AddStream(s2);
+  EXPECT_NE(mgr.Find(s1), mgr.End());
+  EXPECT_NE(mgr.Find(s2), mgr.End());
+
+  mgr.RemoveStreamAndNotifyRanges(it1, Status());
+  EXPECT_EQ(mgr.Find(s1), mgr.End());
+  EXPECT_NE(mgr.Find(s2), mgr.End());
+}
+
 GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
 }  // namespace storage_internal
 }  // namespace cloud