From 10b7ff4738ff62eb291567bff30cd32ee7226406 Mon Sep 17 00:00:00 2001 From: noah Date: Wed, 22 Jun 2022 00:20:45 -0500 Subject: [PATCH] mess with budget --- tokio/src/io/driver/registration.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tokio/src/io/driver/registration.rs b/tokio/src/io/driver/registration.rs index c9393650c20..b4c6fe18232 100644 --- a/tokio/src/io/driver/registration.rs +++ b/tokio/src/io/driver/registration.rs @@ -224,6 +224,9 @@ cfg_io_readiness! { pin!(fut); crate::future::poll_fn(|cx| { + // Keep track of task budget + let coop = ready!(crate::coop::poll_proceed(cx)); + if self.handle.inner.is_shutdown() { return Poll::Ready(Err(io::Error::new( io::ErrorKind::Other, @@ -231,7 +234,11 @@ cfg_io_readiness! { ))); } - Pin::new(&mut fut).poll(cx).map(Ok) + let x = Pin::new(&mut fut).poll(cx).map(Ok); + + coop.made_progress(); + + x }).await }