nuclei/io/
async_fd.rs

1//! [`AsyncFd`]: a caller's own descriptor, driven by readiness. The shard's
2//! driver says when it may be ready (an edge on epoll and kqueue, a poll on
3//! io_uring), and the caller makes its own calls on it, clearing the
4//! readiness when one would block. It is for descriptors the caller
5//! manages (an eventfd, a device, a socket of a kind nuclei does not wrap);
6//! nuclei's sockets and files are completed by the driver instead.
7//!
8//! The descriptor must be in non-blocking mode; its mode is left as it is.
9//! Readiness may be spurious (the side starts ready, and a guard says what
10//! may be ready, not what is): a call that would block answers
11//! `WouldBlock`, and clearing the guard's readiness makes the next wait
12//! last until the driver says so again. A clear that comes after the
13//! driver's next word is not lost: each guard remembers what it saw
14//! (proofs/tla/asyncfd).
15//!
16//! ```
17//! use std::os::unix::net::UnixStream;
18//! use std::io::{Read, Write};
19//! use nuclei::io::unix::AsyncFd;
20//!
21//! # nuclei::block_on(async {
22//! let (a, mut b) = UnixStream::pair()?;
23//! a.set_nonblocking(true)?;
24//! let fd = AsyncFd::new(a)?;
25//! b.write_all(b"ping")?;
26//! let mut buf = [0u8; 4];
27//! let n = fd.async_io(nuclei::io::Interest::READABLE, |mut a| a.read(&mut buf)).await?;
28//! assert_eq!(&buf[..n], b"ping");
29//! # std::io::Result::Ok(())
30//! # }).unwrap();
31//! ```
32
33use core::fmt;
34use core::future::poll_fn;
35use core::task::{Context, Poll, ready};
36use std::io;
37use std::os::fd::AsFd;
38
39use super::caller::{self, Marks};
40use super::source::Kind;
41use super::{Handle, Interest, Ready};
42
43/// A caller's own descriptor, driven by readiness.
44pub struct AsyncFd<T: AsFd> {
45    handle: Handle<T>,
46    interest: Interest,
47}
48
49/// What a readiness wait of an [`AsyncFd`] found: the readiness, and the
50/// edge counts it is of, to clear with.
51pub struct AsyncFdReadyGuard<'a, T: AsFd> {
52    fd: &'a AsyncFd<T>,
53    marks: Marks,
54}
55
56/// As [`AsyncFdReadyGuard`], with the descriptor borrowed mutably.
57pub struct AsyncFdReadyMutGuard<'a, T: AsFd> {
58    fd: &'a mut AsyncFd<T>,
59    marks: Marks,
60}
61
62/// The call of a guard's `try_io` would block: the guard's readiness was
63/// cleared.
64#[derive(Debug)]
65pub struct TryIoError(());
66
67impl<T: AsFd> AsyncFd<T> {
68    /// Drives `inner` by readiness for reading and writing both.
69    ///
70    /// # Errors
71    ///
72    /// Never today: the descriptor registers with the shard that first
73    /// waits on it, and a poller that refuses it (a regular file, on epoll)
74    /// says so at that wait.
75    pub fn new(inner: T) -> io::Result<Self> {
76        Self::with_interest(inner, Interest::READABLE | Interest::WRITABLE)
77    }
78
79    /// Drives `inner` by readiness for what `interest` names.
80    ///
81    /// # Errors
82    ///
83    /// As [`new`](Self::new).
84    pub fn with_interest(inner: T, interest: Interest) -> io::Result<Self> {
85        Ok(Self {
86            handle: Handle::from_nonblocking(inner, Kind::Fd),
87            interest,
88        })
89    }
90
91    /// The descriptor.
92    pub fn get_ref(&self) -> &T {
93        self.handle.get_ref()
94    }
95
96    /// The descriptor, to change.
97    pub fn get_mut(&mut self) -> &mut T {
98        self.handle.get_mut()
99    }
100
101    /// The descriptor back, no longer watched.
102    pub fn into_inner(self) -> T {
103        self.handle.into_inner()
104    }
105
106    /// Polls until the descriptor may be readable.
107    ///
108    /// # Errors
109    ///
110    /// When the descriptor cannot be watched.
111    pub fn poll_read_ready<'a>(
112        &'a self,
113        cx: &mut Context<'_>,
114    ) -> Poll<io::Result<AsyncFdReadyGuard<'a, T>>> {
115        let tick = ready!(self.handle.source().poll_fd_read_ready(cx))?;
116        Poll::Ready(Ok(AsyncFdReadyGuard {
117            fd: self,
118            marks: Marks {
119                ready: Ready::READABLE,
120                read: Some(tick),
121                write: None,
122            },
123        }))
124    }
125
126    /// Polls until the descriptor may be writable.
127    ///
128    /// # Errors
129    ///
130    /// When the descriptor cannot be watched.
131    pub fn poll_write_ready<'a>(
132        &'a self,
133        cx: &mut Context<'_>,
134    ) -> Poll<io::Result<AsyncFdReadyGuard<'a, T>>> {
135        let tick = ready!(self.handle.source().poll_fd_write_ready(cx))?;
136        Poll::Ready(Ok(AsyncFdReadyGuard {
137            fd: self,
138            marks: Marks {
139                ready: Ready::WRITABLE,
140                read: None,
141                write: Some(tick),
142            },
143        }))
144    }
145
146    /// The sides of `interest` the descriptor is watched for; none is an
147    /// error, since a wait for them could never end.
148    fn sides(&self, interest: Interest) -> io::Result<(bool, bool)> {
149        let (read, write) = caller::sides(interest);
150        let (read, write) = (
151            read && self.interest.is_readable(),
152            write && self.interest.is_writable(),
153        );
154        if !read && !write {
155            return Err(io::Error::new(
156                io::ErrorKind::InvalidInput,
157                "the descriptor is not watched for that readiness",
158            ));
159        }
160        Ok((read, write))
161    }
162
163    /// What of `interest` may be ready now, waiting until some is.
164    fn poll_marks(&self, cx: &mut Context<'_>, interest: Interest) -> Poll<io::Result<Marks>> {
165        let (read, write) = self.sides(interest)?;
166        caller::poll_marks(self.handle.source(), cx, read, write)
167    }
168
169    /// Waits until the descriptor may be readable.
170    ///
171    /// # Errors
172    ///
173    /// When the descriptor cannot be watched.
174    pub async fn readable(&self) -> io::Result<AsyncFdReadyGuard<'_, T>> {
175        self.ready(Interest::READABLE).await
176    }
177
178    /// Waits until the descriptor may be writable.
179    ///
180    /// # Errors
181    ///
182    /// When the descriptor cannot be watched.
183    pub async fn writable(&self) -> io::Result<AsyncFdReadyGuard<'_, T>> {
184        self.ready(Interest::WRITABLE).await
185    }
186
187    /// Waits until some of `interest` may be ready.
188    ///
189    /// # Errors
190    ///
191    /// When the descriptor cannot be watched, and `InvalidInput` when it is
192    /// not watched for any of `interest`.
193    pub async fn ready(&self, interest: Interest) -> io::Result<AsyncFdReadyGuard<'_, T>> {
194        let marks = poll_fn(|cx| self.poll_marks(cx, interest)).await?;
195        Ok(AsyncFdReadyGuard { fd: self, marks })
196    }
197
198    /// As [`readable`](Self::readable), with the descriptor borrowed
199    /// mutably.
200    ///
201    /// # Errors
202    ///
203    /// When the descriptor cannot be watched.
204    pub async fn readable_mut(&mut self) -> io::Result<AsyncFdReadyMutGuard<'_, T>> {
205        self.ready_mut(Interest::READABLE).await
206    }
207
208    /// As [`writable`](Self::writable), with the descriptor borrowed
209    /// mutably.
210    ///
211    /// # Errors
212    ///
213    /// When the descriptor cannot be watched.
214    pub async fn writable_mut(&mut self) -> io::Result<AsyncFdReadyMutGuard<'_, T>> {
215        self.ready_mut(Interest::WRITABLE).await
216    }
217
218    /// As [`ready`](Self::ready), with the descriptor borrowed mutably.
219    ///
220    /// # Errors
221    ///
222    /// When the descriptor cannot be watched.
223    pub async fn ready_mut(
224        &mut self,
225        interest: Interest,
226    ) -> io::Result<AsyncFdReadyMutGuard<'_, T>> {
227        let marks = poll_fn(|cx| self.poll_marks(cx, interest)).await?;
228        Ok(AsyncFdReadyMutGuard { fd: self, marks })
229    }
230
231    /// Runs `f` on the descriptor each time some of `interest` may be ready,
232    /// until it answers other than `WouldBlock`, clearing the readiness
233    /// each time it does.
234    ///
235    /// # Errors
236    ///
237    /// What `f` answers, when the descriptor cannot be watched, and
238    /// `InvalidInput` when it is not watched for any of `interest`.
239    pub async fn async_io<R>(
240        &self,
241        interest: Interest,
242        mut f: impl FnMut(&T) -> io::Result<R>,
243    ) -> io::Result<R> {
244        let sides = self.sides(interest)?;
245        caller::async_io(self.handle.source(), sides, || f(self.get_ref())).await
246    }
247
248    /// Runs `f` on the descriptor now if some of `interest` may be ready,
249    /// clearing that readiness if it answers `WouldBlock`.
250    ///
251    /// # Errors
252    ///
253    /// `WouldBlock` when none of `interest` may be ready, `InvalidInput`
254    /// when the descriptor is not watched for any of it, and what `f`
255    /// answers.
256    pub fn try_io<R>(
257        &self,
258        interest: Interest,
259        f: impl FnOnce(&T) -> io::Result<R>,
260    ) -> io::Result<R> {
261        caller::try_io(self.handle.source(), self.sides(interest)?, || {
262            f(self.get_ref())
263        })
264    }
265
266    fn clear(&self, marks: Marks) {
267        caller::clear(self.handle.source(), marks);
268    }
269}
270
271macro_rules! guard {
272    ($guard:ident) => {
273        impl<T: AsFd> $guard<'_, T> {
274            /// What may be ready.
275            pub fn ready(&self) -> Ready {
276                self.marks.ready
277            }
278
279            /// The descriptor would block: the next wait lasts until the
280            /// driver says it may be ready again.
281            pub fn clear_ready(&mut self) {
282                self.fd.clear(self.marks);
283                self.marks.ready = Ready::EMPTY;
284            }
285
286            /// Clears what of `ready` this guard holds.
287            pub fn clear_ready_matching(&mut self, ready: Ready) {
288                let mut marks = self.marks;
289                if !ready.is_readable() {
290                    marks.read = None;
291                }
292                if !ready.is_writable() {
293                    marks.write = None;
294                }
295                self.fd.clear(marks);
296                if ready.is_readable() {
297                    self.marks.read = None;
298                }
299                if ready.is_writable() {
300                    self.marks.write = None;
301                }
302                self.marks.ready = Ready::EMPTY;
303                if self.marks.read.is_some() {
304                    self.marks.ready = self.marks.ready | Ready::READABLE;
305                }
306                if self.marks.write.is_some() {
307                    self.marks.ready = self.marks.ready | Ready::WRITABLE;
308                }
309            }
310
311            /// Keeps the readiness: the next wait may answer at once.
312            pub fn retain_ready(&mut self) {}
313
314            /// The descriptor.
315            pub fn get_inner(&self) -> &T {
316                self.fd.get_ref()
317            }
318        }
319
320        impl<T: AsFd> fmt::Debug for $guard<'_, T> {
321            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
322                f.debug_struct(stringify!($guard))
323                    .field("ready", &self.marks.ready)
324                    .finish_non_exhaustive()
325            }
326        }
327    };
328}
329
330guard!(AsyncFdReadyGuard);
331guard!(AsyncFdReadyMutGuard);
332
333impl<'a, T: AsFd> AsyncFdReadyGuard<'a, T> {
334    /// Runs `f`; if it answers `WouldBlock`, the readiness is cleared and
335    /// `Err(TryIoError)` comes back, else `Ok` with what it answered.
336    ///
337    /// # Errors
338    ///
339    /// [`TryIoError`] when `f` would block.
340    pub fn try_io<R>(
341        &mut self,
342        f: impl FnOnce(&'a AsyncFd<T>) -> io::Result<R>,
343    ) -> Result<io::Result<R>, TryIoError> {
344        match f(self.fd) {
345            Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
346                self.clear_ready();
347                Err(TryIoError(()))
348            }
349            result => Ok(result),
350        }
351    }
352
353    /// The descriptor's `AsyncFd`.
354    pub fn get_ref(&self) -> &'a AsyncFd<T> {
355        self.fd
356    }
357}
358
359impl<T: AsFd> AsyncFdReadyMutGuard<'_, T> {
360    /// Runs `f`; if it answers `WouldBlock`, the readiness is cleared and
361    /// `Err(TryIoError)` comes back, else `Ok` with what it answered.
362    ///
363    /// # Errors
364    ///
365    /// [`TryIoError`] when `f` would block.
366    pub fn try_io<R>(
367        &mut self,
368        f: impl FnOnce(&mut AsyncFd<T>) -> io::Result<R>,
369    ) -> Result<io::Result<R>, TryIoError> {
370        match f(self.fd) {
371            Err(err) if err.kind() == io::ErrorKind::WouldBlock => {
372                self.clear_ready();
373                Err(TryIoError(()))
374            }
375            result => Ok(result),
376        }
377    }
378
379    /// The descriptor's `AsyncFd`.
380    pub fn get_ref(&self) -> &AsyncFd<T> {
381        self.fd
382    }
383
384    /// The descriptor's `AsyncFd`, mutably.
385    pub fn get_mut(&mut self) -> &mut AsyncFd<T> {
386        self.fd
387    }
388
389    /// The descriptor, mutably.
390    pub fn get_inner_mut(&mut self) -> &mut T {
391        self.fd.get_mut()
392    }
393}
394
395impl<T: AsFd + fmt::Debug> fmt::Debug for AsyncFd<T> {
396    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
397        f.debug_struct("AsyncFd")
398            .field("inner", self.get_ref())
399            .finish_non_exhaustive()
400    }
401}
402
403impl fmt::Display for TryIoError {
404    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
405        f.write_str("the call would block")
406    }
407}
408
409impl std::error::Error for TryIoError {}
410
411#[cfg(test)]
412mod tests {
413    use std::io::{Read, Write};
414    use std::os::unix::net::UnixStream;
415    use std::time::{Duration, Instant};
416
417    use super::*;
418
419    /// Whether this thread's shard, which serves the test's descriptor, runs a ring: other tests
420    /// may run shards on the poller meanwhile.
421    fn on_ring() -> bool {
422        crate::Proactor::backend() == crate::IoBackend::IoUring
423    }
424
425    /// An edge that comes between a call that would block and the clear of
426    /// its guard keeps the side ready: the clear is of the look before the
427    /// call (proofs/tla/asyncfd, the no_tick mutation). On a ring the side
428    /// counts only the polls it arms, each after a clear, so the clear
429    /// there comes first and the poll it arms finds the byte.
430    #[test]
431    fn an_edge_after_the_call_survives_the_clear() {
432        crate::block_on(async {
433            let (a, mut b) = UnixStream::pair().unwrap();
434            a.set_nonblocking(true).unwrap();
435            let fd = AsyncFd::with_interest(a, Interest::READABLE).unwrap();
436            let mut buf = [0u8; 8];
437            // A wait that waits registers the descriptor with the driver.
438            let mut guard = fd.readable().await.unwrap();
439            assert!(guard.try_io(|fd| fd.get_ref().read(&mut buf)).is_err());
440            b.write_all(b"r").unwrap();
441            let read = fd.async_io(Interest::READABLE, |mut a| a.read(&mut buf));
442            assert_eq!(read.await.unwrap(), 1);
443            let mut first = fd.readable().await.unwrap();
444            let err = fd.get_ref().read(&mut buf).unwrap_err();
445            assert_eq!(err.kind(), io::ErrorKind::WouldBlock);
446            b.write_all(b"x").unwrap();
447            if on_ring() {
448                first.clear_ready();
449                let n = fd.async_io(Interest::READABLE, |mut a| a.read(&mut buf));
450                assert_eq!(n.await.unwrap(), 1);
451                return;
452            }
453            let deadline = Instant::now() + Duration::from_secs(10);
454            loop {
455                crate::task::yield_now().await;
456                let look = fd.readable().await.unwrap();
457                if look.marks.read != first.marks.read {
458                    break;
459                }
460                assert!(Instant::now() < deadline, "the driver never saw the edge");
461            }
462            first.clear_ready();
463            let mut cx = Context::from_waker(std::task::Waker::noop());
464            assert!(
465                fd.poll_read_ready(&mut cx).is_ready(),
466                "the edge after the call was lost to the clear"
467            );
468            assert_eq!(fd.get_ref().read(&mut buf).unwrap(), 1);
469        });
470    }
471
472    /// Bytes written into `fd` until a write would block.
473    fn fill(fd: &AsyncFd<UnixStream>) -> usize {
474        let chunk = [5u8; 4096];
475        let mut sent = 0;
476        loop {
477            match fd.get_ref().write(&chunk) {
478                Ok(n) => sent += n,
479                Err(err) if err.kind() == io::ErrorKind::WouldBlock => return sent,
480                Err(err) => panic!("{err}"),
481            }
482        }
483    }
484
485    /// As above, on the write side: the peer drains the socket between a
486    /// write that would block and the clear.
487    #[test]
488    fn a_write_edge_after_the_call_survives_the_clear() {
489        crate::block_on(async {
490            let (a, mut b) = UnixStream::pair().unwrap();
491            a.set_nonblocking(true).unwrap();
492            let fd = AsyncFd::with_interest(a, Interest::WRITABLE).unwrap();
493            let mut sink = vec![0u8; 1 << 20];
494            // A wait that waits registers the descriptor with the driver.
495            let mut guard = fd.writable().await.unwrap();
496            let sent = fill(&fd);
497            guard.clear_ready();
498            b.read_exact(&mut sink[..sent]).unwrap();
499            fd.writable().await.unwrap();
500            let mut first = fd.writable().await.unwrap();
501            let sent = fill(&fd);
502            b.read_exact(&mut sink[..sent]).unwrap();
503            if on_ring() {
504                first.clear_ready();
505                let n = fd.async_io(Interest::WRITABLE, |mut a| a.write(b"w"));
506                assert_eq!(n.await.unwrap(), 1);
507                return;
508            }
509            let deadline = Instant::now() + Duration::from_secs(10);
510            loop {
511                crate::task::yield_now().await;
512                let look = fd.writable().await.unwrap();
513                if look.marks.write != first.marks.write {
514                    break;
515                }
516                assert!(Instant::now() < deadline, "the driver never saw the edge");
517            }
518            first.clear_ready();
519            let mut cx = Context::from_waker(std::task::Waker::noop());
520            assert!(
521                fd.poll_write_ready(&mut cx).is_ready(),
522                "the edge after the call was lost to the clear"
523            );
524        });
525    }
526}
nuclei/io/async_fd.rs

nuclei/io/
async_fd.rs
