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;