feat(io): copy-bidirectional (#800)

* feat(io): copy-bidirectional

* feat(io): flush & shutdown after successful copy

Author: George-Miao

compio-fs/src/file.rs

@@ -287,4 +287,13 @@ impl Splittable for &File {
     }
 }
 
+impl Splittable for &mut File {
+    type ReadHalf = File;
+    type WriteHalf = File;
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        (self.clone(), self.clone())
+    }
+}
+
 impl_raw_fd!(File, std::fs::File, inner, file);

compio-io/src/util/copy.rs

@@ -0,0 +1,113 @@
+use futures_util::future::join;
+
+use crate::{
+    AsyncRead, AsyncWrite, AsyncWriteExt, IoResult,
+    util::{DEFAULT_BUF_SIZE, Splittable},
+};
+
+/// Asynchronously copies the entire contents of a reader into a writer.
+///
+/// This function returns a future that will continuously read data from
+/// `reader` and then write it into `writer` in a streaming fashion until
+/// `reader` returns EOF or fails.
+///
+/// On success, the total number of bytes that were copied from `reader` to
+/// `writer` is returned.
+///
+/// This is an asynchronous version of [`std::io::copy`][std].
+///
+/// A heap-allocated copy buffer with 8 KiB is created to take data from the
+/// reader to the writer.
+pub async fn copy<R: AsyncRead, W: AsyncWrite>(reader: &mut R, writer: &mut W) -> IoResult<u64> {
+    copy_with_size(reader, writer, DEFAULT_BUF_SIZE).await
+}
+
+/// Asynchronously copies the entire contents of a reader into a writer with
+/// specified buffer sizes.
+///
+/// This function returns a future that will continuously read data from
+/// `reader` and then write it into `writer` in a streaming fashion until
+/// `reader` returns EOF or fails.
+///
+/// On success, the total number of bytes that were copied from `reader` to
+/// `writer` is returned.
+///
+/// This is an asynchronous version of [`std::io::copy`][std].
+pub async fn copy_with_size<R: AsyncRead, W: AsyncWrite>(
+    reader: &mut R,
+    writer: &mut W,
+    buf_size: usize,
+) -> IoResult<u64> {
+    let mut buf = Vec::with_capacity(buf_size);
+    let mut total = 0u64;
+
+    loop {
+        let res;
+        (res, buf) = reader.read(buf).await.into();
+        match res {
+            Ok(0) => break,
+            Ok(read) => {
+                total += read as u64;
+            }
+            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
+                continue;
+            }
+            Err(e) => return Err(e),
+        }
+        let res;
+        (res, buf) = writer.write_all(buf).await.into();
+        res?;
+        buf.clear();
+    }
+
+    writer.flush().await?;
+    writer.shutdown().await?;
+
+    Ok(total)
+}
+
+/// Asynchronously copies data bidirectionally between two pairs of reader and
+/// writer.
+///
+/// This function takes two `Splittable` objects, `reader` and `writer`, and
+/// splits them into their respective read and write halves. It then
+/// concurrently copies data from the read half of `reader` to the write half of
+/// `writer`, and from the read half of `writer` to the write half of `reader`.
+/// The function returns a tuple containing the results of both copy operations,
+/// which indicate the total number of bytes copied in each direction or any
+/// errors that occurred during the copying process.
+pub async fn copy_bidirectional<A, B>(reader: A, writer: B) -> (IoResult<u64>, IoResult<u64>)
+where
+    A: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
+    B: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
+{
+    let (mut ar, mut aw) = reader.split();
+    let (mut br, mut bw) = writer.split();
+
+    join(copy(&mut ar, &mut bw), copy(&mut br, &mut aw)).await
+}
+
+/// Asynchronously copies data bidirectionally between two pairs of reader and
+/// writer with specified buffer sizes.
+///
+/// This function is like `copy_bidirectional`, but allows you to specify the
+/// buffer sizes for each direction of copying.
+pub async fn copy_bidirectional_with_sizes<A, B>(
+    reader: A,
+    writer: B,
+    a_to_b_size: usize,
+    b_to_a_size: usize,
+) -> (IoResult<u64>, IoResult<u64>)
+where
+    A: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
+    B: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
+{
+    let (mut ar, mut aw) = reader.split();
+    let (mut br, mut bw) = writer.split();
+
+    join(
+        copy_with_size(&mut ar, &mut bw, a_to_b_size),
+        copy_with_size(&mut br, &mut aw, b_to_a_size),
+    )
+    .await
+}

compio-io/src/util/mod.rs

@@ -1,5 +1,6 @@
 //! IO related utilities functions for ease of use.
-use crate::{AsyncRead, AsyncWrite, AsyncWriteExt, IoResult};
+mod copy;
+pub use copy::*;
 
 mod take;
 pub use take::Take;
@@ -15,42 +16,3 @@ pub(crate) use internal::*;
 
 pub mod split;
 pub use split::Splittable;
-
-/// Asynchronously copies the entire contents of a reader into a writer.
-///
-/// This function returns a future that will continuously read data from
-/// `reader` and then write it into `writer` in a streaming fashion until
-/// `reader` returns EOF or fails.
-///
-/// On success, the total number of bytes that were copied from `reader` to
-/// `writer` is returned.
-///
-/// This is an asynchronous version of [`std::io::copy`][std].
-///
-/// A heap-allocated copy buffer with 8 KiB is created to take data from the
-/// reader to the writer.
-pub async fn copy<R: AsyncRead, W: AsyncWrite>(reader: &mut R, writer: &mut W) -> IoResult<u64> {
-    let mut buf = Vec::with_capacity(DEFAULT_BUF_SIZE);
-    let mut total = 0u64;
-
-    loop {
-        let res;
-        (res, buf) = reader.read(buf).await.into();
-        match res {
-            Ok(0) => break,
-            Ok(read) => {
-                total += read as u64;
-            }
-            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
-                continue;
-            }
-            Err(e) => return Err(e),
-        }
-        let res;
-        (res, buf) = writer.write_all(buf).await.into();
-        res?;
-        buf.clear();
-    }
-
-    Ok(total)
-}

compio-net/src/tcp.rs

@@ -605,4 +605,13 @@ impl<'a> Splittable for &'a TcpStream {
     }
 }
 
+impl<'a> Splittable for &'a mut TcpStream {
+    type ReadHalf = ReadHalf<'a, TcpStream>;
+    type WriteHalf = WriteHalf<'a, TcpStream>;
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        crate::split(self)
+    }
+}
+
 impl_raw_fd!(TcpStream, socket2::Socket, inner, socket);

compio-net/src/unix.rs

@@ -524,6 +524,15 @@ impl<'a> Splittable for &'a UnixStream {
     }
 }
 
+impl<'a> Splittable for &'a mut UnixStream {
+    type ReadHalf = ReadHalf<'a, UnixStream>;
+    type WriteHalf = WriteHalf<'a, UnixStream>;
+
+    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
+        crate::split(self)
+    }
+}
+
 impl_raw_fd!(UnixStream, socket2::Socket, inner, socket);
 
 #[cfg(windows)]