Repository: alexcrichton/futures-rs
Pull request #658 (2017-11-24 08:08:43)

test: make sure mpsc::spawn streams get dropped when receiver is dropped
`mpsc::spawn` says:

> The `stream` will be canceled if the `SpawnHandle` is dropped.

But actually the `stream` will only get dropped (my reading of "canceled") after it produced an item that couldn't be sent over the (now dead) channel.

This pull request adds two tests (one for `sync`, one for `unsync`) to make sure the `stream` actually gets dropped; the tests should be failing right now.

alexcrichton: Hm the tests are passing?
stbuehler: That's because I'm stupid :) I wanted to see whether they actually fail and added `#[should_panic]` - removed those lines, so they really should fail now.
stbuehler: I'll try.
stbuehler: At first I though it might be possible to query the "open" state of the sinks in the `Execute` futures. It is, but if the sink isn't parked (due to a pending send), the task won't get notified when the receiver is closed.  It probably isn't a good idea to modify this behavior or clone all those types just to add it.

I still have another idea :)
alexcrichton: I think implementation-wise this may be difficult to do as the same notification mechanism present on oneshots isn't present on channels, but it may still be possible via various trickery.
stbuehler: Sorry for abusing the same branch and pull request for a new feature :)

Basically `job` is a generic way of canceling futures by dropping a handle.  If there is a simpler way to do this I'm happy to look into it.
alexcrichton: Hm ok, thanks for the fix!

I wonder, though, the job idea here seems very similar to the oneshot-style concept, no? Would it be possible to perhaps leverage those existing concepts and have something like this bake on crates.io first before moving into this crate?

Patch:
diff --git a/src/sync/mpsc/mod.rs b/src/sync/mpsc/mod.rs
index 8ef2a5e5..71875796 100644
--- a/src/sync/mpsc/mod.rs
+++ b/src/sync/mpsc/mod.rs
@@ -77,6 +77,7 @@ use std::thread;
 use std::usize;
 
 use sync::mpsc::queue::{Queue, PopResult};
+use sync::oneshot;
 use task::{self, Task};
 use future::Executor;
 use sink::SendAll;
@@ -982,12 +983,14 @@ impl<T> Stream for UnboundedReceiver<T> {
 /// 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>>
+    rx: Receiver<Result<Item, Error>>,
+    _cancel_tx: oneshot::Sender<()>,
 }
 
 /// 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>>, Results<S, SendError<Result<S::Item, S::Error>>>>
+    inner: SendAll<Sender<Result<S::Item, S::Error>>, Results<S, SendError<Result<S::Item, S::Error>>>>,
+    cancel_rx: oneshot::Receiver<()>,
 }
 
 /// Spawns a `stream` onto the instance of `Executor` provided, `executor`,
@@ -1011,12 +1014,15 @@ pub fn spawn<S, E>(stream: S, executor: &E, buffer: usize) -> SpawnHandle<S::Ite
     where S: Stream,
           E: Executor<Execute<S>>
 {
+    let (cancel_tx, cancel_rx) = oneshot::channel();
     let (tx, rx) = channel(buffer);
     executor.execute(Execute {
-        inner: tx.send_all(resultstream::new(stream))
+        inner: tx.send_all(resultstream::new(stream)),
+        cancel_rx: cancel_rx,
     }).expect("failed to spawn stream");
     SpawnHandle {
-        rx: rx
+        rx: rx,
+        _cancel_tx: cancel_tx,
     }
 }
 
@@ -1042,12 +1048,15 @@ pub fn spawn_unbounded<S, E>(stream: S, executor: &E) -> SpawnHandle<S::Item, S:
     where S: Stream,
           E: Executor<Execute<S>>
 {
+    let (cancel_tx, cancel_rx) = oneshot::channel();
     let (tx, rx) = channel2(None);
     executor.execute(Execute {
-        inner: tx.send_all(resultstream::new(stream))
+        inner: tx.send_all(resultstream::new(stream)),
+        cancel_rx: cancel_rx,
     }).expect("failed to spawn stream");
     SpawnHandle {
-        rx: rx
+        rx: rx,
+        _cancel_tx: cancel_tx,
     }
 }
 
