feat(net): add reunite and tests

Author: Berrysoft

compio-io/src/split.rs

@@ -16,7 +16,7 @@ pub fn split<T: AsyncRead + AsyncWrite>(stream: T) -> (ReadHalf<T>, WriteHalf<T>
 #[derive(Debug)]
 pub struct ReadHalf<T>(Arc<Mutex<T>>);
 
-impl<T> ReadHalf<T> {
+impl<T: Unpin> ReadHalf<T> {
     /// Reunites with a previously split [`WriteHalf`].
     ///
     /// # Panics
@@ -26,10 +26,7 @@ impl<T> ReadHalf<T> {
     /// This can be checked ahead of time by comparing the stored pointer
     /// of the two halves.
     #[track_caller]
-    pub fn unsplit(self, w: WriteHalf<T>) -> T
-    where
-        T: Unpin,
-    {
+    pub fn unsplit(self, w: WriteHalf<T>) -> T {
         if Arc::ptr_eq(&self.0, &w.0) {
             drop(w);
             let inner = Arc::try_unwrap(self.0).expect("`Arc::try_unwrap` failed");

compio-net/src/split.rs

@@ -1,4 +1,4 @@
-use std::{io, ops::Deref, sync::Arc};
+use std::{error::Error, fmt, io, ops::Deref, sync::Arc};
 
 use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
 use compio_io::{AsyncRead, AsyncWrite};
@@ -64,6 +64,22 @@ where
 #[derive(Debug)]
 pub struct OwnedReadHalf<T>(Arc<T>);
 
+impl<T: Unpin> OwnedReadHalf<T> {
+    /// Attempts to put the two halves of a `TcpStream` back together and
+    /// recover the original socket. Succeeds only if the two halves
+    /// originated from the same call to `into_split`.
+    pub fn reunite(self, w: OwnedWriteHalf<T>) -> Result<T, ReuniteError<T>> {
+        if Arc::ptr_eq(&self.0, &w.0) {
+            drop(w);
+            Ok(Arc::try_unwrap(self.0)
+                .ok()
+                .expect("`Arc::try_unwrap` failed"))
+        } else {
+            Err(ReuniteError(self, w))
+        }
+    }
+}
+
 impl<T> AsyncRead for OwnedReadHalf<T>
 where
     for<'a> &'a T: AsyncRead,
@@ -101,3 +117,19 @@ where
         self.0.deref().shutdown().await
     }
 }
+
+/// Error indicating that two halves were not from the same socket, and thus
+/// could not be reunited.
+#[derive(Debug)]
+pub struct ReuniteError<T>(pub OwnedReadHalf<T>, pub OwnedWriteHalf<T>);
+
+impl<T> fmt::Display for ReuniteError<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "tried to reunite halves that are not from the same socket"
+        )
+    }
+}
+
+impl<T: fmt::Debug> Error for ReuniteError<T> {}

compio-net/tests/split.rs

@@ -0,0 +1,100 @@
+use std::io::{Read, Write};
+
+use compio_buf::BufResult;
+use compio_io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
+use compio_net::{TcpStream, UnixListener, UnixStream};
+
+#[compio_macros::test]
+async fn tcp_split() {
+    const MSG: &[u8] = b"split";
+
+    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    let handle = compio_runtime::spawn_blocking(move || {
+        let (mut stream, _) = listener.accept().unwrap();
+        stream.write_all(MSG).unwrap();
+
+        let mut read_buf = [0u8; 32];
+        let read_len = stream.read(&mut read_buf).unwrap();
+        assert_eq!(&read_buf[..read_len], MSG);
+    });
+
+    let stream = TcpStream::connect(&addr).await.unwrap();
+    let (mut read_half, mut write_half) = stream.into_split();
+
+    let read_buf = [0u8; 32];
+    let (read_res, buf) = read_half.read(read_buf).await.unwrap();
+    assert_eq!(read_res, MSG.len());
+    assert_eq!(&buf[..MSG.len()], MSG);
+
+    write_half.write_all(MSG).await.unwrap();
+    handle.await;
+}
+
+#[compio_macros::test]
+async fn tcp_unsplit() {
+    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    let handle = compio_runtime::spawn_blocking(move || {
+        drop(listener.accept().unwrap());
+        drop(listener.accept().unwrap());
+    });
+
+    let stream1 = TcpStream::connect(&addr).await.unwrap();
+    let (read1, write1) = stream1.into_split();
+
+    let stream2 = TcpStream::connect(&addr).await.unwrap();
+    let (_, write2) = stream2.into_split();
+
+    let read1 = match read1.reunite(write2) {
+        Ok(_) => panic!("Reunite should not succeed"),
+        Err(err) => err.0,
+    };
+
+    read1.reunite(write1).expect("Reunite should succeed");
+
+    handle.await;
+}
+
+#[compio_macros::test]
+async fn unix_split() {
+    let dir = tempfile::Builder::new()
+        .prefix("compio-uds-split-tests")
+        .tempdir()
+        .unwrap();
+    let sock_path = dir.path().join("connect.sock");
+
+    let listener = UnixListener::bind(&sock_path).unwrap();
+
+    let client = UnixStream::connect(&sock_path).unwrap();
+    let (server, _) = listener.accept().await.unwrap();
+
+    let (mut a_read, mut a_write) = server.into_split();
+    let (mut b_read, mut b_write) = client.into_split();
+
+    let (a_response, b_response) = futures_util::future::try_join(
+        send_recv_all(&mut a_read, &mut a_write, b"A"),
+        send_recv_all(&mut b_read, &mut b_write, b"B"),
+    )
+    .await
+    .unwrap();
+
+    assert_eq!(a_response, b"B");
+    assert_eq!(b_response, b"A");
+}
+
+async fn send_recv_all<R: AsyncRead, W: AsyncWrite>(
+    read: &mut R,
+    write: &mut W,
+    input: &'static [u8],
+) -> std::io::Result<Vec<u8>> {
+    write.write_all(input).await.0?;
+    write.shutdown().await?;
+
+    let output = Vec::with_capacity(2);
+    let BufResult(res, buf) = read.read_exact(output).await;
+    assert_eq!(res.unwrap_err().kind(), std::io::ErrorKind::UnexpectedEof);
+    Ok(buf)
+}