diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index 8858c8a1674..52209a76b22 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -25,11 +25,37 @@ impl RuntimeFlavor { } } +#[derive(Clone, Copy, PartialEq)] +enum UnhandledPanic { + Ignore, + ShutdownRuntime, +} + +impl UnhandledPanic { + fn from_str(s: &str) -> Result { + 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, start_paused: Option, crate_name: Option, + unhandled_panic: Option, } /// 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, + unhandled_panic: Option, } 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); + Ok(()) + } + fn macro_name(&self) -> &'static str { if self.is_test { "tokio::test" @@ -165,6 +213,7 @@ impl Configuration { Ok(FinalConfig { crate_name: self.crate_name.clone(), + unhandled_panic: self.unhandled_panic, flavor, worker_threads, start_paused, @@ -275,9 +324,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 +356,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 +412,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! { diff --git a/tokio/tests/macros_test.rs b/tokio/tests/macros_test.rs index 69ee30b36ce..a91d079b325 100644 --- a/tokio/tests/macros_test.rs +++ b/tokio/tests/macros_test.rs @@ -1,3 +1,4 @@ +#![allow(unexpected_cfgs)] #![cfg(all(feature = "full", not(target_os = "wasi")))] // Wasi doesn't support threading use tokio::test; @@ -92,3 +93,27 @@ pub mod issue_5243 { async fn foo() {} ); } + +#[cfg(tokio_unstable)] +pub mod macro_rt_arg_unhandled_panic { + use tokio_test::assert_err; + + #[tokio::test(unhandled_panic = "shutdown_runtime")] + #[should_panic] + async fn unhandled_panic_shutdown_runtime() { + let rt = tokio::spawn(async { + panic!("This panic should shutdown the runtime."); + }) + .await; + assert_err!(rt); + } + + #[tokio::test(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); + } +}