Do not hang in poll if reactor is destroyed

After channel receiver dropped, sender send returns an error. So
any further `poll` requests fail.

Author: stepancheg

src/channel.rs

@@ -93,13 +93,13 @@ impl<T> Stream for Receiver<T> {
     type Error = io::Error;
 
     fn poll(&mut self) -> Poll<Option<T>, io::Error> {
-        if let Async::NotReady = self.rx.poll_read() {
+        if let Async::NotReady = try!(self.rx.poll_read()) {
             return Ok(Async::NotReady)
         }
         match self.rx.get_ref().try_recv() {
             Ok(t) => Ok(Async::Ready(Some(t))),
             Err(TryRecvError::Empty) => {
-                self.rx.need_read();
+                try!(self.rx.need_read());
                 Ok(Async::NotReady)
             }
             Err(TryRecvError::Disconnected) => Ok(Async::Ready(None)),

src/io/mod.rs

@@ -5,7 +5,7 @@
 
 use std::io::{self, Read, Write};
 
-use futures::{BoxFuture, Async};
+use futures::{BoxFuture, Async, Poll};
 use futures::stream::BoxStream;
 
 /// A convenience typedef around a `Future` whose error component is `io::Error`
@@ -76,8 +76,8 @@ pub trait Io: Read + Write {
     ///
     /// This method is likely to panic if called from outside the context of a
     /// future's task.
-    fn poll_read(&mut self) -> Async<()> {
-        Async::Ready(())
+    fn poll_read(&mut self) -> Poll<(), io::Error> {
+        Ok(Async::Ready(()))
     }
 
     /// Tests to see if this I/O object may be writable.
@@ -98,8 +98,8 @@ pub trait Io: Read + Write {
     ///
     /// This method is likely to panic if called from outside the context of a
     /// future's task.
-    fn poll_write(&mut self) -> Async<()> {
-        Async::Ready(())
+    fn poll_write(&mut self) -> Poll<(), io::Error> {
+        Ok(Async::Ready(()))
     }
 
     /// Helper method for splitting this read/write object into two halves.

src/net/tcp.rs

@@ -76,7 +76,7 @@ impl TcpListener {
     }
 
     /// Test whether this socket is ready to be read or not.
-    pub fn poll_read(&self) -> Async<()> {
+    pub fn poll_read(&self) -> Poll<(), io::Error> {
         self.io.poll_read()
     }
 
@@ -103,13 +103,13 @@ impl TcpListener {
             type Error = io::Error;
 
             fn poll(&mut self) -> Poll<Option<Self::Item>, io::Error> {
-                if let Async::NotReady = self.inner.io.poll_read() {
+                if let Async::NotReady = try!(self.inner.io.poll_read()) {
                     return Ok(Async::NotReady)
                 }
                 match self.inner.io.get_ref().accept() {
                     Ok(pair) => Ok(Async::Ready(Some(pair))),
                     Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                        self.inner.io.need_read();
+                        try!(self.inner.io.need_read());
                         Ok(Async::NotReady)
                     }
                     Err(e) => Err(e)
@@ -128,7 +128,7 @@ impl TcpListener {
                     });
                     tx.complete(res);
                     Ok(())
-                });
+                }).expect("failed to spawn");
                 rx.then(|r| r.expect("shouldn't be canceled"))
             }).boxed(),
         }
@@ -266,7 +266,7 @@ impl TcpStream {
     /// get a notification when the socket does become readable. That is, this
     /// is only suitable for calling in a `Future::poll` method and will
     /// automatically handle ensuring a retry once the socket is readable again.
-    pub fn poll_read(&self) -> Async<()> {
+    pub fn poll_read(&self) -> Poll<(), io::Error> {
         self.io.poll_read()
     }
 
@@ -276,7 +276,7 @@ impl TcpStream {
     /// get a notification when the socket does become writable. That is, this
     /// is only suitable for calling in a `Future::poll` method and will
     /// automatically handle ensuring a retry once the socket is writable again.
-    pub fn poll_write(&self) -> Async<()> {
+    pub fn poll_write(&self) -> Poll<(), io::Error> {
         self.io.poll_write()
     }
 
@@ -379,11 +379,11 @@ impl Write for TcpStream {
 }
 
 impl Io for TcpStream {
-    fn poll_read(&mut self) -> Async<()> {
+    fn poll_read(&mut self) -> Poll<(), io::Error> {
         <TcpStream>::poll_read(self)
     }
 
-    fn poll_write(&mut self) -> Async<()> {
+    fn poll_write(&mut self) -> Poll<(), io::Error> {
         <TcpStream>::poll_write(self)
     }
 }
@@ -405,11 +405,11 @@ impl<'a> Write for &'a TcpStream {
 }
 
 impl<'a> Io for &'a TcpStream {
-    fn poll_read(&mut self) -> Async<()> {
+    fn poll_read(&mut self) -> Poll<(), io::Error> {
         <TcpStream>::poll_read(self)
     }
 
-    fn poll_write(&mut self) -> Async<()> {
+    fn poll_write(&mut self) -> Poll<(), io::Error> {
         <TcpStream>::poll_write(self)
     }
 }
@@ -446,7 +446,7 @@ impl Future for TcpStreamConnect {
             // actually hit an error or not.
             //
             // If all that succeeded then we ship everything on up.
-            if let Async::NotReady = stream.io.poll_write() {
+            if let Async::NotReady = try!(stream.io.poll_write()) {
                 return Ok(Async::NotReady)
             }
             if let Some(e) = try!(stream.io.get_ref().take_error()) {

src/net/udp.rs

@@ -61,7 +61,7 @@ impl UdpSocket {
     /// get a notification when the socket does become readable. That is, this
     /// is only suitable for calling in a `Future::poll` method and will
     /// automatically handle ensuring a retry once the socket is readable again.
-    pub fn poll_read(&self) -> Async<()> {
+    pub fn poll_read(&self) -> Poll<(), io::Error> {
         self.io.poll_read()
     }
 
@@ -71,7 +71,7 @@ impl UdpSocket {
     /// get a notification when the socket does become writable. That is, this
     /// is only suitable for calling in a `Future::poll` method and will
     /// automatically handle ensuring a retry once the socket is writable again.
-    pub fn poll_write(&self) -> Async<()> {
+    pub fn poll_write(&self) -> Poll<(), io::Error> {
         self.io.poll_write()
     }
 
@@ -81,13 +81,13 @@ impl UdpSocket {
     /// Address type can be any implementor of `ToSocketAddrs` trait. See its
     /// documentation for concrete examples.
     pub fn send_to(&self, buf: &[u8], target: &SocketAddr) -> io::Result<usize> {
-        if let Async::NotReady = self.io.poll_write() {
+        if let Async::NotReady = try!(self.io.poll_write()) {
             return Err(mio::would_block())
         }
         match self.io.get_ref().send_to(buf, target) {
             Ok(Some(n)) => Ok(n),
             Ok(None) => {
-                self.io.need_write();
+                try!(self.io.need_write());
                 Err(mio::would_block())
             }
             Err(e) => Err(e),
@@ -97,13 +97,13 @@ impl UdpSocket {
     /// Receives data from the socket. On success, returns the number of bytes
     /// read and the address from whence the data came.
     pub fn recv_from(&self, buf: &mut [u8]) -> io::Result<(usize, SocketAddr)> {
-        if let Async::NotReady = self.io.poll_read() {
+        if let Async::NotReady = try!(self.io.poll_read()) {
             return Err(mio::would_block())
         }
         match self.io.get_ref().recv_from(buf) {
             Ok(Some(n)) => Ok(n),
             Ok(None) => {
-                self.io.need_read();
+                try!(self.io.need_read());
                 Err(mio::would_block())
             }
             Err(e) => Err(e),

src/reactor/channel.rs

@@ -10,25 +10,35 @@ use std::cell::Cell;
 use std::io;
 use std::marker;
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 
 use mio;
 use mio::channel::{ctl_pair, SenderCtl, ReceiverCtl};
 
 use mpsc_queue::{Queue, PopResult};
 
+struct Inner<T> {
+    queue: Queue<T>,
+    receiver_alive: AtomicBool,
+}
+
 pub struct Sender<T> {
     ctl: SenderCtl,
-    inner: Arc<Queue<T>>,
+    inner: Arc<Inner<T>>,
 }
 
 pub struct Receiver<T> {
     ctl: ReceiverCtl,
-    inner: Arc<Queue<T>>,
+    inner: Arc<Inner<T>>,
     _marker: marker::PhantomData<Cell<()>>, // this type is not Sync
 }
 
 pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
-    let inner = Arc::new(Queue::new());
+    let inner = Arc::new(Inner {
+        queue: Queue::new(),
+        receiver_alive: AtomicBool::new(true),
+    });
     let (tx, rx) = ctl_pair();
 
     let tx = Sender {
@@ -45,7 +55,10 @@ pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
 
 impl<T> Sender<T> {
     pub fn send(&self, data: T) -> io::Result<()> {
-        self.inner.push(data);
+        if !self.inner.receiver_alive.load(Ordering::SeqCst) {
+            return Err(io::Error::new(io::ErrorKind::Other, "receiver has been dropped"));
+        }
+        self.inner.queue.push(data);
         self.ctl.inc()
     }
 }
@@ -57,7 +70,7 @@ impl<T> Receiver<T> {
         //
         // We, however, are the only thread with a `Receiver<T>` because this
         // type is not `Sync`. and we never handed out another instance.
-        match unsafe { self.inner.pop() } {
+        match unsafe { self.inner.queue.pop() } {
             PopResult::Data(t) => {
                 try!(self.ctl.dec());
                 Ok(Some(t))
@@ -108,6 +121,12 @@ impl<T> mio::Evented for Receiver<T> {
     }
 }
 
+impl<T> Drop for Receiver<T> {
+    fn drop(&mut self) {
+        self.inner.receiver_alive.store(false, Ordering::SeqCst);
+    }
+}
+
 impl<T> Clone for Sender<T> {
     fn clone(&self) -> Sender<T> {
         Sender {

src/reactor/io_token.rs

@@ -82,8 +82,8 @@ impl IoToken {
     ///
     /// This function will also panic if there is not a currently running future
     /// task.
-    pub fn schedule_read(&self, handle: &Remote) {
-        handle.send(Message::Schedule(self.token, task::park(), Direction::Read));
+    pub fn schedule_read(&self, handle: &Remote) -> io::Result<()> {
+        handle.send(Message::Schedule(self.token, task::park(), Direction::Read))
     }
 
     /// Schedule the current future task to receive a notification when the
@@ -109,8 +109,8 @@ impl IoToken {
     ///
     /// This function will also panic if there is not a currently running future
     /// task.
-    pub fn schedule_write(&self, handle: &Remote) {
-        handle.send(Message::Schedule(self.token, task::park(), Direction::Write));
+    pub fn schedule_write(&self, handle: &Remote) -> io::Result<()> {
+        handle.send(Message::Schedule(self.token, task::park(), Direction::Write))
     }
 
     /// Unregister all information associated with a token on an event loop,
@@ -135,7 +135,7 @@ impl IoToken {
     /// This function will panic if the event loop this handle is associated
     /// with has gone away, or if there is an error communicating with the event
     /// loop.
-    pub fn drop_source(&self, handle: &Remote) {
-        handle.send(Message::DropSource(self.token));
+    pub fn drop_source(&self, handle: &Remote) -> io::Result<()> {
+        handle.send(Message::DropSource(self.token))
     }
 }

src/reactor/mod.rs

@@ -44,7 +44,9 @@ const SLAB_CAPACITY: usize = 1024 * 64;
 pub struct Core {
     events: mio::Events,
     tx: Sender<Message>,
-    rx: Receiver<Message>,
+    // `rx` is `Option` here only because it is needed to be dropped explicitly
+    // in `drop` before other things.
+    rx: Option<Receiver<Message>>,
     inner: Rc<RefCell<Inner>>,
 
     // Used for determining when the future passed to `run` is ready. Once the
@@ -143,7 +145,7 @@ impl Core {
         Ok(Core {
             events: mio::Events::with_capacity(1024),
             tx: tx,
-            rx: rx,
+            rx: Some(rx),
             _future_registration: future_pair.0,
             future_readiness: Arc::new(MySetReadiness(future_pair.1)),
 
@@ -378,8 +380,9 @@ impl Core {
 
     fn consume_queue(&self) {
         debug!("consuming notification queue");
+        let rx = self.rx.as_ref().unwrap();
         // TODO: can we do better than `.unwrap()` here?
-        while let Some(msg) = self.rx.recv().unwrap() {
+        while let Some(msg) = rx.recv().unwrap() {
             self.notify(msg);
         }
     }
@@ -407,6 +410,27 @@ impl Core {
     }
 }
 
+impl Drop for Core {
+    fn drop(&mut self) {
+        // Destroy the receiver, so all schedule operations will be rejected.
+        // Do it explicitly before unparking, to avoid race condition.
+        self.rx.take();
+
+        // Unpark all tasks.
+        // It has no effect for tasks in this event loop,
+        // however tasks in another executors get an error
+        // when they do `poll` right after wakeup.
+        for io in self.inner.borrow_mut().io_dispatch.iter_mut() {
+            if let Some(ref mut reader) = io.reader {
+                reader.unpark();
+            }
+            if let Some(ref mut writer) = io.writer {
+                writer.unpark();
+            }
+        }
+    }
+}
+
 impl Inner {
     fn add_source(&mut self, source: &mio::Evented)
                   -> io::Result<(Arc<AtomicUsize>, usize)> {
@@ -497,26 +521,20 @@ impl Inner {
 }
 
 impl Remote {
-    fn send(&self, msg: Message) {
+    fn send(&self, msg: Message) -> io::Result<()> {
         self.with_loop(|lp| {
             match lp {
                 Some(lp) => {
                     // Need to execute all existing requests first, to ensure
                     // that our message is processed "in order"
                     lp.consume_queue();
                     lp.notify(msg);
+                    Ok(())
                 }
                 None => {
-                    match self.tx.send(msg) {
-                        Ok(()) => {}
-
-                        // This should only happen when there was an error
-                        // writing to the pipe to wake up the event loop,
-                        // hopefully that never happens
-                        Err(e) => {
-                            panic!("error sending message to event loop: {}", e)
-                        }
-                    }
+                    // May return an error if receiver is closed
+                    // or if there was an error writing to the pipe.
+                    self.tx.send(msg)
                 }
             }
         })
@@ -547,15 +565,15 @@ impl Remote {
     ///
     /// Note that while the closure, `F`, requires the `Send` bound as it might
     /// cross threads, the future `R` does not.
-    pub fn spawn<F, R>(&self, f: F)
+    pub fn spawn<F, R>(&self, f: F) -> io::Result<()>
         where F: FnOnce(&Handle) -> R + Send + 'static,
               R: IntoFuture<Item=(), Error=()>,
               R::Future: 'static,
     {
         self.send(Message::Run(Box::new(|lp: &Core| {
             let f = f(&lp.handle());
             lp.inner.borrow_mut().spawn(Box::new(f.into_future()));
-        })));
+        })))
     }
 }