Do not hang in poll if event loop is destroyed

`Loop` destructor marks all loop handles dead, so any further `poll`
requests fail (they currently panic, but should probably return an
error).

Also `Loop` destructor unparks all tasks waiting for IO on self.
It has no effect on tasks bound to self event loop, but tasks running
on different executors call `poll` right afterwards, so they
immediately fail instead of hanging.

Author: stepancheg

src/channel.rs

@@ -91,7 +91,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/event_loop/mod.rs

@@ -3,7 +3,7 @@ use std::io::{self, ErrorKind};
 use std::mem;
 use std::rc::Rc;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicUsize, ATOMIC_USIZE_INIT, Ordering};
+use std::sync::atomic::{AtomicUsize, AtomicBool, ATOMIC_USIZE_INIT, Ordering};
 use std::time::{Instant, Duration};
 
 use futures::{Future, Poll, IntoFuture, Async};
@@ -37,7 +37,7 @@ pub struct Loop {
     id: usize,
     io: mio::Poll,
     events: mio::Events,
-    tx: Arc<Sender<Message>>,
+    inner: Arc<LoopHandleInner>,
     rx: Receiver<Message>,
     io_dispatch: RefCell<Slab<ScheduledIo, usize>>,
     task_dispatch: RefCell<Slab<ScheduledTask, usize>>,
@@ -62,6 +62,31 @@ pub struct Loop {
     timeouts: RefCell<Slab<(Timeout, TimeoutState), usize>>,
 }
 
+impl Drop for Loop {
+    fn drop(&mut self) {
+        // mark event loop as dead, so all schedule operations will be rejected
+        self.inner.alive.store(false, Ordering::SeqCst);
+
+        // 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.io_dispatch.borrow_mut().iter_mut() {
+            if let Some(ref mut reader) = io.reader {
+                reader.unpark();
+            }
+            if let Some(ref mut writer) = io.writer {
+                writer.unpark();
+            }
+        }
+    }
+}
+
+struct LoopHandleInner {
+    tx: Sender<Message>,
+    alive: AtomicBool,
+}
+
 /// Handle to an event loop, used to construct I/O objects, send messages, and
 /// otherwise interact indirectly with the event loop itself.
 ///
@@ -70,7 +95,7 @@ pub struct Loop {
 #[derive(Clone)]
 pub struct LoopHandle {
     id: usize,
-    tx: Arc<Sender<Message>>,
+    inner: Arc<LoopHandleInner>,
 }
 
 /// A non-sendable handle to an event loop, useful for manufacturing instances
@@ -145,7 +170,10 @@ impl Loop {
             id: NEXT_LOOP_ID.fetch_add(1, Ordering::Relaxed),
             io: io,
             events: mio::Events::with_capacity(1024),
-            tx: Arc::new(tx),
+            inner: Arc::new(LoopHandleInner {
+                tx: tx,
+                alive: AtomicBool::new(true),
+            }),
             rx: rx,
             io_dispatch: RefCell::new(Slab::with_capacity(SLAB_CAPACITY)),
             task_dispatch: RefCell::new(Slab::with_capacity(SLAB_CAPACITY)),
@@ -169,7 +197,7 @@ impl Loop {
     pub fn handle(&self) -> LoopHandle {
         LoopHandle {
             id: self.id,
-            tx: self.tx.clone(),
+            inner: self.inner.clone(),
         }
     }
 
@@ -506,18 +534,23 @@ impl Loop {
 }
 
 impl LoopHandle {
-    fn send(&self, msg: Message) {
+    fn send(&self, msg: Message) -> io::Result<()> {
+        if !self.inner.alive.load(Ordering::SeqCst) {
+            return Err(io::Error::new(io::ErrorKind::Other, "loop has been dropped"));
+        }
+
         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(()) => {}
+                    match self.inner.tx.send(msg) {
+                        Ok(()) => Ok(()),
 
                         // This should only happen when there was an error
                         // writing to the pipe to wake up the event loop,
@@ -555,15 +588,15 @@ impl LoopHandle {
     ///
     /// 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(&LoopPin) -> R + Send + 'static,
               R: IntoFuture<Item=(), Error=()>,
               R::Future: 'static,
     {
         self.send(Message::Run(Box::new(|lp: &Loop| {
             let f = f(&lp.pin());
             lp.spawn(Box::new(f.into_future()));
-        })));
+        })))
     }
 }
 
@@ -626,7 +659,7 @@ impl<T, U> LoopFuture<T, U>
                     task.unpark();
                 });
                 self.result = Some((result.clone(), token));
-                self.loop_handle.send(g(data.take().unwrap(), result));
+                try!(self.loop_handle.send(g(data.take().unwrap(), result)));
                 Ok(Async::NotReady)
             }
         }

