itsscb/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2025-09-25 12:00:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

macros: allow unhandled_panic behavior for #[tokio::main] and #[tokio::test] (#6593)

Browse Source
This commit is contained in:
Hai-Hsin 2024-06-07 19:48:56 +08:00 committed by GitHub
parent 126ce89bb4
commit 833ee027d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 227 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tests-build/tests/fail/macros_invalid_input.rs
View File

@ -41,6 +41,9 @@ async fn test_crate_not_path_int() {}
#[tokio::test(crate = "456")]
async fn test_crate_not_path_invalid() {}
#[tokio::test(flavor = "multi_thread", unhandled_panic = "shutdown_runtime")]
async fn test_multi_thread_with_unhandled_panic() {}
#[tokio::test]
#[test]
async fn test_has_second_test_attr() {}

64
tests-build/tests/fail/macros_invalid_input.stderr
View File

@ -1,115 +1,121 @@
error: the `async` keyword is missing from the function declaration
--> $DIR/macros_invalid_input.rs:6:1
--> tests/fail/macros_invalid_input.rs:6:1
|
6 | fn main_is_not_async() {}
| ^^
error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
--> $DIR/macros_invalid_input.rs:8:15
error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.
--> tests/fail/macros_invalid_input.rs:8:15
|
8 | #[tokio::main(foo)]
| ^^^
error: Must have specified ident
--> $DIR/macros_invalid_input.rs:11:15
--> tests/fail/macros_invalid_input.rs:11:15
|
11 | #[tokio::main(threadpool::bar)]
| ^^^^^^^^^^^^^^^
error: the `async` keyword is missing from the function declaration
--> $DIR/macros_invalid_input.rs:15:1
--> tests/fail/macros_invalid_input.rs:15:1
|
15 | fn test_is_not_async() {}
| ^^
error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
--> $DIR/macros_invalid_input.rs:17:15
error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.
--> tests/fail/macros_invalid_input.rs:17:15
|
17 | #[tokio::test(foo)]
| ^^^
error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`
--> $DIR/macros_invalid_input.rs:20:15
error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`
--> tests/fail/macros_invalid_input.rs:20:15
|
20 | #[tokio::test(foo = 123)]
| ^^^^^^^^^
error: Failed to parse value of `flavor` as string.
--> $DIR/macros_invalid_input.rs:23:24
--> tests/fail/macros_invalid_input.rs:23:24
|
23 | #[tokio::test(flavor = 123)]
| ^^^
error: No such runtime flavor `foo`. The runtime flavors are `current_thread` and `multi_thread`.
--> $DIR/macros_invalid_input.rs:26:24
--> tests/fail/macros_invalid_input.rs:26:24
|
26 | #[tokio::test(flavor = "foo")]
| ^^^^^
error: The `start_paused` option requires the `current_thread` runtime flavor. Use `#[tokio::test(flavor = "current_thread")]`
--> $DIR/macros_invalid_input.rs:29:55
--> tests/fail/macros_invalid_input.rs:29:55
|
29 | #[tokio::test(flavor = "multi_thread", start_paused = false)]
| ^^^^^
error: Failed to parse value of `worker_threads` as integer.
--> $DIR/macros_invalid_input.rs:32:57
--> tests/fail/macros_invalid_input.rs:32:57
|
32 | #[tokio::test(flavor = "multi_thread", worker_threads = "foo")]
| ^^^^^
error: The `worker_threads` option requires the `multi_thread` runtime flavor. Use `#[tokio::test(flavor = "multi_thread")]`
--> $DIR/macros_invalid_input.rs:35:59
--> tests/fail/macros_invalid_input.rs:35:59
|
35 | #[tokio::test(flavor = "current_thread", worker_threads = 4)]
| ^
error: Failed to parse value of `crate` as path.
--> $DIR/macros_invalid_input.rs:38:23
--> tests/fail/macros_invalid_input.rs:38:23
|
38 | #[tokio::test(crate = 456)]
| ^^^
error: Failed to parse value of `crate` as path: "456"
--> $DIR/macros_invalid_input.rs:41:23
--> tests/fail/macros_invalid_input.rs:41:23
|
41 | #[tokio::test(crate = "456")]
| ^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> $DIR/macros_invalid_input.rs:45:1
error: The `unhandled_panic` option requires the `current_thread` runtime flavor. Use `#[tokio::test(flavor = "current_thread")]`
--> tests/fail/macros_invalid_input.rs:44:58
|
45 | #[test]
44 | #[tokio::test(flavor = "multi_thread", unhandled_panic = "shutdown_runtime")]
| ^^^^^^^^^^^^^^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> tests/fail/macros_invalid_input.rs:48:1
|
48 | #[test]
| ^^^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> $DIR/macros_invalid_input.rs:49:1
--> tests/fail/macros_invalid_input.rs:52:1
|
49 | #[::core::prelude::v1::test]
52 | #[::core::prelude::v1::test]
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> $DIR/macros_invalid_input.rs:53:1
--> tests/fail/macros_invalid_input.rs:56:1
|
53 | #[core::prelude::rust_2015::test]
56 | #[core::prelude::rust_2015::test]
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> $DIR/macros_invalid_input.rs:57:1
--> tests/fail/macros_invalid_input.rs:60:1
|
57 | #[::std::prelude::rust_2018::test]
60 | #[::std::prelude::rust_2018::test]
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> $DIR/macros_invalid_input.rs:61:1
--> tests/fail/macros_invalid_input.rs:64:1
|
61 | #[std::prelude::rust_2021::test]
64 | #[std::prelude::rust_2021::test]
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
error: second test attribute is supplied, consider removing or changing the order of your test attributes
--> $DIR/macros_invalid_input.rs:64:1
--> tests/fail/macros_invalid_input.rs:67:1
|
64 | #[tokio::test]
67 | #[tokio::test]
| ^^^^^^^^^^^^^^
|
= note: this error originates in the attribute macro `tokio::test` (in Nightly builds, run with -Z macro-backtrace for more info)

