Merge pull request #236 from Berrysoft/dev/async-socket

Berrysoft · web-flow · commit 84ed6e6eb65b · 2024-04-09T18:50:46.000+08:00
feat(driver): add CreateSocket
diff --git a/.github/workflows/ci_test.yml b/.github/workflows/ci_test.yml
@@ -24,7 +24,7 @@ jobs:
         setup:
           - os: "ubuntu-20.04"
           - os: "ubuntu-22.04"
-            features: "io-uring-sqe128,io-uring-cqe32"
+            features: "io-uring-sqe128,io-uring-cqe32,io-uring-socket"
           - os: "ubuntu-20.04"
             features: "polling"
             no_default_features: true
diff --git a/compio-driver/Cargo.toml b/compio-driver/Cargo.toml
@@ -82,6 +82,7 @@ polling = ["dep:polling", "dep:os_pipe"]
 
 io-uring-sqe128 = []
 io-uring-cqe32 = []
+io-uring-socket = []
 
 iocp-global = []
 
diff --git a/compio-driver/src/fusion/mod.rs b/compio-driver/src/fusion/mod.rs
@@ -86,7 +86,11 @@ mod driver_type {
             Close::CODE,
             Shutdown::CODE,
             // Linux kernel 5.19
-            #[cfg(any(feature = "io-uring-seq128", feature = "io-uring-cqe32"))]
+            #[cfg(any(
+                feature = "io-uring-seq128",
+                feature = "io-uring-cqe32",
+                feature = "io-uring-socket"
+            ))]
             Socket::CODE,
         ];
 
diff --git a/compio-driver/src/iour/op.rs b/compio-driver/src/iour/op.rs
@@ -1,4 +1,4 @@
-use std::{ffi::CString, marker::PhantomPinned, os::fd::RawFd, pin::Pin};
+use std::{ffi::CString, io, marker::PhantomPinned, os::fd::RawFd, pin::Pin};
 
 use compio_buf::{
     BufResult, IntoInner, IoBuf, IoBufMut, IoSlice, IoSliceMut, IoVectoredBuf, IoVectoredBufMut,
@@ -12,7 +12,7 @@ use socket2::SockAddr;
 
 use super::OpCode;
 pub use crate::unix::op::*;
-use crate::{op::*, OpEntry};
+use crate::{op::*, syscall, OpEntry};
 
 impl<
     D: std::marker::Send + 'static,
@@ -254,6 +254,22 @@ impl OpCode for HardLink {
     }
 }
 
