Skip to content

Commit cbdceb9

Browse files
authored
io: add AsyncFd::try_io() and try_io_mut() (#6967)
1 parent d4178cf commit cbdceb9

File tree

2 files changed

+153
-1
lines changed

2 files changed

+153
-1
lines changed

tokio/src/io/async_fd.rs

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -872,6 +872,56 @@ impl<T: AsRawFd> AsyncFd<T> {
872872
.async_io(interest, || f(self.inner.as_mut().unwrap()))
873873
.await
874874
}
875+
876+
/// Tries to read or write from the file descriptor using a user-provided IO operation.
877+
///
878+
/// If the file descriptor is ready, the provided closure is called. The closure
879+
/// should attempt to perform IO operation on the file descriptor by manually
880+
/// calling the appropriate syscall. If the operation fails because the
881+
/// file descriptor is not actually ready, then the closure should return a
882+
/// `WouldBlock` error and the readiness flag is cleared. The return value
883+
/// of the closure is then returned by `try_io`.
884+
///
885+
/// If the file descriptor is not ready, then the closure is not called
886+
/// and a `WouldBlock` error is returned.
887+
///
888+
/// The closure should only return a `WouldBlock` error if it has performed
889+
/// an IO operation on the file descriptor that failed due to the file descriptor not being
890+
/// ready. Returning a `WouldBlock` error in any other situation will
891+
/// incorrectly clear the readiness flag, which can cause the file descriptor to
892+
/// behave incorrectly.
893+
///
894+
/// The closure should not perform the IO operation using any of the methods
895+
/// defined on the Tokio `AsyncFd` type, as this will mess with the
896+
/// readiness flag and can cause the file descriptor to behave incorrectly.
897+
///
898+
/// This method is not intended to be used with combined interests.
899+
/// The closure should perform only one type of IO operation, so it should not
900+
/// require more than one ready state. This method may panic or sleep forever
901+
/// if it is called with a combined interest.
902+
pub fn try_io<R>(
903+
&self,
904+
interest: Interest,
905+
f: impl FnOnce(&T) -> io::Result<R>,
906+
) -> io::Result<R> {
907+
self.registration
908+
.try_io(interest, || f(self.inner.as_ref().unwrap()))
909+
}
910+
911+
/// Tries to read or write from the file descriptor using a user-provided IO operation.
912+
///
913+
/// The behavior is the same as [`try_io`], except that the closure can mutate the inner
914+
/// value of the [`AsyncFd`].
915+
///
916+
/// [`try_io`]: AsyncFd::try_io
917+
pub fn try_io_mut<R>(
918+
&mut self,
919+
interest: Interest,
920+
f: impl FnOnce(&mut T) -> io::Result<R>,
921+
) -> io::Result<R> {
922+
self.registration
923+
.try_io(interest, || f(self.inner.as_mut().unwrap()))
924+
}
875925
}
876926

877927
impl<T: AsRawFd> AsRawFd for AsyncFd<T> {

tokio/tests/io_async_fd.rs

Lines changed: 103 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -302,7 +302,7 @@ async fn reregister() {
302302

303303
#[tokio::test]
304304
#[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri.
305-
async fn try_io() {
305+
async fn guard_try_io() {
306306
let (a, mut b) = socketpair();
307307

308308
b.write_all(b"0").unwrap();
@@ -336,6 +336,108 @@ async fn try_io() {
336336
let _ = readable.await.unwrap();
337337
}
338338

339+
#[tokio::test]
340+
#[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri.
341+
async fn try_io_readable() {
342+
let (a, mut b) = socketpair();
343+
let mut afd_a = AsyncFd::new(a).unwrap();
344+
345+
// Give the runtime some time to update bookkeeping.
346+
tokio::task::yield_now().await;
347+
348+
{
349+
let mut called = false;
350+
let _ = afd_a.try_io_mut(Interest::READABLE, |_| {
351+
called = true;
352+
Ok(())
353+
});
354+
assert!(
355+
!called,
356+
"closure should not have been called, since socket should not be readable"
357+
);
358+
}
359+
360+
// Make `a` readable by writing to `b`.
361+
// Give the runtime some time to update bookkeeping.
362+
b.write_all(&[0]).unwrap();
363+
tokio::task::yield_now().await;
364+
365+
{
366+
let mut called = false;
367+
let _ = afd_a.try_io(Interest::READABLE, |_| {
368+
called = true;
369+
Ok(())
370+
});
371+
assert!(
372+
called,
373+
"closure should have been called, since socket should have data available to read"
374+
);
375+
}
376+
377+
{
378+
let mut called = false;
379+
let _ = afd_a.try_io(Interest::READABLE, |_| {
380+
called = true;
381+
io::Result::<()>::Err(ErrorKind::WouldBlock.into())
382+
});
383+
assert!(
384+
called,
385+
"closure should have been called, since socket should have data available to read"
386+
);
387+
}
388+
389+
{
390+
let mut called = false;
391+
let _ = afd_a.try_io(Interest::READABLE, |_| {
392+
called = true;
393+
Ok(())
394+
});
395+
assert!(!called, "closure should not have been called, since socket readable state should have been cleared");
396+
}
397+
}
398+
399+
#[tokio::test]
400+
#[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri.
401+
async fn try_io_writable() {
402+
let (a, _b) = socketpair();
403+
let afd_a = AsyncFd::new(a).unwrap();
404+
405+
// Give the runtime some time to update bookkeeping.
406+
tokio::task::yield_now().await;
407+
408+
{
409+
let mut called = false;
410+
let _ = afd_a.try_io(Interest::WRITABLE, |_| {
411+
called = true;
412+
Ok(())
413+
});
414+
assert!(
415+
called,
416+
"closure should have been called, since socket should still be marked as writable"
417+
);
418+
}
419+
{
420+
let mut called = false;
421+
let _ = afd_a.try_io(Interest::WRITABLE, |_| {
422+
called = true;
423+
io::Result::<()>::Err(ErrorKind::WouldBlock.into())
424+
});
425+
assert!(
426+
called,
427+
"closure should have been called, since socket should still be marked as writable"
428+
);
429+
}
430+
431+
{
432+
let mut called = false;
433+
let _ = afd_a.try_io(Interest::WRITABLE, |_| {
434+
called = true;
435+
Ok(())
436+
});
437+
assert!(!called, "closure should not have been called, since socket writable state should have been cleared");
438+
}
439+
}
440+
339441
#[tokio::test]
340442
#[cfg_attr(miri, ignore)] // No F_GETFL for fcntl in miri.
341443
async fn multiple_waiters() {

0 commit comments

Comments
 (0)