From 8093712604e54a434658e693ca80535d906e26a7 Mon Sep 17 00:00:00 2001
From: Alex Yusiuk <55661041+RRRadicalEdward@users.noreply.github.com>
Date: Mon, 22 Apr 2024 12:03:33 +0300
Subject: [PATCH] io: add `copy_bidirectional_with_sizes` (#6500)

---
 tokio/src/io/mod.rs                     |  2 +-
 tokio/src/io/util/copy.rs               |  6 +--
 tokio/src/io/util/copy_bidirectional.rs | 52 +++++++++++++++++++++++--
 tokio/src/io/util/mod.rs                |  2 +-
 4 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs
index 7dab413ce..b35a20dd3 100644
--- a/tokio/src/io/mod.rs
+++ b/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,
     };
 }
diff --git a/tokio/src/io/util/copy.rs b/tokio/src/io/util/copy.rs
index 56310c86f..47f8d4ebe 100644
--- a/tokio/src/io/util/copy.rs
+++ b/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
     }
 }
diff --git a/tokio/src/io/util/copy_bidirectional.rs b/tokio/src/io/util/copy_bidirectional.rs
index e1a7db127..ce90141e5 100644
--- a/tokio/src/io/util/copy_bidirectional.rs
+++ b/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,
 {
-    let mut a_to_b = TransferState::Running(CopyBuffer::new());
-    let mut b_to_a = TransferState::Running(CopyBuffer::new());
+    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(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)?;
diff --git a/tokio/src/io/util/mod.rs b/tokio/src/io/util/mod.rs
index 47b951f2b..5010fc01d 100644
--- a/tokio/src/io/util/mod.rs
+++ b/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;