Skip to content

Commit

Permalink
macros: allow unhandled_panic behavior for #[tokio::main] and `#[…
Browse files Browse the repository at this point in the history 
…tokio::test]` (#6593)
  • Loading branch information
@name1e5s
name1e5s committed Jun 7, 2024
1 parent 126ce89 commit 833ee02
Show file tree
Hide file tree
Showing 5 changed files with 226 additions and 31 deletions.
3 changes: 3 additions & 0 deletions tests-build/tests/fail/macros_invalid_input.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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() {}
Expand Down
62 changes: 34 additions & 28 deletions tests-build/tests/fail/macros_invalid_input.stderr
View file
Original file line number Diff line number Diff line change
@@ -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)
75 changes: 72 additions & 3 deletions tokio-macros/src/entry.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -38,6 +64,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
worker_threads: None,
start_paused: None,
crate_name: None,
unhandled_panic: None,
};

struct Configuration {
Expand All @@ -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 {
Expand All @@ -63,6 +91,7 @@ impl Configuration {
start_paused: None,
is_test,
crate_name: None,
unhandled_panic: None,
}
}

Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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,
})
}
}
Expand Down Expand Up @@ -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));
Expand All @@ -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));
Expand Down Expand Up @@ -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! {
Expand Down
Loading

0 comments on commit 833ee02

Please sign in to comment.