diff options
author | Jon Gjengset <jon@thesquareplanet.com> | 2020-07-28 19:58:33 -0400 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-07-28 19:58:33 -0400 |
commit | 0366a3e6d1aa4e7bf4a1c717680dd0947589264b (patch) | |
tree | 664ca955e90cad981ddfe68fe9804506543cdf99 /tokio | |
parent | 03b68f4e7564a3a78fa29bacfe5e306c59145aaa (diff) |
Reset coop budget when blocking in block_on (#2711)
Previously, we would fail to reset the coop budget in this case, making
it so that `coop::poll_proceed` would perpetually yield `Poll::Pending`
in nested executers even when run in `block_in_place`.
This is also a further improvement on #2645.
Diffstat (limited to 'tokio')
-rw-r--r-- | tokio/src/runtime/thread_pool/worker.rs | 6 | ||||
-rw-r--r-- | tokio/tests/task_blocking.rs | 67 |
2 files changed, 68 insertions, 5 deletions
diff --git a/tokio/src/runtime/thread_pool/worker.rs b/tokio/src/runtime/thread_pool/worker.rs index abe20da5..c53c9384 100644 --- a/tokio/src/runtime/thread_pool/worker.rs +++ b/tokio/src/runtime/thread_pool/worker.rs @@ -199,7 +199,6 @@ cfg_blocking! { } } - let mut had_core = false; let mut had_entered = false; CURRENT.with(|maybe_cx| { @@ -248,7 +247,6 @@ cfg_blocking! { // // First, move the core back into the worker's shared core slot. cx.worker.core.set(core); - had_core = true; // Next, clone the worker handle and send it to a new thread for // processing. @@ -259,14 +257,12 @@ cfg_blocking! { runtime::spawn_blocking(move || run(worker)); }); - if had_core { + if had_entered { // 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 if had_entered { - crate::runtime::enter::exit(f) } else { f() } diff --git a/tokio/tests/task_blocking.rs b/tokio/tests/task_blocking.rs index 50c070a3..4ca1596e 100644 --- a/tokio/tests/task_blocking.rs +++ b/tokio/tests/task_blocking.rs @@ -176,3 +176,70 @@ fn can_shutdown_now_in_runtime() { rt.shutdown_background(); }); } + +#[test] +fn coop_disabled_in_block_in_place() { + let mut outer = tokio::runtime::Builder::new() + .threaded_scheduler() + .enable_time() + .build() + .unwrap(); + + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + for i in 0..200 { + tx.send(i).unwrap(); + } + drop(tx); + + outer.block_on(async move { + let jh = tokio::spawn(async move { + tokio::task::block_in_place(move || { + futures::executor::block_on(async move { + use tokio::stream::StreamExt; + assert_eq!(rx.fold(0, |n, _| n + 1).await, 200); + }) + }) + }); + + tokio::time::timeout(Duration::from_secs(1), jh) + .await + .expect("timed out (probably hanging)") + .unwrap() + }); +} + +#[test] +fn coop_disabled_in_block_in_place_in_block_on() { + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let done = done_tx.clone(); + thread::spawn(move || { + let mut outer = tokio::runtime::Builder::new() + .threaded_scheduler() + .build() + .unwrap(); + + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + for i in 0..200 { + tx.send(i).unwrap(); + } + drop(tx); + + outer.block_on(async move { + tokio::task::block_in_place(move || { + futures::executor::block_on(async move { + use tokio::stream::StreamExt; + assert_eq!(rx.fold(0, |n, _| n + 1).await, 200); + }) + }) + }); + + let _ = done.send(Ok(())); + }); + + thread::spawn(move || { + thread::sleep(Duration::from_secs(1)); + let _ = done_tx.send(Err("timed out (probably hanging)")); + }); + + done_rx.recv().unwrap().unwrap(); +} |