From addbfb9204be25a8621feb3f20b44a7c1f00edbd Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Thu, 13 Mar 2025 00:18:13 -0700
Subject: [PATCH] rt: skip defer queue in `block_in_place` context (#7216)

---
 .../runtime/scheduler/multi_thread/worker.rs  |  8 +++-
 tokio/tests/rt_handle_block_on.rs             | 11 +++++
 tokio/tests/rt_threaded.rs                    | 45 +++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs
index 8db4fb4ec..e33b9baea 100644
--- a/tokio/src/runtime/scheduler/multi_thread/worker.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs
@@ -782,7 +782,13 @@ impl Context {
     }
 
     pub(crate) fn defer(&self, waker: &Waker) {
-        self.defer.defer(waker);
+        if self.core.borrow().is_none() {
+            // If there is no core, then the worker is currently in a block_in_place. In this case,
+            // we cannot use the defer queue as we aren't really in the current runtime.
+            waker.wake_by_ref();
+        } else {
+            self.defer.defer(waker);
+        }
     }
 
     #[allow(dead_code)]
diff --git a/tokio/tests/rt_handle_block_on.rs b/tokio/tests/rt_handle_block_on.rs
index 4037fb93c..d62815df9 100644
--- a/tokio/tests/rt_handle_block_on.rs
+++ b/tokio/tests/rt_handle_block_on.rs
@@ -126,6 +126,17 @@ fn unbounded_mpsc_channel() {
     })
 }
 
+#[test]
+fn yield_in_block_in_place() {
+    test_with_runtimes(|| {
+        Handle::current().block_on(async {
+            tokio::task::block_in_place(|| {
+                Handle::current().block_on(tokio::task::yield_now());
+            });
+        });
+    })
+}
+
 #[cfg(not(target_os = "wasi"))] // Wasi doesn't support file operations or bind
 rt_test! {
     use tokio::fs;
diff --git a/tokio/tests/rt_threaded.rs b/tokio/tests/rt_threaded.rs
index f0ed8443f..ca2a952fb 100644
--- a/tokio/tests/rt_threaded.rs
+++ b/tokio/tests/rt_threaded.rs
@@ -647,6 +647,51 @@ fn test_nested_block_in_place_with_block_on_between() {
     }
 }
 
+#[test]
+fn yield_now_in_block_in_place() {
+    let rt = runtime::Builder::new_multi_thread()
+        .worker_threads(1)
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        tokio::spawn(async {
+            tokio::task::block_in_place(|| {
+                tokio::runtime::Handle::current().block_on(tokio::task::yield_now());
+            })
+        })
+        .await
+        .unwrap()
+    })
+}
+
+#[test]
+fn mutex_in_block_in_place() {
+    const BUDGET: usize = 128;
+
+    let rt = runtime::Builder::new_multi_thread()
+        .worker_threads(1)
+        .build()
+        .unwrap();
+
+    rt.block_on(async {
+        let lock = tokio::sync::Mutex::new(0);
+
+        tokio::spawn(async move {
+            tokio::task::block_in_place(|| {
+                tokio::runtime::Handle::current().block_on(async move {
+                    for i in 0..(BUDGET + 1) {
+                        let mut guard = lock.lock().await;
+                        *guard = i;
+                    }
+                });
+            })
+        })
+        .await
+        .unwrap();
+    })
+}
+
 // Testing the tuning logic is tricky as it is inherently timing based, and more
 // of a heuristic than an exact behavior. This test checks that the interval
 // changes over time based on load factors. There are no assertions, completion