tokio/io/util/
copy_bidirectional.rs

1use super::copy::CopyBuffer;
2
3use crate::io::{AsyncRead, AsyncWrite};
4
5use std::future::poll_fn;
6use std::io;
7use std::pin::Pin;
8use std::task::{ready, Context, Poll};
9
10enum TransferState {
11    Running(CopyBuffer),
12    ShuttingDown(u64),
13    Done(u64),
14}
15
16fn transfer_one_direction<A, B>(
17    cx: &mut Context<'_>,
18    state: &mut TransferState,
19    r: &mut A,
20    w: &mut B,
21) -> Poll<io::Result<u64>>
22where
23    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
24    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
25{
26    let mut r = Pin::new(r);
27    let mut w = Pin::new(w);
28
29    loop {
30        match state {
31            TransferState::Running(buf) => {
32                let count = ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?;
33                *state = TransferState::ShuttingDown(count);
34            }
35            TransferState::ShuttingDown(count) => {
36                ready!(w.as_mut().poll_shutdown(cx))?;
37
38                *state = TransferState::Done(*count);
39            }
40            TransferState::Done(count) => return Poll::Ready(Ok(*count)),
41        }
42    }
43}
44/// Copies data in both directions between `a` and `b`.
45///
46/// This function returns a future that will read from both streams,
47/// writing any data read to the opposing stream.
48/// This happens in both directions concurrently.
49///
50/// If an EOF is observed on one stream, [`shutdown()`] will be invoked on
51/// the other, and reading from that stream will stop. Copying of data in
52/// the other direction will continue.
53///
54/// The future will complete successfully once both directions of communication has been shut down.
55/// A direction is shut down when the reader reports EOF,
56/// at which point [`shutdown()`] is called on the corresponding writer. When finished,
57/// it will return a tuple of the number of bytes copied from a to b
58/// and the number of bytes copied from b to a, in that order.
59///
60/// It uses two 8 KB buffers for transferring bytes between `a` and `b` by default.
61/// To set your own buffers sizes use [`copy_bidirectional_with_sizes()`].
62///
63/// [`shutdown()`]: crate::io::AsyncWriteExt::shutdown
64///
65/// # Errors
66///
67/// Interrupted reads and writes are retried. The future will immediately
68/// return any other read or write error, or any flush or shutdown error. Some
69/// data read from either stream may be lost (not written to the other stream)
70/// in this case.
71///
72/// # Return value
73///
74/// Returns a tuple of bytes copied `a` to `b` and bytes copied `b` to `a`.
75#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
76pub async fn copy_bidirectional<A, B>(a: &mut A, b: &mut B) -> io::Result<(u64, u64)>
77where
78    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
79    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
80{
81    copy_bidirectional_impl(
82        a,
83        b,
84        CopyBuffer::new(super::DEFAULT_BUF_SIZE),
85        CopyBuffer::new(super::DEFAULT_BUF_SIZE),
86    )
87    .await
88}
89
90/// Copies data in both directions between `a` and `b` using buffers of the specified size.
91///
92/// This method is the same as the [`copy_bidirectional()`], except that it allows you to set the
93/// size of the internal buffers used when copying data.
94///
95/// # Panics
96///
97/// Panics if either buffer size is zero.
98#[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
99pub async fn copy_bidirectional_with_sizes<A, B>(
100    a: &mut A,
101    b: &mut B,
102    a_to_b_buf_size: usize,
103    b_to_a_buf_size: usize,
104) -> io::Result<(u64, u64)>
105where
106    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
107    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
108{
109    assert!(
110        a_to_b_buf_size > 0,
111        "`a_to_b_buf_size` must be greater than 0"
112    );
113    assert!(
114        b_to_a_buf_size > 0,
115        "`b_to_a_buf_size` must be greater than 0"
116    );
117
118    copy_bidirectional_impl(
119        a,
120        b,
121        CopyBuffer::new(a_to_b_buf_size),
122        CopyBuffer::new(b_to_a_buf_size),
123    )
124    .await
125}
126
127async fn copy_bidirectional_impl<A, B>(
128    a: &mut A,
129    b: &mut B,
130    a_to_b_buffer: CopyBuffer,
131    b_to_a_buffer: CopyBuffer,
132) -> io::Result<(u64, u64)>
133where
134    A: AsyncRead + AsyncWrite + Unpin + ?Sized,
135    B: AsyncRead + AsyncWrite + Unpin + ?Sized,
136{
137    let mut a_to_b = TransferState::Running(a_to_b_buffer);
138    let mut b_to_a = TransferState::Running(b_to_a_buffer);
139    poll_fn(|cx| {
140        let a_to_b = transfer_one_direction(cx, &mut a_to_b, a, b)?;
141        let b_to_a = transfer_one_direction(cx, &mut b_to_a, b, a)?;
142
143        // It is not a problem if ready! returns early because transfer_one_direction for the
144        // other direction will keep returning TransferState::Done(count) in future calls to poll
145        let a_to_b = ready!(a_to_b);
146        let b_to_a = ready!(b_to_a);
147
148        Poll::Ready(Ok((a_to_b, b_to_a)))
149    })
150    .await
151}
tokio/io/util/copy_bidirectional.rs

tokio/io/util/
copy_bidirectional.rs
