From 8d517d390f6ee5fe558339c287686b831da8cc1f 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 | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tokio/src/io/poll_evented.rs b/tokio/src/io/poll_evented.rs index b8c0be2c09f..bebabc23bcd 100644 --- a/tokio/src/io/poll_evented.rs +++ b/tokio/src/io/poll_evented.rs @@ -188,7 +188,26 @@ 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 write only part of our buffer, this is sufficient on unix to show + // that the socket buffer is full + 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)), + } + } } #[cfg(feature = "net")]