feat(storage): implement multi stream manager for async downloads #15858
Conversation
product-auto-label
bot
added
the
api: storage
Summary of ChangesHello @v-pratap, I'm Gemini Code Assist1! 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 Highlights
🧠 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 AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a MultiStreamManager to support multiple gRPC streams for asynchronous downloads, aiming to improve performance through connection pooling and reuse. The implementation is well-structured, encapsulating stream management logic effectively. The changes include proactive stream creation and reuse of idle streams, with corresponding updates to ObjectDescriptorImpl. The test suite has been significantly enhanced to cover the new multi-stream functionality, including scenarios for proactive creation, stream reuse, and various error conditions. I've found one critical issue regarding unhandled exceptions that could lead to application termination.
| stream_future.then([w = WeakFromThis()](auto f) { | ||
| auto self = w.lock(); | ||
| if (!self) return; | ||
|
|
||
| auto stream_result = f.get(); | ||
| if (!stream_result) { | ||
| // Stream creation failed. | ||
| // The next call to AssurePendingStreamQueued will retry creation. | ||
| return; | ||
| } | ||
|
|
||
| std::unique_lock<std::mutex> lk(self->mu_); | ||
| if (self->cancelled_) return; | ||
|
|
||
| auto read_stream = | ||
| std::make_shared<ReadStream>(std::move(stream_result->stream), | ||
| self->resume_policy_prototype_->clone()); | ||
|
|
||
| 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); | ||
|
|
||
| lk.unlock(); | ||
| self->OnRead(new_it, std::move(stream_result->first_response)); | ||
| }); |
There was a problem hiding this comment.
The call to f.get() on line 106 may throw an unhandled exception if the stream_future is cancelled, which would terminate the application. The ObjectDescriptorImpl::Cancel() method does cancel pending_stream_, which is where stream_future comes from. The continuation lambda should handle this potential exception, for example by wrapping the f.get() call in a try-catch block.
stream_future.then([w = WeakFromThis()](auto f) {
auto self = w.lock();
if (!self) return;
StatusOr<OpenStreamResult> stream_result;
try {
stream_result = f.get();
} catch (...) {
// The future was cancelled, which is expected if the ObjectDescriptor is
// destroyed. This is not an error. The next call to
// AssurePendingStreamQueued() will create a new stream if needed.
return;
}
if (!stream_result) {
// Stream creation failed.
// The next call to AssurePendingStreamQueued will retry creation.
return;
}
std::unique_lock<std::mutex> lk(self->mu_);
if (self->cancelled_) return;
auto read_stream =
std::make_shared<ReadStream>(std::move(stream_result->stream),
self->resume_policy_prototype_->clone());
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);
lk.unlock();
self->OnRead(new_it, std::move(stream_result->first_response));
});
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #15858 +/- ##
==========================================
+ Coverage 92.95% 92.96% +0.01%
==========================================
Files 2458 2460 +2
Lines 227590 228211 +621
==========================================
+ Hits 211548 212158 +610
- Misses 16042 16053 +11 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| // 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_(), {}}); |
There was a problem hiding this comment.
nit: here and other places where push_back or push_front is called while constructing the element at the same time are good opportunities to use emplace_back or emplace_front instead.
| } | ||
|
|
||
| bool Empty() const { return streams_.empty(); } | ||
| StreamIterator End() { return streams_.end(); } |
2 participants
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),
Bechmarking the changes by reading 500 GB data (thread count = 30),
Benchmarking by changing the MaximumRangeSizeOption values,
Code accurate Control flow diagram for this snippet,