From c03187281afe92b672d101e278bf1b91e315e417 Mon Sep 17 00:00:00 2001
From: Noah Kennedy <nkennedy@cloudflare.com>
Date: Fri, 20 Oct 2023 12:07:44 -0500
Subject: [PATCH] add extra apis

---
 tokio/src/io/async_fd.rs             | 27 +++++++++++++++++++++++++++
 tokio/src/io/poll_evented.rs         |  2 +-
 tokio/src/net/tcp/listener.rs        |  1 -
 tokio/src/net/unix/listener.rs       | 11 +++++++++++
 tokio/src/runtime/io/registration.rs |  2 +-
 5 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/tokio/src/io/async_fd.rs b/tokio/src/io/async_fd.rs
index b27c60bf6..6f54d6387 100644
--- a/tokio/src/io/async_fd.rs
+++ b/tokio/src/io/async_fd.rs
@@ -243,6 +243,18 @@ impl<T: AsRawFd> AsyncFd<T> {
         Self::new_with_handle_and_interest(inner, scheduler::Handle::current(), interest)
     }
 
+    /// Create a new AsyncFd with the provided raw epoll flags for registration.
+    ///
+    /// These flags replace any epoll flags would normally set when registering the fd.
+    #[track_caller]
+    #[cfg(all(target_os = "linux", tokio_unstable))]
+    pub fn with_flags(inner: T, flags: u32) -> io::Result<Self>
+    where
+        T: AsRawFd,
+    {
+        Self::new_with_handle_and_flags(inner, scheduler::Handle::current(), flags)
+    }
+
     #[track_caller]
     pub(crate) fn new_with_handle_and_interest(
         inner: T,
@@ -260,6 +272,21 @@ impl<T: AsRawFd> AsyncFd<T> {
         })
     }
 
+    #[track_caller]
+    #[cfg(all(target_os = "linux", tokio_unstable))]
+    pub(crate) fn new_with_handle_and_flags(
+        mut inner: T,
+        handle: scheduler::Handle,
+        flags: u32,
+    ) -> io::Result<Self> {
+        let registration = Registration::new_with_flags_and_handle(&mut inner, flags, handle)?;
+
+        Ok(AsyncFd {
+            registration,
+            inner: Some(inner),
+        })
+    }
+
     /// Returns a shared reference to the backing object of this [`AsyncFd`].
     #[inline]
     pub fn get_ref(&self) -> &T {
diff --git a/tokio/src/io/poll_evented.rs b/tokio/src/io/poll_evented.rs
index f2f5442b0..71076eac6 100644
--- a/tokio/src/io/poll_evented.rs
+++ b/tokio/src/io/poll_evented.rs
@@ -122,7 +122,7 @@ impl<E: std::os::unix::io::AsRawFd + Source> PollEvented<E> {
         flags: u32,
         handle: scheduler::Handle,
     ) -> io::Result<Self> {
-        let registration = Registration::new_with_interest_and_handle_raw(&mut io, flags, handle)?;
+        let registration = Registration::new_with_flags_and_handle(&mut io, flags, handle)?;
         Ok(Self {
             io: Some(io),
             registration,
diff --git a/tokio/src/net/tcp/listener.rs b/tokio/src/net/tcp/listener.rs
index 281ced04c..0c344dd1a 100644
--- a/tokio/src/net/tcp/listener.rs
+++ b/tokio/src/net/tcp/listener.rs
@@ -243,7 +243,6 @@ impl TcpListener {
     /// Create a new TcpListener with the provided raw epoll flags.
     ///
     /// These flags replace any epoll flags would normally set when registering the fd.
-    #[cfg(all(target_os = "linux"))]
     #[track_caller]
     #[cfg(all(target_os = "linux", tokio_unstable))]
     pub fn from_std_with_flags(listener: net::TcpListener, flags: u32) -> io::Result<TcpListener> {
diff --git a/tokio/src/net/unix/listener.rs b/tokio/src/net/unix/listener.rs
index a7e9115ea..5a937b25c 100644
--- a/tokio/src/net/unix/listener.rs
+++ b/tokio/src/net/unix/listener.rs
@@ -114,6 +114,17 @@ impl UnixListener {
         Ok(UnixListener { io })
     }
 
+    /// Create a new UnixListener with the provided raw epoll flags.
+    ///
+    /// These flags replace any epoll flags would normally set when registering the fd.
+    #[track_caller]
+    #[cfg(all(target_os = "linux", tokio_unstable))]
+    pub fn from_std_with_flags(listener: net::UnixListener, flags: u32) -> io::Result<UnixListener> {
+        let io = mio::net::UnixListener::from_std(listener);
+        let io = PollEvented::new_raw(io, flags)?;
+        Ok(UnixListener { io })
+    }
+
     /// Turns a [`tokio::net::UnixListener`] into a [`std::os::unix::net::UnixListener`].
     ///
     /// The returned [`std::os::unix::net::UnixListener`] will have nonblocking mode
diff --git a/tokio/src/runtime/io/registration.rs b/tokio/src/runtime/io/registration.rs
index 636feb11b..e6f27e044 100644
--- a/tokio/src/runtime/io/registration.rs
+++ b/tokio/src/runtime/io/registration.rs
@@ -90,7 +90,7 @@ impl Registration {
     /// - `Ok` if the registration happened successfully
     /// - `Err` if an error was encountered during registration
     #[cfg(all(target_os = "linux", tokio_unstable))]
-    pub(crate) fn new_with_interest_and_handle_raw(
+    pub(crate) fn new_with_flags_and_handle(
         io: &mut impl std::os::unix::io::AsRawFd,
         flags: u32,
         handle: scheduler::Handle,