feat(driver): make thread safety optional

Author: Berrysoft

compio-dispatcher/tests/listener.rs

@@ -1,6 +1,6 @@
 use std::num::NonZeroUsize;
 
-use compio_buf::arrayvec::ArrayVec;
+use compio_buf::{IntoInner, arrayvec::ArrayVec};
 use compio_dispatcher::Dispatcher;
 use compio_io::{AsyncReadExt, AsyncWriteExt};
 use compio_net::{TcpListener, TcpStream};
@@ -27,9 +27,11 @@ async fn listener_dispatch() {
     });
     let mut handles = FuturesUnordered::new();
     for _i in 0..CLIENT_NUM {
-        let (mut srv, _) = listener.accept().await.unwrap();
+        let (srv, _) = listener.accept().await.unwrap();
+        let srv = srv.try_into_sync().unwrap();
         let handle = dispatcher
             .dispatch(move || async move {
+                let mut srv = srv.into_inner();
                 let (_, buf) = srv.read_exact(ArrayVec::<u8, 12>::new()).await.unwrap();
                 assert_eq!(buf.as_slice(), b"Hello world!");
             })

compio-driver/src/fd.rs

@@ -3,35 +3,35 @@ use std::os::fd::FromRawFd;
 #[cfg(windows)]
 use std::os::windows::io::{FromRawHandle, FromRawSocket, RawHandle, RawSocket};
 use std::{
+    cell::{Cell, RefCell},
     future::{Future, poll_fn},
     mem::ManuallyDrop,
     ops::Deref,
     panic::RefUnwindSafe,
-    sync::{
-        Arc,
-        atomic::{AtomicBool, Ordering},
-    },
-    task::Poll,
+    rc::Rc,
+    task::{Poll, Waker},
 };
 
-use futures_util::task::AtomicWaker;
-
 use crate::{AsFd, AsRawFd, BorrowedFd, RawFd};
 
 #[derive(Debug)]
 struct Inner<T> {
     fd: T,
     // whether there is a future waiting
-    waits: AtomicBool,
-    waker: AtomicWaker,
+    waits: Cell<bool>,
+    waker: RefCell<Option<Waker>>,
 }
 
 impl<T> RefUnwindSafe for Inner<T> {}
 
 /// A shared fd. It is passed to the operations to make sure the fd won't be
 /// closed before the operations complete.
 #[derive(Debug)]
-pub struct SharedFd<T>(Arc<Inner<T>>);
+pub struct SharedFd<T>(Rc<Inner<T>>);
+// We use `Rc` internally to avoid the overhead of `Arc`. It is not `Send` or
+// `Sync`, but we have to access it in another thread when processing blocking
+// operations. We ensure that the access is safe because there will be only
+// one thread accessing it at a time.
 
 impl<T: AsFd> SharedFd<T> {
     /// Create the shared fd from an owned fd.
@@ -46,10 +46,10 @@ impl<T> SharedFd<T> {
     /// # Safety
     /// * T should own the fd.
     pub unsafe fn new_unchecked(fd: T) -> Self {
-        Self(Arc::new(Inner {
+        Self(Rc::new(Inner {
             fd,
-            waits: AtomicBool::new(false),
-            waker: AtomicWaker::new(),
+            waits: Cell::new(false),
+            waker: RefCell::new(None),
         }))
     }
 
@@ -67,7 +67,7 @@ impl<T> SharedFd<T> {
     unsafe fn try_unwrap_inner(this: &ManuallyDrop<Self>) -> Option<T> {
         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)) {
+        match Rc::try_unwrap(ManuallyDrop::into_inner(ptr)) {
             Ok(inner) => Some(inner.fd),
             Err(ptr) => {
                 std::mem::forget(ptr);
@@ -80,13 +80,13 @@ impl<T> SharedFd<T> {
     pub fn take(self) -> impl Future<Output = Option<T>> {
         let this = ManuallyDrop::new(self);
         async move {
-            if !this.0.waits.swap(true, Ordering::AcqRel) {
+            if !this.0.waits.replace(true) {
                 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());
+                    this.0.waker.borrow_mut().replace(cx.waker().clone());
 
                     if let Some(fd) = unsafe { Self::try_unwrap_inner(&this) } {
                         Poll::Ready(Some(fd))
@@ -105,8 +105,10 @@ impl<T> SharedFd<T> {
 impl<T> Drop for SharedFd<T> {
     fn drop(&mut self) {
         // It's OK to wake multiple times.
-        if Arc::strong_count(&self.0) == 2 && self.0.waits.load(Ordering::Acquire) {
-            self.0.waker.wake()
+        if Rc::strong_count(&self.0) == 2 && self.0.waits.get() {
+            if let Some(waker) = self.0.waker.borrow_mut().take() {
+                waker.wake()
+            }
         }
     }
 }

compio-driver/src/iocp/mod.rs

@@ -44,6 +44,17 @@ pub enum OwnedFd {
     Socket(OwnedSocket),
 }
 
+impl OwnedFd {
+    /// Creates a new [`OwnedFd`] instance that shares the same underlying
+    /// object as the existing [`OwnedFd`] instance.
+    pub fn try_clone(&self) -> io::Result<Self> {
+        match self {
+            Self::File(fd) => fd.try_clone().map(OwnedFd::File),
+            Self::Socket(s) => s.try_clone().map(OwnedFd::Socket),
+        }
+    }
+}
+
 impl AsRawFd for OwnedFd {
     fn as_raw_fd(&self) -> RawFd {
         match self {
@@ -152,6 +163,17 @@ pub enum BorrowedFd<'a> {
     Socket(BorrowedSocket<'a>),
 }
 
+impl<'a> BorrowedFd<'a> {
+    /// Creates a new [`OwnedFd`] instance that shares the same underlying
+    /// object as the existing [`BorrowedFd`] instance.
+    pub fn try_clone_to_owned(&self) -> io::Result<OwnedFd> {
+        match self {
+            Self::File(fd) => fd.try_clone_to_owned().map(OwnedFd::File),
+            Self::Socket(s) => s.try_clone_to_owned().map(OwnedFd::Socket),
+        }
+    }
+}
+
 impl AsRawFd for BorrowedFd<'_> {
     fn as_raw_fd(&self) -> RawFd {
         match self {

compio-fs/src/async_fd.rs

@@ -44,6 +44,18 @@ impl<T: AsFd> AsyncFd<T> {
             inner: Attacher::new_unchecked(source),
         }
     }
+
+    /// Create [`AsyncFd`] from a shared file descriptor without attaching
+    /// the source.
+    ///
+    /// # Safety
+    ///
+    /// See [`AsyncFd::new_unchecked`].
+    pub unsafe fn from_shared_fd_unchecked(fd: SharedFd<T>) -> Self {
+        Self {
+            inner: Attacher::from_shared_fd_unchecked(fd),
+        }
+    }
 }
 
 impl<T: AsFd + 'static> AsyncRead for AsyncFd<T> {

compio-fs/src/file.rs

@@ -59,6 +59,12 @@ impl File {
         })
     }
 
+    pub(crate) unsafe fn from_std_unchecked(file: std::fs::File) -> Self {
+        Self {
+            inner: Attacher::new_unchecked(file),
+        }
+    }
+
     /// Attempts to open a file in read-only mode.
     ///
     /// See the [`OpenOptions::open`] method for more details.
@@ -105,7 +111,7 @@ impl File {
     /// Queries metadata about the underlying file.
     #[cfg(windows)]
     pub async fn metadata(&self) -> io::Result<Metadata> {
-        let file = self.inner.clone();
+        let file = self.inner.try_clone()?;
         compio_runtime::spawn_blocking(move || file.metadata().map(Metadata::from_std))
             .await
             .unwrap_or_else(|e| resume_unwind(e))
@@ -122,7 +128,7 @@ impl File {
     /// Changes the permissions on the underlying file.
     #[cfg(windows)]
     pub async fn set_permissions(&self, perm: Permissions) -> io::Result<()> {
-        let file = self.inner.clone();
+        let file = self.inner.try_clone()?;
         compio_runtime::spawn_blocking(move || file.set_permissions(perm.0))
             .await
             .unwrap_or_else(|e| resume_unwind(e))
@@ -172,6 +178,21 @@ impl File {
     pub async fn sync_data(&self) -> io::Result<()> {
         self.sync_impl(true).await
     }
+
+    /// Attempts to clone the file handle.
+    pub fn try_clone(&self) -> io::Result<SyncFile> {
+        self.inner.try_clone().map(SyncFile::new)
+    }
+
+    /// Try to convert into a thread safe file handle.
+    pub fn try_into_sync(self) -> Result<SyncFile, Self> {
+        match self.inner.into_inner().try_unwrap() {
+            Ok(inner) => Ok(SyncFile::new(inner)),
+            Err(fd) => Err(Self {
+                inner: unsafe { Attacher::from_shared_fd_unchecked(fd) },
+            }),
+        }
+    }
 }
 
 impl AsyncReadAt for File {
@@ -267,3 +288,29 @@ impl Splittable for &File {
 }
 
 impl_raw_fd!(File, std::fs::File, inner, file);
+
+/// A thread safe file handle.
+#[derive(Debug)]
+pub struct SyncFile {
+    file: std::fs::File,
+}
+
+impl SyncFile {
+    pub(crate) fn new(file: std::fs::File) -> Self {
+        Self { file }
+    }
+
+    /// Attempts to clone the file handle.
+    pub fn try_clone(&self) -> io::Result<Self> {
+        self.file.try_clone().map(Self::new)
+    }
+}
+
+impl IntoInner for SyncFile {
+    type Inner = File;
+
+    fn into_inner(self) -> Self::Inner {
+        // Safety: the handle is cloned from an existing File.
+        unsafe { File::from_std_unchecked(self.file) }
+    }
+}

compio-fs/src/named_pipe.rs

@@ -6,7 +6,7 @@
 use std::ptr::null_mut;
 use std::{ffi::OsStr, io, os::windows::io::FromRawHandle, ptr::null};
 
-use compio_buf::{BufResult, IoBuf, IoBufMut};
+use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut};
 use compio_driver::{AsRawFd, RawFd, ToSharedFd, impl_raw_fd, op::ConnectNamedPipe, syscall};
 use compio_io::{
     AsyncRead, AsyncReadAt, AsyncReadManaged, AsyncReadManagedAt, AsyncWrite, AsyncWriteAt,
@@ -29,7 +29,7 @@ use windows_sys::Win32::{
     },
 };
 
-use crate::{AsyncFd, File, OpenOptions};
+use crate::{AsyncFd, File, OpenOptions, SyncFile};
 
 /// A [Windows named pipe] server.
 ///
@@ -180,6 +180,21 @@ impl NamedPipeServer {
         syscall!(BOOL, DisconnectNamedPipe(self.as_raw_fd() as _))?;
         Ok(())
     }
+
+    /// Attempts to clone the named pipe server.
+    pub fn try_clone(&self) -> io::Result<SyncNamedPipeServer> {
+        self.handle.try_clone().map(SyncNamedPipeServer::new)
+    }
+
+    /// Try to convert the named pipe server into a thread safe handle.
+    pub fn try_into_sync(self) -> Result<SyncNamedPipeServer, Self> {
+        match self.handle.into_inner().try_unwrap() {
+            Ok(file) => Ok(SyncNamedPipeServer::new(file)),
+            Err(fd) => Err(Self {
+                handle: unsafe { AsyncFd::from_shared_fd_unchecked(fd) },
+            }),
+        }
+    }
 }
 
 impl AsyncRead for NamedPipeServer {
@@ -277,6 +292,33 @@ impl Splittable for &NamedPipeServer {
 
 impl_raw_fd!(NamedPipeServer, std::fs::File, handle, file);
 
+/// A thread safe named pipe server.
+#[derive(Debug)]
+pub struct SyncNamedPipeServer {
+    inner: std::fs::File,
+}
+
+impl SyncNamedPipeServer {
+    pub(crate) fn new(inner: std::fs::File) -> Self {
+        Self { inner }
+    }
+
+    /// Attempts to clone the named pipe server.
+    pub fn try_clone(&self) -> io::Result<Self> {
+        self.inner.try_clone().map(Self::new)
+    }
+}
+
+impl IntoInner for SyncNamedPipeServer {
+    type Inner = NamedPipeServer;
+
+    fn into_inner(self) -> Self::Inner {
+        NamedPipeServer {
+            handle: unsafe { AsyncFd::new_unchecked(self.inner) },
+        }
+    }
+}
+
 /// A [Windows named pipe] client.
 ///
 /// Constructed using [`ClientOptions::open`].
@@ -344,6 +386,19 @@ impl NamedPipeClient {
         // Safety: we're ensuring the lifetime of the named pipe.
         unsafe { named_pipe_info(self.as_raw_fd()) }
     }
+
+    /// Attempts to clone the named pipe client.
+    pub fn try_clone(&self) -> io::Result<SyncNamedPipeClient> {
+        self.handle.try_clone().map(SyncNamedPipeClient::new)
+    }
+
+    /// Try to convert the named pipe client into a thread safe handle.
+    pub fn try_into_sync(self) -> Result<SyncNamedPipeClient, Self> {
+        match self.handle.try_into_sync() {
+            Ok(file) => Ok(SyncNamedPipeClient::new(file)),
+            Err(handle) => Err(Self { handle }),
+        }
+    }
 }
 
 impl AsyncRead for NamedPipeClient {
@@ -443,6 +498,33 @@ impl Splittable for &NamedPipeClient {
 
 impl_raw_fd!(NamedPipeClient, std::fs::File, handle, file);
 
+/// A thread safe named pipe client.
+#[derive(Debug)]
+pub struct SyncNamedPipeClient {
+    inner: SyncFile,
+}
+
+impl SyncNamedPipeClient {
+    pub(crate) fn new(inner: SyncFile) -> Self {
+        Self { inner }
+    }
+
+    /// Attempts to clone the named pipe client.
+    pub fn try_clone(&self) -> io::Result<Self> {
+        self.inner.try_clone().map(Self::new)
+    }
+}
+
+impl IntoInner for SyncNamedPipeClient {
+    type Inner = NamedPipeClient;
+
+    fn into_inner(self) -> Self::Inner {
+        NamedPipeClient {
+            handle: self.inner.into_inner(),
+        }
+    }
+}
+
 /// A builder structure for construct a named pipe with named pipe-specific
 /// options. This is required to use for named pipe servers who wants to modify
 /// pipe-related options.