From eacd28219bb9a16adb1d0ebfe0ed6ce42c6971b1 Mon Sep 17 00:00:00 2001
From: Masahiro Kozuka <masa.koz@kozuka.jp>
Date: Wed, 23 Mar 2022 20:06:34 +0900
Subject: [PATCH] tokio: make try_io methods call mio's try_io internally

A user defined I/O closure should be called by mio's try_io to ensure
that the I/O receives more events if it hits a WouldBlock error.

Fixes: #4510
---
 tokio/src/net/tcp/stream.rs           | 2 +-
 tokio/src/net/udp.rs                  | 2 +-
 tokio/src/net/unix/datagram/socket.rs | 2 +-
 tokio/src/net/unix/stream.rs          | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tokio/src/net/tcp/stream.rs b/tokio/src/net/tcp/stream.rs
index ebb67b84d16..ab0b51e33f2 100644
--- a/tokio/src/net/tcp/stream.rs
+++ b/tokio/src/net/tcp/stream.rs
@@ -968,7 +968,7 @@ impl TcpStream {
         interest: Interest,
         f: impl FnOnce() -> io::Result<R>,
     ) -> io::Result<R> {
-        self.io.registration().try_io(interest, f)
+        self.io.registration().try_io(interest, || self.io.try_io(f))
     }
 
     /// Receives data on the socket from the remote address to which it is
diff --git a/tokio/src/net/udp.rs b/tokio/src/net/udp.rs
index 12af5152c28..baf53456857 100644
--- a/tokio/src/net/udp.rs
+++ b/tokio/src/net/udp.rs
@@ -1272,7 +1272,7 @@ impl UdpSocket {
         interest: Interest,
         f: impl FnOnce() -> io::Result<R>,
     ) -> io::Result<R> {
-        self.io.registration().try_io(interest, f)
+        self.io.registration().try_io(interest, || self.io.try_io(f))
     }
 
     /// Receives data from the socket, without removing it from the input queue.
diff --git a/tokio/src/net/unix/datagram/socket.rs b/tokio/src/net/unix/datagram/socket.rs
index d5b618663dc..0f763536b51 100644
--- a/tokio/src/net/unix/datagram/socket.rs
+++ b/tokio/src/net/unix/datagram/socket.rs
@@ -1241,7 +1241,7 @@ impl UnixDatagram {
         interest: Interest,
         f: impl FnOnce() -> io::Result<R>,
     ) -> io::Result<R> {
-        self.io.registration().try_io(interest, f)
+        self.io.registration().try_io(interest, || self.io.try_io(f))
     }
 
     /// Returns the local address that this socket is bound to.
diff --git a/tokio/src/net/unix/stream.rs b/tokio/src/net/unix/stream.rs
index 4e7ef87b416..95996d94fee 100644
--- a/tokio/src/net/unix/stream.rs
+++ b/tokio/src/net/unix/stream.rs
@@ -685,7 +685,7 @@ impl UnixStream {
         interest: Interest,
         f: impl FnOnce() -> io::Result<R>,
     ) -> io::Result<R> {
-        self.io.registration().try_io(interest, f)
+        self.io.registration().try_io(interest, || self.io.try_io(f))
     }
 
     /// Creates new `UnixStream` from a `std::os::unix::net::UnixStream`.