From bbb96fc61a95cfb1bcedf805651cb4fa286fbdc3 Mon Sep 17 00:00:00 2001 From: Noah Kennedy Date: Fri, 2 Sep 2022 14:20:43 -0500 Subject: [PATCH] io: reduce syscalls in `poll_write` This applies the same optimization made in #4840 to writes. --- tokio/src/io/poll_evented.rs | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/tokio/src/io/poll_evented.rs b/tokio/src/io/poll_evented.rs index b8c0be2c09f..cc5db039d58 100644 --- a/tokio/src/io/poll_evented.rs +++ b/tokio/src/io/poll_evented.rs @@ -188,7 +188,28 @@ feature! { &'a E: io::Write + 'a, { use std::io::Write; - self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write(buf)) + + loop { + let evt = ready!(self.registration.poll_write_ready(cx))?; + + match self.io.as_ref().unwrap().write(buf) { + 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 < buf.len()) { + self.registration.clear_readiness(evt); + } + + return Poll::Ready(Ok(n)); + }, + Err(e) if e.kind() == io::ErrorKind::WouldBlock => { + self.registration.clear_readiness(evt); + } + Err(e) => return Poll::Ready(Err(e)), + } + } + + // self.registration.poll_write_io(cx, || self.io.as_ref().unwrap().write(buf)) } #[cfg(feature = "net")]