From 5db10f538b683fe88d699dfd11be31d193db011c Mon Sep 17 00:00:00 2001 From: Timo Date: Sun, 12 Apr 2026 15:06:06 +0200 Subject: [PATCH] net: add 'try_io' function to 'unix::pipe' sender and receiver types (#8030) --- tokio/src/net/unix/pipe.rs | 64 ++++++++++++++++++++++ tokio/tests/net_unix_pipe.rs | 102 +++++++++++++++++++++++++++++++++++ 2 files changed, 166 insertions(+) diff --git a/tokio/src/net/unix/pipe.rs b/tokio/src/net/unix/pipe.rs index 65bd51ad8..f605043a9 100644 --- a/tokio/src/net/unix/pipe.rs +++ b/tokio/src/net/unix/pipe.rs @@ -722,6 +722,38 @@ impl Sender { .try_io(Interest::WRITABLE, || (&*self.io).write_vectored(buf)) } + /// Tries to write from the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation on the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `pipe::Sender` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// Usually, [`writable()`] or [`ready()`] is used with this function. + /// + /// [`writable()`]: Self::writable() + /// [`ready()`]: Self::ready() + pub fn try_io(&self, f: impl FnOnce() -> io::Result) -> io::Result { + self.io + .registration() + .try_io(Interest::WRITABLE, || self.io.try_io(f)) + } + /// Converts the pipe into an [`OwnedFd`] in blocking mode. /// /// This function will deregister this pipe end from the event loop, set @@ -1237,6 +1269,38 @@ impl Receiver { .try_io(Interest::READABLE, || (&*self.io).read_vectored(bufs)) } + /// Tries to read to the socket using a user-provided IO operation. + /// + /// If the socket is ready, the provided closure is called. The closure + /// should attempt to perform IO operation on the socket by manually + /// calling the appropriate syscall. If the operation fails because the + /// socket is not actually ready, then the closure should return a + /// `WouldBlock` error and the readiness flag is cleared. The return value + /// of the closure is then returned by `try_io`. + /// + /// If the socket is not ready, then the closure is not called + /// and a `WouldBlock` error is returned. + /// + /// The closure should only return a `WouldBlock` error if it has performed + /// an IO operation on the socket that failed due to the socket not being + /// ready. Returning a `WouldBlock` error in any other situation will + /// incorrectly clear the readiness flag, which can cause the socket to + /// behave incorrectly. + /// + /// The closure should not perform the IO operation using any of the methods + /// defined on the Tokio `pipe::Receiver` type, as this will mess with the + /// readiness flag and can cause the socket to behave incorrectly. + /// + /// Usually, [`readable()`] or [`ready()`] is used with this function. + /// + /// [`readable()`]: Self::readable() + /// [`ready()`]: Self::ready() + pub fn try_io(&self, f: impl FnOnce() -> io::Result) -> io::Result { + self.io + .registration() + .try_io(Interest::READABLE, || self.io.try_io(f)) + } + cfg_io_util! { /// Tries to read data from the pipe into the provided buffer, advancing the /// buffer's internal cursor, returning how many bytes were read. diff --git a/tokio/tests/net_unix_pipe.rs b/tokio/tests/net_unix_pipe.rs index c35b1c81f..066c678e8 100644 --- a/tokio/tests/net_unix_pipe.rs +++ b/tokio/tests/net_unix_pipe.rs @@ -544,3 +544,105 @@ async fn anon_pipe_into_blocking_fd() -> std::io::Result<()> { Ok(()) } + +#[tokio::test] +async fn try_io_writable() -> std::io::Result<()> { + let (tx, _rx) = pipe::pipe()?; + + // Give the runtime some time to update bookkeeping. + tokio::task::yield_now().await; + + { + let mut called = false; + let _ = tx.try_io(|| { + called = true; + Ok(()) + }); + assert!( + called, + "closure should have been called, since socket should still be marked as writable" + ); + } + { + let mut called = false; + let _ = tx.try_io(|| { + called = true; + io::Result::<()>::Err(io::ErrorKind::WouldBlock.into()) + }); + assert!( + called, + "closure should have been called, since socket should still be marked as writable" + ); + } + + { + let mut called = false; + let _ = tx.try_io(|| { + called = true; + Ok(()) + }); + assert!(!called, "closure should not have been called, since socket writable state should have been cleared"); + } + + Ok(()) +} + +#[tokio::test] +async fn try_io_readable() -> io::Result<()> { + let (mut tx, rx) = pipe::pipe()?; + + // Give the runtime some time to update bookkeeping. + tokio::task::yield_now().await; + + { + let mut called = false; + let _ = rx.try_io(|| { + called = true; + Ok(()) + }); + assert!( + !called, + "closure should not have been called, since socket should not be readable" + ); + } + + // Make `a` readable by writing to `b`. + // Give the runtime some time to update bookkeeping. + tx.write_all(&[0]).await?; + tokio::task::yield_now().await; + + { + let mut called = false; + let _ = rx.try_io(|| { + called = true; + Ok(()) + }); + assert!( + called, + "closure should have been called, since socket should have data available to read" + ); + } + + { + let mut called = false; + let _ = rx.try_io(|| { + called = true; + io::Result::<()>::Err(io::ErrorKind::WouldBlock.into()) + }); + assert!( + called, + "closure should have been called, since socket should have data available to read" + ); + } + + { + let mut called = false; + let _ = rx.try_io(|| { + called = true; + Ok(()) + }); + assert!(!called, "closure should not have been called, since socket readable state should have been cleared"); + } + + Ok(()) +}