src/event_loop/source.rs

@@ -74,8 +74,8 @@ impl LoopHandle {
     ///
     /// This function will also panic if there is not a currently running future
     /// task.
-    pub fn schedule_read(&self, tok: &IoToken) {
-        self.send(Message::Schedule(tok.token, task::park(), Direction::Read));
+    pub fn schedule_read(&self, tok: &IoToken) -> io::Result<()> {
+        self.send(Message::Schedule(tok.token, task::park(), Direction::Read))
     }
 
     /// Schedule the current future task to receive a notification when the
@@ -101,8 +101,8 @@ impl LoopHandle {
     ///
     /// This function will also panic if there is not a currently running future
     /// task.
-    pub fn schedule_write(&self, tok: &IoToken) {
-        self.send(Message::Schedule(tok.token, task::park(), Direction::Write));
+    pub fn schedule_write(&self, tok: &IoToken) -> io::Result<()> {
+        self.send(Message::Schedule(tok.token, task::park(), Direction::Write))
     }
 
     /// Unregister all information associated with a token on an event loop,
@@ -127,8 +127,8 @@ impl LoopHandle {
     /// 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, tok: &IoToken) {
-        self.send(Message::DropSource(tok.token));
+    pub fn drop_source(&self, tok: &IoToken) -> io::Result<()> {
+        self.send(Message::DropSource(tok.token))
     }
 }
 

src/event_loop/timeout.rs

@@ -25,7 +25,7 @@ impl LoopHandle {
     ///
     /// This method will panic if the timeout specified was not created by this
     /// loop handle's `add_timeout` method.
-    pub fn update_timeout(&self, timeout: &TimeoutToken) {
+    pub fn update_timeout(&self, timeout: &TimeoutToken) -> io::Result<()> {
         self.send(Message::UpdateTimeout(timeout.token, task::park()))
     }
 
@@ -35,7 +35,7 @@ impl LoopHandle {
     ///
     /// This method will panic if the timeout specified was not created by this
     /// loop handle's `add_timeout` method.
-    pub fn cancel_timeout(&self, timeout: &TimeoutToken) {
+    pub fn cancel_timeout(&self, timeout: &TimeoutToken) -> io::Result<()> {
         debug!("cancel timeout {}", timeout.token);
         self.send(Message::CancelTimeout(timeout.token))
     }

src/readiness_stream.rs

@@ -66,7 +66,7 @@ impl<E> ReadinessStream<E> {
         if self.readiness.load(Ordering::SeqCst) & 1 != 0 {
             Ok(Async::Ready(()))
         } else {
-            self.handle.schedule_read(&self.token);
+            try!(self.handle.schedule_read(&self.token));
             Ok(Async::NotReady)
         }
     }
@@ -86,7 +86,7 @@ impl<E> ReadinessStream<E> {
         if self.readiness.load(Ordering::SeqCst) & 2 != 0 {
             Ok(Async::Ready(()))
         } else {
-            self.handle.schedule_write(&self.token);
+            try!(self.handle.schedule_write(&self.token));
             Ok(Async::NotReady)
         }
     }
@@ -102,9 +102,9 @@ impl<E> ReadinessStream<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.handle.schedule_read(&self.token);
+        self.handle.schedule_read(&self.token)
     }
 
     /// Indicates to this source of events that the corresponding I/O object is
@@ -118,9 +118,9 @@ impl<E> ReadinessStream<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.handle.schedule_write(&self.token);
+        self.handle.schedule_write(&self.token)
     }
 
     /// Returns a reference to the event loop handle that this readiness stream
@@ -161,6 +161,7 @@ impl<E> Future for ReadinessStreamNew<E>
 
 impl<E> Drop for ReadinessStream<E> {
     fn drop(&mut self) {
-        self.handle.drop_source(&self.token)
+        // Ignore error
+        drop(self.handle.drop_source(&self.token))
     }
 }

src/tcp.rs

@@ -96,7 +96,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)
@@ -391,7 +391,7 @@ impl<'a> Read for &'a TcpStream {
         }
         let r = self.io.get_ref().read(buf);
         if is_wouldblock(&r) {
-            self.io.need_read();
+            try!(self.io.need_read());
         }
         return r
     }
@@ -404,7 +404,7 @@ impl<'a> Write for &'a TcpStream {
         }
         let r = self.io.get_ref().write(buf);
         if is_wouldblock(&r) {
-            self.io.need_write();
+            try!(self.io.need_write());
         }
         return r
     }
@@ -415,7 +415,7 @@ impl<'a> Write for &'a TcpStream {
         }
         let r = self.io.get_ref().flush();
         if is_wouldblock(&r) {
-            self.io.need_write();
+            try!(self.io.need_write());
         }
         return r
     }

src/timeout.rs

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

src/udp.rs

@@ -90,7 +90,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(io::Error::new(io::ErrorKind::WouldBlock, "would block"))
             }
             Err(e) => Err(e),
@@ -106,7 +106,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(io::Error::new(io::ErrorKind::WouldBlock, "would block"))
             }
             Err(e) => Err(e),

