Explicitly deregister I/O resources on drop (#189)

Mio will be requiring `deregister` to be called explicitly in order to
guarantee that Poll releases any state associated with the I/O resource.
See carllerche/mio#753.

This patch adds an explicit `deregister` function to `Registration` and
updates `PollEvented` to call this function on drop.

`Registration::deregister` is also called on `PollEvented::into_inner`.

Closes #168

src/reactor/poll_evented.rs

@@ -10,6 +10,7 @@
 
 use std::fmt;
 use std::io::{self, Read, Write};
+use std::sync::Mutex;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::Relaxed;
 
@@ -29,7 +30,7 @@ pub struct PollEvented<E> {
 }
 
 struct Inner {
-    registration: Registration,
+    registration: Mutex<Registration>,
 
     /// Currently visible read readiness
     read_readiness: AtomicUsize,
@@ -58,7 +59,7 @@ impl<E> PollEvented<E> {
         Ok(PollEvented {
             io: io,
             inner: Inner {
-                registration,
+                registration: Mutex::new(registration),
                 read_readiness: AtomicUsize::new(0),
                 write_readiness: AtomicUsize::new(0),
             },
@@ -89,12 +90,14 @@ impl<E> PollEvented<E> {
     }
 
     fn poll_read2(&self) -> Async<Ready> {
+        let r = self.inner.registration.lock().unwrap();
+
         // Load the cached readiness
         match self.inner.read_readiness.load(Relaxed) {
             0 => {}
             mut n => {
                 // Check what's new with the reactor.
-                if let Some(ready) = self.inner.registration.take_read_ready().unwrap() {
+                if let Some(ready) = r.take_read_ready().unwrap() {
                     n |= ready2usize(ready);
                     self.inner.read_readiness.store(n, Relaxed);
                 }
@@ -103,7 +106,7 @@ impl<E> PollEvented<E> {
             }
         }
 
-        let ready = match self.inner.registration.poll_read_ready().unwrap() {
+        let ready = match r.poll_read_ready().unwrap() {
             Async::Ready(r) => r,
             _ => return Async::NotReady,
         };
@@ -129,11 +132,13 @@ impl<E> PollEvented<E> {
     /// This function will panic if called outside the context of a future's
     /// task.
     pub fn poll_write(&mut self) -> Async<()> {
+        let r = self.inner.registration.lock().unwrap();
+
         match self.inner.write_readiness.load(Relaxed) {
             0 => {}
             mut n => {
                 // Check what's new with the reactor.
-                if let Some(ready) = self.inner.registration.take_write_ready().unwrap() {
+                if let Some(ready) = r.take_write_ready().unwrap() {
                     n |= ready2usize(ready);
                     self.inner.write_readiness.store(n, Relaxed);
                 }
@@ -142,7 +147,7 @@ impl<E> PollEvented<E> {
             }
         }
 
-        let ready = match self.inner.registration.poll_write_ready().unwrap() {
+        let ready = match r.poll_write_ready().unwrap() {
             Async::Ready(r) => r,
             _ => return Async::NotReady,
         };
@@ -331,9 +336,8 @@ impl<E> PollEvented<E> {
     pub fn deregister(&self) -> io::Result<()>
         where E: Evented,
     {
-        // Nothing has to happen here anymore as I/O objects are explicitly
-        // deregistered before dropped.
-        Ok(())
+        self.inner.registration.lock().unwrap()
+            .deregister(&self.io)
     }
 }
 

tokio-reactor/src/lib.rs

@@ -559,6 +559,11 @@ impl Inner {
         Ok(key)
     }
 
+    /// Deregisters an I/O resource from the reactor.
+    fn deregister_source(&self, source: &Evented) -> io::Result<()> {
+        self.io.deregister(source)
+    }
+
     fn drop_source(&self, token: usize) {
         debug!("dropping I/O source: {}", token);
         self.io_dispatch.write().unwrap().remove(token);

tokio-reactor/src/poll_evented.rs

@@ -82,8 +82,8 @@ use std::sync::atomic::Ordering::Relaxed;
 /// [`mio::Evented`]: https://docs.rs/mio/0.6/mio/trait.Evented.html
 /// [`Registration`]: struct.Registration.html
 /// [`TcpListener`]: ../net/struct.TcpListener.html
-pub struct PollEvented<E> {
-    io: E,
+pub struct PollEvented<E: Evented> {
+    io: Option<E>,
     inner: Inner,
 }
 
@@ -105,7 +105,7 @@ where E: Evented
     /// Creates a new `PollEvented` associated with the default reactor.
     pub fn new(io: E) -> PollEvented<E> {
         PollEvented {
-            io: io,
+            io: Some(io),
             inner: Inner {
                 registration: Registration::new(),
                 read_readiness: AtomicUsize::new(0),
@@ -117,10 +117,36 @@ where E: Evented
     /// Creates a new `PollEvented` associated with the specified reactor.
     pub fn new_with_handle(io: E, handle: &Handle) -> io::Result<Self> {
         let ret = PollEvented::new(io);
-        ret.inner.registration.register_with(&ret.io, handle)?;
+        ret.inner.registration.register_with(ret.io.as_ref().unwrap(), handle)?;
         Ok(ret)
     }
 
+    /// Returns a shared reference to the underlying I/O object this readiness
+    /// stream is wrapping.
+    pub fn get_ref(&self) -> &E {
+        self.io.as_ref().unwrap()
+    }
+
+    /// Returns a mutable reference to the underlying I/O object this readiness
+    /// stream is wrapping.
+    pub fn get_mut(&mut self) -> &mut E {
+        self.io.as_mut().unwrap()
+    }
+
+    /// Consumes self, returning the inner I/O object
+    ///
+    /// This function will deregister the I/O resource from the reactor before
+    /// returning. If the deregistration operation fails, an error is returned.
+    ///
+    /// Note that deregistering does not guarantee that the I/O resource can be
+    /// registered with a different reactor. Some I/O resource types can only be
+    /// associated with a single reactor instance for their lifetime.
+    pub fn into_inner(mut self) -> io::Result<E> {
+        let io = self.io.take().unwrap();
+        self.inner.registration.deregister(&io)?;
+        Ok(io)
+    }
+
     /// Check the I/O resource's read readiness state.
     ///
     /// If the resource is not ready for a read then `Async::NotReady` is
@@ -241,30 +267,11 @@ where E: Evented
 
     /// Ensure that the I/O resource is registered with the reactor.
     fn register(&self) -> io::Result<()> {
-        self.inner.registration.register(&self.io)?;
+        self.inner.registration.register(self.io.as_ref().unwrap())?;
         Ok(())
     }
 }
 
-impl<E> PollEvented<E> {
-    /// Returns a shared reference to the underlying I/O object this readiness
-    /// stream is wrapping.
-    pub fn get_ref(&self) -> &E {
-        &self.io
-    }
-
-    /// Returns a mutable reference to the underlying I/O object this readiness
-    /// stream is wrapping.
-    pub fn get_mut(&mut self) -> &mut E {
-        &mut self.io
-    }
-
-    /// Consumes self, returning the inner I/O object
-    pub fn into_inner(self) -> E {
-        self.io
-    }
-}
-
 // ===== Read / Write impls =====
 
 impl<E> Read for PollEvented<E>
@@ -403,10 +410,19 @@ fn is_wouldblock<T>(r: &io::Result<T>) -> bool {
 }
 
 
-impl<E: fmt::Debug> fmt::Debug for PollEvented<E> {
+impl<E: Evented + fmt::Debug> fmt::Debug for PollEvented<E> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         f.debug_struct("PollEvented")
          .field("io", &self.io)
          .finish()
     }
 }
+
+impl<E: Evented> Drop for PollEvented<E> {
+    fn drop(&mut self) {
+        if let Some(io) = self.io.as_ref() {
+            // Ignore errors
+            let _ = self.inner.registration.deregister(io);
+        }
+    }
+}

tokio-reactor/src/registration.rs

@@ -104,6 +104,8 @@ impl Registration {
     /// the first call will establish the registration. Subsequent calls will be
     /// no-ops.
     ///
+    /// # Return
+    ///
     /// If the registration happened successfully, `Ok(true)` is returned.
     ///
     /// If an I/O resource has previously been successfully registered,
@@ -116,6 +118,35 @@ impl Registration {
         self.register2(io, || Handle::try_current())
     }
 
+    /// Deregister the I/O resource from the reactor it is associatd with.
+    ///
+    /// This function must be called before the I/O resource associated with the
+    /// registration is dropped.
+    ///
+    /// Note that deregistering does not guarantee that the I/O resource can be
+    /// registered with a different reactor. Some I/O resource types can only be
+    /// associated with a single reactor instance for their lifetime.
+    ///
+    /// # Return
+    ///
+    /// If the deregistration was successful, `Ok` is returned. Any calls to
+    /// `Reactor::turn` that happen after a successful call to `deregister` will
+    /// no longer result in notifications getting sent for this registration.
+    ///
+    /// `Err` is returned if an error is encountered.
+    pub fn deregister<T>(&mut self, io: &T) -> io::Result<()>
+    where T: Evented,
+    {
+        // The state does not need to be checked and coordination is not
+        // necessary as this function takes `&mut self`. This guarantees a
+        // single thread is accessing the instance.
+        if let Some(inner) = unsafe { (*self.inner.get()).as_ref() } {
+            inner.deregister(io)?;
+        }
+
+        Ok(())
+    }
+
     /// Register the I/O resource with the specified reactor.
     ///
     /// This function is safe to call concurrently and repeatedly. However, only
@@ -424,6 +455,19 @@ impl Inner {
         inner.register(self.token, direction, task);
     }
 
+    fn deregister<E: Evented>(&self, io: &E) -> io::Result<()> {
+        if self.token == ERROR {
+            return Err(io::Error::new(io::ErrorKind::Other, "failed to associate with reactor"));
+        }
+
+        let inner = match self.handle.inner() {
+            Some(inner) => inner,
+            None => return Err(io::Error::new(io::ErrorKind::Other, "reactor gone")),
+        };
+
+        inner.deregister_source(io)
+    }
+
     fn poll_ready(&self, direction: Direction, notify: bool)
         -> io::Result<Option<mio::Ready>>
     {