1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
use core::{future::poll_fn, pin::Pin, slice};
use std::{io, net::Shutdown};
use tokio::{
io::{AsyncWrite, Interest},
net::{TcpStream, UnixStream},
};
use crate::buf::{BoundedBuf, BoundedBufMut};
use super::{async_buf_read::AsyncBufRead, async_buf_write::AsyncBufWrite};
macro_rules! trait_impl {
($ty: ty) => {
impl AsyncBufRead for $ty {
#[allow(unsafe_code)]
async fn read<B>(&self, mut buf: B) -> (io::Result<usize>, B)
where
B: BoundedBufMut,
{
let init = buf.bytes_init();
let total = buf.bytes_total();
// Safety: construct a mutable slice over the spare capacity.
// try_read writes contiguously from the start of the slice
// and returns the exact byte count written on Ok(n).
let spare = unsafe { slice::from_raw_parts_mut(buf.stable_mut_ptr().add(init), total - init) };
let mut written = 0;
let res = loop {
if written == spare.len() {
break Ok(written);
}
match self.try_read(&mut spare[written..]) {
Ok(0) => break Ok(written),
Ok(n) => written += n,
Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
if written > 0 {
break Ok(written);
}
if let Err(e) = self.ready(Interest::READABLE).await {
break Err(e);
}
}
Err(e) => break Err(e),
}
};
// SAFETY: TcpStream::try_read has put written bytes into buf.
unsafe {
buf.set_init(init + written);
}
(res, buf)
}
}
impl AsyncBufWrite for $ty {
async fn write<B>(&self, buf: B) -> (io::Result<usize>, B)
where
B: BoundedBuf,
{
let data = buf.chunk();
let mut written = 0;
let res = loop {
if written == data.len() {
break Ok(written);
}
match self.try_write(&data[written..]) {
Ok(0) => break Ok(written),
Ok(n) => written += n,
Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
if written > 0 {
break Ok(written);
}
if let Err(e) = self.ready(Interest::WRITABLE).await {
break Err(e);
}
}
Err(e) => break Err(e),
}
};
(res, buf)
}
async fn shutdown(mut self, _: Shutdown) -> io::Result<()> {
poll_fn(|cx| Pin::new(&mut self).poll_shutdown(cx)).await
}
}
};
}
trait_impl!(TcpStream);
trait_impl!(UnixStream);