fix: spawn tasks in Stream::buffered and Stream::buffer_unordered to max concurrency

Author: andylokandy

futures-util/src/stream/stream/buffer_unordered.rs

@@ -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,28 @@ 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 {
+            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 we have any ready outputs, return the first one.
+        if let Some(output) = this.ready_queue.pop() {
+            // If there are still ready outputs, wake the task to poll again.
+            if !this.ready_queue.is_empty() {
+                cx.waker().wake_by_ref();
+            }
+
+            return Poll::Ready(Some(output));
         }
 
-        // If more values are still coming from the stream, we're not done yet
-        if this.stream.is_done() {
+        if this.stream.is_done() && this.in_progress_queue.is_empty() {
             Poll::Ready(None)
         } else {
             Poll::Pending
@@ -98,10 +117,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 +129,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;
 

futures-util/src/stream/stream/buffered.rs

@@ -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,
     {
         #[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,28 @@ 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 we have any ready outputs, return the first one.
+        if let Some(output) = this.ready_queue.pop_front() {
+            // If there are still ready outputs, wake the task to poll again.
+            if !this.ready_queue.is_empty() {
+                cx.waker().wake_by_ref();
+            }
+
+            return Poll::Ready(Some(output));
         }
 
-        // If more values are still coming from the stream, we're not done yet
-        if this.stream.is_done() {
+        if this.stream.is_done() && this.in_progress_queue.is_empty() {
             Poll::Ready(None)
         } else {
             Poll::Pending
@@ -100,10 +116,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 +128,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;
 

futures-util/src/stream/stream/mod.rs

@@ -1487,10 +1487,11 @@ pub trait StreamExt: Stream {
     /// library is activated, and it is activated by default.
     #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
     #[cfg(feature = "alloc")]
-    fn buffered(self, n: impl Into<Option<usize>>) -> Buffered<Self>
+    fn buffered<F>(self, n: impl Into<Option<usize>>) -> Buffered<Self, F>
     where
-        Self::Item: Future,
         Self: Sized,
+        Self: Stream<Item = F>,
+        F: Future,
     {
         assert_stream::<<Self::Item as Future>::Output, _>(Buffered::new(self, n.into()))
     }
@@ -1536,10 +1537,11 @@ pub trait StreamExt: Stream {
     /// ```
     #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
     #[cfg(feature = "alloc")]
-    fn buffer_unordered(self, n: impl Into<Option<usize>>) -> BufferUnordered<Self>
+    fn buffer_unordered<F>(self, n: impl Into<Option<usize>>) -> BufferUnordered<Self, F>
     where
-        Self::Item: Future,
         Self: Sized,
+        Self: Stream<Item = F>,
+        F: Future,
     {
         assert_stream::<<Self::Item as Future>::Output, _>(BufferUnordered::new(self, n.into()))
     }

futures/tests/auto_traits.rs

@@ -1112,24 +1112,27 @@ mod stream {
     assert_not_impl!(AndThen<PhantomPinned, (), ()>: Unpin);
     assert_not_impl!(AndThen<(), PhantomPinned, ()>: Unpin);
 
-    assert_impl!(BufferUnordered<SendStream<()>>: Send);
-    assert_not_impl!(BufferUnordered<SendStream>: Send);
-    assert_not_impl!(BufferUnordered<LocalStream>: Send);
-    assert_impl!(BufferUnordered<SyncStream<()>>: Sync);
-    assert_not_impl!(BufferUnordered<SyncStream>: Sync);
-    assert_not_impl!(BufferUnordered<LocalStream>: Sync);
-    assert_impl!(BufferUnordered<UnpinStream>: Unpin);
-    assert_not_impl!(BufferUnordered<PinnedStream>: Unpin);
-
-    assert_impl!(Buffered<SendStream<SendFuture<()>>>: Send);
-    assert_not_impl!(Buffered<SendStream<SendFuture>>: Send);
-    assert_not_impl!(Buffered<SendStream<LocalFuture>>: Send);
-    assert_not_impl!(Buffered<LocalStream<SendFuture<()>>>: Send);
-    assert_impl!(Buffered<SyncStream<SendSyncFuture<()>>>: Sync);
-    assert_not_impl!(Buffered<SyncStream<SyncFuture<()>>>: Sync);
-    assert_not_impl!(Buffered<LocalStream<SendSyncFuture<()>>>: Sync);
-    assert_impl!(Buffered<UnpinStream<PinnedFuture>>: Unpin);
-    assert_not_impl!(Buffered<PinnedStream<PinnedFuture>>: Unpin);
+    assert_impl!(BufferUnordered<SendStream<SendFuture<()>>, SendFuture<()>>: Send);
+    assert_not_impl!(BufferUnordered<SendStream<SendFuture>, SendFuture>: Send);
+    assert_not_impl!(BufferUnordered<SendStream<LocalFuture>, LocalFuture>: Send);
+    assert_not_impl!(BufferUnordered<LocalStream<LocalFuture>, LocalFuture>: Send);
+    assert_impl!(BufferUnordered<SyncStream<SendSyncFuture<()>>, SendSyncFuture<()>>: Sync);
+    assert_not_impl!(BufferUnordered<SyncStream<SyncFuture<()>>, SyncFuture<()>>: Sync);
+    assert_not_impl!(BufferUnordered<SyncStream<LocalFuture>, LocalFuture>: Sync);
+    assert_not_impl!(BufferUnordered<LocalStream<LocalFuture>, LocalFuture>: Sync);
+    assert_impl!(BufferUnordered<UnpinStream<UnpinFuture>, UnpinFuture>: Unpin);
+    assert_not_impl!(BufferUnordered<PinnedStream<PinnedFuture>, PinnedFuture>: Unpin);
+
+    assert_impl!(Buffered<SendStream<SendFuture<()>>, SendFuture<()>>: Send);
+    assert_not_impl!(Buffered<SendStream<SendFuture>, SendFuture>: Send);
+    assert_not_impl!(Buffered<SendStream<LocalFuture>, LocalFuture>: Send);
+    assert_not_impl!(Buffered<LocalStream<LocalFuture>, LocalFuture>: Send);
+    assert_impl!(Buffered<SyncStream<SendSyncFuture<()>>, SendSyncFuture<()>>: Sync);
+    assert_not_impl!(Buffered<SyncStream<SyncFuture<()>>, SyncFuture<()>>: Sync);
+    assert_not_impl!(Buffered<SyncStream<LocalFuture>, LocalFuture>: Sync);
+    assert_not_impl!(Buffered<LocalStream<LocalFuture>, LocalFuture>: Sync);
+    assert_impl!(Buffered<UnpinStream<UnpinFuture>, UnpinFuture>: Unpin);
+    assert_not_impl!(Buffered<PinnedStream<PinnedFuture>, PinnedFuture>: Unpin);
 
     assert_impl!(CatchUnwind<SendStream>: Send);
     assert_not_impl!(CatchUnwind<LocalStream>: Send);