Add std::os::unix::{register, unregister, Interest}

Closes async-rs#293, based on async-rs#293 (comment)

Author: leo60228

Diff for: src/net/driver/mod.rs

@@ -1,5 +1,6 @@
 use std::fmt;
 use std::sync::{Arc, Mutex};
+use std::os::raw::c_int;
 
 use mio::{self, Evented};
 use once_cell::sync::Lazy;
@@ -11,19 +12,22 @@ use crate::utils::abort_on_panic;
 
 /// Data associated with a registered I/O handle.
 #[derive(Debug)]
-struct Entry {
+pub struct Entry {
     /// A unique identifier.
     token: mio::Token,
 
+    /// File descriptor.
+    fd: Option<c_int>,
+
     /// Tasks that are blocked on reading from this I/O handle.
-    readers: Mutex<Vec<Waker>>,
+    pub readers: Mutex<Vec<Waker>>,
 
     /// Thasks that are blocked on writing to this I/O handle.
-    writers: Mutex<Vec<Waker>>,
+    pub writers: Mutex<Vec<Waker>>,
 }
 
 /// The state of a networking driver.
-struct Reactor {
+pub struct Reactor {
     /// A mio instance that polls for new events.
     poller: mio::Poll,
 
@@ -51,14 +55,14 @@ impl Reactor {
         };
 
         // Register a dummy I/O handle for waking up the polling thread.
-        let entry = reactor.register(&reactor.notify_reg.0)?;
+        let entry = reactor.register(&reactor.notify_reg.0, None)?;
         reactor.notify_token = entry.token;
 
         Ok(reactor)
     }
 
     /// Registers an I/O event source and returns its associated entry.
-    fn register(&self, source: &dyn Evented) -> io::Result<Arc<Entry>> {
+    pub fn register(&self, source: &dyn Evented, fd: Option<c_int>) -> io::Result<Arc<Entry>> {
         let mut entries = self.entries.lock().unwrap();
 
         // Reserve a vacant spot in the slab and use its key as the token value.
@@ -68,6 +72,7 @@ impl Reactor {
         // Allocate an entry and insert it into the slab.
         let entry = Arc::new(Entry {
             token,
+            fd,
             readers: Mutex::new(Vec::new()),
             writers: Mutex::new(Vec::new()),
         });
@@ -82,7 +87,7 @@ impl Reactor {
     }
 
     /// Deregisters an I/O event source associated with an entry.
-    fn deregister(&self, source: &dyn Evented, entry: &Entry) -> io::Result<()> {
+    pub fn deregister(&self, source: &dyn Evented, entry: &Entry) -> io::Result<()> {
         // Deregister the I/O object from the mio instance.
         self.poller.deregister(source)?;
 
@@ -92,6 +97,17 @@ impl Reactor {
         Ok(())
     }
 
+    /// Deregisters an I/O event source associated with a file descriptor.
+    pub fn deregister_fd(&self, source: &dyn Evented, fd: c_int) -> io::Result<()> {
+        // Deregister the I/O object from the mio instance.
+        self.poller.deregister(source)?;
+
+        // Remove the entry associated with the I/O object.
+        self.entries.lock().unwrap().retain(|_, e| e.fd != Some(fd));
+
+        Ok(())
+    }
+
     // fn notify(&self) {
     //     self.notify_reg
     //         .1
@@ -101,7 +117,7 @@ impl Reactor {
 }
 
 /// The state of the global networking driver.
-static REACTOR: Lazy<Reactor> = Lazy::new(|| {
+pub static REACTOR: Lazy<Reactor> = Lazy::new(|| {
     // Spawn a thread that waits on the poller for new events and wakes up tasks blocked on I/O
     // handles.
     std::thread::Builder::new()
@@ -181,7 +197,7 @@ impl<T: Evented> Watcher<T> {
     pub fn new(source: T) -> Watcher<T> {
         Watcher {
             entry: REACTOR
-                .register(&source)
+                .register(&source, None)
                 .expect("cannot register an I/O event source"),
             source: Some(source),
         }

Diff for: src/os/unix/io.rs

@@ -1,7 +1,39 @@
 //! Unix-specific I/O extensions.
 
+use crate::task::Context;
+use crate::io;
+
 cfg_not_docs! {
     pub use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
+    use mio::unix::EventedFd;
+    use crate::net::driver::REACTOR;
+
+    /// Registers an I/O handle so that the current task gets woken up when it becomes ready.
+    #[cfg(feature = "unstable")]
+    pub fn register(cx: &mut Context<'_>, fd: RawFd, interest: Interest) -> io::Result<()> {
+        let evented = EventedFd(&fd);
+        let entry = REACTOR.register(&evented, Some(fd))?;
+        if interest == Interest::Read || interest == Interest::Both {
+            let mut list = entry.readers.lock().unwrap();
+            if list.iter().all(|w| !w.will_wake(cx.waker())) {
+                list.push(cx.waker().clone());
+            }
+        }
+        if interest == Interest::Write || interest == Interest::Both {
+            let mut list = entry.writers.lock().unwrap();
+            if list.iter().all(|w| !w.will_wake(cx.waker())) {
+                list.push(cx.waker().clone());
+            }
+        }
+        Ok(())
+    }
+
+    /// Unregisters an I/O handle.
+    #[cfg(feature = "unstable")]
+    pub fn unregister(fd: RawFd) -> io::Result<()> {
+        let evented = EventedFd(&fd);
+        REACTOR.deregister_fd(&evented, fd)
+    }
 }
 
 cfg_docs! {
@@ -51,4 +83,30 @@ cfg_docs! {
         /// and must close the descriptor once it's no longer needed.
         fn into_raw_fd(self) -> RawFd;
     }
+
+    /// Registers an I/O handle so that the current task gets woken up when it becomes ready.
+    #[doc(cfg(unstable))]
+    pub fn register(cx: &mut Context<'_>, fd: RawFd, interest: Interest) -> io::Result<()> {
+        unreachable!()
+    }
+
+    /// Unregisters an I/O handle.
+    #[doc(cfg(unstable))]
+    pub fn unregister(fd: RawFd) -> io::Result<()> {
+        unreachable!()
+    }
+}
+
+cfg_unstable! {
+    /// Decides which possibility a task is interested in.
+    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+    pub enum Interest {
+        /// A task is interested in reading from a file descriptor.
+        Read,
+        /// A task is interested in writing to a file descriptor.
+        Write,
+        /// A task is interested in either being able to write or being able to read from a file
+        /// descriptor.
+        Both,
+    }
 }

Diff for: tests/register_rawfd.rs

@@ -0,0 +1,41 @@
+#![cfg(all(unix, feature = "unstable"))]
+
+use async_std::os::unix::io::{register, unregister, Interest, IntoRawFd, FromRawFd};
+use async_std::future::poll_fn;
+use async_std::net::{TcpListener, TcpStream};
+use async_std::task::{self, block_on, Poll};
+
+#[test]
+fn register_stream() {
+    block_on(async {
+        for i in 0..3 {
+            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+            let addr = listener.local_addr().unwrap();
+            let t = task::spawn(async move { listener.accept().await });
+
+            let stream = TcpStream::connect(&addr).await.unwrap().into_raw_fd();
+
+            let mut woken = false;
+            poll_fn(|cx| {
+                if woken {
+                    Poll::Ready(())
+                } else {
+                    match i {
+                        0 => register(cx, stream, Interest::Read).unwrap(),
+                        1 => register(cx, stream, Interest::Write).unwrap(),
+                        2 => register(cx, stream, Interest::Both).unwrap(),
+                        _ => unreachable!(),
+                    }
+
+                    unregister(stream).unwrap();
+                    woken = true;
+
+                    Poll::Pending
+                }
+            }).await;
+
+            t.await.unwrap();
+            unsafe { TcpStream::from_raw_fd(stream); }
+        }
+    });
+}