Change need_read and need_write to return an error

This commit is targeted at solving tokio-rs/tokio-core#12 and incorporates the
solution from tokio-rs/tokio-core#17. Namely the `need_read` and `need_write`
functions on `PollEvented` now return an error when the connected reactor has
gone away and the task cannot be blocked. This will typically naturally
translate to errors being returned by various connected I/O objects and should
help tear down the world in a clean-ish fashion.

Author: alexcrichton

src/net/tcp.rs

@@ -64,7 +64,7 @@ impl TcpListener {
         match self.io.get_ref().accept() {
             Err(e) => {
                 if e.kind() == io::ErrorKind::WouldBlock {
-                    self.io.need_read();
+                    self.io.need_read()?;
                 }
                 Err(e)
             },
@@ -576,7 +576,7 @@ impl<'a> AsyncRead for &'a TcpStream {
                 Ok(Async::Ready(n))
             }
             Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                self.io.need_read();
+                self.io.need_read()?;
                 Ok(Async::NotReady)
             }
             Err(e) => Err(e),
@@ -614,7 +614,7 @@ impl<'a> AsyncWrite for &'a TcpStream {
                 Ok(Async::Ready(n))
             }
             Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
-                self.io.need_write();
+                self.io.need_write()?;
                 Ok(Async::NotReady)
             }
             Err(e) => Err(e),

src/net/udp/mod.rs

