feat(net): add split and into_split for TCP & UNIX stream

Author: Berrysoft

compio-io/src/split.rs

@@ -13,6 +13,7 @@ pub fn split<T: AsyncRead + AsyncWrite>(stream: T) -> (ReadHalf<T>, WriteHalf<T>
 }
 
 /// The readable half of a value returned from [`split`].
+#[derive(Debug)]
 pub struct ReadHalf<T>(Arc<Mutex<T>>);
 
 impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
@@ -26,6 +27,7 @@ impl<T: AsyncRead> AsyncRead for ReadHalf<T> {
 }
 
 /// The writable half of a value returned from [`split`].
+#[derive(Debug)]
 pub struct WriteHalf<T>(Arc<Mutex<T>>);
 
 impl<T: AsyncWrite> AsyncWrite for WriteHalf<T> {

compio-net/src/lib.rs

@@ -7,13 +7,15 @@
 
 mod resolve;
 mod socket;
+pub(crate) mod split;
 mod tcp;
 mod udp;
 mod unix;
 
 pub use resolve::ToSocketAddrsAsync;
 pub(crate) use resolve::{each_addr, first_addr_buf};
 pub(crate) use socket::*;
+pub use split::*;
 pub use tcp::*;
 pub use udp::*;
 pub use unix::*;

compio-net/src/split.rs

@@ -0,0 +1,103 @@
+use std::{io, ops::Deref, sync::Arc};
+
+use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
+use compio_io::{AsyncRead, AsyncWrite};
+
+pub(crate) fn split<T>(stream: &T) -> (ReadHalf<T>, WriteHalf<T>)
+where
+    for<'a> &'a T: AsyncRead + AsyncWrite,
+{
+    (ReadHalf(stream), WriteHalf(stream))
+}
+
+/// Borrowed read half.
+#[derive(Debug)]
+pub struct ReadHalf<'a, T>(&'a T);
+
+impl<T> AsyncRead for ReadHalf<'_, T>
+where
+    for<'a> &'a T: AsyncRead,
+{
+    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
+        self.0.read(buf).await
+    }
+
+    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
+        self.0.read_vectored(buf).await
+    }
+}
+
+/// Borrowed write half.
+#[derive(Debug)]
+pub struct WriteHalf<'a, T>(&'a T);
+
+impl<T> AsyncWrite for WriteHalf<'_, T>
+where
+    for<'a> &'a T: AsyncWrite,
+{
+    async fn write<B: IoBuf>(&mut self, buf: B) -> BufResult<usize, B> {
+        self.0.write(buf).await
+    }
+
+    async fn write_vectored<B: IoVectoredBuf>(&mut self, buf: B) -> BufResult<usize, B> {
+        self.0.write_vectored(buf).await
+    }
+
+    async fn flush(&mut self) -> io::Result<()> {
+        self.0.flush().await
+    }
+
+    async fn shutdown(&mut self) -> io::Result<()> {
+        self.0.shutdown().await
+    }
+}
+
+pub(crate) fn into_split<T>(stream: T) -> (OwnedReadHalf<T>, OwnedWriteHalf<T>)
+where
+    for<'a> &'a T: AsyncRead + AsyncWrite,
+{
+    let stream = Arc::new(stream);
+    (OwnedReadHalf(stream.clone()), OwnedWriteHalf(stream))
+}
+
+/// Owned read half.
+#[derive(Debug)]
+pub struct OwnedReadHalf<T>(Arc<T>);
+
+impl<T> AsyncRead for OwnedReadHalf<T>
+where
+    for<'a> &'a T: AsyncRead,
+{
+    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
+        self.0.deref().read(buf).await
+    }
+
+    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
+        self.0.deref().read_vectored(buf).await
+    }
+}
+
+/// Owned write half.
+#[derive(Debug)]
+pub struct OwnedWriteHalf<T>(Arc<T>);
+
+impl<T> AsyncWrite for OwnedWriteHalf<T>
+where
+    for<'a> &'a T: AsyncWrite,
+{
+    async fn write<B: IoBuf>(&mut self, buf: B) -> BufResult<usize, B> {
+        self.0.deref().write(buf).await
+    }
+
+    async fn write_vectored<B: IoVectoredBuf>(&mut self, buf: B) -> BufResult<usize, B> {
+        self.0.deref().write_vectored(buf).await
+    }
+
+    async fn flush(&mut self) -> io::Result<()> {
+        self.0.deref().flush().await
+    }
+
+    async fn shutdown(&mut self) -> io::Result<()> {
+        self.0.deref().shutdown().await
+    }
+}

