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

@@ -114,7 +114,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

@@ -73,7 +73,7 @@ impl TcpListener {
             match self.io.get_ref().accept() {
                 Err(e) => {
                     if e.kind() == io::ErrorKind::WouldBlock {
-                        self.io.need_read();
+                        try!(self.io.need_read());
                     }
                     return Err(e)
                 },
@@ -87,7 +87,7 @@ impl TcpListener {
                             });
                         tx.complete(res);
                         Ok(())
-                    });
+                    }).expect("failed to spawn");
                     self.pending_accept = Some(rx);
                     // continue to polling the `rx` at the beginning of the loop
                 }

src/net/udp/mod.rs

@@ -106,7 +106,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),
@@ -149,7 +149,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/interval.rs

@@ -66,18 +66,19 @@ impl Stream for Interval {
         let now = Instant::now();
         if self.next <= now {
             self.next = next_interval(self.next, now, self.interval);
-            self.token.reset_timeout(self.next, &self.handle);
+            try!(self.token.reset_timeout(self.next, &self.handle));
             Ok(Async::Ready(Some(())))
         } else {
-            self.token.update_timeout(&self.handle);
+            try!(self.token.update_timeout(&self.handle));
             Ok(Async::NotReady)
         }
     }
 }
 
 impl Drop for Interval {
     fn drop(&mut self) {
-        self.token.cancel_timeout(&self.handle);
+        // Ignore error
+        drop(self.token.cancel_timeout(&self.handle));
     }
 }
 

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

@@ -412,6 +412,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)> {
@@ -519,26 +540,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)
                 }
             }
         })
@@ -569,15 +584,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()));
-        })));
+        })))
     }
 }