feat(io,net): zerocopy API

Author: George-Miao

compio-io/src/write/mod.rs

@@ -1,6 +1,6 @@
 #[cfg(feature = "allocator_api")]
 use std::alloc::Allocator;
-use std::io::Cursor;
+use std::{future::ready, io::Cursor};
 
 use compio_buf::{BufResult, IntoInner, IoBuf, IoVectoredBuf, buf_try, t_alloc};
 
@@ -35,6 +35,32 @@ pub trait AsyncWrite {
         loop_write_vectored!(buf, iter, self.write(iter))
     }
 
+    /// Like `write`, except that it attempts to write the buffer without
+    /// copying in kernel.
+    ///
+    /// The returned future can be used to wait for the buffer to be release by
+    /// the kernel. By default this falls back to `write`, and the returned
+    /// future will be ready immediately.
+    async fn write_zerocopy<T: IoBuf>(
+        &mut self,
+        buf: T,
+    ) -> BufResult<usize, impl Future<Output = T> + use<'_, Self, T>> {
+        self.write(buf).await.map_buffer(ready)
+    }
+
+    /// Like `write_zerocopy`, except that it attempts to write the buffers
+    /// without copying in kernel.
+    ///
+    /// The returned future can be used to wait for the buffer to be release by
+    /// the kernel. By default this falls back to `write_vectored`, and the
+    /// returned future will be ready immediately.
+    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
+        &mut self,
+        buf: T,
+    ) -> BufResult<usize, impl Future<Output = T> + use<'_, Self, T>> {
+        self.write_vectored(buf).await.map_buffer(ready)
+    }
+
     /// Attempts to flush the object, ensuring that any buffered data reach
     /// their destination.
     async fn flush(&mut self) -> IoResult<()>;

compio-net/Cargo.toml

@@ -24,6 +24,7 @@ cfg-if = { workspace = true }
 either = "1.9.0"
 once_cell = { workspace = true }
 socket2 = { workspace = true }
+futures-util = { workspace = true }
 
 [target.'cfg(windows)'.dependencies]
 widestring = { workspace = true }
@@ -40,7 +41,6 @@ libc = { workspace = true }
 # Shared dev dependencies for all platforms
 [dev-dependencies]
 compio-macros = { workspace = true }
-futures-util = { workspace = true }
 tempfile = { workspace = true }
 
 [features]

compio-net/src/lib.rs

@@ -4,7 +4,7 @@
 
 #![cfg_attr(docsrs, feature(doc_cfg))]
 #![cfg_attr(feature = "once_cell_try", feature(once_cell_try))]
-#![allow(unused_features)]
+#![allow(unused_features, refining_impl_trait)]
 #![warn(missing_docs)]
 #![deny(rustdoc::broken_intra_doc_links)]
 #![doc(

compio-net/src/socket.rs

@@ -8,15 +8,17 @@ use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBuf, IoVectore
 #[cfg(unix)]
 use compio_driver::op::CreateSocket;
 use compio_driver::{
-    AsRawFd, ToSharedFd, impl_raw_fd,
+    AsRawFd, OpCode, ToSharedFd, impl_raw_fd,
     op::{
         Accept, BufResultExt, CloseSocket, Connect, Recv, RecvFrom, RecvFromManaged,
         RecvFromVectored, RecvManaged, RecvMsg, RecvResultExt, RecvVectored, ResultTakeBuffer,
-        Send, SendMsg, SendTo, SendToVectored, SendVectored, ShutdownSocket, VecBufResultExt,
+        Send, SendMsg, SendMsgZc, SendTo, SendToVectored, SendVectored, SendVectoredZc, SendZc,
+        ShutdownSocket, VecBufResultExt,
     },
     syscall,
 };
 use compio_runtime::{Attacher, BorrowedBuffer, BufferPool, fd::PollFd};
+use futures_util::StreamExt;
 use socket2::{Domain, Protocol, SockAddr, Socket as Socket2, Type};
 
 #[derive(Debug, Clone)]
@@ -208,6 +210,22 @@ impl Socket {
         compio_runtime::submit(op).await.into_inner()
     }
 
