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

[v-pratap]

This PR implements scalable multi-stream reading for Zonal Buckets by introducing a MultiStreamManager to pool and reuse gRPC connections. It mitigates performance regressions from frequent stream creation by reusing idle streams and load-balancing requests, while maintaining a proactive background connection to minimize latency during ramp-up.
Class UML diagrams of the final architecture: Link
Benchmarking script used: Link

Bechmarking the changes by reading 500 GB data (thread count = 1),

Single steam: 1.14 GB/s
Multi stream with Reuse (poc code by vaibhav): 1.13 GB/s
Multi stream with Reuse (current PR code): 1.14 GB/s
Bechmarking the changes by reading 500 GB data (thread count = 30),

Single steam: 1.09 GB/s
Multi stream with Reuse (poc code by vaibhav): 1.48 GB/s
Multi stream with Reuse (current PR code): 1.49 GB/s
Benchmarking by changing the MaximumRangeSizeOption values,

64 MB size: 1.52 GB/s
128 MB size: 1.59 GB/s
256 MB size: 1.53 GB/s
Code accurate Control flow diagram for this snippet,

auto async = MakeAsyncClient(project_id);
google::storage::v2::BidiReadObjectSpec spec; spec.set_bucket("projects/_/buckets/" + bucket); spec.set_object(object);
auto descriptor = *async.Open(spec).get();
auto [reader, token] = descriptor.Read(/*offset=*/0, /*limit=*/400 * 1024 * 1024);
for (auto t = std::move(token); t.valid(); ) { auto r = reader.Read(std::move(t)).get(); t = std::get<1>(*r); /* process std::get<0>(*r) */ }

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Summary of Changes

Hello @v-pratap, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the asynchronous download capabilities for Google Cloud Storage by introducing a MultiStreamManager. This manager intelligently pools and reuses gRPC streams, which drastically improves performance by reducing the overhead of establishing new connections. The ObjectDescriptorImpl has been updated to leverage this new system, enabling more efficient resource utilization, better load distribution across streams, and a proactive approach to stream availability, ultimately leading to faster and more reliable data transfers, especially for large objects and concurrent operations.

Highlights

• MultiStreamManager Implementation: Introduced a new MultiStreamManager class to efficiently pool and reuse gRPC connections for asynchronous downloads, specifically targeting Zonal Buckets. This manager handles stream lifecycle, load-balancing, and proactive stream creation.
• Performance Improvement: Benchmarking results demonstrate significant performance gains, particularly under high concurrency (e.g., 1.09 GB/s to 1.49 GB/s with 30 threads), by mitigating overhead from frequent stream creation.
• ObjectDescriptorImpl Refactoring: The ObjectDescriptorImpl class has been refactored to integrate with the MultiStreamManager, centralizing stream management and enabling features like reusing idle streams and proactive background connection establishment.
• Proactive Stream Creation: The system now proactively creates background streams to minimize latency during ramp-up, ensuring a ready stream is often available for new read requests.
• Enhanced Test Coverage: New unit tests for MultiStreamManager and updated tests for ObjectDescriptorImpl have been added to validate the new multi-streaming logic, including scenarios for proactive stream creation, stream reuse, and error handling.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a MultiStreamManager to implement scalable multi-stream reading for async downloads, which is a significant feature for improving performance. The implementation includes pooling and reusing gRPC connections, load-balancing requests, and maintaining a proactive background connection. The changes are well-structured, with the new logic encapsulated in MultiStreamManager and integrated into ObjectDescriptorImpl. The accompanying tests are thorough and cover the new complex asynchronous logic well. My feedback focuses on a few opportunities to simplify the API of the new MultiStreamManager class for better clarity and maintainability.

[gemini-code-assist]

The MoveActiveRanges function does not appear to be used anywhere in ObjectDescriptorImpl or the associated tests. It seems to be a leftover from a previous design. To improve maintainability and simplify the API, it's best to remove unused code.

[gemini-code-assist]

The stream_factory_ member seems to be unused in the primary use case within ObjectDescriptorImpl. The responsibility of creating streams lies with ObjectDescriptorImpl, which then adds them to the manager. This suggests that stream creation is not a core responsibility of MultiStreamManager in this design.

Consider removing stream_factory_ and simplifying the constructors. For example:

// In multi_stream_manager.h

// Remove the factory-based constructor.
// explicit MultiStreamManager(StreamFactory stream_factory) ...

// Simplify this constructor to not take a factory.
explicit MultiStreamManager(std::shared_ptr<StreamT> initial_stream) {
  streams_.push_back(Stream{std::move(initial_stream), {}});
}

// ... private members ...
// std::list<Stream> streams_;
// StreamFactory stream_factory_; // <-- remove this

The test code would need a small adjustment:

// In multi_stream_manager_test.cc

static Manager MakeManager() {
  return Manager(std::make_shared<FakeStream>());
}

This change would make the class's responsibility clearer: it only manages streams, it doesn't create them.

[gemini-code-assist]

This no-op lambda returning nullptr highlights that the stream_factory parameter in MultiStreamManager's constructor is not used in this context. This is related to my suggestion on google/cloud/storage/internal/async/multi_stream_manager.h to simplify the MultiStreamManager API, which would remove the need for this parameter.