Repository: alexcrichton/futures-rs
Pull request #536 (2017-07-20 06:16:13)

Add mpsc::spawn function.
Add a function to spawn a `Stream` in an `Executor`.

This is related to issue #470


Patch:
diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs
index 670d7384..b7e1eadc 100644
--- a/src/sync/mpsc/mod.rs
+++ b/src/sync/mpsc/mod.rs
@@ -78,7 +78,9 @@ use std::usize;
 
 use sync::mpsc::queue::{Queue, PopResult};
 use task::{self, Task};
-use {Async, AsyncSink, Poll, StartSend, Sink, Stream};
+use future::Executor;
+use sink::SendAll;
+use {Async, AsyncSink, Future, Poll, StartSend, Sink, Stream};
 
 mod queue;
 
@@ -824,6 +826,141 @@ impl<T> Stream for UnboundedReceiver<T> {
     }
 }
 
+/// Handle returned from the `spawn` function.
+///
+/// This handle is a stream that proxies a stream on a seperate `Executor`.
+/// Created through the `mpsc::spawn` function, this handle will produce
+/// the same values as the proxied stream, as they are produced in the executor,
+/// and uses a limited buffer to exert back-pressure on the remote stream.
+///
+/// If this handle is dropped, then the stream will no longer be polled and is
+/// scheduled to be dropped.
+pub struct SpawnHandle<Item, Error> {
+    rx: Receiver<Result<Item, Error>>
+}
+
+/// Type of future which `Executor` instances must be able to execute for `spawn`.
+pub struct Execute<S: Stream> {
+    inner: SendAll<Sender<Result<S::Item, S::Error>>, ResultStream<S>>
+}
+
+struct ResultStream<S: Stream>(S);
+
+/// Spawns a `stream` onto the instance of `Executor` provided, `executor`,
+/// returning a handle representing the remote stream.
+///
+/// The `stream` will be canceled if the `SpawnHandle` is dropped.
+///
+/// The `SpawnHandle` returned is a stream that is a proxy for `stream` itself.
+/// When `stream` has additional items available, then the `SpawnHandle`
+/// will have those same items available.
+///
+/// At most `buffer + 1` elements will be buffered at a time. If the buffer
+/// is full, then `stream` will stop progressing until more space is available.
+/// This allows the `SpawnHandle` to exert backpressure on the `stream`.
+///
+/// # Panics
+///
+/// This function will panic if `executor` is unable spawn a `Future` containing
+/// the entirety of the `stream`.
+pub fn spawn<S, E>(stream: S, executor: &E, buffer: usize) -> SpawnHandle<S::Item, S::Error>
+    where S: Stream,
+          E: Executor<Execute<S>>
+{
+    let (tx, rx) = channel(buffer);
+    executor.execute(Execute {
+        inner: tx.send_all(ResultStream(stream))
+    }).expect("failed to spawn stream");
+    SpawnHandle {
+        rx: rx
+    }
+}
+
+/// Spawns a `stream` onto the instance of `Executor` provided, `executor`,
+/// returning a handle representing the remote stream, with unbounded buffering.
+///
+/// The `stream` will be canceled if the `SpawnHandle` is dropped.
+///
+/// The `SpawnHandle` returned is a stream that is a proxy for `stream` itself.
+/// When `stream` has additional items available, then the `SpawnHandle`
+/// will have those same items available.
+///
+/// An unbounded buffer is used, which means that values will be buffered as
+/// fast as `stream` can produce them, without any backpressure. Therefore, if
+/// `stream` is an infinite stream, it can use an unbounded amount of memory, and
+/// potentially hog CPU resources.
+///
+/// # Panics
+///
+/// This function will panic if `executor` is unable spawn a `Future` containing
+/// the entirety of the `stream`.
+pub fn spawn_unbounded<S, E>(stream: S, executor: &E) -> SpawnHandle<S::Item, S::Error>
+    where S: Stream,
+          E: Executor<Execute<S>>
+{
+    let (tx, rx) = channel2(None);
+    executor.execute(Execute {
+        inner: tx.send_all(ResultStream(stream))
+    }).expect("failed to spawn stream");
+    SpawnHandle {
+        rx: rx
+    }
+}
+
+impl<I, E> Stream for SpawnHandle<I, E> {
+    type Item = I;
+    type Error = E;
+
+    fn poll(&mut self) -> Poll<Option<I>, E> {
+        match self.rx.poll() {
+            Ok(Async::Ready(Some(Ok(t)))) => Ok(Async::Ready(Some(t.into()))),
+            Ok(Async::Ready(Some(Err(e)))) => Err(e),
+            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            Err(_) => unreachable!("mpsc::Receiver should never return Err"),
+        }
+    }
+}
+
+impl<I, E> fmt::Debug for SpawnHandle<I, E> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("SpawnHandle")
+         .finish()
+    }
+}
+
+impl<S: Stream> Future for Execute<S> {
+    type Item = ();
+    type Error = ();
+
+    fn poll(&mut self) -> Poll<(), ()> {
+        match self.inner.poll() {
+            Ok(Async::NotReady) => Ok(Async::NotReady),
+            _ => Ok(Async::Ready(()))
+        }
+    }
+}
+
+impl<S: Stream> fmt::Debug for Execute<S> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("Execute")
+         .finish()
+    }
+}
+
+impl<S: Stream> Stream for ResultStream<S> {
+    type Item = Result<S::Item, S::Error>;
+    type Error = SendError<Self::Item>;
+    fn poll(&mut self) -> Poll<Option<Result<S::Item, S::Error>>, Self::Error> {
+        match self.0.poll() {
+            Ok(Async::Ready(Some(item))) => Ok(Async::Ready(Some(Ok(item)))),
+            Err(e) => Ok(Async::Ready(Some(Err(e)))),
+            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
+            Ok(Async::NotReady) => Ok(Async::NotReady)
+        }
+    }
+}
+
 /*
  *
  * ===== impl Inner =====


Test patch:
diff --git a/tests/mpsc.rs b/tests/mpsc.rs
index a411bcef..a4e114dc 100644
--- a/tests/mpsc.rs
+++ b/tests/mpsc.rs
@@ -3,7 +3,8 @@
 extern crate futures;
 
 use futures::{Future, Stream, Sink, Async, AsyncSink};
-use futures::future::lazy;
+use futures::future::{lazy, ok};
+use futures::stream::unfold;
 use futures::sync::mpsc;
 
 use std::time::Duration;
@@ -133,6 +134,15 @@ fn tx_close_gets_none() {
     }).wait().unwrap();
 }
 
+#[test]
+fn spawn_sends_items() {
+    let core = local_executor::Core::new();
+    let stream = unfold(0, |i| Some(ok::<_,u8>((i, i + 1))));
+    let rx = mpsc::spawn(stream, &core, 1);
+    assert_eq!(core.run(rx.take(4).collect()).unwrap(),
+               [0, 1, 2, 3]);
+}
+
 #[test]
 fn stress_shared_unbounded() {
     const AMT: u32 = 10000;