io: add copy_bidirectional_with_size

This is the same as `copy_bidirectional` but accepts the `a` -> `b` and `b` -> `a` buffers sizes.
`copy_bidirectional` uses 8Kb buffers under the hood which isn't
configurable. `copy_bidirectional_with_size` fixes it by allowing an
user to set its own sizes for underlying buffers.

Fixes: #6454.
Refs: #3572.

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_size, 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 8Kb buffers for transferring bytes between `a` and `b` by default.
+/// To set your own buffers sizes use [`copy_bidirectional_with_size()`].
+///
 /// [`shutdown()`]: crate::io::AsyncWriteExt::shutdown
 ///
 /// # Errors
@@ -74,8 +77,26 @@ 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
+}
+
+/// The same as the [`copy_bidirectional()`], but allows to set the underlying `a` to `b` and `b` to `a` buffers sizes.
+#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
+pub async fn copy_bidirectional_with_size<A, B>(a: &mut A, b: &mut B, a_to_b_buf_size: usize, b_to_a_buf_size: usize) -> Result<(u64, u64), std::io::Error>
+ 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) -> Result<(u64, u64), std::io::Error>
+ 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)?;
@@ -87,5 +108,5 @@ where
 
  Poll::Ready(Ok((a_to_b, b_to_a)))
  })
- .await
-}
+  .await
+}

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_size};
 
  mod copy_buf;
  pub use copy_buf::copy_buf;