@@ -93,7 +93,7 @@ impl UdpSocket {
             Ok(n) => Ok(n),
             Err(e) => {
                 if e.kind() == io::ErrorKind::WouldBlock {
-                    self.io.need_write();
+                    self.io.need_write()?;
                 }
                 Err(e)
             }
@@ -115,7 +115,7 @@ impl UdpSocket {
             Ok(n) => Ok(n),
             Err(e) => {
                 if e.kind() == io::ErrorKind::WouldBlock {
-                    self.io.need_read();
+                    self.io.need_read()?;
                 }
                 Err(e)
             }
@@ -163,7 +163,7 @@ impl UdpSocket {
             Ok(n) => Ok(n),
             Err(e) => {
                 if e.kind() == io::ErrorKind::WouldBlock {
-                    self.io.need_write();
+                    self.io.need_write()?;
                 }
                 Err(e)
             }
@@ -205,7 +205,7 @@ impl UdpSocket {
             Ok(n) => Ok(n),
             Err(e) => {
                 if e.kind() == io::ErrorKind::WouldBlock {
-                    self.io.need_read();
+                    self.io.need_read()?;
                 }
                 Err(e)
             }

src/reactor/io_token.rs

@@ -92,13 +92,14 @@ impl IoToken {
     ///
     /// This function will also panic if there is not a currently running future
     /// task.
-    pub fn schedule_read(&self) {
+    pub fn schedule_read(&self) -> io::Result<()> {
         let inner = match self.handle.inner.upgrade() {
             Some(inner) => inner,
-            None => return,
+            None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")),
         };
 
         inner.schedule(self.token, Direction::Read);
+        Ok(())
     }
 
     /// Schedule the current future task to receive a notification when the
@@ -124,13 +125,14 @@ impl IoToken {
     ///
     /// This function will also panic if there is not a currently running future
     /// task.
-    pub fn schedule_write(&self) {
+    pub fn schedule_write(&self) -> io::Result<()> {
         let inner = match self.handle.inner.upgrade() {
             Some(inner) => inner,
-            None => return,
+            None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")),
         };
 
         inner.schedule(self.token, Direction::Write);
+        Ok(())
     }
 
     /// Unregister all information associated with a token on an event loop,

src/reactor/mod.rs

@@ -231,6 +231,19 @@ impl fmt::Debug for Core {
     }
 }
 
+impl Drop for Inner {
+    fn drop(&mut self) {
+        // When a reactor is dropped it needs to wake up all blocked tasks as
+        // they'll never receive a notification, and all connected I/O objects
+        // will start returning errors pretty quickly.
+        let io = self.io_dispatch.read().unwrap();
+        for (_, io) in io.iter() {
+            io.writer.notify();
+            io.reader.notify();
+        }
+    }
+}
+
 impl Inner {
     /// Register an I/O resource with the reactor.
     ///

src/reactor/poll_evented.rs

@@ -184,9 +184,13 @@ impl<E> PollEvented<E> {
         match self.readiness.load(Ordering::SeqCst) & bits {
             0 => {
                 if mask.is_writable() {
-                    self.need_write();
+                    if self.need_write().is_err() {
+                        return Async::Ready(mask)
+                    }
                 } else {
-                    self.need_read();
+                    if self.need_read().is_err() {
+                        return Async::Ready(mask)
+                    }
                 }
                 Async::NotReady
             }
@@ -213,14 +217,22 @@ impl<E> PollEvented<E> {
     /// previously indicated that the object is readable. That is, this function
     /// must always be paired with calls to `poll_read` previously.
     ///
+    /// # Errors
+    ///
+    /// This function will return an error if the `Core` that this `PollEvented`
+    /// is associated with has gone away (been destroyed). The error means that
+    /// the ambient futures task could not be scheduled to receive a
+    /// notification and typically means that the error should be propagated
+    /// outwards.
+    ///
     /// # Panics
     ///
     /// This function will panic if called outside the context of a future's
     /// task.
-    pub fn need_read(&self) {
+    pub fn need_read(&self) -> io::Result<()> {
         let bits = super::ready2usize(super::read_ready());
         self.readiness.fetch_and(!bits, Ordering::SeqCst);
-        self.token.schedule_read();
+        self.token.schedule_read()
     }
 
     /// Indicates to this source of events that the corresponding I/O object is
@@ -239,14 +251,22 @@ impl<E> PollEvented<E> {
     /// previously indicated that the object is writable. That is, this function
     /// must always be paired with calls to `poll_write` previously.
     ///
+    /// # Errors
+    ///
+    /// This function will return an error if the `Core` that this `PollEvented`
+    /// is associated with has gone away (been destroyed). The error means that
+    /// the ambient futures task could not be scheduled to receive a
+    /// notification and typically means that the error should be propagated
+    /// outwards.
+    ///
     /// # Panics
     ///
     /// This function will panic if called outside the context of a future's
     /// task.
-    pub fn need_write(&self) {
+    pub fn need_write(&self) -> io::Result<()> {
         let bits = super::ready2usize(Ready::writable());
         self.readiness.fetch_and(!bits, Ordering::SeqCst);
-        self.token.schedule_write();
+        self.token.schedule_write()
     }
 
     /// Returns a reference to the event loop handle that this readiness stream
@@ -275,7 +295,7 @@ impl<E: Read> Read for PollEvented<E> {
         }
         let r = self.get_mut().read(buf);
         if is_wouldblock(&r) {
-            self.need_read();
+            self.need_read()?;
         }
         return r
     }
@@ -288,7 +308,7 @@ impl<E: Write> Write for PollEvented<E> {
         }
         let r = self.get_mut().write(buf);
         if is_wouldblock(&r) {
-            self.need_write();
+            self.need_write()?;
         }
         return r
     }
@@ -299,7 +319,7 @@ impl<E: Write> Write for PollEvented<E> {
         }
         let r = self.get_mut().flush();
         if is_wouldblock(&r) {
-            self.need_write();
+            self.need_write()?;
         }
         return r
     }
@@ -323,7 +343,7 @@ impl<'a, E> Read for &'a PollEvented<E>
         }
         let r = self.get_ref().read(buf);
         if is_wouldblock(&r) {
-            self.need_read();
+            self.need_read()?;
         }
         return r
     }
@@ -338,7 +358,7 @@ impl<'a, E> Write for &'a PollEvented<E>
         }
         let r = self.get_ref().write(buf);
         if is_wouldblock(&r) {
-            self.need_write();
+            self.need_write()?;
         }
         return r
     }
@@ -349,7 +369,7 @@ impl<'a, E> Write for &'a PollEvented<E>
         }
         let r = self.get_ref().flush();
         if is_wouldblock(&r) {
-            self.need_write();
+            self.need_write()?;
         }
         return r
     }

tests/drop-core.rs

@@ -0,0 +1,39 @@
+extern crate tokio;
+extern crate futures;
+
+use std::thread;
+
+use futures::future;
+use futures::prelude::*;
+use futures::sync::oneshot;
+use tokio::net::TcpListener;
+use tokio::reactor::Core;
+
+#[test]
+fn tcp_doesnt_block() {
+    let core = Core::new().unwrap();
+    let handle = core.handle();
+    let listener = TcpListener::bind(&"127.0.0.1:0".parse().unwrap(), &handle).unwrap();
+    drop(core);
+    assert!(listener.incoming().wait().next().unwrap().is_err());
+}
+
+#[test]
+fn drop_wakes() {
+    let core = Core::new().unwrap();
+    let handle = core.handle();
+    let listener = TcpListener::bind(&"127.0.0.1:0".parse().unwrap(), &handle).unwrap();
+    let (tx, rx) = oneshot::channel::<()>();
+    let t = thread::spawn(move || {
+        let incoming = listener.incoming();
+        let new_socket = incoming.into_future().map_err(|_| ());
+        let drop_tx = future::lazy(|| {
+            drop(tx);
+            future::ok(())
+        });
+        assert!(new_socket.join(drop_tx).wait().is_err());
+    });
+    drop(rx.wait());
+    drop(core);
+    t.join().unwrap();
+}