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

@@ -108,7 +108,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)
@@ -127,7 +127,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

@@ -79,7 +79,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),
@@ -95,7 +95,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 closed"));
+        }
+        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))
@@ -85,6 +98,13 @@ impl<T> Receiver<T> {
     }
 }
 
+// Close receiver, so further send operations would fail.
+// This function is used internally in `Core` and is not exposed as public API.
+pub fn close_receiver<T>(receiver: &Receiver<T>) {
+    receiver.inner.as_ref().receiver_alive.store(false, Ordering::SeqCst);
+}
+
+
 // Just delegate everything to `self.ctl`
 impl<T> mio::Evented for Receiver<T> {
     fn register(&self,
@@ -108,6 +128,12 @@ impl<T> mio::Evented for Receiver<T> {
     }
 }
 
+impl<T> Drop for Receiver<T> {
+    fn drop(&mut self) {
+        close_receiver(self);
+    }
+}
+
 impl<T> Clone for Sender<T> {
     fn clone(&self) -> Sender<T> {
         Sender {

src/reactor/io_token.rs

@@ -71,6 +71,8 @@ impl IoToken {
     /// receive further notifications it will need to call `schedule_read`
     /// again.
     ///
+    /// This function returns an error if reactor is destroyed.
+    ///
     /// > **Note**: This method should generally not be used directly, but
     /// >           rather the `ReadinessStream` type should be used instead.
     ///
@@ -82,8 +84,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
@@ -98,6 +100,8 @@ impl IoToken {
     /// receive further notifications it will need to call `schedule_write`
     /// again.
     ///
+    /// This function returns an error if reactor is destroyed.
+    ///
     /// > **Note**: This method should generally not be used directly, but
     /// >           rather the `ReadinessStream` type should be used instead.
     ///
@@ -109,8 +113,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,
@@ -127,6 +131,8 @@ impl IoToken {
     /// ensure that the callbacks are **not** invoked, so pending scheduled
     /// callbacks cannot be relied upon to get called.
     ///
+    /// This function returns an error if reactor is destroyed.
+    ///
     /// > **Note**: This method should generally not be used directly, but
     /// >           rather the `ReadinessStream` type should be used instead.
     ///
@@ -135,7 +141,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

@@ -408,6 +408,27 @@ impl Core {
     }
 }
 
+impl Drop for Core {
+    fn drop(&mut self) {
+        // Close the receiver, so all schedule operations will be rejected.
+        // Do it explicitly before unparking to avoid race condition.
+        channel::close_receiver(&self.rx);
+
+        // 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 +519,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 +563,17 @@ 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)
+    ///
+    /// This function returns an error if reactor is destroyed.
+    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(()),
+            }
         }
     }
 
@@ -111,7 +118,9 @@ impl<E> PollEvented<E> {
     /// Note that it is also only valid to call this method if `poll_read`
     /// previously indicated that the object is readable. That is, this function
     /// must always be paired with calls to `poll_read` previously.
-    pub fn need_read(&self) {
+    ///
+    /// This function returns an error if reactor is destroyed.
+    pub fn need_read(&self) -> io::Result<()> {
         self.readiness.fetch_and(!1, Ordering::SeqCst);
         self.token.schedule_read(&self.handle)
     }
@@ -131,7 +140,9 @@ impl<E> PollEvented<E> {
     /// Note that it is also only valid to call this method if `poll_write`
     /// previously indicated that the object is writeable. That is, this function
     /// must always be paired with calls to `poll_write` previously.
-    pub fn need_write(&self) {
+    ///
+    /// This function returns an error if reactor is destroyed.
+    pub fn need_write(&self) -> io::Result<()> {
         self.readiness.fetch_and(!2, Ordering::SeqCst);
         self.token.schedule_write(&self.handle)
     }
@@ -162,7 +173,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
     }
@@ -175,7 +186,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
     }
@@ -186,7 +197,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
     }
@@ -211,7 +222,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
     }
@@ -226,7 +237,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
     }
@@ -237,7 +248,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
     }
@@ -264,6 +275,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));
     }
 }