tokio/io/
blocking.rs

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    /// `true` if operations are spawned with `spawn_mandatory_blocking`.
20    mandatory: bool,
21}
22
23#[derive(Debug)]
24pub(crate) struct Buf {
25    buf: Vec<u8>,
26    pos: usize,
27}
28
29pub(crate) const DEFAULT_MAX_BUF_SIZE: usize = 2 * 1024 * 1024;
30
31#[derive(Debug)]
32enum State<T> {
33    /// Holds `None` if an operation could not be spawned, normally because the
34    /// runtime is shutting down. `inner` is gone in that case, so all further
35    /// operations fail.
36    Idle(Option<Buf>),
37    Busy(sys::Blocking<(io::Result<usize>, Buf, T)>),
38}
39
40cfg_io_blocking! {
41    impl<T> Blocking<T> {
42        /// # Safety
43        ///
44        /// The `Read` implementation of `inner` must never read from the buffer
45        /// it is borrowing and must correctly report the length of the data
46        /// written into the buffer.
47        #[cfg_attr(feature = "fs", allow(dead_code))]
48        pub(crate) unsafe fn new(inner: T) -> Blocking<T> {
49            Blocking {
50                inner: Some(inner),
51                state: State::Idle(Some(Buf::with_capacity(0))),
52                need_flush: false,
53                mandatory: false,
54            }
55        }
56
57        /// Like `new`, but operations are spawned with
58        /// `spawn_mandatory_blocking`, so an operation that was started before
59        /// the runtime began shutting down still runs.
60        ///
61        /// # Safety
62        ///
63        /// Same as `new`.
64        #[cfg_attr(not(feature = "io-std"), allow(dead_code))]
65        pub(crate) unsafe fn new_mandatory(inner: T) -> Blocking<T> {
66            Blocking {
67                inner: Some(inner),
68                state: State::Idle(Some(Buf::with_capacity(0))),
69                need_flush: false,
70                mandatory: true,
71            }
72        }
73    }
74}
75
76/// Runs `f` on the blocking pool.
77///
78/// Returns `None` if `mandatory` is set and the operation could not be
79/// spawned, normally because the runtime is shutting down.
80fn spawn<F, R>(mandatory: bool, f: F) -> Option<sys::Blocking<R>>
81where
82    F: FnOnce() -> R + Send + 'static,
83    R: Send + 'static,
84{
85    if mandatory {
86        sys::spawn_mandatory_blocking(f)
87    } else {
88        Some(sys::run(f))
89    }
90}
91
92impl<T> AsyncRead for Blocking<T>
93where
94    T: Read + Unpin + Send + 'static,
95{
96    fn poll_read(
97        mut self: Pin<&mut Self>,
98        cx: &mut Context<'_>,
99        dst: &mut ReadBuf<'_>,
100    ) -> Poll<io::Result<()>> {
101        loop {
102            match self.state {
103                State::Idle(ref mut buf_cell) => {
104                    let Some(mut buf) = buf_cell.take() else {
105                        return Poll::Ready(Err(gone()));
106                    };
107
108                    if !buf.is_empty() {
109                        buf.copy_to(dst);
110                        *buf_cell = Some(buf);
111                        return Poll::Ready(Ok(()));
112                    }
113
114                    let mut inner = self.inner.take().unwrap();
115
116                    let max_buf_size = cmp::min(dst.remaining(), DEFAULT_MAX_BUF_SIZE);
117                    let rx = spawn(self.mandatory, move || {
118                        // SAFETY: the caller of `Blocking::new` or `Blocking::new_mandatory`
119                        // promised that `inner` meets the requirements of `read_from`.
120                        let res = unsafe { buf.read_from(&mut inner, max_buf_size) };
121                        (res, buf, inner)
122                    });
123
124                    let Some(rx) = rx else {
125                        return Poll::Ready(Err(gone()));
126                    };
127                    self.state = State::Busy(rx);
128                }
129                State::Busy(ref mut rx) => {
130                    let (res, mut buf, inner) = ready!(Pin::new(rx).poll(cx))?;
131                    self.inner = Some(inner);
132
133                    match res {
134                        Ok(_) => {
135                            buf.copy_to(dst);
136                            self.state = State::Idle(Some(buf));
137                            return Poll::Ready(Ok(()));
138                        }
139                        Err(e) => {
140                            assert!(buf.is_empty());
141
142                            self.state = State::Idle(Some(buf));
143                            return Poll::Ready(Err(e));
144                        }
145                    }
146                }
147            }
148        }
149    }
150}
151
152impl<T> AsyncWrite for Blocking<T>
153where
154    T: Write + Unpin + Send + 'static,
155{
156    fn poll_write(
157        mut self: Pin<&mut Self>,
158        cx: &mut Context<'_>,
159        src: &[u8],
160    ) -> Poll<io::Result<usize>> {
161        loop {
162            match self.state {
163                State::Idle(ref mut buf_cell) => {
164                    let Some(mut buf) = buf_cell.take() else {
165                        return Poll::Ready(Err(gone()));
166                    };
167
168                    assert!(buf.is_empty());
169
170                    let n = buf.copy_from(src, DEFAULT_MAX_BUF_SIZE);
171                    let mut inner = self.inner.take().unwrap();
172
173                    let rx = spawn(self.mandatory, move || {
174                        let n = buf.len();
175                        let res = buf.write_to(&mut inner).map(|()| n);
176
177                        (res, buf, inner)
178                    });
179
180                    let Some(rx) = rx else {
181                        return Poll::Ready(Err(gone()));
182                    };
183                    self.state = State::Busy(rx);
184                    self.need_flush = true;
185
186                    return Poll::Ready(Ok(n));
187                }
188                State::Busy(ref mut rx) => {
189                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?;
190                    self.state = State::Idle(Some(buf));
191                    self.inner = Some(inner);
192
193                    // If error, return
194                    res?;
195                }
196            }
197        }
198    }
199
200    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
201        loop {
202            let need_flush = self.need_flush;
203            match self.state {
204                // The buffer is not used here
205                State::Idle(ref mut buf_cell) => {
206                    if need_flush {
207                        let Some(buf) = buf_cell.take() else {
208                            return Poll::Ready(Err(gone()));
209                        };
210                        let mut inner = self.inner.take().unwrap();
211
212                        let rx = spawn(self.mandatory, move || {
213                            let res = inner.flush().map(|()| 0);
214                            (res, buf, inner)
215                        });
216
217                        let Some(rx) = rx else {
218                            return Poll::Ready(Err(gone()));
219                        };
220                        self.state = State::Busy(rx);
221
222                        self.need_flush = false;
223                    } else {
224                        return Poll::Ready(Ok(()));
225                    }
226                }
227                State::Busy(ref mut rx) => {
228                    let (res, buf, inner) = ready!(Pin::new(rx).poll(cx))?;
229                    self.state = State::Idle(Some(buf));
230                    self.inner = Some(inner);
231
232                    // If error, return
233                    res?;
234                }
235            }
236        }
237    }
238
239    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
240        Poll::Ready(Ok(()))
241    }
242}
243
244/// Repeats operations that are interrupted.
245macro_rules! uninterruptibly {
246    ($e:expr) => {{
247        loop {
248            match $e {
249                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
250                res => break res,
251            }
252        }
253    }};
254}
255
256impl Buf {
257    pub(crate) fn with_capacity(n: usize) -> Buf {
258        Buf {
259            buf: Vec::with_capacity(n),
260            pos: 0,
261        }
262    }
263
264    pub(crate) fn is_empty(&self) -> bool {
265        self.len() == 0
266    }
267
268    pub(crate) fn len(&self) -> usize {
269        self.buf.len() - self.pos
270    }
271
272    pub(crate) fn copy_to(&mut self, dst: &mut ReadBuf<'_>) -> usize {
273        let n = cmp::min(self.len(), dst.remaining());
274        dst.put_slice(&self.bytes()[..n]);
275        self.pos += n;
276
277        if self.pos == self.buf.len() {
278            self.buf.clear();
279            self.pos = 0;
280        }
281
282        n
283    }
284
285    pub(crate) fn copy_from(&mut self, src: &[u8], max_buf_size: usize) -> usize {
286        assert!(self.is_empty());
287
288        let n = cmp::min(src.len(), max_buf_size);
289
290        self.buf.extend_from_slice(&src[..n]);
291        n
292    }
293
294    pub(crate) fn bytes(&self) -> &[u8] {
295        &self.buf[self.pos..]
296    }
297
298    /// # Safety
299    ///
300    /// `rd` must not read from the buffer `read` is borrowing and must correctly
301    /// report the length of the data written into the buffer.
302    pub(crate) unsafe fn read_from<T: Read>(
303        &mut self,
304        rd: &mut T,
305        max_buf_size: usize,
306    ) -> io::Result<usize> {
307        assert!(self.is_empty());
308        self.buf.reserve(max_buf_size);
309
310        let buf = &mut self.buf.spare_capacity_mut()[..max_buf_size];
311        // SAFETY: The memory may be uninitialized, but `rd.read` will only write to the buffer.
312        let buf = unsafe { &mut *(buf as *mut [MaybeUninit<u8>] as *mut [u8]) };
313        let res = uninterruptibly!(rd.read(buf));
314
315        if let Ok(n) = res {
316            // SAFETY: the caller promises that `rd.read` initializes
317            // a section of `buf` and correctly reports that length.
318            // The `self.is_empty()` assertion verifies that `n`
319            // equals the length of the `buf` capacity that was written
320            // to (and that `buf` isn't being shrunk).
321            unsafe { self.buf.set_len(n) }
322        } else {
323            self.buf.clear();
324        }
325
326        assert_eq!(self.pos, 0);
327
328        res
329    }
330
331    pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> {
332        assert_eq!(self.pos, 0);
333
334        // `write_all` already ignores interrupts
335        let res = wr.write_all(&self.buf);
336        self.buf.clear();
337        res
338    }
339}
340
341cfg_io_uring! {
342    impl Buf {
343        /// Prepare the internal buffer for an io-uring read operation.
344        ///
345        /// Returns a pointer to the spare capacity and the length available
346        /// for the kernel to write into.
347        pub(crate) fn prepare_uring_read(&mut self, max_buf_size: usize) -> (*mut u8, u32) {
348            assert!(self.is_empty());
349            let max_buf_size = cmp::min(max_buf_size, u32::MAX as usize);
350            self.buf.reserve(max_buf_size);
351            let spare = self.buf.spare_capacity_mut();
352            let len = cmp::min(spare.len(), max_buf_size);
353            let ptr = spare.as_mut_ptr().cast::<u8>();
354            (ptr, len as u32)
355        }
356
357        /// Complete an io-uring read operation.
358        ///
359        /// # Safety
360        ///
361        /// The caller must ensure that the kernel wrote exactly `n` bytes
362        /// into the buffer that was returned by `prepare_uring_read`.
363        pub(crate) unsafe fn complete_uring_read(&mut self, n: usize) {
364            assert_eq!(self.pos, 0);
365            // SAFETY: `prepare_uring_read` handed out a pointer to
366            // `self.buf.spare_capacity_mut()` after asserting it's empty.
367            // The caller guarantees the kernel initialised exactly `n` bytes
368            // starting at that pointer, so bytes `0..n` are now initialised and
369            // it is sound to set the Vec length to `n`.
370            unsafe { self.buf.set_len(n) };
371        }
372    }
373}
374
375cfg_fs! {
376    impl Buf {
377        pub(crate) fn discard_read(&mut self) -> i64 {
378            let ret = -(self.bytes().len() as i64);
379            self.pos = 0;
380            self.buf.clear();
381            ret
382        }
383
384        pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>], max_buf_size: usize) -> usize {
385            assert!(self.is_empty());
386
387            let mut rem = max_buf_size;
388            for buf in bufs {
389                if rem == 0 {
390                    break
391                }
392
393                let len = buf.len().min(rem);
394                self.buf.extend_from_slice(&buf[..len]);
395                rem -= len;
396            }
397
398            max_buf_size - rem
399        }
400    }
401}
402
403fn gone() -> io::Error {
404    io::Error::other(crate::util::error::RUNTIME_SHUTTING_DOWN_ERROR)
405}
tokio/io/blocking.rs

tokio/io/
blocking.rs
