rt: skip defer queue in block_in_place context (#7216)

carllerche · web-flow · commit addbfb9204be · 2025-03-13T08:18:13.000+01:00
diff --git 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
@@ -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
@@ -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

Original file line number	Diff line number	Diff line change
`@@ -782,7 +782,13 @@ impl Context {`
`782`	`782`	`}`
`783`	`783`
`784`	`784`	`pub(crate) fn defer(&self, waker: &Waker) {`
`785`		`- self.defer.defer(waker);`
	`785`	`+ if self.core.borrow().is_none() {`
	`786`	`+ // If there is no core, then the worker is currently in a block_in_place. In this case,`
	`787`	`+ // we cannot use the defer queue as we aren't really in the current runtime.`
	`788`	`+ waker.wake_by_ref();`
	`789`	`+ } else {`
	`790`	`+ self.defer.defer(waker);`
	`791`	`+ }`
`786`	`792`	`}`
`787`	`793`
`788`	`794`	`#[allow(dead_code)]`