Apply async_io to all IO assets

tokio/src/net/tcp/stream.rs

@@ -1002,6 +1002,48 @@ impl TcpStream {
             .registration()
             .try_io(interest, || self.io.try_io(f))
     }
+
+    /// Reads or writes from the socket using a user-provided IO operation.
+    ///
+    /// The readiness of the socket is awaited and when 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, the readiness
+    /// flag is cleared and the socket readiness is awaited again. This loop
+    /// repeated until the closure returns an `Ok` or an error that doesn't
+    /// have the `WouldBlock` value.
+    ///
+    /// 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 `UdpSocket` type, as this will mess with the
+    /// readiness flag and can cause the socket to behave incorrectly.
+    ///
+    /// This method is not intended to be used with combined interests.
+    /// The closure should perform only one type of IO operation, so it should not
+    /// require more than one ready state. This method may panic or sleep forever
+    /// if it is called with a combined interest.
+    ///
+    /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function.
+    ///
+    /// [`readable()`]: UdpSocket::readable()
+    /// [`writable()`]: UdpSocket::writable()
+    /// [`ready()`]: UdpSocket::ready()
+    pub async fn async_io<R>(
+        &self, 
+        interest: Interest, 
+        mut f: impl FnMut() -> io::Result<R>,
+    ) -> io::Result<R> {
+        self.io
+            .registration()
+            .async_io(interest, || self.io.try_io(&mut f) )
+            .await
+    }
 
     /// Receives data on the socket from the remote address to which it is
     /// connected, without removing that data from the queue. On success,

tokio/src/net/unix/datagram/socket.rs

@@ -1252,6 +1252,48 @@ impl UnixDatagram {
             .registration()
             .try_io(interest, || self.io.try_io(f))
     }
+
+    /// Reads or writes from the socket using a user-provided IO operation.
+    ///
+    /// The readiness of the socket is awaited and when 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, the readiness
+    /// flag is cleared and the socket readiness is awaited again. This loop
+    /// repeated until the closure returns an `Ok` or an error that doesn't
+    /// have the `WouldBlock` value.
+    ///
+    /// 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 `UdpSocket` type, as this will mess with the
+    /// readiness flag and can cause the socket to behave incorrectly.
+    ///
+    /// This method is not intended to be used with combined interests.
+    /// The closure should perform only one type of IO operation, so it should not
+    /// require more than one ready state. This method may panic or sleep forever
+    /// if it is called with a combined interest.
+    ///
+    /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function.
+    ///
+    /// [`readable()`]: UdpSocket::readable()
+    /// [`writable()`]: UdpSocket::writable()
+    /// [`ready()`]: UdpSocket::ready()
+    pub async fn async_io<R>(
+        &self, 
+        interest: Interest, 
+        mut f: impl FnMut() -> io::Result<R>,
+    ) -> io::Result<R> {
+        self.io
+            .registration()
+            .async_io(interest, || self.io.try_io(&mut f) )
+            .await
+    }
 
     /// Returns the local address that this socket is bound to.
     ///

tokio/src/net/unix/stream.rs

@@ -699,6 +699,48 @@ impl UnixStream {
             .registration()
             .try_io(interest, || self.io.try_io(f))
     }
+
+    /// Reads or writes from the socket using a user-provided IO operation.
+    ///
+    /// The readiness of the socket is awaited and when 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, the readiness
+    /// flag is cleared and the socket readiness is awaited again. This loop
+    /// repeated until the closure returns an `Ok` or an error that doesn't
+    /// have the `WouldBlock` value.
+    ///
+    /// 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 `UdpSocket` type, as this will mess with the
+    /// readiness flag and can cause the socket to behave incorrectly.
+    ///
+    /// This method is not intended to be used with combined interests.
+    /// The closure should perform only one type of IO operation, so it should not
+    /// require more than one ready state. This method may panic or sleep forever
+    /// if it is called with a combined interest.
+    ///
+    /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function.
+    ///
+    /// [`readable()`]: UdpSocket::readable()
+    /// [`writable()`]: UdpSocket::writable()
+    /// [`ready()`]: UdpSocket::ready()
+    pub async fn async_io<R>(
+        &self, 
+        interest: Interest, 
+        mut f: impl FnMut() -> io::Result<R>,
+    ) -> io::Result<R> {
+        self.io
+            .registration()
+            .async_io(interest, || self.io.try_io(&mut f) )
+            .await
+    }
 
     /// Creates new `UnixStream` from a `std::os::unix::net::UnixStream`.
     ///

tokio/src/net/windows/named_pipe.rs

@@ -845,6 +845,48 @@ impl NamedPipeServer {
     ) -> io::Result<R> {
         self.io.registration().try_io(interest, f)
     }
+
+    /// Reads or writes from the socket using a user-provided IO operation.
+    ///
+    /// The readiness of the socket is awaited and when 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, the readiness
+    /// flag is cleared and the socket readiness is awaited again. This loop
+    /// repeated until the closure returns an `Ok` or an error that doesn't
+    /// have the `WouldBlock` value.
+    ///
+    /// 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 `UdpSocket` type, as this will mess with the
+    /// readiness flag and can cause the socket to behave incorrectly.
+    ///
+    /// This method is not intended to be used with combined interests.
+    /// The closure should perform only one type of IO operation, so it should not
+    /// require more than one ready state. This method may panic or sleep forever
+    /// if it is called with a combined interest.
+    ///
+    /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function.
+    ///
+    /// [`readable()`]: UdpSocket::readable()
+    /// [`writable()`]: UdpSocket::writable()
+    /// [`ready()`]: UdpSocket::ready()
+    pub async fn async_io<R>(
+        &self, 
+        interest: Interest, 
+        mut f: impl FnMut() -> io::Result<R>,
+    ) -> io::Result<R> {
+        self.io
+            .registration()
+            .async_io(interest, || self.io.try_io(&mut f) )
+            .await
+    }
 }
 
 impl AsyncRead for NamedPipeServer {
@@ -1589,6 +1631,48 @@ impl NamedPipeClient {
     ) -> io::Result<R> {
         self.io.registration().try_io(interest, f)
     }
+
+    /// Reads or writes from the socket using a user-provided IO operation.
+    ///
+    /// The readiness of the socket is awaited and when 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, the readiness
+    /// flag is cleared and the socket readiness is awaited again. This loop
+    /// repeated until the closure returns an `Ok` or an error that doesn't
+    /// have the `WouldBlock` value.
+    ///
+    /// 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 `UdpSocket` type, as this will mess with the
+    /// readiness flag and can cause the socket to behave incorrectly.
+    ///
+    /// This method is not intended to be used with combined interests.
+    /// The closure should perform only one type of IO operation, so it should not
+    /// require more than one ready state. This method may panic or sleep forever
+    /// if it is called with a combined interest.
+    ///
+    /// Usually, [`readable()`], [`writable()`] or [`ready()`] is used with this function.
+    ///
+    /// [`readable()`]: UdpSocket::readable()
+    /// [`writable()`]: UdpSocket::writable()
+    /// [`ready()`]: UdpSocket::ready()
+    pub async fn async_io<R>(
+        &self, 
+        interest: Interest, 
+        mut f: impl FnMut() -> io::Result<R>,
+    ) -> io::Result<R> {
+        self.io
+            .registration()
+            .async_io(interest, || self.io.try_io(&mut f) )
+            .await
+    }
 }
 
 impl AsyncRead for NamedPipeClient {