From bff21aba6c1568f260fdf48d2eb3c0566e293f5a Mon Sep 17 00:00:00 2001 From: Carl Lerche Date: Thu, 7 May 2020 16:25:04 -0700 Subject: [PATCH] rt: set task budget after block_in_place call (#2502) In some cases, when a call to `block_in_place` completes, the runtime is reinstated on the thread. In this case, the task budget must also be set in order to avoid starving other tasks on the worker. --- tokio/src/coop.rs | 28 +++++++++--- tokio/src/runtime/thread_pool/worker.rs | 34 +++++++------- tokio/tests/rt_threaded.rs | 59 +++++++++++++++++++++++++ 3 files changed, 95 insertions(+), 26 deletions(-) diff --git a/tokio/src/coop.rs b/tokio/src/coop.rs index 4d57161fa..2aae093af 100644 --- a/tokio/src/coop.rs +++ b/tokio/src/coop.rs @@ -92,10 +92,20 @@ cfg_rt_threaded! { /// Run the given closure with a cooperative task budget. When the function /// returns, the budget is reset to the value prior to calling the function. #[inline(always)] -pub(crate) fn budget(f: F) -> R -where - F: FnOnce() -> R, -{ +pub(crate) fn budget(f: impl FnOnce() -> R) -> R { + with_budget(Budget::initial(), f) +} + +cfg_rt_threaded! { + /// Set the current task's budget + #[cfg(feature = "blocking")] + pub(crate) fn set(budget: Budget) { + CURRENT.with(|cell| cell.set(budget)) + } +} + +#[inline(always)] +fn with_budget(budget: Budget, f: impl FnOnce() -> R) -> R { struct ResetGuard<'a> { cell: &'a Cell, prev: Budget, @@ -110,7 +120,7 @@ where CURRENT.with(move |cell| { let prev = cell.get(); - cell.set(Budget::initial()); + cell.set(budget); let _guard = ResetGuard { cell, prev }; @@ -127,10 +137,14 @@ cfg_rt_threaded! { cfg_blocking_impl! { /// Forcibly remove the budgeting constraints early. - pub(crate) fn stop() { + /// + /// Returns the remaining budget + pub(crate) fn stop() -> Budget { CURRENT.with(|cell| { + let prev = cell.get(); cell.set(Budget::unconstrained()); - }); + prev + }) } } diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index e31f237cc..40cd69e7f 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -4,6 +4,7 @@ //! "core" is handed off to a new thread allowing the scheduler to continue to //! make progress while the originating thread blocks. +use crate::coop; use crate::loom::rand::seed; use crate::loom::sync::{Arc, Mutex}; use crate::park::{Park, Unpark}; @@ -179,31 +180,27 @@ cfg_blocking! { F: FnOnce() -> R, { // Try to steal the worker core back - struct Reset(bool); + struct Reset(coop::Budget); impl Drop for Reset { fn drop(&mut self) { CURRENT.with(|maybe_cx| { - if !self.0 { - // We were not the ones to give away the core, - // so we do not get to restore it either. - // This is necessary so that with a nested - // block_in_place, the inner block_in_place - // does not restore the core. - return; - } - if let Some(cx) = maybe_cx { let core = cx.worker.core.take(); let mut cx_core = cx.core.borrow_mut(); assert!(cx_core.is_none()); *cx_core = core; + + // Reset the task budget as we are re-entering the + // runtime. + coop::set(self.0); } }); } } let mut had_core = false; + CURRENT.with(|maybe_cx| { match (crate::runtime::enter::context(), maybe_cx.is_some()) { (EnterContext::Entered { .. }, true) => { @@ -231,16 +228,12 @@ cfg_blocking! { return; } } + let cx = maybe_cx.expect("no .is_some() == false cases above should lead here"); // Get the worker core. If none is set, then blocking is fine! let core = match cx.core.borrow_mut().take() { - Some(core) => { - // We are effectively leaving the executor, so we need to - // forcibly end budgeting. - crate::coop::stop(); - core - }, + Some(core) => core, None => return, }; @@ -263,9 +256,12 @@ cfg_blocking! { runtime::spawn_blocking(move || run(worker)); }); - let _reset = Reset(had_core); if had_core { + // Unset the current task's budget. Blocking sections are not + // constrained by task budgets. + let _reset = Reset(coop::stop()); + crate::runtime::enter::exit(f) } else { f() @@ -349,7 +345,7 @@ impl Context { *self.core.borrow_mut() = Some(core); // Run the task - crate::coop::budget(|| { + coop::budget(|| { task.run(); // As long as there is budget remaining and a task exists in the @@ -368,7 +364,7 @@ impl Context { None => return Ok(core), }; - if crate::coop::has_budget_remaining() { + if coop::has_budget_remaining() { // Run the LIFO task, then loop *self.core.borrow_mut() = Some(core); task.run(); diff --git a/tokio/tests/rt_threaded.rs b/tokio/tests/rt_threaded.rs index ad063348f..b5ec96dec 100644 --- a/tokio/tests/rt_threaded.rs +++ b/tokio/tests/rt_threaded.rs @@ -7,6 +7,7 @@ use tokio::runtime::{self, Runtime}; use tokio::sync::oneshot; use tokio_test::{assert_err, assert_ok}; +use futures::future::poll_fn; use std::future::Future; use std::pin::Pin; use std::sync::atomic::AtomicUsize; @@ -322,6 +323,64 @@ fn multi_threadpool() { done_rx.recv().unwrap(); } +// When `block_in_place` returns, it attempts to reclaim the yielded runtime +// worker. In this case, the remainder of the task is on the runtime worker and +// must take part in the cooperative task budgeting system. +// +// The test ensures that, when this happens, attempting to consume from a +// channel yields occasionally even if there are values ready to receive. +#[test] +fn coop_and_block_in_place() { + use tokio::sync::mpsc; + + let mut rt = tokio::runtime::Builder::new() + .threaded_scheduler() + // Setting max threads to 1 prevents another thread from claiming the + // runtime worker yielded as part of `block_in_place` and guarantees the + // same thread will reclaim the worker at the end of the + // `block_in_place` call. + .max_threads(1) + .build() + .unwrap(); + + rt.block_on(async move { + let (mut tx, mut rx) = mpsc::channel(1024); + + // Fill the channel + for _ in 0..1024 { + tx.send(()).await.unwrap(); + } + + drop(tx); + + tokio::spawn(async move { + // Block in place without doing anything + tokio::task::block_in_place(|| {}); + + // Receive all the values, this should trigger a `Pending` as the + // coop limit will be reached. + poll_fn(|cx| { + while let Poll::Ready(v) = { + tokio::pin! { + let fut = rx.recv(); + } + + Pin::new(&mut fut).poll(cx) + } { + if v.is_none() { + panic!("did not yield"); + } + } + + Poll::Ready(()) + }) + .await + }) + .await + .unwrap(); + }); +} + // Testing this does not panic #[test] fn max_threads() {