diff --git a/tests-build/tests/fail/macros_invalid_input.stderr b/tests-build/tests/fail/macros_invalid_input.stderr index e6e80e8114a..8bb59a8f318 100644 --- a/tests-build/tests/fail/macros_invalid_input.stderr +++ b/tests-build/tests/fail/macros_invalid_input.stderr @@ -4,7 +4,7 @@ error: the `async` keyword is missing from the function declaration 6 | fn main_is_not_async() {} | ^^ -error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`. +error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `thread_name`. --> tests/fail/macros_invalid_input.rs:8:15 | 8 | #[tokio::main(foo)] @@ -22,13 +22,13 @@ error: the `async` keyword is missing from the function declaration 15 | fn test_is_not_async() {} | ^^ -error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`. +error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `thread_name`. --> 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`, `unhandled_panic` +error: Unknown attribute foo is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `thread_name` --> tests/fail/macros_invalid_input.rs:20:15 | 20 | #[tokio::test(foo = 123)] diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index 184718784e7..67597511adc 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -56,6 +56,7 @@ struct FinalConfig { start_paused: Option, crate_name: Option, unhandled_panic: Option, + thread_name: Option, } /// Config used in case of the attribute not being able to build a valid config @@ -65,6 +66,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig { start_paused: None, crate_name: None, unhandled_panic: None, + thread_name: None, }; struct Configuration { @@ -76,6 +78,7 @@ struct Configuration { is_test: bool, crate_name: Option, unhandled_panic: Option<(UnhandledPanic, Span)>, + thread_name: Option, } impl Configuration { @@ -92,6 +95,7 @@ impl Configuration { is_test, crate_name: None, unhandled_panic: None, + thread_name: None, } } @@ -165,6 +169,16 @@ impl Configuration { Ok(()) } + fn set_thread_name(&mut self, thread_name: syn::Lit, span: Span) -> Result<(), syn::Error> { + if self.thread_name.is_some() { + return Err(syn::Error::new(span, "`thread_name` set multiple times.")); + } + + let thread_name = parse_string(thread_name, span, "thread_name")?; + self.thread_name = Some(thread_name); + Ok(()) + } + fn macro_name(&self) -> &'static str { if self.is_test { "tokio::test" @@ -229,6 +243,7 @@ impl Configuration { worker_threads, start_paused, unhandled_panic, + thread_name: self.thread_name.clone(), }) } } @@ -340,9 +355,12 @@ fn build_config( config .set_unhandled_panic(lit.clone(), syn::spanned::Spanned::span(lit))?; } + "thread_name" => { + config.set_thread_name(lit.clone(), syn::spanned::Spanned::span(lit))?; + } name => { let msg = format!( - "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`", + "Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `thread_name`", name, ); return Err(syn::Error::new_spanned(namevalue, msg)); @@ -372,7 +390,7 @@ fn build_config( format!("The `{}` attribute requires an argument.", name) } name => { - format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.", name) + format!("Unknown attribute {} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `thread_name`.", name) } }; return Err(syn::Error::new_spanned(path, msg)); @@ -428,6 +446,9 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt let unhandled_panic = v.into_tokens(&crate_path); rt = quote_spanned! {last_stmt_start_span=> #rt.unhandled_panic(#unhandled_panic) }; } + if let Some(v) = config.thread_name { + rt = quote_spanned! {last_stmt_start_span=> #rt.thread_name(#v) }; + } let generated_attrs = if is_test { quote! { diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index 29ea2867cff..e68be9db82f 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -150,6 +150,30 @@ use proc_macro::TokenStream; /// } /// ``` /// +/// ### Configure name of spawned threads +/// +/// ```rust +/// #[tokio::main(thread_name = "foo")] +/// async fn main() { +/// println!("Hello world"); +/// } +/// ``` +/// +/// Equivalent code not using `#[tokio::main]` +/// +/// ```rust +/// fn main() { +/// tokio::runtime::Builder::new_multi_thread() +/// .thread_name("foo") +/// .enable_all() +/// .build() +/// .unwrap() +/// .block_on(async { +/// println!("Hello world"); +/// }) +/// } +/// ``` +/// /// ### Configure the runtime to start with time paused /// /// ```rust