From f8948ea021bd4b5626bbd8dee89bee51ba766326 Mon Sep 17 00:00:00 2001
From: quininer <quininer@live.com>
Date: Mon, 2 Dec 2024 20:52:00 +0800
Subject: [PATCH] runtime: do not defer `yield_now` inside `block_in_place`
 (#6999)

---
 tokio/src/runtime/context.rs  |  9 ++++++++-
 tokio/tests/task_yield_now.rs | 10 +++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs
index 76918114b..12f6edf13 100644
--- a/tokio/src/runtime/context.rs
+++ b/tokio/src/runtime/context.rs
@@ -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))
     }
 
diff --git a/tokio/tests/task_yield_now.rs b/tokio/tests/task_yield_now.rs
index 3cb8cb16e..3b462e922 100644
--- a/tokio/tests/task_yield_now.rs
+++ b/tokio/tests/task_yield_now.rs
@@ -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;
@@ -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();
+}