Merge pull request #230 from Berrysoft/refactor/shared-fd

refactor(driver): add SharedFd & OwnedFd and require them in operations

Author: Berrysoft

compio-driver/Cargo.toml

@@ -35,6 +35,7 @@ compio-log = { workspace = true }
 # Utils
 cfg-if = { workspace = true }
 crossbeam-channel = { workspace = true }
+futures-util = { workspace = true }
 slab = { workspace = true }
 socket2 = { workspace = true }
 

compio-driver/src/fd.rs

@@ -0,0 +1,172 @@
+#[cfg(unix)]
+use std::os::fd::FromRawFd;
+#[cfg(windows)]
+use std::os::windows::io::{
+    FromRawHandle, FromRawSocket, OwnedHandle, OwnedSocket, RawHandle, RawSocket,
+};
+use std::{
+    future::{poll_fn, Future},
+    mem::ManuallyDrop,
+    panic::RefUnwindSafe,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    task::Poll,
+};
+
+use futures_util::task::AtomicWaker;
+
+use crate::{AsRawFd, OwnedFd, RawFd};
+
+#[derive(Debug)]
+struct Inner {
+    fd: OwnedFd,
+    // whether there is a future waiting
+    waits: AtomicBool,
+    waker: AtomicWaker,
+}
+
+impl RefUnwindSafe for Inner {}
+
+/// A shared fd. It is passed to the operations to make sure the fd won't be
+/// closed before the operations complete.
+#[derive(Debug, Clone)]
+pub struct SharedFd(Arc<Inner>);
+
+impl SharedFd {
+    /// Create the shared fd from an owned fd.
+    pub fn new(fd: impl Into<OwnedFd>) -> Self {
+        Self(Arc::new(Inner {
+            fd: fd.into(),
+            waits: AtomicBool::new(false),
+            waker: AtomicWaker::new(),
+        }))
+    }
+
+    /// Try to take the inner owned fd.
+    pub fn try_unwrap(self) -> Result<OwnedFd, Self> {
+        let this = ManuallyDrop::new(self);
+        if let Some(fd) = unsafe { Self::try_unwrap_inner(&this) } {
+            Ok(fd)
+        } else {
+            Err(ManuallyDrop::into_inner(this))
+        }
+    }
+
+    // SAFETY: if `Some` is returned, the method should not be called again.
+    unsafe fn try_unwrap_inner(this: &ManuallyDrop<Self>) -> Option<OwnedFd> {
+        let ptr = ManuallyDrop::new(std::ptr::read(&this.0));
+        // The ptr is duplicated without increasing the strong count, should forget.
+        match Arc::try_unwrap(ManuallyDrop::into_inner(ptr)) {
+            Ok(inner) => Some(inner.fd),
+            Err(ptr) => {
+                std::mem::forget(ptr);
+                None
+            }
+        }
+    }
+
+    /// Wait and take the inner owned fd.
+    pub fn take(self) -> impl Future<Output = Option<OwnedFd>> {
+        let this = ManuallyDrop::new(self);
+        async move {
+            if !this.0.waits.swap(true, Ordering::AcqRel) {
+                poll_fn(move |cx| {
+                    if let Some(fd) = unsafe { Self::try_unwrap_inner(&this) } {
+                        return Poll::Ready(Some(fd));
+                    }
+
+                    this.0.waker.register(cx.waker());
+
+                    if let Some(fd) = unsafe { Self::try_unwrap_inner(&this) } {
+                        Poll::Ready(Some(fd))
+                    } else {
+                        Poll::Pending
+                    }
+                })
+                .await
+            } else {
+                None
+            }
+        }
+    }
+}
+
+impl Drop for SharedFd {
+    fn drop(&mut self) {
+        // It's OK to wake multiple times.
+        if Arc::strong_count(&self.0) == 2 {
+            self.0.waker.wake()
+        }
+    }
+}
+
+#[cfg(windows)]
+#[doc(hidden)]
+impl SharedFd {
+    pub unsafe fn to_file(&self) -> ManuallyDrop<std::fs::File> {
+        ManuallyDrop::new(std::fs::File::from_raw_handle(self.as_raw_fd() as _))
+    }
+
+    pub unsafe fn to_socket(&self) -> ManuallyDrop<socket2::Socket> {
+        ManuallyDrop::new(socket2::Socket::from_raw_socket(self.as_raw_fd() as _))
+    }
+}
+
+#[cfg(unix)]
+#[doc(hidden)]
+impl SharedFd {
+    pub unsafe fn to_file(&self) -> ManuallyDrop<std::fs::File> {
+        ManuallyDrop::new(std::fs::File::from_raw_fd(self.as_raw_fd() as _))
+    }
+
+    pub unsafe fn to_socket(&self) -> ManuallyDrop<socket2::Socket> {
+        ManuallyDrop::new(socket2::Socket::from_raw_fd(self.as_raw_fd() as _))
+    }
+}
+
+impl AsRawFd for SharedFd {
+    fn as_raw_fd(&self) -> RawFd {
+        self.0.fd.as_raw_fd()
+    }
+}
+
+#[cfg(windows)]
+impl FromRawHandle for SharedFd {
+    unsafe fn from_raw_handle(handle: RawHandle) -> Self {
+        Self::new(OwnedFd::File(OwnedHandle::from_raw_handle(handle)))
+    }
+}
+
+#[cfg(windows)]
+impl FromRawSocket for SharedFd {
+    unsafe fn from_raw_socket(sock: RawSocket) -> Self {
+        Self::new(OwnedFd::Socket(OwnedSocket::from_raw_socket(sock)))
+    }
+}
+
+#[cfg(unix)]
+impl FromRawFd for SharedFd {
+    unsafe fn from_raw_fd(fd: RawFd) -> Self {
+        Self::new(OwnedFd::from_raw_fd(fd))
+    }
+}
+
+impl From<OwnedFd> for SharedFd {
+    fn from(value: OwnedFd) -> Self {
+        Self::new(value)
+    }
+}
+
+/// Get a clone of [`SharedFd`].
+pub trait ToSharedFd {
+    /// Return a cloned [`SharedFd`].
+    fn to_shared_fd(&self) -> SharedFd;
+}
+
+impl ToSharedFd for SharedFd {
+    fn to_shared_fd(&self) -> SharedFd {
+        self.clone()
+    }
+}

