From 28ce4eeab2fe983a159b6f6d5070c16e6ff705b9 Mon Sep 17 00:00:00 2001
From: Noah Kennedy <nkennedy@cloudflare.com>
Date: Tue, 26 Jul 2022 00:52:30 -0500
Subject: [PATCH] io: reduce syscalls in poll_read (#4840)

As the [epoll documentation points out](https://man7.org/linux/man-pages/man7/epoll.7.html), a read that only partially fills a buffer is sufficient to show that the socket buffer has been drained.

We can take advantage of this by clearing readiness in this case, which seems to significantly improve performance under certain conditions.

Co-authored-by: Carl Lerche <me@carllerche.com>
---
 tokio/src/io/driver/registration.rs |  1 +
 tokio/src/io/poll_evented.rs        | 34 +++++++++++++++++++++--------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/tokio/src/io/driver/registration.rs b/tokio/src/io/driver/registration.rs
index c9393650c20..dbc34592e3a 100644
--- a/tokio/src/io/driver/registration.rs
+++ b/tokio/src/io/driver/registration.rs
@@ -115,6 +115,7 @@ impl Registration {
 
     // Uses the poll path, requiring the caller to ensure mutual exclusion for
     // correctness. Only the last task to call this function is notified.
+    #[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
     pub(crate) fn poll_read_io<R>(
         &self,
         cx: &mut Context<'_>,
diff --git a/tokio/src/io/poll_evented.rs b/tokio/src/io/poll_evented.rs
index 25cece62764..f495fddd138 100644
--- a/tokio/src/io/poll_evented.rs
+++ b/tokio/src/io/poll_evented.rs
@@ -153,16 +153,32 @@ feature! {
         {
             use std::io::Read;
 
-            let n = ready!(self.registration.poll_read_io(cx, || {
+            loop {
+                let evt = ready!(self.registration.poll_read_ready(cx))?;
+
                 let b = &mut *(buf.unfilled_mut() as *mut [std::mem::MaybeUninit<u8>] as *mut [u8]);
-                self.io.as_ref().unwrap().read(b)
-            }))?;
-
-            // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the
-            // buffer.
-            buf.assume_init(n);
-            buf.advance(n);
-            Poll::Ready(Ok(()))
+                let len = b.len();
+
+                match self.io.as_ref().unwrap().read(b) {
+                    Ok(n) => {
+                        // if we read a partially full buffer, this is sufficient on unix to show
+                        // that the socket buffer has been drained
+                        if n > 0 && (!cfg!(windows) && n < len) {
+                            self.registration.clear_readiness(evt);
+                        }
+
+                        // Safety: We trust `TcpStream::read` to have filled up `n` bytes in the
+                        // buffer.
+                        buf.assume_init(n);
+                        buf.advance(n);
+                        return Poll::Ready(Ok(()));
+                    },
+                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                        self.registration.clear_readiness(evt);
+                    }
+                    Err(e) => return Poll::Ready(Err(e)),
+                }
+            }
         }
 
         pub(crate) fn poll_write<'a>(&'a self, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>>