tokio/io/
blocking.rs

Help
1use crate::io::sys;
2use crate::io::{AsyncRead, AsyncWrite, ReadBuf};
3
4use std::cmp;
5use std::future::Future;
6use std::io;
7use std::io::prelude::*;
8use std::mem::MaybeUninit;
9use std::pin::Pin;
10use std::task::{ready, Context, Poll};
11
12/// `T` should not implement _both_ Read and Write.
13#[derive(Debug)]
14pub(crate) struct Blocking<T> {
15    inner: Option<T>,
16    state: State<T>,
17    /// `true` if the lower IO layer needs flushing.
18    need_flush: bool,
19}
20
21#[derive(Debug)]
22pub(crate) struct Buf {
23    buf: Vec<u8>,
24    pos: usize,
25}
26
27pub(crate) const DEFAULT_MAX_BUF_SIZE: usize = 2 * 1024 * 1024;
28
29#[derive(Debug)]
30enum State<T> {
31    Idle(Option<Buf>),
32    Busy(sys::Blocking<(io::Result<usize>, Buf, T)>),
33}
34
35cfg_io_blocking! {
36    impl<T> Blocking<T> {
37        /// # Safety
38        ///
39        /// The `Read` implementation of `inner` must never read from the buffer
40        /// it is borrowing and must correctly report the length of the data
41        /// written into the buffer.
42        #[cfg_attr(feature = "fs", allow(dead_code))]
43        pub(crate) unsafe fn new(inner: T) -> Blocking<T> {
44            Blocking {
45                inner: Some(inner),
46                state: State::Idle(Some(Buf::with_capacity(0))),
47                need_flush: false,
48            }
49        }
50    }
51}
52
53impl<T> AsyncRead for Blocking<T>
54where
55    T: Read + Unpin + Send + 'static,
56{
57    fn poll_read(
58        mut self: Pin<&mut Self>,
59        cx: &mut Context<'_>,
60        dst: &mut ReadBuf<'_>,
61    ) -> Poll<io::Result<()>> {
62        loop {
63            match self.state {
64                State::Idle(ref mut buf_cell) => {
65                    let mut buf = buf_cell.take().unwrap();
66
67                    if !buf.is_empty() {
68                        buf.copy_to(dst);
69                        *buf_cell = Some(buf);
70                        return Poll::Ready(Ok(()));
71                    }
72
73                    let mut inner = self.inner.take().unwrap();
74
75                    let max_buf_size = cmp::min(dst.remaining(), DEFAULT_MAX_BUF_SIZE);
76                    self.state = State::Busy(sys::run(move || {
77                        // SAFETY: the requirements are satisfied by `Blocking::new`.
78                        let res = unsafe { buf.read_from(&mut inner, max_buf_size) };
79                        (res, buf, inner)
80                    }));
81                }
82                State::Busy(ref mut rx) => {
83                    let (res, mut buf, inner) = ready!(Pin::new(rx).poll(cx))?;
84                    self.inner = Some(inner);
85
86                    match res {
87                        Ok(_) => {
88                            buf.copy_to(dst);
89                            self.state = State::Idle(Some(buf));
90                            return Poll::Ready(Ok(()));
91                        }
92                        Err(e) => {
93                            assert!(buf.is_empty());
94
95                            self.state = State::Idle(Some(buf));
96                            return Poll::Ready(Err(e));
97                        }
98                    }
99                }
100            }
101        }
102    }
103}
104
105impl<T> AsyncWrite for Blocking<T>
106where
107    T: Write + Unpin + Send + 'static,
108{
109    fn poll_write(
110        mut self: Pin<&mut Self>,
111        cx: &mut Context<'_>,
112        src: &[u8],
113    ) -> Poll<io::Result<usize>> {
114        loop {
115            match self.state {
116                State::Idle(ref mut buf_cell) => {
117                    let mut buf = buf_cell.take().unwrap();
118
119                    assert!(buf.is_empty());
120
121                    let n = buf.copy_from(src, DEFAULT_MAX_BUF_SIZE);
122                    let mut inner = self.inner.take().unwrap();
123
124                    self.state = State::Busy(sys::run(move || {
125                        let n = buf.len();
126                        let res = buf.write_to(&mut inner).map(|()| n);
127
128                        (res, buf, inner)
129                    }));
130                    self.need_flush = true;
131
132                    return Poll::Ready(Ok(n));
133                }
134                State::Busy(ref mut rx) => {
135                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?;
136                    self.state = State::Idle(Some(buf));
137                    self.inner = Some(inner);
138
139                    // If error, return
140                    res?;
141                }
142            }
143        }
144    }
145
146    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
147        loop {
148            let need_flush = self.need_flush;
149            match self.state {
150                // The buffer is not used here
151                State::Idle(ref mut buf_cell) => {
152                    if need_flush {
153                        let buf = buf_cell.take().unwrap();
154                        let mut inner = self.inner.take().unwrap();
155
156                        self.state = State::Busy(sys::run(move || {
157                            let res = inner.flush().map(|()| 0);
158                            (res, buf, inner)
159                        }));
160
161                        self.need_flush = false;
162                    } else {
163                        return Poll::Ready(Ok(()));
164                    }
165                }
166                State::Busy(ref mut rx) => {
167                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?;
168                    self.state = State::Idle(Some(buf));
169                    self.inner = Some(inner);
170
171                    // If error, return
172                    res?;
173                }
174            }
175        }
176    }
177
178    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
179        Poll::Ready(Ok(()))
180    }
181}
182
183/// Repeats operations that are interrupted.
184macro_rules! uninterruptibly {
185    ($e:expr) => {{
186        loop {
187            match $e {
188                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
189                res => break res,
190            }
191        }
192    }};
193}
194
195impl Buf {
196    pub(crate) fn with_capacity(n: usize) -> Buf {
197        Buf {
198            buf: Vec::with_capacity(n),
199            pos: 0,
200        }
201    }
202
203    pub(crate) fn is_empty(&self) -> bool {
204        self.len() == 0
205    }
206
207    pub(crate) fn len(&self) -> usize {
208        self.buf.len() - self.pos
209    }
210
211    pub(crate) fn copy_to(&mut self, dst: &mut ReadBuf<'_>) -> usize {
212        let n = cmp::min(self.len(), dst.remaining());
213        dst.put_slice(&self.bytes()[..n]);
214        self.pos += n;
215
216        if self.pos == self.buf.len() {
217            self.buf.truncate(0);
218            self.pos = 0;
219        }
220
221        n
222    }
223
224    pub(crate) fn copy_from(&mut self, src: &[u8], max_buf_size: usize) -> usize {
225        assert!(self.is_empty());
226
227        let n = cmp::min(src.len(), max_buf_size);
228
229        self.buf.extend_from_slice(&src[..n]);
230        n
231    }
232
233    pub(crate) fn bytes(&self) -> &[u8] {
234        &self.buf[self.pos..]
235    }
236
237    /// # Safety
238    ///
239    /// `rd` must not read from the buffer `read` is borrowing and must correctly
240    /// report the length of the data written into the buffer.
241    pub(crate) unsafe fn read_from<T: Read>(
242        &mut self,
243        rd: &mut T,
244        max_buf_size: usize,
245    ) -> io::Result<usize> {
246        assert!(self.is_empty());
247        self.buf.reserve(max_buf_size);
248
249        let buf = &mut self.buf.spare_capacity_mut()[..max_buf_size];
250        // SAFETY: The memory may be uninitialized, but `rd.read` will only write to the buffer.
251        let buf = unsafe { &mut *(buf as *mut [MaybeUninit<u8>] as *mut [u8]) };
252        let res = uninterruptibly!(rd.read(buf));
253
254        if let Ok(n) = res {
255            // SAFETY: the caller promises that `rd.read` initializes
256            // a section of `buf` and correctly reports that length.
257            // The `self.is_empty()` assertion verifies that `n`
258            // equals the length of the `buf` capacity that was written
259            // to (and that `buf` isn't being shrunk).
260            unsafe { self.buf.set_len(n) }
261        } else {
262            self.buf.clear();
263        }
264
265        assert_eq!(self.pos, 0);
266
267        res
268    }
269
270    pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> {
271        assert_eq!(self.pos, 0);
272
273        // `write_all` already ignores interrupts
274        let res = wr.write_all(&self.buf);
275        self.buf.clear();
276        res
277    }
278}
279
280cfg_fs! {
281    impl Buf {
282        pub(crate) fn discard_read(&mut self) -> i64 {
283            let ret = -(self.bytes().len() as i64);
284            self.pos = 0;
285            self.buf.truncate(0);
286            ret
287        }
288
289        pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>], max_buf_size: usize) -> usize {
290            assert!(self.is_empty());
291
292            let mut rem = max_buf_size;
293            for buf in bufs {
294                if rem == 0 {
295                    break
296                }
297
298                let len = buf.len().min(rem);
299                self.buf.extend_from_slice(&buf[..len]);
300                rem -= len;
301            }
302
303            max_buf_size - rem
304        }
305    }
306}
tokio/io/blocking.rs

tokio/io/
blocking.rs
