From 7096a8007502526b23ee1707a6cb37c68c4f0a84 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Tue, 20 Sep 2022 22:55:51 +0200 Subject: [PATCH] task: add track_caller to `block_in_place` and `spawn_local` (#5034) Functions that may panic can be annotated with `#[track_caller]` so that in the event of a panic, the function where the user called the panicking function is shown instead of the file and line within Tokio source. This change adds `#[track_caller]` to two public APIs in tokio task module which weren't added in #4848. * `tokio::task::block_in_place` * `tokio::task::spawn_local` These APIs had call stacks that went through closures, which is a use case not supported by `#[track_caller]`. These two cases have been refactored so that the `panic!` call is no longer inside a closure. Tests have been added for these two new cases. Refs: #4413 --- .../runtime/scheduler/multi_thread/worker.rs | 20 +++++++---- tokio/src/task/blocking.rs | 1 + tokio/src/task/local.rs | 10 +++--- tokio/tests/task_panic.rs | 34 +++++++++++++++++-- 4 files changed, 51 insertions(+), 14 deletions(-) diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 7da4a3f0f56..34ef0d9f126 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -249,6 +249,7 @@ pub(super) fn create( (handle, launch) } +#[track_caller] pub(crate) fn block_in_place(f: F) -> R where F: FnOnce() -> R, @@ -275,7 +276,7 @@ where let mut had_entered = false; - CURRENT.with(|maybe_cx| { + let setup_result = CURRENT.with(|maybe_cx| { match (crate::runtime::enter::context(), maybe_cx.is_some()) { (EnterContext::Entered { .. }, true) => { // We are on a thread pool runtime thread, so we just need to @@ -288,22 +289,24 @@ where // method: if allow_blocking { had_entered = true; - return; + return Ok(()); } else { // This probably means we are on the current_thread runtime or in a // LocalSet, where it is _not_ okay to block. - panic!("can call blocking only when running on the multi-threaded runtime"); + return Err( + "can call blocking only when running on the multi-threaded runtime", + ); } } (EnterContext::NotEntered, true) => { // This is a nested call to block_in_place (we already exited). // All the necessary setup has already been done. - return; + return Ok(()); } (EnterContext::NotEntered, false) => { // We are outside of the tokio runtime, so blocking is fine. // We can also skip all of the thread pool blocking setup steps. - return; + return Ok(()); } } @@ -312,7 +315,7 @@ where // Get the worker core. If none is set, then blocking is fine! let core = match cx.core.borrow_mut().take() { Some(core) => core, - None => return, + None => return Ok(()), }; // The parker should be set here @@ -331,8 +334,13 @@ where // steal the core back. let worker = cx.worker.clone(); runtime::spawn_blocking(move || run(worker)); + Ok(()) }); + if let Err(panic_message) = setup_result { + panic!("{}", panic_message); + } + if had_entered { // Unset the current task's budget. Blocking sections are not // constrained by task budgets. diff --git a/tokio/src/task/blocking.rs b/tokio/src/task/blocking.rs index bcebbf5edc9..46756a9b302 100644 --- a/tokio/src/task/blocking.rs +++ b/tokio/src/task/blocking.rs @@ -70,6 +70,7 @@ cfg_rt_multi_thread! { /// This function panics if called from a [`current_thread`] runtime. /// /// [`current_thread`]: fn@crate::runtime::Builder::new_current_thread + #[track_caller] pub fn block_in_place(f: F) -> R where F: FnOnce() -> R, diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 4ad80649a57..0ac28e67fc2 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -314,12 +314,10 @@ cfg_rt! { where F: Future + 'static, F::Output: 'static { - CURRENT.with(|maybe_cx| { - match maybe_cx.get() { - None => panic!("`spawn_local` called from outside of a `task::LocalSet`"), - Some(cx) => cx.spawn(future, name) - } - }) + match CURRENT.with(|maybe_cx| maybe_cx.get()) { + None => panic!("`spawn_local` called from outside of a `task::LocalSet`"), + Some(cx) => cx.spawn(future, name) + } } } diff --git a/tokio/tests/task_panic.rs b/tokio/tests/task_panic.rs index 451243d2e6b..126195222e5 100644 --- a/tokio/tests/task_panic.rs +++ b/tokio/tests/task_panic.rs @@ -3,13 +3,43 @@ use futures::future; use std::error::Error; -use tokio::{runtime::Builder, spawn, task}; +use tokio::runtime::Builder; +use tokio::task::{self, block_in_place}; mod support { pub mod panic; } use support::panic::test_panic; +#[test] +fn block_in_place_panic_caller() -> Result<(), Box> { + let panic_location_file = test_panic(|| { + let rt = Builder::new_current_thread().enable_all().build().unwrap(); + rt.block_on(async { + block_in_place(|| {}); + }); + }); + + // The panic location should be in this file + assert_eq!(&panic_location_file.unwrap(), file!()); + + Ok(()) +} + +#[test] +fn local_set_spawn_local_panic_caller() -> Result<(), Box> { + let panic_location_file = test_panic(|| { + let _local = task::LocalSet::new(); + + let _ = task::spawn_local(async {}); + }); + + // The panic location should be in this file + assert_eq!(&panic_location_file.unwrap(), file!()); + + Ok(()) +} + #[test] fn local_set_block_on_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { @@ -30,7 +60,7 @@ fn local_set_block_on_panic_caller() -> Result<(), Box> { #[test] fn spawn_panic_caller() -> Result<(), Box> { let panic_location_file = test_panic(|| { - spawn(future::pending::<()>()); + tokio::spawn(future::pending::<()>()); }); // The panic location should be in this file