-
Notifications
You must be signed in to change notification settings - Fork 659
fix: spawn tasks in Stream::buffered and Stream::buffer_unordered to max concurrency #2962
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,5 @@ | ||
use crate::stream::{Fuse, FuturesUnordered, StreamExt}; | ||
use alloc::vec::Vec; | ||
use core::fmt; | ||
use core::num::NonZeroUsize; | ||
use core::pin::Pin; | ||
|
@@ -13,20 +14,23 @@ pin_project! { | |
/// Stream for the [`buffer_unordered`](super::StreamExt::buffer_unordered) | ||
/// method. | ||
#[must_use = "streams do nothing unless polled"] | ||
pub struct BufferUnordered<St> | ||
pub struct BufferUnordered<St, F> | ||
where | ||
St: Stream, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
#[pin] | ||
stream: Fuse<St>, | ||
in_progress_queue: FuturesUnordered<St::Item>, | ||
ready_queue: Vec<F::Output>, | ||
max: Option<NonZeroUsize>, | ||
} | ||
} | ||
|
||
impl<St> fmt::Debug for BufferUnordered<St> | ||
impl<St, F> fmt::Debug for BufferUnordered<St, F> | ||
where | ||
St: Stream + fmt::Debug, | ||
St: Stream<Item = F> + fmt::Debug, | ||
F: Future, | ||
{ | ||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { | ||
f.debug_struct("BufferUnordered") | ||
|
@@ -37,26 +41,27 @@ where | |
} | ||
} | ||
|
||
impl<St> BufferUnordered<St> | ||
impl<St, F> BufferUnordered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
pub(super) fn new(stream: St, n: Option<usize>) -> Self { | ||
Self { | ||
stream: super::Fuse::new(stream), | ||
in_progress_queue: FuturesUnordered::new(), | ||
ready_queue: Vec::new(), | ||
max: n.and_then(NonZeroUsize::new), | ||
} | ||
} | ||
|
||
delegate_access_inner!(stream, St, (.)); | ||
} | ||
|
||
impl<St> Stream for BufferUnordered<St> | ||
impl<St, F> Stream for BufferUnordered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
type Item = <St::Item as Future>::Output; | ||
|
||
|
@@ -72,14 +77,21 @@ where | |
} | ||
} | ||
|
||
// Attempt to pull the next value from the in_progress_queue | ||
match this.in_progress_queue.poll_next_unpin(cx) { | ||
x @ Poll::Pending | x @ Poll::Ready(Some(_)) => return x, | ||
Poll::Ready(None) => {} | ||
// Try to poll all ready futures in the in_progress_queue. | ||
loop { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I believe this fix could be incorrect, as it assumes that tasks are finished in order. And it introduces extra costs. I have similar issues in opendal, and my final solution is to find a way to track the status of futures and subtract the already finished ones from the maximum concurrent limit. Do you think it's a good idea to implement similiar thing in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Would you like to elaborate more why you think it assumes that tasks are finished in order? |
||
match this.in_progress_queue.poll_next_unpin(cx) { | ||
Poll::Ready(Some(output)) => { | ||
this.ready_queue.push(output); | ||
} | ||
Poll::Ready(None) => break, | ||
Poll::Pending => break, | ||
} | ||
} | ||
|
||
// If more values are still coming from the stream, we're not done yet | ||
if this.stream.is_done() { | ||
if let Some(output) = this.ready_queue.pop() { | ||
// If we have any ready outputs, return the first one. | ||
Poll::Ready(Some(output)) | ||
} else if this.stream.is_done() && this.in_progress_queue.is_empty() { | ||
Poll::Ready(None) | ||
} else { | ||
Poll::Pending | ||
|
@@ -98,10 +110,10 @@ where | |
} | ||
} | ||
|
||
impl<St> FusedStream for BufferUnordered<St> | ||
impl<St, F> FusedStream for BufferUnordered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
fn is_terminated(&self) -> bool { | ||
self.in_progress_queue.is_terminated() && self.stream.is_terminated() | ||
|
@@ -110,10 +122,10 @@ where | |
|
||
// Forwarding impl of Sink from the underlying stream | ||
#[cfg(feature = "sink")] | ||
impl<S, Item> Sink<Item> for BufferUnordered<S> | ||
impl<S, F, Item> Sink<Item> for BufferUnordered<S, F> | ||
where | ||
S: Stream + Sink<Item>, | ||
S::Item: Future, | ||
S: Stream<Item = F> + Sink<Item>, | ||
F: Future, | ||
{ | ||
type Error = S::Error; | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,9 +1,9 @@ | ||
use crate::stream::{Fuse, FusedStream, FuturesOrdered, StreamExt}; | ||
use alloc::collections::VecDeque; | ||
use core::fmt; | ||
use core::num::NonZeroUsize; | ||
use core::pin::Pin; | ||
use futures_core::future::Future; | ||
use futures_core::ready; | ||
use futures_core::stream::Stream; | ||
use futures_core::task::{Context, Poll}; | ||
#[cfg(feature = "sink")] | ||
|
@@ -13,22 +13,23 @@ use pin_project_lite::pin_project; | |
pin_project! { | ||
/// Stream for the [`buffered`](super::StreamExt::buffered) method. | ||
#[must_use = "streams do nothing unless polled"] | ||
pub struct Buffered<St> | ||
pub struct Buffered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
Comment on lines
+16
to
+19
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Seems like a breaking change? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It seems ok since we are on 0.4-alpha |
||
{ | ||
#[pin] | ||
stream: Fuse<St>, | ||
in_progress_queue: FuturesOrdered<St::Item>, | ||
ready_queue: VecDeque<F::Output>, | ||
max: Option<NonZeroUsize>, | ||
} | ||
} | ||
|
||
impl<St> fmt::Debug for Buffered<St> | ||
impl<St, F> fmt::Debug for Buffered<St, F> | ||
where | ||
St: Stream + fmt::Debug, | ||
St::Item: Future, | ||
St: Stream<Item = F> + fmt::Debug, | ||
F: Future, | ||
{ | ||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { | ||
f.debug_struct("Buffered") | ||
|
@@ -39,26 +40,27 @@ where | |
} | ||
} | ||
|
||
impl<St> Buffered<St> | ||
impl<St, F> Buffered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
pub(super) fn new(stream: St, n: Option<usize>) -> Self { | ||
Self { | ||
stream: super::Fuse::new(stream), | ||
in_progress_queue: FuturesOrdered::new(), | ||
ready_queue: VecDeque::new(), | ||
max: n.and_then(NonZeroUsize::new), | ||
} | ||
} | ||
|
||
delegate_access_inner!(stream, St, (.)); | ||
} | ||
|
||
impl<St> Stream for Buffered<St> | ||
impl<St, F> Stream for Buffered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
type Item = <St::Item as Future>::Output; | ||
|
||
|
@@ -74,14 +76,21 @@ where | |
} | ||
} | ||
|
||
// Attempt to pull the next value from the in_progress_queue | ||
let res = this.in_progress_queue.poll_next_unpin(cx); | ||
if let Some(val) = ready!(res) { | ||
return Poll::Ready(Some(val)); | ||
// Try to poll all ready futures in the in_progress_queue. | ||
loop { | ||
match this.in_progress_queue.poll_next_unpin(cx) { | ||
Poll::Ready(Some(output)) => { | ||
this.ready_queue.push_back(output); | ||
} | ||
Poll::Ready(None) => break, | ||
Poll::Pending => break, | ||
} | ||
} | ||
|
||
// If more values are still coming from the stream, we're not done yet | ||
if this.stream.is_done() { | ||
if let Some(output) = this.ready_queue.pop_front() { | ||
// If we have any ready outputs, return the first one. | ||
Poll::Ready(Some(output)) | ||
} else if this.stream.is_done() && this.in_progress_queue.is_empty() { | ||
Poll::Ready(None) | ||
} else { | ||
Poll::Pending | ||
|
@@ -100,10 +109,10 @@ where | |
} | ||
} | ||
|
||
impl<St> FusedStream for Buffered<St> | ||
impl<St, F> FusedStream for Buffered<St, F> | ||
where | ||
St: Stream, | ||
St::Item: Future, | ||
St: Stream<Item = F>, | ||
F: Future, | ||
{ | ||
fn is_terminated(&self) -> bool { | ||
self.stream.is_done() && self.in_progress_queue.is_terminated() | ||
|
@@ -112,10 +121,10 @@ where | |
|
||
// Forwarding impl of Sink from the underlying stream | ||
#[cfg(feature = "sink")] | ||
impl<S, Item> Sink<Item> for Buffered<S> | ||
impl<S, F, Item> Sink<Item> for Buffered<S, F> | ||
where | ||
S: Stream + Sink<Item>, | ||
S::Item: Future, | ||
S: Stream<Item = F> + Sink<Item>, | ||
F: Future, | ||
{ | ||
type Error = S::Error; | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can we fix this in
FuturesUnordered
? This could actually be a breaking change.Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I've tried to implement it in andylokandy@1d9fe21.
The short answer is no -- many components in futures -rs relies on the feature of
FuturesUnordered
not greedily poll internal futures. Instead,buffered
is supposed to provide the ability to maximize the concurrency respecting the limit.