io: support PRIORITY epoll events (#5566)

Add support for epoll priority events. The commit adds `Interest::PRIORITY`, `ready`, and `ready_mut` functions to `AsyncFd`.

Closes #4885

Author: folkertdev

tokio/Cargo.toml

@@ -103,7 +103,7 @@ pin-project-lite = "0.2.0"
 
 # Everything else is optional...
 bytes = { version = "1.0.0", optional = true }
-mio = { version = "0.8.4", optional = true, default-features = false }
+mio = { version = "0.8.6", optional = true, default-features = false }
 num_cpus = { version = "1.8.0", optional = true }
 parking_lot = { version = "0.12.0", optional = true }
 

tokio/src/io/async_fd.rs

@@ -44,6 +44,11 @@ impl Interest {
     /// Writable interest includes write-closed events.
     pub const WRITABLE: Interest = Interest(mio::Interest::WRITABLE);
 
+    /// Returns a `Interest` set representing priority completion interests.
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))]
+    pub const PRIORITY: Interest = Interest(mio::Interest::PRIORITY);
+
     /// Returns true if the value includes readable interest.
     ///
     /// # Examples
@@ -78,6 +83,25 @@ impl Interest {
         self.0.is_writable()
     }
 
+    /// Returns true if the value includes priority interest.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::io::Interest;
+    ///
+    /// assert!(!Interest::READABLE.is_priority());
+    /// assert!(Interest::PRIORITY.is_priority());
+    ///
+    /// let both = Interest::READABLE | Interest::PRIORITY;
+    /// assert!(both.is_priority());
+    /// ```
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))]
+    pub const fn is_priority(self) -> bool {
+        self.0.is_priority()
+    }
+
     /// Add together two `Interest` values.
     ///
     /// This function works from a `const` context.
@@ -104,6 +128,8 @@ impl Interest {
         match self {
             Interest::READABLE => Ready::READABLE | Ready::READ_CLOSED,
             Interest::WRITABLE => Ready::WRITABLE | Ready::WRITE_CLOSED,
+            #[cfg(any(target_os = "linux", target_os = "android"))]
+            Interest::PRIORITY => Ready::PRIORITY | Ready::READ_CLOSED,
             _ => Ready::EMPTY,
         }
     }

tokio/src/io/interest.rs

@@ -73,6 +73,9 @@ cfg_io_driver! {
 impl<E: Source> PollEvented<E> {
     /// Creates a new `PollEvented` associated with the default reactor.
     ///
+    /// The returned `PollEvented` has readable and writable interests. For more control, use
+    /// [`Self::new_with_interest`].
+    ///
     /// # Panics
     ///
     /// This function panics if thread-local runtime is not set.

tokio/src/io/poll_evented.rs

@@ -7,6 +7,8 @@ const READABLE: usize = 0b0_01;
 const WRITABLE: usize = 0b0_10;
 const READ_CLOSED: usize = 0b0_0100;
 const WRITE_CLOSED: usize = 0b0_1000;
+#[cfg(any(target_os = "linux", target_os = "android"))]
+const PRIORITY: usize = 0b1_0000;
 
 /// Describes the readiness state of an I/O resources.
 ///
@@ -31,7 +33,17 @@ impl Ready {
     /// Returns a `Ready` representing write closed readiness.
     pub const WRITE_CLOSED: Ready = Ready(WRITE_CLOSED);
 
+    /// Returns a `Ready` representing priority readiness.
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))]
+    pub const PRIORITY: Ready = Ready(PRIORITY);
+
+    /// Returns a `Ready` representing readiness for all operations.
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED | PRIORITY);
+
     /// Returns a `Ready` representing readiness for all operations.
+    #[cfg(not(any(target_os = "linux", target_os = "android")))]
     pub const ALL: Ready = Ready(READABLE | WRITABLE | READ_CLOSED | WRITE_CLOSED);
 
     // Must remain crate-private to avoid adding a public dependency on Mio.
@@ -65,6 +77,13 @@ impl Ready {
             ready |= Ready::WRITE_CLOSED;
         }
 
+        #[cfg(any(target_os = "linux", target_os = "android"))]
+        {
+            if event.is_priority() {
+                ready |= Ready::PRIORITY;
+            }
+        }
+
         ready
     }
 
@@ -144,6 +163,23 @@ impl Ready {
         self.contains(Ready::WRITE_CLOSED)
     }
 
+    /// Returns `true` if the value includes priority `readiness`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::io::Ready;
+    ///
+    /// assert!(!Ready::EMPTY.is_priority());
+    /// assert!(!Ready::WRITABLE.is_priority());
+    /// assert!(Ready::PRIORITY.is_priority());
+    /// ```
+    #[cfg(any(target_os = "linux", target_os = "android"))]
+    #[cfg_attr(docsrs, doc(cfg(any(target_os = "linux", target_os = "android"))))]
+    pub fn is_priority(self) -> bool {
+        self.contains(Ready::PRIORITY)
+    }
+
     /// Returns true if `self` is a superset of `other`.
     ///
     /// `other` may represent more than one readiness operations, in which case
