macros: allow setting unhandled_panic behavior in tokio::{main, test}

tests-build/tests/fail/macros_invalid_input.rs

@@ -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() {}

tests-build/tests/fail/macros_invalid_input.stderr

@@ -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: The `unhandled_panic` option requires the `current_thread` runtime flavor. Use `#[tokio::test(flavor = "current_thread")]`
+  --> tests/fail/macros_invalid_input.rs:44:58
+   |
+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
-  --> $DIR/macros_invalid_input.rs:45:1
+  --> tests/fail/macros_invalid_input.rs:48:1
    |
-45 | #[test]
+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)

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<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! {