Skip to content

Commit

Permalink
runtime: do not defer yield_now inside block_in_place (#6999)
Browse files Browse the repository at this point in the history
  • Loading branch information
quininer authored Dec 2, 2024
1 parent bce9780 commit f8948ea
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 2 deletions.
9 changes: 8 additions & 1 deletion tokio/src/runtime/context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,14 @@ cfg_rt! {
#[track_caller]
pub(super) fn with_scheduler<R>(f: impl FnOnce(Option<&scheduler::Context>) -> R) -> R {
let mut f = Some(f);
CONTEXT.try_with(|c| c.scheduler.with(f.take().unwrap()))
CONTEXT.try_with(|c| {
let f = f.take().unwrap();
if matches!(c.runtime.get(), EnterRuntime::Entered { .. }) {
c.scheduler.with(f)
} else {
f(None)
}
})
.unwrap_or_else(|_| (f.take().unwrap())(None))
}

Expand Down
10 changes: 9 additions & 1 deletion tokio/tests/task_yield_now.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#![allow(unknown_lints, unexpected_cfgs)]
#![cfg(all(feature = "full", tokio_unstable))]
#![cfg(all(feature = "full", not(target_os = "wasi"), tokio_unstable))]

use tokio::task;
use tokio_test::task::spawn;
Expand All @@ -15,3 +15,11 @@ fn yield_now_outside_of_runtime() {
assert!(task.is_woken());
assert!(task.poll().is_ready());
}

#[tokio::test(flavor = "multi_thread")]
async fn yield_now_external_executor_and_block_in_place() {
let j = tokio::spawn(async {
task::block_in_place(|| futures::executor::block_on(task::yield_now()));
});
j.await.unwrap();
}

0 comments on commit f8948ea

Please sign in to comment.