Add for_each_concurrent and try_for_each_concurrent

Author: cramertj

futures-util/src/stream/for_each_concurrent.rs

@@ -0,0 +1,106 @@
+use crate::stream::{FuturesUnordered, StreamExt};
+use core::marker::Unpin;
+use core::mem::PinMut;
+use core::num::NonZeroUsize;
+use futures_core::future::Future;
+use futures_core::stream::Stream;
+use futures_core::task::{self, Poll};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// A stream combinator which executes a unit closure over each item on a
+/// stream concurrently.
+///
+/// This structure is returned by the
+/// [`StreamExt::for_each_concurrent`](super::StreamExt::for_each_concurrent)
+/// method.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct ForEachConcurrent<St, Fut, F> {
+    stream: Option<St>,
+    f: F,
+    futures: FuturesUnordered<Fut>,
+    limit: Option<NonZeroUsize>,
+}
+
+impl<St, Fut, F> Unpin for ForEachConcurrent<St, Fut, F>
+where St: Unpin,
+      Fut: Unpin,
+{}
+
+impl<St, Fut, F> ForEachConcurrent<St, Fut, F>
+where St: Stream,
+      F: FnMut(St::Item) -> Fut,
+      Fut: Future<Output = ()>,
+{
+    unsafe_pinned!(stream: Option<St>);
+    unsafe_unpinned!(f: F);
+    unsafe_unpinned!(futures: FuturesUnordered<Fut>);
+    unsafe_unpinned!(limit: Option<NonZeroUsize>);
+
+    pub(super) fn new(stream: St, limit: Option<usize>, f: F) -> ForEachConcurrent<St, Fut, F> {
+        ForEachConcurrent {
+            stream: Some(stream),
+            // Note: `limit` = 0 gets ignored.
+            limit: limit.and_then(NonZeroUsize::new),
+            f,
+            futures: FuturesUnordered::new(),
+        }
+    }
+}
+
+impl<St, Fut, F> Future for ForEachConcurrent<St, Fut, F>
+    where St: Stream,
+          F: FnMut(St::Item) -> Fut,
+          Fut: Future<Output = ()>,
+{
+    type Output = ();
+
+    fn poll(mut self: PinMut<Self>, cx: &mut task::Context) -> Poll<()> {
+        loop {
+            let mut made_progress_this_iter = false;
+
+            // Try and pull an item from the stream
+            let current_len = self.futures().len();
+            // Check if we've already created a number of futures greater than `limit`
+            if self.limit().map(|limit| limit.get() > current_len).unwrap_or(true) {
+                let mut stream_completed = false;
+                let elem = if let Some(stream) = self.stream().as_pin_mut() {
+                    match stream.poll_next(cx) {
+                        Poll::Ready(Some(elem)) => {
+                            made_progress_this_iter = true;
+                            Some(elem)
+                        },
+                        Poll::Ready(None) => {
+                            stream_completed = true;
+                            None
+                        }
+                        Poll::Pending => None,
+                    }
+                } else {
+                    None
+                };
+                if stream_completed {
+                    PinMut::set(self.stream(), None);
+                }
+                if let Some(elem) = elem {
+                    let next_future = (self.f())(elem);
+                    self.futures().push(next_future);
+                }
+            }
+
+            match self.futures().poll_next_unpin(cx) {
+                Poll::Ready(Some(())) => made_progress_this_iter = true,
+                Poll::Ready(None) => {
+                    if self.stream().is_none() {
+                        return Poll::Ready(())
+                    }
+                },
+                Poll::Pending => {}
+            }
+
+            if !made_progress_this_iter {
+                return Poll::Pending;
+            }
+        }
+    }
+}

futures-util/src/stream/mod.rs