@@ -1078,6 +1087,10 @@ impl<S: Stream> Future for Execute<S> {
     type Error = ();
 
     fn poll(&mut self) -> Poll<(), ()> {
+        match self.cancel_rx.poll() {
+            Ok(Async::NotReady) => (),
+            _ => return Ok(Async::Ready(())),
+        }
         match self.inner.poll() {
             Ok(Async::NotReady) => Ok(Async::NotReady),
             _ => Ok(Async::Ready(()))
diff --git a/src/unsync/mpsc.rs b/src/unsync/mpsc.rs
index 0a5bf457..0cff01b1 100644
--- a/src/unsync/mpsc.rs
+++ b/src/unsync/mpsc.rs
@@ -16,6 +16,7 @@ use task::{self, Task};
 use future::Executor;
 use sink::SendAll;
 use resultstream::{self, Results};
+use unsync::oneshot;
 use {Async, AsyncSink, Future, Poll, StartSend, Sink, Stream};
 
 /// Creates a bounded in-memory channel with buffered storage.
@@ -355,12 +356,14 @@ impl<T> SendError<T> {
 /// If this handle is dropped, then the stream will no longer be polled and is
 /// scheduled to be dropped.
 pub struct SpawnHandle<Item, Error> {
-    inner: Receiver<Result<Item, Error>>
+    inner: Receiver<Result<Item, Error>>,
+    _cancel_tx: oneshot::Sender<()>,
 }
 
 /// 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>>, Results<S, SendError<Result<S::Item, S::Error>>>>
+    inner: SendAll<Sender<Result<S::Item, S::Error>>, Results<S, SendError<Result<S::Item, S::Error>>>>,
+    cancel_rx: oneshot::Receiver<()>,
 }
 
 /// Spawns a `stream` onto the instance of `Executor` provided, `executor`,
@@ -384,12 +387,15 @@ pub fn spawn<S, E>(stream: S, executor: &E, buffer: usize) -> SpawnHandle<S::Ite
     where S: Stream,
           E: Executor<Execute<S>>
 {
+    let (cancel_tx, cancel_rx) = oneshot::channel();
     let (tx, rx) = channel(buffer);
     executor.execute(Execute {
-        inner: tx.send_all(resultstream::new(stream))
+        inner: tx.send_all(resultstream::new(stream)),
+        cancel_rx: cancel_rx,
     }).expect("failed to spawn stream");
     SpawnHandle {
-        inner: rx
+        inner: rx,
+        _cancel_tx: cancel_tx,
     }
 }
 
@@ -417,12 +423,15 @@ pub fn spawn_unbounded<S,E>(stream: S, executor: &E) -> SpawnHandle<S::Item, S::
     where S: Stream,
           E: Executor<Execute<S>>
 {
+    let (cancel_tx, cancel_rx) = oneshot::channel();
     let (tx, rx) = channel_(None);
     executor.execute(Execute {
-        inner: tx.send_all(resultstream::new(stream))
+        inner: tx.send_all(resultstream::new(stream)),
+        cancel_rx: cancel_rx,
     }).expect("failed to spawn stream");
     SpawnHandle {
-        inner: rx
+        inner: rx,
+        _cancel_tx: cancel_tx,
     }
 }
 
@@ -453,6 +462,10 @@ impl<S: Stream> Future for Execute<S> {
     type Error = ();
 
     fn poll(&mut self) -> Poll<(), ()> {
+        match self.cancel_rx.poll() {
+            Ok(Async::NotReady) => (),
+            _ => return Ok(Async::Ready(())),
+        }
         match self.inner.poll() {
             Ok(Async::NotReady) => Ok(Async::NotReady),
             _ => Ok(Async::Ready(()))


Test patch:
diff --git a/tests/mpsc.rs b/tests/mpsc.rs
index a64ad9e0..8df98d49 100644
--- a/tests/mpsc.rs
+++ b/tests/mpsc.rs
@@ -151,6 +151,64 @@ fn spawn_sends_items() {
                [0, 1, 2, 3]);
 }
 
+#[test]
+fn spawn_kill_dead_stream() {
+    use std::thread;
+    use std::time::Duration;
+    use futures::future::Either;
+    use futures::sync::oneshot;
+
+    // a stream which never returns anything (maybe a remote end isn't
+    // responding), but dropping it leads to observable side effects
+    // (like closing connections, releasing limited resources, ...)
+    #[derive(Debug)]
+    struct Dead {
+        // when dropped you should get Err(oneshot::Canceled) on the
+        // receiving end
+        done: oneshot::Sender<()>,
+    }
+    impl Stream for Dead {
+        type Item = ();
+        type Error = ();
+
+        fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+            Ok(Async::NotReady)
+        }
+    }
+
+    // need to implement a timeout for the test, as it would hang
+    // forever right now
+    let (timeout_tx, timeout_rx) = oneshot::channel();
+    thread::spawn(move || {
+        thread::sleep(Duration::from_millis(1000));
+        let _ = timeout_tx.send(());
+    });
+
+    let core = local_executor::Core::new();
+    let (done_tx, done_rx) = oneshot::channel();
+    let stream = Dead{done: done_tx};
+    let rx = mpsc::spawn(stream, &core, 1);
+    let res = core.run(
+        Ok::<_, ()>(())
+        .into_future()
+        .then(move |_| {
+            // now drop the spawned stream: maybe some timeout exceeded,
+            // or some connection on this end was closed by the remote
+            // end.
+            drop(rx);
+            // and wait for the spawned stream to release its resources
+            done_rx
+        })
+        .select2(timeout_rx)
+    );
+    match res {
+        Err(Either::A((oneshot::Canceled, _))) => (),
+        _ => {
+            panic!("dead stream wasn't canceled");
+        },
+    }
+}
+
 #[test]
 fn stress_shared_unbounded() {
     const AMT: u32 = 10000;
diff --git a/tests/unsync.rs b/tests/unsync.rs
index 1ad0ed31..e86b95ca 100644
--- a/tests/unsync.rs
+++ b/tests/unsync.rs
@@ -121,3 +121,60 @@ fn spawn_sends_items() {
     assert_eq!(core.run(rx.take(4).collect()).unwrap(),
                [0, 1, 2, 3]);
 }
+
+#[test]
+fn spawn_kill_dead_stream() {
+    use std::thread;
+    use std::time::Duration;
+    use futures::future::Either;
+
+    // a stream which never returns anything (maybe a remote end isn't
+    // responding), but dropping it leads to observable side effects
+    // (like closing connections, releasing limited resources, ...)
+    #[derive(Debug)]
+    struct Dead {
+        // when dropped you should get Err(oneshot::Canceled) on the
+        // receiving end
+        done: oneshot::Sender<()>,
+    }
+    impl Stream for Dead {
+        type Item = ();
+        type Error = ();
+
+        fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+            Ok(Async::NotReady)
+        }
+    }
+
+    // need to implement a timeout for the test, as it would hang
+    // forever right now
+    let (timeout_tx, timeout_rx) = futures::sync::oneshot::channel();
+    thread::spawn(move || {
+        thread::sleep(Duration::from_millis(1000));
+        let _ = timeout_tx.send(());
+    });
+
+    let core = Core::new();
+    let (done_tx, done_rx) = oneshot::channel();
+    let stream = Dead{done: done_tx};
+    let rx = mpsc::spawn(stream, &core, 1);
+    let res = core.run(
+        Ok::<_, ()>(())
+        .into_future()
+        .then(move |_| {
+            // now drop the spawned stream: maybe some timeout exceeded,
+            // or some connection on this end was closed by the remote
+            // end.
+            drop(rx);
+            // and wait for the spawned stream to release its resources
+            done_rx
+        })
+        .select2(timeout_rx)
+    );
+    match res {
+        Err(Either::A((oneshot::Canceled, _))) => (),
+        _ => {
+            panic!("dead stream wasn't canceled");
+        },
+    }
+}