@@ -191,6 +227,12 @@ cfg_io_readiness! {
                 ready |= Ready::WRITE_CLOSED;
             }
 
+            #[cfg(any(target_os = "linux", target_os = "android"))]
+            if interest.is_priority() {
+                ready |= Ready::PRIORITY;
+                ready |= Ready::READ_CLOSED;
+            }
+
             ready
         }
 
@@ -240,11 +282,16 @@ impl ops::Sub<Ready> for Ready {
 
 impl fmt::Debug for Ready {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("Ready")
-            .field("is_readable", &self.is_readable())
+        let mut fmt = fmt.debug_struct("Ready");
+
+        fmt.field("is_readable", &self.is_readable())
             .field("is_writable", &self.is_writable())
             .field("is_read_closed", &self.is_read_closed())
-            .field("is_write_closed", &self.is_write_closed())
-            .finish()
+            .field("is_write_closed", &self.is_write_closed());
+
+        #[cfg(any(target_os = "linux", target_os = "android"))]
+        fmt.field("is_priority", &self.is_priority());
+
+        fmt.finish()
     }
 }

tokio/src/io/ready.rs

@@ -63,6 +63,18 @@ pub(crate) struct ReadyEvent {
     is_shutdown: bool,
 }
 
+cfg_net_unix!(
+    impl ReadyEvent {
+        pub(crate) fn with_ready(&self, ready: Ready) -> Self {
+            Self {
+                ready,
+                tick: self.tick,
+                is_shutdown: self.is_shutdown,
+            }
+        }
+    }
+);
+
 struct IoDispatcher {
     allocator: slab::Allocator<ScheduledIo>,
     is_shutdown: bool,

tokio/src/runtime/io/mod.rs

@@ -599,3 +599,89 @@ fn driver_shutdown_wakes_poll_race() {
         assert_err!(futures::executor::block_on(poll_writable(&afd_a)));
     }
 }
+
+#[tokio::test]
+#[cfg(any(target_os = "linux", target_os = "android"))]
+async fn priority_event_on_oob_data() {
+    use std::net::SocketAddr;
+
+    use tokio::io::Interest;
+
+    let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
+
+    let listener = std::net::TcpListener::bind(addr).unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    let client = std::net::TcpStream::connect(addr).unwrap();
+    let client = AsyncFd::with_interest(client, Interest::PRIORITY).unwrap();
+
+    let (stream, _) = listener.accept().unwrap();
+
+    // Sending out of band data should trigger priority event.
+    send_oob_data(&stream, b"hello").unwrap();
+
+    let _ = client.ready(Interest::PRIORITY).await.unwrap();
+}
+
+#[cfg(any(target_os = "linux", target_os = "android"))]
+fn send_oob_data<S: AsRawFd>(stream: &S, data: &[u8]) -> io::Result<usize> {
+    unsafe {
+        let res = libc::send(
+            stream.as_raw_fd(),
+            data.as_ptr().cast(),
+            data.len(),
+            libc::MSG_OOB,
+        );
+        if res == -1 {
+            Err(io::Error::last_os_error())
+        } else {
+            Ok(res as usize)
+        }
+    }
+}
+
+#[tokio::test]
+async fn clear_ready_matching_clears_ready() {
+    use tokio::io::{Interest, Ready};
+
+    let (a, mut b) = socketpair();
+
+    let afd_a = AsyncFd::new(a).unwrap();
+    b.write_all(b"0").unwrap();
+
+    let mut guard = afd_a
+        .ready(Interest::READABLE | Interest::WRITABLE)
+        .await
+        .unwrap();
+
+    assert_eq!(guard.ready(), Ready::READABLE | Ready::WRITABLE);
+
+    guard.clear_ready_matching(Ready::READABLE);
+    assert_eq!(guard.ready(), Ready::WRITABLE);
+
+    guard.clear_ready_matching(Ready::WRITABLE);
+    assert_eq!(guard.ready(), Ready::EMPTY);
+}
+
+#[tokio::test]
+async fn clear_ready_matching_clears_ready_mut() {
+    use tokio::io::{Interest, Ready};
+
+    let (a, mut b) = socketpair();
+
+    let mut afd_a = AsyncFd::new(a).unwrap();
+    b.write_all(b"0").unwrap();
+
+    let mut guard = afd_a
+        .ready_mut(Interest::READABLE | Interest::WRITABLE)
+        .await
+        .unwrap();
+
+    assert_eq!(guard.ready(), Ready::READABLE | Ready::WRITABLE);
+
+    guard.clear_ready_matching(Ready::READABLE);
+    assert_eq!(guard.ready(), Ready::WRITABLE);
+
+    guard.clear_ready_matching(Ready::WRITABLE);
+    assert_eq!(guard.ready(), Ready::EMPTY);
+}