Skip to content

Commit

Permalink
macros: support setting name of spawned threads
Browse files Browse the repository at this point in the history
Support setting the names of spawned tokio threads.
  • Loading branch information
cgzones committed Oct 18, 2024
1 parent 32e0b43 commit 572d41f
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 5 deletions.
6 changes: 3 additions & 3 deletions tests-build/tests/fail/macros_invalid_input.stderr
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand All @@ -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)]
Expand Down
25 changes: 23 additions & 2 deletions tokio-macros/src/entry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ struct FinalConfig {
start_paused: Option<bool>,
crate_name: Option<Path>,
unhandled_panic: Option<UnhandledPanic>,
thread_name: Option<String>,
}

/// Config used in case of the attribute not being able to build a valid config
Expand All @@ -65,6 +66,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig {
start_paused: None,
crate_name: None,
unhandled_panic: None,
thread_name: None,
};

struct Configuration {
Expand All @@ -76,6 +78,7 @@ struct Configuration {
is_test: bool,
crate_name: Option<Path>,
unhandled_panic: Option<(UnhandledPanic, Span)>,
thread_name: Option<String>,
}

impl Configuration {
Expand All @@ -92,6 +95,7 @@ impl Configuration {
is_test,
crate_name: None,
unhandled_panic: None,
thread_name: None,
}
}

Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -229,6 +243,7 @@ impl Configuration {
worker_threads,
start_paused,
unhandled_panic,
thread_name: self.thread_name.clone(),
})
}
}
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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! {
Expand Down
24 changes: 24 additions & 0 deletions tokio-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 572d41f

Please sign in to comment.