Add TryStreamExt::try_buffer_unordered

Author: Greg Bonik

futures-util/src/try_stream/mod.rs

@@ -28,6 +28,9 @@ pub use self::try_for_each::TryForEach;
 mod try_filter_map;
 pub use self::try_filter_map::TryFilterMap;
 
+mod try_buffer_unordered;
+pub use self::try_buffer_unordered::TryBufferUnordered;
+
 if_std! {
     mod try_collect;
     pub use self::try_collect::TryCollect;
@@ -292,4 +295,68 @@ pub trait TryStreamExt: TryStream {
     {
         TryFilterMap::new(self, f)
     }
+
+    /// Attempt to execute several futures from a stream concurrently.
+    ///
+    /// This stream's `Ok` type must be a [`TryFuture`] with an `Error` type
+    /// that matches the stream's `Error` type.
+    ///
+    /// This adaptor will buffer up to `n` futures and then return their
+    /// outputs in the order in which they complete. If the underlying stream
+    /// returns an error, it will be immediately propagated.
+    ///
+    /// The returned stream will be a stream of results, each containing either
+    /// an error or a future's output. An error can be produced either by the
+    /// underlying stream itself or by one of the futures it yielded.
+    ///
+    /// # Examples
+    ///
+    /// Results are returned in the order of completion:
+    /// ```
+    /// #![feature(async_await, await_macro)]
+    /// # futures::executor::block_on(async {
+    /// use futures::channel::oneshot;
+    /// use futures::stream::{self, StreamExt, TryStreamExt};
+    ///
+    /// let (send_one, recv_one) = oneshot::channel();
+    /// let (send_two, recv_two) = oneshot::channel();
+    ///
+    /// let stream_of_futures = stream::iter(vec![Ok(recv_one), Ok(recv_two)]);
+    ///
+    /// let mut buffered = stream_of_futures.try_buffer_unordered(10);
+    ///
+    /// send_two.send(2i32);
+    /// assert_eq!(await!(buffered.next()), Some(Ok(2i32)));
+    ///
+    /// send_one.send(1i32);
+    /// assert_eq!(await!(buffered.next()), Some(Ok(1i32)));
+    ///
+    /// assert_eq!(await!(buffered.next()), None);
+    /// # })
+    /// ```
+    ///
+    /// Errors from the underlying stream itself are propagated:
+    /// ```
+    /// #![feature(async_await, await_macro)]
+    /// # futures::executor::block_on(async {
+    /// use futures::channel::mpsc;
+    /// use futures::future;
+    /// use futures::stream::{StreamExt, TryStreamExt};
+    ///
+    /// let (sink, stream_of_futures) = mpsc::unbounded();
+    /// let mut buffered = stream_of_futures.try_buffer_unordered(10);
+    ///
+    /// sink.unbounded_send(Ok(future::ready(Ok(7i32))));
+    /// assert_eq!(await!(buffered.next()), Some(Ok(7i32)));
+    ///
+    /// sink.unbounded_send(Err("error in the stream"));
+    /// assert_eq!(await!(buffered.next()), Some(Err("error in the stream")));
+    /// # })
+    /// ```
+    fn try_buffer_unordered(self, n: usize) -> TryBufferUnordered<Self>
+        where Self::Ok: TryFuture<Error = Self::Error>,
+              Self: Sized
+    {
+        TryBufferUnordered::new(self, n)
+    }
 }

futures-util/src/try_stream/try_buffer_unordered.rs

@@ -0,0 +1,98 @@
+use crate::stream::{Fuse, FuturesUnordered, StreamExt};
+use crate::try_future::{IntoFuture, TryFutureExt};
+use crate::try_stream::IntoStream;
+use futures_core::future::TryFuture;
+use futures_core::stream::{Stream, TryStream};
+use futures_core::task::{self, Poll};
+use std::marker::Unpin;
+use std::mem::PinMut;
+
+/// A stream returned by the
+/// [`try_buffer_unordered`](super::TryStreamExt::try_buffer_unordered) method
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct TryBufferUnordered<St>
+    where St: TryStream
+{
+    stream: Fuse<IntoStream<St>>,
+    in_progress_queue: FuturesUnordered<IntoFuture<St::Ok>>,
+    max: usize,
+}
+
+impl<St> Unpin for TryBufferUnordered<St>
+    where St: TryStream + Unpin
+{}
+
+impl<St> TryBufferUnordered<St>
+    where St: TryStream,
+          St::Ok: TryFuture,
+{
+    unsafe_pinned!(stream: Fuse<IntoStream<St>>);
+    unsafe_unpinned!(in_progress_queue: FuturesUnordered<IntoFuture<St::Ok>>);
+
+    pub(super) fn new(stream: St, n: usize) -> Self {
+        TryBufferUnordered {
+            stream: IntoStream::new(stream).fuse(),
+            in_progress_queue: FuturesUnordered::new(),
+            max: n,
+        }
+    }
+
+    /// Acquires a reference to the underlying stream that this combinator is
+    /// pulling from.
+    pub fn get_ref(&self) -> &St {
+        self.stream.get_ref().get_ref()
+    }
+
+    /// Acquires a mutable reference to the underlying stream that this
+    /// combinator is pulling from.
+    ///
+    /// Note that care must be taken to avoid tampering with the state of the
+    /// stream which may otherwise confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        self.stream.get_mut().get_mut()
+    }
+
+    /// Consumes this combinator, returning the underlying stream.
+    ///
+    /// Note that this may discard intermediate state of this combinator, so
+    /// care should be taken to avoid losing resources when this is called.
+    pub fn into_inner(self) -> St {
+        self.stream.into_inner().into_inner()
+    }
+}
+
+impl<St> Stream for TryBufferUnordered<St>
+    where St: TryStream,
+          St::Ok: TryFuture<Error = St::Error>,
+{
+    type Item = Result<<St::Ok as TryFuture>::Ok, St::Error>;
+
+    fn poll_next(
+        mut self: PinMut<Self>,
+        cx: &mut task::Context,
+    ) -> Poll<Option<Self::Item>> {
+        // First up, try to spawn off as many futures as possible by filling up
+        // our slab of futures. Propagate errors from the stream immediately.
+        while self.in_progress_queue.len() < self.max {
+            match self.stream().poll_next(cx) {
+                Poll::Ready(Some(Ok(fut))) => self.in_progress_queue().push(fut.into_future()),
+                Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(e))),
+                Poll::Ready(None) | Poll::Pending => break,
+            }
+        }
+
+        // Attempt to pull the next value from the in_progress_queue
+        match PinMut::new(self.in_progress_queue()).poll_next(cx) {
+            x @ Poll::Pending | x @ Poll::Ready(Some(_)) => return x,
+            Poll::Ready(None) => {}
+        }
+
+        // If more values are still coming from the stream, we're not done yet
+        if self.stream.is_done() {
+            Poll::Ready(None)
+        } else {
+            Poll::Pending
+        }
+    }
+}