tests/channel.rs

@@ -0,0 +1,34 @@
+extern crate tokio_core;
+extern crate futures;
+
+use tokio_core::Loop;
+
+use futures::stream::Stream;
+
+
+#[test]
+fn recv_after_loop_drop() {
+    let mut lp: Loop = Loop::new().unwrap();
+
+    let (_sender, receiver) = lp.handle().channel::<u32>();
+
+    let receiver = lp.run(receiver).unwrap();
+
+    drop(lp);
+
+    assert!(receiver.wait().next().unwrap().is_err());
+}
+
+#[test]
+fn recv_after_loop_and_sender_drop() {
+    let mut lp: Loop = Loop::new().unwrap();
+
+    let (sender, receiver) = lp.handle().channel::<u32>();
+
+    let receiver = lp.run(receiver).unwrap();
+
+    drop(lp);
+    drop(sender);
+
+    assert!(receiver.wait().next().unwrap().is_err());
+}

tests/poll.rs

@@ -52,7 +52,7 @@ fn poll_after_ready() {
 
         fn poll(&mut self) -> Poll<(), ()> {
             if self.n == 0 {
-                self.handle.schedule_read(&self.token);
+                self.handle.schedule_read(&self.token).expect("schedule_read");
                 self.n += 1;
                 Ok(Async::NotReady)
             } else {

tests/spawn.rs

@@ -21,7 +21,7 @@ fn simple() {
             tx2.complete(2);
             Ok(())
         })
-    });
+    }).expect("spawn");
 
     assert_eq!(lp.run(rx1.join(rx2)).unwrap(), (1, 2));
 }
@@ -41,7 +41,7 @@ fn spawn_in_poll() {
                 tx2.complete(2);
                 Ok(())
             })
-        });
+        }).expect("spawn");
         Ok(())
     }));