+impl OpCode for CreateSocket {
+    fn create_entry(self: Pin<&mut Self>) -> OpEntry {
+        if cfg!(feature = "io-uring-socket") {
+            opcode::Socket::new(self.domain, self.socket_type, self.protocol)
+                .build()
+                .into()
+        } else {
+            OpEntry::Blocking
+        }
+    }
+
+    fn call_blocking(self: Pin<&mut Self>) -> io::Result<usize> {
+        Ok(syscall!(libc::socket(self.domain, self.socket_type, self.protocol))? as _)
+    }
+}
+
 impl OpCode for ShutdownSocket {
     fn create_entry(self: Pin<&mut Self>) -> OpEntry {
         opcode::Shutdown::new(Fd(self.fd), self.how())
diff --git a/compio-driver/src/op.rs b/compio-driver/src/op.rs
@@ -16,8 +16,8 @@ pub use crate::sys::op::{
 };
 #[cfg(unix)]
 pub use crate::sys::op::{
-    CreateDir, FileStat, HardLink, OpenFile, PathStat, ReadVectoredAt, Rename, Symlink, Unlink,
-    WriteVectoredAt,
+    CreateDir, CreateSocket, FileStat, HardLink, OpenFile, PathStat, ReadVectoredAt, Rename,
+    Symlink, Unlink, WriteVectoredAt,
 };
 use crate::sys::{sockaddr_storage, socklen_t, RawFd};
 
diff --git a/compio-driver/src/poll/op.rs b/compio-driver/src/poll/op.rs
@@ -347,6 +347,18 @@ impl OpCode for HardLink {
     }
 }
 
+impl OpCode for CreateSocket {
+    fn pre_submit(self: Pin<&mut Self>) -> io::Result<Decision> {
+        Ok(Decision::blocking_dummy())
+    }
+
+    fn on_event(self: Pin<&mut Self>, _: &Event) -> Poll<io::Result<usize>> {
+        Poll::Ready(Ok(
+            syscall!(libc::socket(self.domain, self.socket_type, self.protocol))? as _,
+        ))
+    }
+}
+
 impl OpCode for ShutdownSocket {
     fn pre_submit(self: Pin<&mut Self>) -> io::Result<Decision> {
         Ok(Decision::blocking_dummy())
diff --git a/compio-driver/src/unix/op.rs b/compio-driver/src/unix/op.rs
@@ -209,6 +209,24 @@ impl HardLink {
     }
 }
 
+/// Create a socket.
+pub struct CreateSocket {
+    pub(crate) domain: i32,
+    pub(crate) socket_type: i32,
+    pub(crate) protocol: i32,
+}
+
+impl CreateSocket {
+    /// Create [`CreateSocket`].
+    pub fn new(domain: i32, socket_type: i32, protocol: i32) -> Self {
+        Self {
+            domain,
+            socket_type,
+            protocol,
+        }
+    }
+}
+
 impl ShutdownSocket {
     pub(crate) fn how(&self) -> i32 {
         match self.how {
diff --git a/compio-net/src/socket.rs b/compio-net/src/socket.rs
@@ -1,6 +1,8 @@
 use std::{future::Future, io, mem::ManuallyDrop};
 
 use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
+#[cfg(unix)]
+use compio_driver::op::CreateSocket;
 use compio_driver::{
     impl_raw_fd,
     op::{
@@ -32,8 +34,62 @@ impl Socket {
         self.socket.local_addr()
     }
 
-    pub fn new(domain: Domain, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
-        let socket = Socket2::new(domain, ty, protocol)?;
+    #[cfg(windows)]
+    pub async fn new(domain: Domain, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
+        let socket =
+            compio_runtime::spawn_blocking(move || Socket2::new(domain, ty, protocol)).await?;
+        Self::from_socket2(socket)
+    }
+
+    #[cfg(unix)]
+    pub async fn new(domain: Domain, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
+        use compio_driver::FromRawFd;
+
+        #[allow(unused_mut)]
+        let mut ty: i32 = ty.into();
+        #[cfg(any(
+            target_os = "android",
+            target_os = "dragonfly",
+            target_os = "freebsd",
+            target_os = "fuchsia",
+            target_os = "hurd",
+            target_os = "illumos",
+            target_os = "linux",
+            target_os = "netbsd",
+            target_os = "openbsd",
+        ))]
+        {
+            ty |= libc::SOCK_CLOEXEC;
+        }
+
+        let op = CreateSocket::new(
+            domain.into(),
+            ty,
+            protocol.map(|p| p.into()).unwrap_or_default(),
+        );
+        let BufResult(res, _) = Runtime::current().submit(op).await;
+        let socket = unsafe { Socket2::from_raw_fd(res? as _) };
+        #[cfg(not(any(
+            target_os = "android",
+            target_os = "dragonfly",
+            target_os = "freebsd",
+            target_os = "fuchsia",
+            target_os = "hurd",
+            target_os = "illumos",
+            target_os = "linux",
+            target_os = "netbsd",
+            target_os = "openbsd",
+            target_os = "espidf",
+            target_os = "vita",
+        )))]
+        socket.set_cloexec(true)?;
+        #[cfg(any(
+            target_os = "ios",
+            target_os = "macos",
+            target_os = "tvos",
+            target_os = "watchos",
+        ))]
+        socket.set_nosigpipe(true)?;
         // On Linux we use blocking socket
         // Newer kernels have the patch that allows to arm io_uring poll mechanism for
         // non blocking socket when there is no connections in listen queue
