Do not hang in poll if reactor is destroyed

If channel is dropped, receiver may still return EOF, and if channel is
alive, receiver produces an error.

Author: stepancheg

src/channel.rs

@@ -99,7 +99,7 @@ impl<T> Stream for Receiver<T> {
         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/net/tcp.rs

@@ -109,7 +109,7 @@ impl TcpListener {
                 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(),
         }

src/net/udp.rs

@@ -81,7 +81,7 @@ impl UdpSocket {
         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,7 +97,7 @@ impl UdpSocket {
         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)),
 
@@ -379,8 +381,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);
         }
     }
@@ -408,6 +411,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)> {
@@ -498,26 +522,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)
                 }
             }
         })
@@ -548,15 +566,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()));
-        })));
+        })))
     }
 }
 

src/reactor/poll_evented.rs

@@ -71,8 +71,11 @@ impl<E> PollEvented<E> {
         if self.readiness.load(Ordering::SeqCst) & 1 != 0 {
             Async::Ready(())
         } else {
-            self.token.schedule_read(&self.handle);
-            Async::NotReady
+            match self.token.schedule_read(&self.handle) {
+                Ok(()) => Async::NotReady,
+                // next read will return error
+                Err(_) => Async::Ready(()),
+            }
         }
     }
 
@@ -91,8 +94,12 @@ impl<E> PollEvented<E> {
         if self.readiness.load(Ordering::SeqCst) & 2 != 0 {
             Async::Ready(())
         } else {
-            self.token.schedule_write(&self.handle);
-            Async::NotReady
+            // ignore error, will be handled in need_write
+            match self.token.schedule_write(&self.handle) {
+                Ok(()) => Async::NotReady,
+                // next read will return error
+                Err(_) => Async::Ready(()),
+            }
         }
     }
 
@@ -107,7 +114,7 @@ impl<E> PollEvented<E> {
     /// The flag indicating that this stream is readable is unset and the
     /// current task is scheduled to receive a notification when the stream is
     /// then again readable.
-    pub fn need_read(&self) {
+    pub fn need_read(&self) -> io::Result<()> {
         self.readiness.fetch_and(!1, Ordering::SeqCst);
         self.token.schedule_read(&self.handle)
     }
@@ -123,7 +130,7 @@ impl<E> PollEvented<E> {
     /// The flag indicating that this stream is writable is unset and the
     /// current task is scheduled to receive a notification when the stream is
     /// then again writable.
-    pub fn need_write(&self) {
+    pub fn need_write(&self) -> io::Result<()> {
         self.readiness.fetch_and(!2, Ordering::SeqCst);
         self.token.schedule_write(&self.handle)
     }
@@ -154,7 +161,7 @@ impl<E: Read> Read for PollEvented<E> {
         }
         let r = self.get_mut().read(buf);
         if is_wouldblock(&r) {
-            self.need_read();
+            try!(self.need_read());
         }
         return r
     }
@@ -167,7 +174,7 @@ impl<E: Write> Write for PollEvented<E> {
         }
         let r = self.get_mut().write(buf);
         if is_wouldblock(&r) {
-            self.need_write();
+            try!(self.need_write());
         }
         return r
     }
@@ -178,7 +185,7 @@ impl<E: Write> Write for PollEvented<E> {
         }
         let r = self.get_mut().flush();
         if is_wouldblock(&r) {
-            self.need_write();
+            try!(self.need_write());
         }
         return r
     }
@@ -203,7 +210,7 @@ impl<'a, E> Read for &'a PollEvented<E>
         }
         let r = self.get_ref().read(buf);
         if is_wouldblock(&r) {
-            self.need_read();
+            try!(self.need_read());
         }
         return r
     }
@@ -218,7 +225,7 @@ impl<'a, E> Write for &'a PollEvented<E>
         }
         let r = self.get_ref().write(buf);
         if is_wouldblock(&r) {
-            self.need_write();
+            try!(self.need_write());
         }
         return r
     }
@@ -229,7 +236,7 @@ impl<'a, E> Write for &'a PollEvented<E>
         }
         let r = self.get_ref().flush();
         if is_wouldblock(&r) {
-            self.need_write();
+            try!(self.need_write());
         }
         return r
     }
@@ -256,6 +263,7 @@ fn is_wouldblock<T>(r: &io::Result<T>) -> bool {
 
 impl<E> Drop for PollEvented<E> {
     fn drop(&mut self) {
-        self.token.drop_source(&self.handle);
+        // Ignore error
+        drop(self.token.drop_source(&self.handle));
     }
 }

src/reactor/timeout.rs

@@ -56,14 +56,15 @@ impl Future for Timeout {
         if *self.token.when() <= now {
             Ok(Async::Ready(()))
         } else {
-            self.token.update_timeout(&self.handle);
+            try!(self.token.update_timeout(&self.handle));
             Ok(Async::NotReady)
         }
     }
 }
 
 impl Drop for Timeout {
     fn drop(&mut self) {
-        self.token.cancel_timeout(&self.handle);
+        // Ignore error
+        drop(self.token.cancel_timeout(&self.handle));
     }
 }