75
tokio-macros/src/entry.rs
View File

@ -25,11 +25,37 @@ impl RuntimeFlavor {
}
}
#[derive(Clone, Copy, PartialEq)]
enum UnhandledPanic {
Ignore,
ShutdownRuntime,
}
impl UnhandledPanic {
fn from_str(s: &str) -> Result<UnhandledPanic, String> {
match s {
"ignore" => Ok(UnhandledPanic::Ignore),
"shutdown_runtime" => Ok(UnhandledPanic::ShutdownRuntime),
_ => Err(format!("No such unhandled panic behavior `{}`. The unhandled panic behaviors are `ignore` and `shutdown_runtime`.", s)),
}
}
fn into_tokens(self, crate_path: &TokenStream) -> TokenStream {
match self {
UnhandledPanic::Ignore => quote! { #crate_path::runtime::UnhandledPanic::Ignore },
UnhandledPanic::ShutdownRuntime => {
quote! { #crate_path::runtime::UnhandledPanic::ShutdownRuntime }
}
}
}
}
struct FinalConfig {
flavor: RuntimeFlavor,
worker_threads: Option<usize>,
start_paused: Option<bool>,
crate_name: Option<Path>,
unhandled_panic: Option<UnhandledPanic>,
}
/// Config used in case of the attribute not being able to build a valid config
@ -38,6 +64,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
worker_threads: None,
start_paused: None,
crate_name: None,
unhandled_panic: None,
};
struct Configuration {
@ -48,6 +75,7 @@ struct Configuration {
start_paused: Option<(bool, Span)>,
is_test: bool,
crate_name: Option<Path>,
unhandled_panic: Option<(UnhandledPanic, Span)>,
}
impl Configuration {
@ -63,6 +91,7 @@ impl Configuration {
start_paused: None,
is_test,
crate_name: None,
unhandled_panic: None,
}
}
@ -117,6 +146,25 @@ impl Configuration {
Ok(())
}
fn set_unhandled_panic(
&mut self,
unhandled_panic: syn::Lit,
span: Span,
) -> Result<(), syn::Error> {
if self.unhandled_panic.is_some() {
return Err(syn::Error::new(
span,
"`unhandled_panic` set multiple times.",
));
}
let unhandled_panic = parse_string(unhandled_panic, span, "unhandled_panic")?;
let unhandled_panic =
UnhandledPanic::from_str(&unhandled_panic).map_err(|err| syn::Error::new(span, err))?;
self.unhandled_panic = Some((unhandled_panic, span));
Ok(())
}
fn macro_name(&self) -> &'static str {
if self.is_test {
"tokio::test"
@ -163,11 +211,24 @@ impl Configuration {
(_, None) => None,
};
let unhandled_panic = match (flavor, self.unhandled_panic) {
(F::Threaded, Some((_, unhandled_panic_span))) => {
let msg = format!(
"The `unhandled_panic` option requires the `current_thread` runtime flavor. Use `#[{}(flavor = \"current_thread\")]`",
self.macro_name(),
);
return Err(syn::Error::new(unhandled_panic_span, msg));
}
(F::CurrentThread, Some((unhandled_panic, _))) => Some(unhandled_panic),
(_, None) => None,
};
Ok(FinalConfig {
crate_name: self.crate_name.clone(),
flavor,
worker_threads,
start_paused,
unhandled_panic,
})
}
}
@ -275,9 +336,13 @@ fn build_config(
"crate" => {
config.set_crate_name(lit.clone(), syn::spanned::Spanned::span(lit))?;
}
"unhandled_panic" => {
config
.set_unhandled_panic(lit.clone(), syn::spanned::Spanned::span(lit))?;
}
name => {
let msg = format!(
"Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`",
"Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`",
name,
);
return Err(syn::Error::new_spanned(namevalue, msg));
@ -303,11 +368,11 @@ fn build_config(
macro_name
)
}
"flavor" | "worker_threads" | "start_paused" => {
"flavor" | "worker_threads" | "start_paused" | "crate" | "unhandled_panic" => {
format!("The `{}` attribute requires an argument.", name)
}
name => {
format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`", name)
format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.", name)
}
};
return Err(syn::Error::new_spanned(path, msg));
@ -359,6 +424,10 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
if let Some(v) = config.start_paused {
rt = quote_spanned! {last_stmt_start_span=> #rt.start_paused(#v) };
}
if let Some(v) = config.unhandled_panic {
let unhandled_panic = v.into_tokens(&crate_path);
rt = quote_spanned! {last_stmt_start_span=> #rt.unhandled_panic(#unhandled_panic) };
}
let generated_attrs = if is_test {
quote! {

93
tokio-macros/src/lib.rs
View File

@ -202,6 +202,52 @@ use proc_macro::TokenStream;
/// })
/// }
/// ```
///
/// ### Configure unhandled panic behavior
///
/// Available options are `shutdown_runtime` and `ignore`. For more details, see
/// [`Builder::unhandled_panic`].
///
/// This option is only compatible with the `current_thread` runtime.
///
/// ```no_run
/// #[cfg(tokio_unstable)]
/// #[tokio::main(flavor = "current_thread", unhandled_panic = "shutdown_runtime")]
/// async fn main() {
/// let _ = tokio::spawn(async {
/// panic!("This panic will shutdown the runtime.");
/// }).await;
/// }
/// # #[cfg(not(tokio_unstable))]
/// # fn main() { }
/// ```
///
/// Equivalent code not using `#[tokio::main]`
///
/// ```no_run
/// #[cfg(tokio_unstable)]
/// fn main() {
/// tokio::runtime::Builder::new_current_thread()
/// .enable_all()
/// .unhandled_panic(UnhandledPanic::ShutdownRuntime)
/// .build()
/// .unwrap()
/// .block_on(async {
/// let _ = tokio::spawn(async {
/// panic!("This panic will shutdown the runtime.");
/// }).await;
/// })
/// }
/// # #[cfg(not(tokio_unstable))]
/// # fn main() { }
/// ```
///
/// **Note**: This option depends on Tokio's [unstable API][unstable]. See [the
/// documentation on unstable features][unstable] for details on how to enable
/// Tokio's unstable features.
///
/// [`Builder::unhandled_panic`]: ../tokio/runtime/struct.Builder.html#method.unhandled_panic
/// [unstable]: ../tokio/index.html#unstable-features
#[proc_macro_attribute]
pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
entry::main(args.into(), item.into(), true).into()
@ -423,6 +469,53 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream {
/// println!("Hello world");
/// }
/// ```
///
/// ### Configure unhandled panic behavior
///
/// Available options are `shutdown_runtime` and `ignore`. For more details, see
/// [`Builder::unhandled_panic`].
///
/// This option is only compatible with the `current_thread` runtime.
///
/// ```no_run
/// #[cfg(tokio_unstable)]
/// #[tokio::test(flavor = "current_thread", unhandled_panic = "shutdown_runtime")]
/// async fn my_test() {
/// let _ = tokio::spawn(async {
/// panic!("This panic will shutdown the runtime.");
/// }).await;
/// }
/// # #[cfg(not(tokio_unstable))]
/// # fn main() { }
/// ```
///
/// Equivalent code not using `#[tokio::test]`
///
/// ```no_run
/// #[cfg(tokio_unstable)]
/// #[test]
/// fn my_test() {
/// tokio::runtime::Builder::new_current_thread()
/// .enable_all()
/// .unhandled_panic(UnhandledPanic::ShutdownRuntime)
/// .build()
/// .unwrap()
/// .block_on(async {
/// let _ = tokio::spawn(async {
/// panic!("This panic will shutdown the runtime.");
/// }).await;
/// })
/// }
/// # #[cfg(not(tokio_unstable))]
/// # fn main() { }
/// ```
///
/// **Note**: This option depends on Tokio's [unstable API][unstable]. See [the
/// documentation on unstable features][unstable] for details on how to enable
/// Tokio's unstable features.
///
/// [`Builder::unhandled_panic`]: ../tokio/runtime/struct.Builder.html#method.unhandled_panic
/// [unstable]: ../tokio/index.html#unstable-features
#[proc_macro_attribute]
pub fn test(args: TokenStream, item: TokenStream) -> TokenStream {
entry::test(args.into(), item.into(), true).into()

24
tokio/tests/macros_test.rs
View File

@ -1,3 +1,4 @@
#![allow(unknown_lints, unexpected_cfgs)]
#![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threading
use tokio::test;
@ -92,3 +93,26 @@ pub mod issue_5243 {
async fn foo() {}
);
}
#[cfg(tokio_unstable)]
pub mod macro_rt_arg_unhandled_panic {
use tokio_test::assert_err;
#[tokio::test(flavor = "current_thread", unhandled_panic = "shutdown_runtime")]
#[should_panic]
async fn unhandled_panic_shutdown_runtime() {
let _ = tokio::spawn(async {
panic!("This panic should shutdown the runtime.");
})
.await;
}
#[tokio::test(flavor = "current_thread", unhandled_panic = "ignore")]
async fn unhandled_panic_ignore() {
let rt = tokio::spawn(async {
panic!("This panic should be forwarded to rt as an error.");
})
.await;
assert_err!(rt);
}
}