@@ -48,8 +104,8 @@ impl Socket {
         Self::from_socket2(socket)
     }
 
-    pub fn bind(addr: &SockAddr, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
-        let socket = Self::new(addr.domain(), ty, protocol)?;
+    pub async fn bind(addr: &SockAddr, ty: Type, protocol: Option<Protocol>) -> io::Result<Self> {
+        let socket = Self::new(addr.domain(), ty, protocol).await?;
         socket.socket.bind(addr)?;
         Ok(socket)
     }
@@ -97,19 +153,16 @@ impl Socket {
 
     #[cfg(windows)]
     pub async fn accept(&self) -> io::Result<(Self, SockAddr)> {
-        let local_addr = self.local_addr()?;
-        // We should allow users sending this accepted socket to a new thread.
-        let accept_sock = Socket2::new(
-            local_addr.domain(),
-            self.socket.r#type()?,
-            self.socket.protocol()?,
-        )?;
+        let domain = self.local_addr()?.domain();
+        let ty = self.socket.r#type()?;
+        let protocol = self.socket.protocol()?;
+        let accept_sock = Self::new(domain, ty, protocol).await?;
         let op = Accept::new(self.as_raw_fd(), accept_sock.as_raw_fd() as _);
         let BufResult(res, op) = Runtime::current().submit(op).await;
         res?;
         op.update_context()?;
         let addr = op.into_addr()?;
-        Ok((Self::from_socket2(accept_sock)?, addr))
+        Ok((accept_sock, addr))
     }
 
     pub fn close(self) -> impl Future<Output = io::Result<()>> {
diff --git a/compio-net/src/tcp.rs b/compio-net/src/tcp.rs
@@ -55,7 +55,8 @@ impl TcpListener {
     /// to this listener.
     pub async fn bind(addr: impl ToSocketAddrsAsync) -> io::Result<Self> {
         super::each_addr(addr, |addr| async move {
-            let socket = Socket::bind(&SockAddr::from(addr), Type::STREAM, Some(Protocol::TCP))?;
+            let socket =
+                Socket::bind(&SockAddr::from(addr), Type::STREAM, Some(Protocol::TCP)).await?;
             socket.listen(128)?;
             Ok(Self { inner: socket })
         })
@@ -157,9 +158,9 @@ impl TcpStream {
                         "Unsupported address domain.",
                     ));
                 };
-                Socket::bind(&bind_addr, Type::STREAM, Some(Protocol::TCP))?
+                Socket::bind(&bind_addr, Type::STREAM, Some(Protocol::TCP)).await?
             } else {
-                Socket::new(addr2.domain(), Type::STREAM, Some(Protocol::TCP))?
+                Socket::new(addr2.domain(), Type::STREAM, Some(Protocol::TCP)).await?
             };
             socket.connect_async(&addr2).await?;
             Ok(Self { inner: socket })
diff --git a/compio-net/src/udp.rs b/compio-net/src/udp.rs
@@ -96,7 +96,8 @@ impl UdpSocket {
     pub async fn bind(addr: impl ToSocketAddrsAsync) -> io::Result<Self> {
         super::each_addr(addr, |addr| async move {
             Ok(Self {
-                inner: Socket::bind(&SockAddr::from(addr), Type::DGRAM, Some(Protocol::UDP))?,
+                inner: Socket::bind(&SockAddr::from(addr), Type::DGRAM, Some(Protocol::UDP))
+                    .await?,
             })
         })
         .await
diff --git a/compio-net/src/unix.rs b/compio-net/src/unix.rs
@@ -24,7 +24,7 @@ use crate::{OwnedReadHalf, OwnedWriteHalf, ReadHalf, Socket, WriteHalf};
 /// let sock_file = dir.path().join("unix-server.sock");
 ///
 /// # compio_runtime::Runtime::new().unwrap().block_on(async move {
-/// let listener = UnixListener::bind(&sock_file).unwrap();
+/// let listener = UnixListener::bind(&sock_file).await.unwrap();
 ///
 /// let (mut tx, (mut rx, _)) =
 ///     futures_util::try_join!(UnixStream::connect(&sock_file), listener.accept()).unwrap();
@@ -45,22 +45,22 @@ impl UnixListener {
     /// Creates a new [`UnixListener`], which will be bound to the specified
     /// file path. The file path cannot yet exist, and will be cleaned up
     /// upon dropping [`UnixListener`]
-    pub fn bind(path: impl AsRef<Path>) -> io::Result<Self> {
-        Self::bind_addr(&SockAddr::unix(path)?)
+    pub async fn bind(path: impl AsRef<Path>) -> io::Result<Self> {
+        Self::bind_addr(&SockAddr::unix(path)?).await
     }
 
     /// Creates a new [`UnixListener`] with [`SockAddr`], which will be bound to
     /// the specified file path. The file path cannot yet exist, and will be
     /// cleaned up upon dropping [`UnixListener`]
-    pub fn bind_addr(addr: &SockAddr) -> io::Result<Self> {
+    pub async fn bind_addr(addr: &SockAddr) -> io::Result<Self> {
         if !addr.is_unix() {
             return Err(io::Error::new(
                 io::ErrorKind::InvalidInput,
                 "addr is not unix socket address",
             ));
         }
 
-        let socket = Socket::bind(addr, Type::STREAM, None)?;
+        let socket = Socket::bind(addr, Type::STREAM, None).await?;
         socket.listen(1024)?;
         Ok(UnixListener { inner: socket })
     }
@@ -138,12 +138,12 @@ impl UnixStream {
         #[cfg(windows)]
         let socket = {
             let new_addr = empty_unix_socket();
-            Socket::bind(&new_addr, Type::STREAM, None)?
+            Socket::bind(&new_addr, Type::STREAM, None).await?
         };
         #[cfg(unix)]
         let socket = {
             use socket2::Domain;
-            Socket::new(Domain::UNIX, Type::STREAM, None)?
+            Socket::new(Domain::UNIX, Type::STREAM, None).await?
         };
         socket.connect_async(addr).await?;
         let unix_stream = UnixStream { inner: socket };
diff --git a/compio-net/tests/split.rs b/compio-net/tests/split.rs
@@ -66,7 +66,7 @@ async fn unix_split() {
         .unwrap();
     let sock_path = dir.path().join("connect.sock");
 
-    let listener = UnixListener::bind(&sock_path).unwrap();
+    let listener = UnixListener::bind(&sock_path).await.unwrap();
 
     let (client, (server, _)) =
         futures_util::try_join!(UnixStream::connect(&sock_path), listener.accept()).unwrap();
diff --git a/compio-net/tests/unix_stream.rs b/compio-net/tests/unix_stream.rs
@@ -9,7 +9,7 @@ async fn accept_read_write() -> std::io::Result<()> {
         .unwrap();
     let sock_path = dir.path().join("connect.sock");
 
-    let listener = UnixListener::bind(&sock_path)?;
+    let listener = UnixListener::bind(&sock_path).await?;
 
     let (mut client, (mut server, _)) =
         futures_util::try_join!(UnixStream::connect(&sock_path), listener.accept()).unwrap();
@@ -33,7 +33,7 @@ async fn shutdown() -> std::io::Result<()> {
         .unwrap();
     let sock_path = dir.path().join("connect.sock");
 
-    let listener = UnixListener::bind(&sock_path)?;
+    let listener = UnixListener::bind(&sock_path).await?;
 
     let (mut client, (mut server, _)) =
         futures_util::try_join!(UnixStream::connect(&sock_path), listener.accept()).unwrap();
diff --git a/compio/examples/unix.rs b/compio/examples/unix.rs
@@ -8,7 +8,7 @@ use tempfile::tempdir;
 async fn main() {
     let dir = tempdir().unwrap();
     let path = dir.path().join("unix-example.sock");
-    let listener = UnixListener::bind(&path).unwrap();
+    let listener = UnixListener::bind(&path).await.unwrap();
 
     let addr = listener.local_addr().unwrap();
 
