feat(storage): implement multi stream manager for async downloads

google/cloud/storage/async/object_descriptor.cc

@@ -27,14 +27,11 @@ absl::optional<google::storage::v2::Object> ObjectDescriptor::metadata() const {
 
 std::pair<AsyncReader, AsyncToken> ObjectDescriptor::Read(std::int64_t offset,
                                                           std::int64_t limit) {
-  // TODO(15340): This change is causing performance regression. We need to
-  // revisit it after benchmarking our code.
-
-  // std::int64_t max_range =
-  // impl_->options().get<storage_experimental::MaximumRangeSizeOption>();
-  // if (limit > max_range) {
-  //   impl_->MakeSubsequentStream();
-  // }
+  std::int64_t max_range =
+      impl_->options().get<storage_experimental::MaximumRangeSizeOption>();
+  if (limit > max_range) {
+    impl_->MakeSubsequentStream();
+  }
   auto reader = impl_->Read({offset, limit});
   auto token = storage_internal::MakeAsyncToken(reader.get());
   return {AsyncReader(std::move(reader)), std::move(token)};

google/cloud/storage/async/object_descriptor_test.cc

@@ -149,7 +149,6 @@ TEST(ObjectDescriptor, ReadLast) {
 }
 
 TEST(ObjectDescriptor, ReadExceedsMaxRange) {
-  GTEST_SKIP();
   auto mock = std::make_shared<MockAsyncObjectDescriptorConnection>();
   auto constexpr kMaxRange = 1024;
   EXPECT_CALL(*mock, options)

google/cloud/storage/google_cloud_cpp_storage_grpc.bzl

@@ -43,6 +43,7 @@ google_cloud_cpp_storage_grpc_hdrs = [
     "internal/async/default_options.h",
     "internal/async/handle_redirect_error.h",
     "internal/async/insert_object.h",
+    "internal/async/multi_stream_manager.h",
     "internal/async/object_descriptor_connection_tracing.h",
     "internal/async/object_descriptor_impl.h",
     "internal/async/object_descriptor_reader.h",
@@ -120,6 +121,7 @@ google_cloud_cpp_storage_grpc_srcs = [
     "internal/async/default_options.cc",
     "internal/async/handle_redirect_error.cc",
     "internal/async/insert_object.cc",
+    "internal/async/multi_stream_manager.cc",
     "internal/async/object_descriptor_connection_tracing.cc",
     "internal/async/object_descriptor_impl.cc",
     "internal/async/object_descriptor_reader.cc",

google/cloud/storage/google_cloud_cpp_storage_grpc.cmake

@@ -110,6 +110,8 @@ add_library(
     internal/async/handle_redirect_error.h
     internal/async/insert_object.cc
     internal/async/insert_object.h
+    internal/async/multi_stream_manager.cc
+    internal/async/multi_stream_manager.h
     internal/async/object_descriptor_connection_tracing.cc
     internal/async/object_descriptor_connection_tracing.h
     internal/async/object_descriptor_impl.cc
@@ -442,6 +444,7 @@ set(storage_client_grpc_unit_tests
     internal/async/default_options_test.cc
     internal/async/handle_redirect_error_test.cc
     internal/async/insert_object_test.cc
+    internal/async/multi_stream_manager_test.cc
     internal/async/object_descriptor_connection_tracing_test.cc
     internal/async/object_descriptor_impl_test.cc
     internal/async/object_descriptor_reader_test.cc

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

@@ -44,6 +44,8 @@ using ::google::cloud::testing_util::IsProtoEqual;
 using ::google::cloud::testing_util::MockCompletionQueueImpl;
 using ::google::cloud::testing_util::StatusIs;
 using ::google::protobuf::TextFormat;
+using ::testing::InvokeWithoutArgs;
+using ::testing::NiceMock;
 using ::testing::NotNull;
 using ::testing::Optional;
 
@@ -183,10 +185,29 @@ TEST(AsyncConnectionImplTest, OpenSimple) {
               [](auto) { return Status{}; });
         });
 
+        return std::unique_ptr<BidiReadStream>(std::move(stream));
+      })
+      .WillRepeatedly([](CompletionQueue const&,
+                         std::shared_ptr<grpc::ClientContext> const&,
+                         google::cloud::internal::ImmutableOptions const&) {
+        auto stream = std::make_unique<NiceMock<MockStream>>();
+        ON_CALL(*stream, Start).WillByDefault(InvokeWithoutArgs([] {
+          return make_ready_future(false);
+        }));
+        ON_CALL(*stream, Finish).WillByDefault(InvokeWithoutArgs([] {
+          return make_ready_future(Status{});
+        }));
+        ON_CALL(*stream, Cancel).WillByDefault([] {});
         return std::unique_ptr<BidiReadStream>(std::move(stream));
       });
 
   auto mock_cq = std::make_shared<MockCompletionQueueImpl>();
+  EXPECT_CALL(*mock_cq, MakeRelativeTimer)
+      .WillRepeatedly([](std::chrono::nanoseconds) {
+        return make_ready_future(
+            StatusOr<std::chrono::system_clock::time_point>(
+                std::chrono::system_clock::now()));
+      });
   auto connection = std::make_shared<AsyncConnectionImpl>(
       CompletionQueue(mock_cq), std::shared_ptr<GrpcChannelRefresh>(), mock,
       TestOptions());

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

@@ -0,0 +1,29 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "google/cloud/storage/internal/async/multi_stream_manager.h"
+#include "google/cloud/storage/internal/async/object_descriptor_impl.h"
+
+namespace google {
+namespace cloud {
+namespace storage_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+// Explicit instantiation for ObjectDescriptorImpl usage.
+template class MultiStreamManager<ReadStream, ReadRange>;
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace storage_internal
+}  // namespace cloud
+}  // namespace google

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

@@ -0,0 +1,175 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_STORAGE_INTERNAL_ASYNC_MULTI_STREAM_MANAGER_H
+#define GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_STORAGE_INTERNAL_ASYNC_MULTI_STREAM_MANAGER_H
+
+#include "google/cloud/status.h"
+#include "google/cloud/version.h"
+#include <cstdint>
+#include <functional>
+#include <list>
+#include <memory>
+#include <unordered_map>
+
+namespace google {
+namespace cloud {
+namespace storage_internal {
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_BEGIN
+
+// Defines the interface contract that any stream type (e.g., ReadStream,
+// WriteStream) managed by MultiStreamManager must implement. This explicit base
+// class ensures we have a clear, enforceable interface for operations like
+// CancelAll().
+class StreamBase {
+ public:
+  virtual ~StreamBase() = default;
+  virtual void Cancel() = 0;
+};
+
+// Manages a collection of streams.
+//
+// This class implements the "Subsequent Stream" logic where idle streams
+// are moved to the front of the queue for reuse.
+//
+// THREAD SAFETY:
+// This class is NOT thread-safe. The owner (e.g. ObjectDescriptorImpl
+// or AsyncWriterImpl etc) must serialize access, typically by holding
+// an external mutex while calling these methods.
+//
+// EXAMPLE USAGE:
+//   class MyOwner {
+//     std::mutex mu_;
+//     MultiStreamManager<MyStream, MyRange> manager_;
+//
+//     void StartRead() {
+//        std::unique_lock<std::mutex> lk(mu_);
+//        auto it = manager_.GetLeastBusyStream();
+//     }
+//   };
+template <typename StreamT, typename RangeT>
+class MultiStreamManager {
+ public:
+  struct Stream {
+    std::shared_ptr<StreamT> stream;
+    std::unordered_map<std::int64_t, std::shared_ptr<RangeT>> active_ranges;
+  };
+
+  using StreamIterator = typename std::list<Stream>::iterator;
+  using StreamFactory = std::function<std::shared_ptr<StreamT>()>;
+
+  // Constructor creates the first stream using the factory immediately.
+  explicit MultiStreamManager(StreamFactory stream_factory)
+      : stream_factory_(std::move(stream_factory)) {
+    streams_.push_back(Stream{stream_factory_(), {}});
+  }
+
+  // Constructor accepts an already-created initial stream.
+  // This is required by ObjectDescriptorImpl which receives an OpenStream.
+  MultiStreamManager(StreamFactory stream_factory,
+                     std::shared_ptr<StreamT> initial_stream)
+      : stream_factory_(std::move(stream_factory)) {
+    streams_.push_back(Stream{std::move(initial_stream), {}});
+  }
+
+  StreamIterator GetFirstStream() {
+    if (streams_.empty()) return streams_.end();
+    return streams_.begin();
+  }
+
+  StreamIterator GetLeastBusyStream() {
+    if (streams_.empty()) return streams_.end();
+    auto least_busy_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_it->active_ranges.size();
+    if (min_ranges == 0) return least_busy_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();
+      if (size < min_ranges) {
+        least_busy_it = it;
+        min_ranges = size;
+        if (min_ranges == 0) return least_busy_it;
+      }
+    }
+    return least_busy_it;
+  }
+
+  StreamIterator AddStream(std::shared_ptr<StreamT> stream) {
+    streams_.push_front(Stream{std::move(stream), {}});
+    return streams_.begin();
+  }
+
+  void CancelAll() {
+    for (auto& s : streams_) {
+      if (s.stream) s.stream->Cancel();
+    }
+  }
+
+  void RemoveStreamAndNotifyRanges(StreamIterator it, Status const& status) {
+    auto ranges = std::move(it->active_ranges);
+    streams_.erase(it);
+    for (auto const& kv : ranges) {
+      kv.second->OnFinish(status);
+    }
+  }
+
+  void MoveActiveRanges(StreamIterator from, StreamIterator to) {
+    to->active_ranges = std::move(from->active_ranges);
+  }
+
+  void CleanupDoneRanges(StreamIterator it) {
+    auto& active_ranges = it->active_ranges;
+    for (auto i = active_ranges.begin(); i != active_ranges.end();) {
+      if (i->second->IsDone()) {
+        i = active_ranges.erase(i);
+      } else {
+        ++i;
+      }
+    }
+  }
+
+  template <typename Pred>
+  bool ReuseIdleStreamToFront(Pred pred) {
+    for (auto it = streams_.begin(); it != streams_.end(); ++it) {
+      if (!pred(*it)) continue;
+
+      // If the idle stream is already at the front, we don't
+      // need to move it. Otherwise splice to the front in O(1).
+      if (it != streams_.begin()) {
+        streams_.splice(streams_.begin(), streams_, it);
+      }
+      return true;
+    }
+    return false;
+  }
+
+  bool Empty() const { return streams_.empty(); }
+  StreamIterator End() { return streams_.end(); }
+  std::size_t Size() const { return streams_.size(); }
+
+ private:
+  std::list<Stream> streams_;
+  StreamFactory stream_factory_;
+};
+
+GOOGLE_CLOUD_CPP_INLINE_NAMESPACE_END
+}  // namespace storage_internal
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_GOOGLE_CLOUD_STORAGE_INTERNAL_ASYNC_MULTI_STREAM_MANAGER_H