Merge pull request rust-lang#134 from Stebalien/unordered

[RFC] Add an unordered buffer stream adapter.

Author: alexcrichton

src/stream/buffer_unordered.rs

@@ -0,0 +1,85 @@
+use std::prelude::v1::*;
+
+use {Async, IntoFuture, Poll, Future};
+use stream::{Stream, Fuse};
+
+/// An adaptor for a stream of futures to execute the futures concurrently, if
+/// possible, delivering results as they become available.
+///
+/// This adaptor will buffer up a list of pending futures, and then return their
+/// results in the order that they complete. This is created by the
+/// `Stream::buffer_unordered` method.
+#[must_use = "streams do nothing unless polled"]
+pub struct BufferUnordered<S>
+    where S: Stream,
+          S::Item: IntoFuture,
+{
+    stream: Fuse<S>,
+    futures: Vec<Option<<S::Item as IntoFuture>::Future>>,
+    cur: usize,
+}
+
+pub fn new<S>(s: S, amt: usize) -> BufferUnordered<S>
+    where S: Stream,
+          S::Item: IntoFuture<Error=<S as Stream>::Error>,
+{
+    BufferUnordered {
+        stream: super::fuse::new(s),
+        futures: (0..amt).map(|_| None).collect(),
+        cur: 0,
+    }
+}
+
+impl<S> Stream for BufferUnordered<S>
+    where S: Stream,
+          S::Item: IntoFuture<Error=<S as Stream>::Error>,
+{
+    type Item = <S::Item as IntoFuture>::Item;
+    type Error = <S as Stream>::Error;
+
+    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        // First, try to fill in all the futures
+        for future in &mut self.futures {
+            if future.is_none() {
+                match try!(self.stream.poll()) {
+                    Async::Ready(Some(s)) => {
+                        *future = Some(s.into_future());
+                    }
+                    Async::Ready(None) => break,
+                    Async::NotReady => break,
+                }
+            }
+        }
+
+        // Next, try and step futures forward until we find a ready one.
+        // Always start at `cur` for fairness.
+        let mut waiting = false;
+        for i in 0..self.futures.len() {
+            let mut idx = self.cur + i;
+            if idx >= self.futures.len() {
+                idx -= self.futures.len();
+            }
+            let future = &mut self.futures[idx];
+            let result = match *future {
+                Some(ref mut s) => match s.poll() {
+                    Ok(Async::NotReady) => {
+                        waiting = true;
+                        continue
+                    },
+                    Ok(Async::Ready(e)) => Ok(Async::Ready(Some(e))),
+                    Err(e) => Err(e),
+                },
+                None => continue,
+            };
+            self.cur = i + 1;
+            *future = None;
+            return result;
+        }
+
+        Ok(if waiting || !self.stream.is_done() {
+            Async::NotReady
+        } else {
+            Async::Ready(None)
+        })
+    }
+}

src/stream/mod.rs

@@ -57,10 +57,12 @@ pub use self::peek::Peekable;
 
 if_std! {
     mod buffered;
+    mod buffer_unordered;
     mod channel;
     mod collect;
     mod wait;
     pub use self::buffered::Buffered;
+    pub use self::buffer_unordered::BufferUnordered;
     pub use self::channel::{channel, Sender, Receiver, FutureSender};
     pub use self::collect::Collect;
     pub use self::wait::Wait;
@@ -650,6 +652,24 @@ pub trait Stream {
         buffered::new(self, amt)
     }
 
+    /// An adaptor for creating a buffered list of pending futures (unordered).
+    ///
+    /// If this stream's item can be converted into a future, then this adaptor
+    /// will buffer up to `amt` futures and then return results in the order
+    /// in which they complete. No more than `amt` futures will be buffered at
+    /// any point in time, and less than `amt` may also be buffered depending on
+    /// the state of each future.
+    ///
+    /// The returned stream will be a stream of each future's result, with
+    /// errors passed through whenever they occur.
+    #[cfg(feature = "use_std")]
+    fn buffer_unordered(self, amt: usize) -> BufferUnordered<Self>
+        where Self::Item: IntoFuture<Error = <Self as Stream>::Error>,
+              Self: Sized
+    {
+        buffer_unordered::new(self, amt)
+    }
+
     /// An adapter for merging the output of two streams.
     ///
     /// The merged stream produces items from one or both of the underlying

tests/stream.rs

@@ -155,6 +155,45 @@ fn buffered() {
     assert_eq!(rx.next(), None);
 }
 
+#[test]
+fn unordered() {
+    let (tx, rx) = channel::<_, u32>();
+    let (a, b) = oneshot::<u32>();
+    let (c, d) = oneshot::<u32>();
+
+    tx.send(Ok(b.map_err(|_| 2).boxed()))
+      .and_then(|tx| tx.send(Ok(d.map_err(|_| 4).boxed())))
+      .forget();
+
+    let mut rx = rx.buffer_unordered(2);
+    sassert_empty(&mut rx);
+    let mut rx = rx.wait();
+    c.complete(3);
+    assert_eq!(rx.next(), Some(Ok(3)));
+    a.complete(5);
+    assert_eq!(rx.next(), Some(Ok(5)));
+    assert_eq!(rx.next(), None);
+
+    let (tx, rx) = channel::<_, u32>();
+    let (a, b) = oneshot::<u32>();
+    let (c, d) = oneshot::<u32>();
+
+    tx.send(Ok(b.map_err(|_| 2).boxed()))
+      .and_then(|tx| tx.send(Ok(d.map_err(|_| 4).boxed())))
+      .forget();
+
+    // We don't even get to see `c` until `a` completes.
+    let mut rx = rx.buffer_unordered(1);
+    sassert_empty(&mut rx);
+    c.complete(3);
+    sassert_empty(&mut rx);
+    a.complete(5);
+    let mut rx = rx.wait();
+    assert_eq!(rx.next(), Some(Ok(5)));
+    assert_eq!(rx.next(), Some(Ok(3)));
+    assert_eq!(rx.next(), None);
+}
+
 #[test]
 fn zip() {
     assert_done(|| list().zip(list()).collect(),