diff --git a/tokio/src/io/mod.rs b/tokio/src/io/mod.rs index 7dab413ceb6..b35a20dd35b 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 56310c86f59..47f8d4ebec6 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 e1a7db127a7..ce90141e5a5 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: &mut A, b: &mut B) -> Result<(u64, u64), std::io::Error> +pub async fn copy_bidirectional(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: &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: &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)?; diff --git a/tokio/src/io/util/mod.rs b/tokio/src/io/util/mod.rs index 47b951f2b83..5010fc01d29 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;