io: add copy_bidirectional_with_sizes (#6500)

tokio/src/io/mod.rs

@@ -271,7 +271,7 @@ cfg_io_util! {
     pub(crate) mod seek;
     pub(crate) mod util;
     pub use util::{
-        copy, copy_bidirectional, copy_buf, duplex, empty, repeat, sink, AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt,
+        copy, copy_bidirectional, copy_bidirectional_with_sizes, copy_buf, duplex, empty, repeat, sink, AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt,
         BufReader, BufStream, BufWriter, DuplexStream, Empty, Lines, Repeat, Sink, Split, Take,
     };
 }

tokio/src/io/util/copy.rs

@@ -16,14 +16,14 @@ pub(super) struct CopyBuffer {
 }
 
 impl CopyBuffer {
-    pub(super) fn new() -> Self {
+    pub(super) fn new(buf_size: usize) -> Self {
         Self {
             read_done: false,
             need_flush: false,
             pos: 0,
             cap: 0,
             amt: 0,
-            buf: vec![0; super::DEFAULT_BUF_SIZE].into_boxed_slice(),
+            buf: vec![0; buf_size].into_boxed_slice(),
         }
     }
 
@@ -269,7 +269,7 @@ cfg_io_util! {
         Copy {
             reader,
             writer,
-            buf: CopyBuffer::new()
+            buf: CopyBuffer::new(super::DEFAULT_BUF_SIZE)
         }.await
     }
 }

tokio/src/io/util/copy_bidirectional.rs

@@ -57,6 +57,9 @@ where
 /// it will return a tuple of the number of bytes copied from a to b
 /// and the number of bytes copied from b to a, in that order.
 ///
+/// It uses two 8 KB buffers for transferring bytes between `a` and `b` by default.
+/// To set your own buffers sizes use [`copy_bidirectional_with_sizes()`].
+///
 /// [`shutdown()`]: crate::io::AsyncWriteExt::shutdown
 ///
 /// # Errors
@@ -69,13 +72,56 @@ where
 ///
 /// Returns a tuple of bytes copied `a` to `b` and bytes copied `b` to `a`.
 #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
-pub async fn copy_bidirectional<A, B>(a: &mut A, b: &mut B) -> Result<(u64, u64), std::io::Error>
+pub async fn copy_bidirectional<A, B>(a: &mut A, b: &mut B) -> io::Result<(u64, u64)>
+where
+    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
+    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
+{
+    copy_bidirectional_impl(
+        a,
+        b,
+        CopyBuffer::new(super::DEFAULT_BUF_SIZE),
+        CopyBuffer::new(super::DEFAULT_BUF_SIZE),
+    )
+    .await
+}
+
+/// Copies data in both directions between `a` and `b` using buffers of the specified size.
+///
+/// This method is the same as the [`copy_bidirectional()`], except that it allows you to set the
+/// size of the internal buffers used when copying data.
+#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
+pub async fn copy_bidirectional_with_sizes<A, B>(
+    a: &mut A,
+    b: &mut B,
+    a_to_b_buf_size: usize,
+    b_to_a_buf_size: usize,
+) -> io::Result<(u64, u64)>
+where
+    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
+    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
+{
+    copy_bidirectional_impl(
+        a,
+        b,
+        CopyBuffer::new(a_to_b_buf_size),
+        CopyBuffer::new(b_to_a_buf_size),
+    )
+    .await
+}
+
+async fn copy_bidirectional_impl<A, B>(
+    a: &mut A,
+    b: &mut B,
+    a_to_b_buffer: CopyBuffer,
+    b_to_a_buffer: CopyBuffer,
+) -> io::Result<(u64, u64)>
 where
     A: AsyncRead + AsyncWrite + Unpin + ?Sized,
     B: AsyncRead + AsyncWrite + Unpin + ?Sized,
 {
-    let mut a_to_b = TransferState::Running(CopyBuffer::new());
-    let mut b_to_a = TransferState::Running(CopyBuffer::new());
+    let mut a_to_b = TransferState::Running(a_to_b_buffer);
+    let mut b_to_a = TransferState::Running(b_to_a_buffer);
     poll_fn(|cx| {
         let a_to_b = transfer_one_direction(cx, &mut a_to_b, a, b)?;
         let b_to_a = transfer_one_direction(cx, &mut b_to_a, b, a)?;

tokio/src/io/util/mod.rs

@@ -28,7 +28,7 @@ cfg_io_util! {
     pub use copy::copy;
 
     mod copy_bidirectional;
-    pub use copy_bidirectional::copy_bidirectional;
+    pub use copy_bidirectional::{copy_bidirectional, copy_bidirectional_with_sizes};
 
     mod copy_buf;
     pub use copy_buf::copy_buf;