Add for_each_concurrent

Author: cramertj

futures-util/src/stream/for_each_concurrent.rs

@@ -0,0 +1,77 @@
+use futures_core::{Async, Future, IntoFuture, Poll, Stream};
+use futures_core::task;
+
+use super::futures_unordered::FuturesUnordered;
+
+/// A stream combinator which executes a unit closure over each item on a
+/// stream concurrently.
+///
+/// This structure is returned by the `Stream::for_each` method.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct ForEachConcurrent<S, U, F> where U: IntoFuture {
+    stream: Option<S>,
+    stream_done: bool,
+    f: F,
+    futures: FuturesUnordered<U::Future>,
+}
+
+pub fn new<S, U, F>(s: S, f: F) -> ForEachConcurrent<S, U, F>
+    where S: Stream,
+          F: FnMut(S::Item) -> U,
+          U: IntoFuture<Item = (), Error = S::Error>,
+{
+    ForEachConcurrent {
+        stream: Some(s),
+        stream_done: false,
+        f: f,
+        futures: FuturesUnordered::new(),
+    }
+}
+
+impl<S, U, F> Future for ForEachConcurrent<S, U, F>
+    where S: Stream,
+          F: FnMut(S::Item) -> U,
+          U: IntoFuture<Item= (), Error = S::Error>,
+{
+    type Item = S;
+    type Error = S::Error;
+
+    fn poll(&mut self, cx: &mut task::Context) -> Poll<S, S::Error> {
+        loop {
+            let mut made_progress_this_iter = false;
+
+            // Try and pull an item off of the stream
+            if !self.stream_done {
+                // `unwrap` is valid because the stream is only taken after `stream_done` is set
+                match self.stream.as_mut().unwrap().poll_next(cx)? {
+                    Async::Ready(Some(x)) => {
+                        self.futures.push((self.f)(x).into_future());
+                        made_progress_this_iter = true;
+                    }
+                    // The stream completed, so it shouldn't be polled
+                    // anymore.
+                    Async::Ready(None) => self.stream_done = true,
+                    Async::Pending => {},
+                }
+            }
+
+            match self.futures.poll_next(cx)? {
+                Async::Ready(Some(())) => made_progress_this_iter = true,
+                Async::Ready(None) if self.stream_done => {
+                    // We've processed all of self.futures and self.stream,
+                    // so return self.stream
+                    return Ok(Async::Ready(self.stream.take().expect(
+                        "polled for_each_concurrent after completion"
+                    )));
+                }
+                Async::Ready(None)
+                | Async::Pending => {}
+            }
+
+            if !made_progress_this_iter {
+                return Ok(Async::Pending);
+            }
+        }
+    }
+}

futures-util/src/stream/mod.rs

@@ -85,6 +85,7 @@ if_std! {
     mod catch_unwind;
     mod chunks;
     mod collect;
+    mod for_each_concurrent;
     mod select_all;
     mod split;
     mod futures_unordered;
@@ -96,6 +97,7 @@ if_std! {
     pub use self::collect::Collect;
     pub use self::select_all::{select_all, SelectAll};
     pub use self::split::{SplitStream, SplitSink, ReuniteError};
+    pub use self::for_each_concurrent::ForEachConcurrent;
     pub use self::futures_unordered::{futures_unordered, FuturesUnordered};
     pub use self::futures_ordered::{futures_ordered, FuturesOrdered};
 }
@@ -584,10 +586,10 @@ pub trait StreamExt: Stream {
     /// to successfully, producing a future. That future will then be executed
     /// to completion before moving on to the next item.
     ///
-    /// The returned value is a `Future` where the `Item` type is `()` and
-    /// errors are otherwise threaded through. Any error on the stream or in the
-    /// closure will cause iteration to be halted immediately and the future
-    /// will resolve to that error.
+    /// The returned value is a `Future` where the `Item` type is the completed
+    /// stream, and errors are otherwise threaded through. Any error on the
+    /// stream or in the provided future will cause iteration to be halted
+    /// immediately and the future will resolve to that error.
     ///
     /// To process each item in the stream and produce another stream instead
     /// of a single future, use `and_then` instead.
@@ -599,6 +601,30 @@ pub trait StreamExt: Stream {
         for_each::new(self, f)
     }
 
+    /// Runs this stream to completion, executing the provided closure for each
+    /// element on the stream. This is similar to `for_each` but may begin
+    /// processing an element while previous elements are still being processed.
+    ///
+    /// When this stream successfully resolves to an item, the closure will be
+    /// called to produce a future. That future will then be added to
+    /// the set of futures to resolve.
+    ///
+    /// The returned value is a `Future` where the `Item` type is the completed
+    /// stream, and errors are otherwise threaded through. Any error on the
+    /// stream or in the provided future will cause iteration to be halted
+    /// immediately and the future will resolve to that error.
+    ///
+    /// To process each item in the stream and produce another stream instead
+    /// of a single future, use `and_then` instead.
+    #[cfg(feature = "std")]
+    fn for_each_concurrent<U, F>(self, f: F) -> ForEachConcurrent<Self, U, F>
+        where F: FnMut(Self::Item) -> U,
+              U: IntoFuture<Item=(), Error = Self::Error>,
+              Self: Sized
+    {
+        for_each_concurrent::new(self, f)
+    }
+
     /// Map this stream's error to a different type using the `Into` trait.
     ///
     /// This function does for streams what `try!` does for `Result`,

futures/tests/stream.rs

@@ -86,6 +86,28 @@ fn fold() {
     assert_done(|| err_list().fold(0, |a, b| ok::<i32, u32>(a + b)), Err(3));
 }
 
+#[test]
+fn for_each_concurrent() {
+    let (sender, receiver) = oneshot::channel::<()>();
+    let (sender, receiver) = (&mut Some(sender), &mut Some(receiver));
+    let fut = list().for_each_concurrent(move |num| {
+        match num {
+            // The first future is added
+            1 => receiver.take().unwrap().map_err(|_| 0).left_future(),
+            // Second future is added and completes immediately
+            2 => ok::<(), _>(()).right_future(),
+            // Third future is added, which when run completes the first
+            3 => {
+                sender.take().unwrap().send(()).unwrap();
+                ok::<(), _>(()).right_future()
+            }
+            _ => panic!(),
+        }
+    }).map(|_| ());
+
+    assert_done(|| fut, Ok(()));
+}
+
 #[test]
 fn filter() {
     assert_done(|| list().filter(|a| ok(*a % 2 == 0)).collect(), Ok(vec![2]));