compio-driver/src/fusion/mod.rs

@@ -7,7 +7,7 @@ mod iour;
 pub(crate) mod op;
 
 #[cfg_attr(all(doc, docsrs), doc(cfg(all())))]
-pub use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
+pub use std::os::fd::{AsRawFd, OwnedFd, RawFd};
 use std::{io, task::Poll, time::Duration};
 
 pub use driver_type::DriverType;

compio-driver/src/fusion/op.rs

@@ -5,6 +5,7 @@ use socket2::SockAddr;
 
 use super::*;
 pub use crate::unix::op::*;
+use crate::SharedFd;
 
 macro_rules! op {
     (<$($ty:ident: $trait:ident),* $(,)?> $name:ident( $($arg:ident: $arg_t:ident),* $(,)? )) => {
@@ -91,9 +92,9 @@ mod iour { pub use crate::sys::iour::{op::*, OpCode}; }
 #[rustfmt::skip]
 mod poll { pub use crate::sys::poll::{op::*, OpCode}; }
 
-op!(<T: IoBufMut> RecvFrom(fd: RawFd, buffer: T));
-op!(<T: IoBuf> SendTo(fd: RawFd, buffer: T, addr: SockAddr));
-op!(<T: IoVectoredBufMut> RecvFromVectored(fd: RawFd, buffer: T));
-op!(<T: IoVectoredBuf> SendToVectored(fd: RawFd, buffer: T, addr: SockAddr));
-op!(<> FileStat(fd: RawFd));
+op!(<T: IoBufMut> RecvFrom(fd: SharedFd, buffer: T));
+op!(<T: IoBuf> SendTo(fd: SharedFd, buffer: T, addr: SockAddr));
+op!(<T: IoVectoredBufMut> RecvFromVectored(fd: SharedFd, buffer: T));
+op!(<T: IoVectoredBuf> SendToVectored(fd: SharedFd, buffer: T, addr: SockAddr));
+op!(<> FileStat(fd: SharedFd));
 op!(<> PathStat(path: CString, follow_symlink: bool));

compio-driver/src/iocp/cp/global.rs

@@ -74,7 +74,7 @@ fn iocp_start() -> io::Result<()> {
                     )
                 ) {
                     error!(
-                        "fail to dispatch entry ({}, {}, {:p}) to driver {:p}: {:?}",
+                        "fail to dispatch entry ({}, {}, {:p}) to driver {:x}: {:?}",
                         entry.dwNumberOfBytesTransferred,
                         entry.lpCompletionKey,
                         entry.lpOverlapped,

compio-driver/src/iocp/cp/mod.rs

@@ -158,7 +158,7 @@ impl CompletionPort {
                         )
                     ) {
                         error!(
-                            "fail to repost entry ({}, {}, {:p}) to driver {:p}: {:?}",
+                            "fail to repost entry ({}, {}, {:p}) to driver {:x}: {:?}",
                             entry.dwNumberOfBytesTransferred,
                             entry.lpCompletionKey,
                             entry.lpOverlapped,

compio-driver/src/iocp/cp/multi.rs

@@ -28,7 +28,7 @@ impl Port {
     }
 
     pub fn poll(&self, timeout: Option<Duration>) -> io::Result<impl Iterator<Item = Entry> + '_> {
-        let current_id = self.as_raw_handle();
+        let current_id = self.as_raw_handle() as _;
         self.port.poll(timeout, Some(current_id))
     }
 }

compio-driver/src/iocp/mod.rs

@@ -4,9 +4,9 @@ use std::{
     mem::ManuallyDrop,
     os::{
         raw::c_void,
-        windows::prelude::{
-            AsRawHandle, AsRawSocket, FromRawHandle, FromRawSocket, IntoRawHandle, IntoRawSocket,
-            RawHandle,
+        windows::{
+            io::{OwnedHandle, OwnedSocket},
+            prelude::{AsRawHandle, AsRawSocket},
         },
     },
     pin::Pin,
@@ -44,66 +44,53 @@ pub(crate) use windows_sys::Win32::Networking::WinSock::{
 /// On windows, handle and socket are in the same size.
 /// Both of them could be attached to an IOCP.
 /// Therefore, both could be seen as fd.
-pub type RawFd = RawHandle;
+pub type RawFd = isize;
 
 /// Extracts raw fds.
 pub trait AsRawFd {
     /// Extracts the raw fd.
     fn as_raw_fd(&self) -> RawFd;
 }
 
-/// Construct IO objects from raw fds.
-pub trait FromRawFd {
-    /// Constructs an IO object from the specified raw fd.
-    ///
-    /// # Safety
-    ///
-    /// The `fd` passed in must:
-    ///   - be a valid open handle or socket,
-    ///   - be opened with `FILE_FLAG_OVERLAPPED` if it's a file handle,
-    ///   - have not been attached to a driver.
-    unsafe fn from_raw_fd(fd: RawFd) -> Self;
-}
-
-/// Consumes an object and acquire ownership of its raw fd.
-pub trait IntoRawFd {
-    /// Consumes this object, returning the raw underlying fd.
-    fn into_raw_fd(self) -> RawFd;
+/// Owned handle or socket on Windows.
+#[derive(Debug)]
+pub enum OwnedFd {
+    /// Win32 handle.
+    File(OwnedHandle),
+    /// Windows socket handle.
+    Socket(OwnedSocket),
 }
 
-impl AsRawFd for std::fs::File {
+impl AsRawFd for OwnedFd {
     fn as_raw_fd(&self) -> RawFd {
-        self.as_raw_handle()
-    }
-}
-
-impl AsRawFd for socket2::Socket {
-    fn as_raw_fd(&self) -> RawFd {
-        self.as_raw_socket() as _
+        match self {
+            Self::File(fd) => fd.as_raw_handle() as _,
+            Self::Socket(s) => s.as_raw_socket() as _,
+        }
     }
 }
 
-impl FromRawFd for std::fs::File {
-    unsafe fn from_raw_fd(fd: RawFd) -> Self {
-        Self::from_raw_handle(fd)
+impl From<OwnedHandle> for OwnedFd {
+    fn from(value: OwnedHandle) -> Self {
+        Self::File(value)
     }
 }
 
-impl FromRawFd for socket2::Socket {
-    unsafe fn from_raw_fd(fd: RawFd) -> Self {
-        Self::from_raw_socket(fd as _)
+impl From<std::fs::File> for OwnedFd {
+    fn from(value: std::fs::File) -> Self {
+        Self::File(OwnedHandle::from(value))
     }
 }
 
-impl IntoRawFd for std::fs::File {
-    fn into_raw_fd(self) -> RawFd {
-        self.into_raw_handle()
+impl From<OwnedSocket> for OwnedFd {
+    fn from(value: OwnedSocket) -> Self {
+        Self::Socket(value)
     }
 }
 
-impl IntoRawFd for socket2::Socket {
-    fn into_raw_fd(self) -> RawFd {
-        self.into_raw_socket() as _
+impl From<socket2::Socket> for OwnedFd {
+    fn from(value: socket2::Socket) -> Self {
+        Self::Socket(OwnedSocket::from(value))
     }
 }
 
@@ -298,7 +285,7 @@ impl Driver {
 
 impl AsRawFd for Driver {
     fn as_raw_fd(&self) -> RawFd {
-        self.port.as_raw_handle()
+        self.port.as_raw_handle() as _
     }
 }