+    pub async fn send_zerocopy<T: IoBuf>(
+        &self,
+        buf: T,
+        flags: i32,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        submit_zc(SendZc::new(self.to_shared_fd(), buf, flags)).await
+    }
+
+    pub async fn send_zerocopy_vectored<T: IoVectoredBuf>(
+        &self,
+        buf: T,
+        flags: i32,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        submit_zc(SendVectoredZc::new(self.to_shared_fd(), buf, flags)).await
+    }
+
     pub async fn recv_from<T: IoBufMut>(
         &self,
         buffer: T,
@@ -275,6 +293,26 @@ impl Socket {
         compio_runtime::submit(op).await.into_inner()
     }
 
+    pub async fn send_to_zerocopy<T: IoBuf>(
+        &self,
+        buffer: T,
+        addr: &SockAddr,
+        flags: i32,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        let op = SendTo::new(self.to_shared_fd(), buffer, addr.clone(), flags);
+        submit_zc(op).await
+    }
+
+    pub async fn send_to_zerocopy_vectored<T: IoVectoredBuf>(
+        &self,
+        buffer: T,
+        addr: &SockAddr,
+        flags: i32,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        let op = SendToVectored::new(self.to_shared_fd(), buffer, addr.clone(), flags);
+        submit_zc(op).await
+    }
+
     pub async fn send_msg<T: IoBuf, C: IoBuf>(
         &self,
         buffer: T,
@@ -299,6 +337,33 @@ impl Socket {
         compio_runtime::submit(op).await.into_inner()
     }
 
+    pub async fn send_msg_zerocopy<T: IoBuf, C: IoBuf>(
+        &self,
+        buffer: T,
+        control: C,
+        addr: Option<&SockAddr>,
+        flags: i32,
+    ) -> BufResult<usize, impl Future<Output = (T, C)> + use<T, C>> {
+        self.send_msg_vectoed_zerocopy([buffer], control, addr, flags)
+            .await
+            .map_buffer(|fut| async move {
+                let ([buffer], control) = fut.await;
+                (buffer, control)
+            })
+    }
+
+    pub async fn send_msg_vectoed_zerocopy<T: IoVectoredBuf, C: IoBuf>(
+        &self,
+        buffer: T,
+        control: C,
+        addr: Option<&SockAddr>,
+        flags: i32,
+    ) -> BufResult<usize, impl Future<Output = (T, C)> + use<T, C>> {
+        let fd = self.to_shared_fd();
+        let op = SendMsgZc::new(fd, buffer, control, addr.cloned(), flags);
+        submit_zc(op).await
+    }
+
     #[cfg(unix)]
     pub unsafe fn get_socket_option<T: Copy>(&self, level: i32, name: i32) -> io::Result<T> {
         let mut value: MaybeUninit<T> = MaybeUninit::uninit();
@@ -377,3 +442,27 @@ impl Socket {
 }
 
 impl_raw_fd!(Socket, Socket2, socket, socket);
+
+async fn submit_zc<T: OpCode + IntoInner + 'static>(
+    op: T,
+) -> BufResult<usize, impl Future<Output = T::Inner> + use<T>> {
+    let mut stream = compio_runtime::submit_multi(op);
+    let res = stream
+        .next()
+        .await
+        .expect("SubmitMulti should yield at least one item")
+        .0;
+
+    let fut = async move {
+        loop {
+            match stream.try_take() {
+                Ok(op) => return op.into_inner(),
+                Err(st) => stream = st,
+            }
+
+            _ = stream.next().await;
+        }
+    };
+
+    BufResult(res, fut)
+}

compio-net/src/split.rs

@@ -37,19 +37,33 @@ where
     for<'a> &'a T: AsyncWrite,
 {
     async fn write<B: IoBuf>(&mut self, buf: B) -> BufResult<usize, B> {
-        self.0.write(buf).await
+        (self.0).write(buf).await
     }
 
     async fn write_vectored<B: IoVectoredBuf>(&mut self, buf: B) -> BufResult<usize, B> {
-        self.0.write_vectored(buf).await
+        (self.0).write_vectored(buf).await
+    }
+
+    async fn write_zerocopy<B: IoBuf>(
+        &mut self,
+        buf: B,
+    ) -> BufResult<usize, impl Future<Output = B> + use<B, T>> {
+        (self.0).write_zerocopy(buf).await
+    }
+
+    async fn write_zerocopy_vectored<B: IoVectoredBuf>(
+        &mut self,
+        buf: B,
+    ) -> BufResult<usize, impl Future<Output = B> + use<B, T>> {
+        (self.0).write_zerocopy_vectored(buf).await
     }
 
     async fn flush(&mut self) -> io::Result<()> {
-        self.0.flush().await
+        (self.0).flush().await
     }
 
     async fn shutdown(&mut self) -> io::Result<()> {
-        self.0.shutdown().await
+        (self.0).shutdown().await
     }
 }
 
@@ -108,6 +122,20 @@ where
         (&self.0).write_vectored(buf).await
     }
 
+    async fn write_zerocopy<B: IoBuf>(
+        &mut self,
+        buf: B,
+    ) -> BufResult<usize, impl Future<Output = B> + use<'_, B, T>> {
+        (&self.0).write_zerocopy(buf).await
+    }
+
+    async fn write_zerocopy_vectored<B: IoVectoredBuf>(
+        &mut self,
+        buf: B,
+    ) -> BufResult<usize, impl Future<Output = B> + use<'_, B, T>> {
+        (&self.0).write_zerocopy_vectored(buf).await
+    }
+
     async fn flush(&mut self) -> io::Result<()> {
         (&self.0).flush().await
     }

compio-net/src/tcp.rs

@@ -391,6 +391,22 @@ impl AsyncWrite for TcpStream {
         (&*self).write_vectored(buf).await
     }
 
+    #[inline]
+    async fn write_zerocopy<T: IoBuf>(
+        &mut self,
+        buf: T,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        self.inner.send_zerocopy(buf, 0).await
+    }
+
+    #[inline]
+    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
+        &mut self,
+        buf: T,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        self.inner.send_zerocopy_vectored(buf, 0).await
+    }
+
     #[inline]
     async fn flush(&mut self) -> io::Result<()> {
         (&*self).flush().await
@@ -413,6 +429,22 @@ impl AsyncWrite for &TcpStream {
         self.inner.send_vectored(buf, 0).await
     }
 
+    #[inline]
+    async fn write_zerocopy<T: IoBuf>(
+        &mut self,
+        buf: T,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        self.inner.send_zerocopy(buf, 0).await
+    }
+
+    #[inline]
+    async fn write_zerocopy_vectored<T: IoVectoredBuf>(
+        &mut self,
+        buf: T,
+    ) -> BufResult<usize, impl Future<Output = T> + use<T>> {
+        self.inner.send_zerocopy_vectored(buf, 0).await
+    }
+
     #[inline]
     async fn flush(&mut self) -> io::Result<()> {
         Ok(())