net: add 'try_io' function to 'unix::pipe' sender and receiver types (#8030)

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<R>(&self, f: impl FnOnce() -> io::Result<R>) -> io::Result<R> {
+        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<R>(&self, f: impl FnOnce() -> io::Result<R>) -> io::Result<R> {
+        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.

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(())
+}