compio_io/util/
copy.rs

1use futures_util::future::join;
2
3use crate::{
4    AsyncRead, AsyncWrite, AsyncWriteExt, IoResult,
5    util::{DEFAULT_BUF_SIZE, Splittable},
6};
7
8/// Asynchronously copies the entire contents of a reader into a writer.
9///
10/// This function returns a future that will continuously read data from
11/// `reader` and then write it into `writer` in a streaming fashion until
12/// `reader` returns EOF or fails.
13///
14/// On success, the total number of bytes that were copied from `reader` to
15/// `writer` is returned.
16///
17/// This is an asynchronous version of [`std::io::copy`][std].
18///
19/// A heap-allocated copy buffer with 8 KiB is created to take data from the
20/// reader to the writer.
21pub async fn copy<R: AsyncRead, W: AsyncWrite>(reader: &mut R, writer: &mut W) -> IoResult<u64> {
22    copy_with_size(reader, writer, DEFAULT_BUF_SIZE).await
23}
24
25/// Asynchronously copies the entire contents of a reader into a writer with
26/// specified buffer sizes.
27///
28/// This function returns a future that will continuously read data from
29/// `reader` and then write it into `writer` in a streaming fashion until
30/// `reader` returns EOF or fails.
31///
32/// On success, the total number of bytes that were copied from `reader` to
33/// `writer` is returned.
34///
35/// This is an asynchronous version of [`std::io::copy`][std].
36pub async fn copy_with_size<R: AsyncRead, W: AsyncWrite>(
37    reader: &mut R,
38    writer: &mut W,
39    buf_size: usize,
40) -> IoResult<u64> {
41    let mut buf = Vec::with_capacity(buf_size);
42    let mut total = 0u64;
43
44    loop {
45        let res;
46        (res, buf) = reader.read(buf).await.into();
47        match res {
48            Ok(0) => break,
49            Ok(read) => {
50                total += read as u64;
51            }
52            Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
53                continue;
54            }
55            Err(e) => return Err(e),
56        }
57        let res;
58        (res, buf) = writer.write_all(buf).await.into();
59        res?;
60        buf.clear();
61    }
62
63    writer.flush().await?;
64    writer.shutdown().await?;
65
66    Ok(total)
67}
68
69/// Asynchronously copies data bidirectionally between two pairs of reader and
70/// writer.
71///
72/// This function takes two `Splittable` objects, `reader` and `writer`, and
73/// splits them into their respective read and write halves. It then
74/// concurrently copies data from the read half of `reader` to the write half of
75/// `writer`, and from the read half of `writer` to the write half of `reader`.
76/// The function returns a tuple containing the results of both copy operations,
77/// which indicate the total number of bytes copied in each direction or any
78/// errors that occurred during the copying process.
79pub async fn copy_bidirectional<A, B>(reader: A, writer: B) -> (IoResult<u64>, IoResult<u64>)
80where
81    A: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
82    B: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
83{
84    let (mut ar, mut aw) = reader.split();
85    let (mut br, mut bw) = writer.split();
86
87    join(copy(&mut ar, &mut bw), copy(&mut br, &mut aw)).await
88}
89
90/// Asynchronously copies data bidirectionally between two pairs of reader and
91/// writer with specified buffer sizes.
92///
93/// This function is like `copy_bidirectional`, but allows you to specify the
94/// buffer sizes for each direction of copying.
95pub async fn copy_bidirectional_with_sizes<A, B>(
96    reader: A,
97    writer: B,
98    a_to_b_size: usize,
99    b_to_a_size: usize,
100) -> (IoResult<u64>, IoResult<u64>)
101where
102    A: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
103    B: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
104{
105    let (mut ar, mut aw) = reader.split();
106    let (mut br, mut bw) = writer.split();
107
108    join(
109        copy_with_size(&mut ar, &mut bw, a_to_b_size),
110        copy_with_size(&mut br, &mut aw, b_to_a_size),
111    )
112    .await
113}
compio_io/util/copy.rs

compio_io/util/
copy.rs