compio-net/src/tcp.rs

@@ -5,7 +5,7 @@ use compio_io::{AsyncRead, AsyncWrite};
 use compio_runtime::{impl_attachable, impl_try_as_raw_fd};
 use socket2::{Protocol, SockAddr, Type};
 
-use crate::{Socket, ToSocketAddrsAsync};
+use crate::{OwnedReadHalf, OwnedWriteHalf, ReadHalf, Socket, ToSocketAddrsAsync, WriteHalf};
 
 /// A TCP socket server, listening for connections.
 ///
@@ -203,6 +203,25 @@ impl TcpStream {
             .local_addr()
             .map(|addr| addr.as_socket().expect("should be SocketAddr"))
     }
+
+    /// Splits a [`TcpStream`] into a read half and a write half, which can be
+    /// used to read and write the stream concurrently.
+    ///
+    /// This method is more efficient than
+    /// [`into_split`](TcpStream::into_split), but the halves cannot
+    /// be moved into independently spawned tasks.
+    pub fn split(&self) -> (ReadHalf<Self>, WriteHalf<Self>) {
+        crate::split(self)
+    }
+
+    /// Splits a [`TcpStream`] into a read half and a write half, which can be
+    /// used to read and write the stream concurrently.
+    ///
+    /// Unlike [`split`](TcpStream::split), the owned halves can be moved to
+    /// separate tasks, however this comes at the cost of a heap allocation.
+    pub fn into_split(self) -> (OwnedReadHalf<Self>, OwnedWriteHalf<Self>) {
+        crate::into_split(self)
+    }
 }
 
 impl AsyncRead for TcpStream {

compio-net/src/unix.rs

@@ -5,7 +5,7 @@ use compio_io::{AsyncRead, AsyncWrite};
 use compio_runtime::{impl_attachable, impl_try_as_raw_fd};
 use socket2::{Domain, SockAddr, Type};
 
-use crate::Socket;
+use crate::{OwnedReadHalf, OwnedWriteHalf, ReadHalf, Socket, WriteHalf};
 
 /// A Unix socket server, listening for connections.
 ///
@@ -159,6 +159,25 @@ impl UnixStream {
     pub fn local_addr(&self) -> io::Result<SockAddr> {
         self.inner.local_addr()
     }
+
+    /// Splits a [`UnixStream`] into a read half and a write half, which can be
+    /// used to read and write the stream concurrently.
+    ///
+    /// This method is more efficient than
+    /// [`into_split`](UnixStream::into_split), but the halves cannot
+    /// be moved into independently spawned tasks.
+    pub fn split(&self) -> (ReadHalf<Self>, WriteHalf<Self>) {
+        crate::split(self)
+    }
+
+    /// Splits a [`UnixStream`] into a read half and a write half, which can be
+    /// used to read and write the stream concurrently.
+    ///
+    /// Unlike [`split`](UnixStream::split), the owned halves can be moved to
+    /// separate tasks, however this comes at the cost of a heap allocation.
+    pub fn into_split(self) -> (OwnedReadHalf<Self>, OwnedWriteHalf<Self>) {
+        crate::into_split(self)
+    }
 }
 
 impl AsyncRead for UnixStream {