use futures_util::future::join;
use crate::{
AsyncRead, AsyncWrite, AsyncWriteExt, IoResult,
util::{DEFAULT_BUF_SIZE, Splittable},
};
pub async fn copy<R: AsyncRead, W: AsyncWrite>(reader: &mut R, writer: &mut W) -> IoResult<u64> {
copy_with_size(reader, writer, DEFAULT_BUF_SIZE).await
}
pub async fn copy_with_size<R: AsyncRead, W: AsyncWrite>(
reader: &mut R,
writer: &mut W,
buf_size: usize,
) -> IoResult<u64> {
let mut buf = Vec::with_capacity(buf_size);
let mut total = 0u64;
loop {
let res;
(res, buf) = reader.read(buf).await.into();
match res {
Ok(0) => break,
Ok(read) => {
total += read as u64;
}
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {
continue;
}
Err(e) => return Err(e),
}
let res;
(res, buf) = writer.write_all(buf).await.into();
res?;
buf.clear();
}
writer.flush().await?;
writer.shutdown().await?;
Ok(total)
}
pub async fn copy_bidirectional<A, B>(reader: A, writer: B) -> (IoResult<u64>, IoResult<u64>)
where
A: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
B: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
{
let (mut ar, mut aw) = reader.split();
let (mut br, mut bw) = writer.split();
join(copy(&mut ar, &mut bw), copy(&mut br, &mut aw)).await
}
pub async fn copy_bidirectional_with_sizes<A, B>(
reader: A,
writer: B,
a_to_b_size: usize,
b_to_a_size: usize,
) -> (IoResult<u64>, IoResult<u64>)
where
A: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
B: Splittable<ReadHalf: AsyncRead, WriteHalf: AsyncWrite>,
{
let (mut ar, mut aw) = reader.split();
let (mut br, mut bw) = writer.split();
join(
copy_with_size(&mut ar, &mut bw, a_to_b_size),
copy_with_size(&mut br, &mut aw, b_to_a_size),
)
.await
}