@@ -111,6 +111,9 @@ if_std! {
     mod collect;
     pub use self::collect::Collect;
 
+    mod for_each_concurrent;
+    pub use self::for_each_concurrent::ForEachConcurrent;
+
     mod futures_ordered;
     pub use self::futures_ordered::{futures_ordered, FuturesOrdered};
 
@@ -552,6 +555,65 @@ pub trait StreamExt: Stream {
         ForEach::new(self, f)
     }
 
+    /// Runs this stream to completion, executing the provided asynchronous
+    /// closure for each element on the stream concurrently as elements become
+    /// available.
+    ///
+    /// This is similar to [`StreamExt::for_each`], but the futures
+    /// produced by the closure are run concurrently (but not in parallel--
+    /// this combinator does not introduce any threads).
+    ///
+    /// The closure provided will be called for each item this stream produces,
+    /// yielding a future. That future will then be executed to completion
+    /// concurrently with the other futures produced by the closure.
+    ///
+    /// The first argument is an optional limit on the number of concurrent
+    /// futures. If this limit is not `None`, no more than `limit` futures
+    /// will be run concurrently. The `limit` argument is of type
+    /// `Into<Option<usize>>`, and so can be provided as either `None`,
+    /// `Some(10)`, or just `10`. Note: a limit of zero is interpreted as
+    /// no limit at all, and will have the same result as passing in `None`.
+    ///
+    /// This method is only available when the `std` feature of this
+    /// library is activated, and it is activated by default.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await, await_macro)]
+    /// # futures::executor::block_on(async {
+    /// use futures::channel::oneshot;
+    /// use futures::stream::{self, StreamExt};
+    ///
+    /// let (tx1, rx1) = oneshot::channel();
+    /// let (tx2, rx2) = oneshot::channel();
+    /// let (tx3, rx3) = oneshot::channel();
+    ///
+    /// let fut = stream::iter(vec![rx1, rx2, rx3]).for_each_concurrent(
+    ///     /* limit */ 2,
+    ///     async move |rx| {
+    ///         await!(rx).unwrap();
+    ///     }
+    /// );
+    /// tx1.send(()).unwrap();
+    /// tx2.send(()).unwrap();
+    /// tx3.send(()).unwrap();
+    /// await!(fut);
+    /// # })
+    /// ```
+    #[cfg(feature = "std")]
+    fn for_each_concurrent<Fut, F>(
+        self,
+        limit: impl Into<Option<usize>>,
+        f: F,
+    ) -> ForEachConcurrent<Self, Fut, F>
+        where F: FnMut(Self::Item) -> Fut,
+              Fut: Future<Output = ()>,
+              Self: Sized,
+    {
+        ForEachConcurrent::new(self, limit.into(), f)
+    }
+
     /// Creates a new stream of at most `n` items of the underlying stream.
     ///
     /// Once `n` items have been yielded from this stream then it will always

futures-util/src/try_stream/mod.rs

@@ -46,6 +46,10 @@ if_std! {
 
     mod try_collect;
     pub use self::try_collect::TryCollect;
+
+    mod try_for_each_concurrent;
+    pub use self::try_for_each_concurrent::TryForEachConcurrent;
+    use futures_core::future::Future;
 }
 
 impl<S: TryStream> TryStreamExt for S {}
@@ -255,6 +259,61 @@ pub trait TryStreamExt: TryStream {
         TrySkipWhile::new(self, f)
     }
 
+    /// Attempts to run this stream to completion, executing the provided asynchronous
+    /// closure for each element on the stream concurrently as elements become
+    /// available, exiting as soon as an error occurs.
+    ///
+    /// This is similar to
+    /// [`StreamExt::for_each_concurrent`](super::StreamExt::for_each_concurrent),
+    /// but will resolve to an error immediately if the underlying stream or the provided
+    /// closure return an error.
+    ///
+    /// This method is only available when the `std` feature of this
+    /// library is activated, and it is activated by default.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(async_await, await_macro)]
+    /// # futures::executor::block_on(async {
+    /// use futures::channel::oneshot;
+    /// use futures::stream::{self, StreamExt, TryStreamExt};
+    ///
+    /// let (tx1, rx1) = oneshot::channel();
+    /// let (tx2, rx2) = oneshot::channel();
+    /// let (_tx3, rx3) = oneshot::channel();
+    ///
+    /// let stream = stream::iter(vec![rx1, rx2, rx3]);
+    /// let fut = stream.map(Ok).try_for_each_concurrent(
+    ///     /* limit */ 2,
+    ///     async move |rx| {
+    ///         let res: Result<(), oneshot::Canceled> = await!(rx);
+    ///         res
+    ///     }
+    /// );
+    ///
+    /// tx1.send(()).unwrap();
+    /// // Drop the second sender so that `rx2` resolves to `Canceled`.
+    /// drop(tx2);
+    ///
+    /// // The final result is an error because the second future
+    /// // resulted in an error.
+    /// assert_eq!(Err(oneshot::Canceled), await!(fut));
+    /// # })
+    /// ```
+    #[cfg(feature = "std")]
+    fn try_for_each_concurrent<Fut, F>(
+        self,
+        limit: impl Into<Option<usize>>,
+        f: F,
+    ) -> TryForEachConcurrent<Self, Fut, F>
+        where F: FnMut(Self::Ok) -> Fut,
+              Fut: Future<Output = Result<(), Self::Error>>,
+              Self: Sized,
+    {
+        TryForEachConcurrent::new(self, limit.into(), f)
+    }
+
     /// Attempt to Collect all of the values of this stream into a vector,
     /// returning a future representing the result of that computation.
     ///

futures-util/src/try_stream/try_for_each_concurrent.rs

@@ -0,0 +1,106 @@
+use crate::stream::{FuturesUnordered, StreamExt};
+use core::marker::Unpin;
+use core::mem::PinMut;
+use core::num::NonZeroUsize;
+use futures_core::future::Future;
+use futures_core::stream::TryStream;
+use futures_core::task::{self, Poll};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// A stream combinator which executes a unit closure over each item on a
+/// stream concurrently.
+///
+/// This structure is returned by the
+/// [`TryStreamExt::try_for_each_concurrent`](super::TryStreamExt::try_for_each_concurrent)
+/// method.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct TryForEachConcurrent<St, Fut, F> {
+    stream: Option<St>,
+    f: F,
+    futures: FuturesUnordered<Fut>,
+    limit: Option<NonZeroUsize>,
+}
+
+impl<St, Fut, F> Unpin for TryForEachConcurrent<St, Fut, F>
+where St: Unpin,
+      Fut: Unpin,
+{}
+
+impl<St, Fut, F> TryForEachConcurrent<St, Fut, F>
+where St: TryStream,
+      F: FnMut(St::Ok) -> Fut,
+      Fut: Future<Output = Result<(), St::Error>>,
+{
+    unsafe_pinned!(stream: Option<St>);
+    unsafe_unpinned!(f: F);
+    unsafe_unpinned!(futures: FuturesUnordered<Fut>);
+    unsafe_unpinned!(limit: Option<NonZeroUsize>);
+
+    pub(super) fn new(stream: St, limit: Option<usize>, f: F) -> TryForEachConcurrent<St, Fut, F> {
+        TryForEachConcurrent {
+            stream: Some(stream),
+            // Note: `limit` = 0 gets ignored.
+            limit: limit.and_then(NonZeroUsize::new),
+            f,
+            futures: FuturesUnordered::new(),
+        }
+    }
+}
+
+impl<St, Fut, F> Future for TryForEachConcurrent<St, Fut, F>
+    where St: TryStream,
+          F: FnMut(St::Ok) -> Fut,
+          Fut: Future<Output = Result<(), St::Error>>,
+{
+    type Output = Result<(), St::Error>;
+
+    fn poll(mut self: PinMut<Self>, cx: &mut task::Context) -> Poll<Self::Output> {
+        loop {
+            let mut made_progress_this_iter = false;
+
+            // Try and pull an item from the stream
+            let current_len = self.futures().len();
+            // Check if we've already created a number of futures greater than `limit`
+            if self.limit().map(|limit| limit.get() > current_len).unwrap_or(true) {
+                let mut stream_completed = false;
+                let elem = if let Some(stream) = self.stream().as_pin_mut() {
+                    match stream.try_poll_next(cx)? {
+                        Poll::Ready(Some(elem)) => {
+                            made_progress_this_iter = true;
+                            Some(elem)
+                        },
+                        Poll::Ready(None) => {
+                            stream_completed = true;
+                            None
+                        }
+                        Poll::Pending => None,
+                    }
+                } else {
+                    None
+                };
+                if stream_completed {
+                    PinMut::set(self.stream(), None);
+                }
+                if let Some(elem) = elem {
+                    let next_future = (self.f())(elem);
+                    self.futures().push(next_future);
+                }
+            }
+
+            match self.futures().poll_next_unpin(cx)? {
+                Poll::Ready(Some(())) => made_progress_this_iter = true,
+                Poll::Ready(None) => {
+                    if self.stream().is_none() {
+                        return Poll::Ready(Ok(()))
+                    }
+                },
+                Poll::Pending => {}
+            }
+
+            if !made_progress_this_iter {
+                return Poll::Pending;
+            }
+        }
+    }
+}