diff --git a/tests-build/tests/fail/macros_invalid_input.rs b/tests-build/tests/fail/macros_invalid_input.rs index 425b274ff7f..5e89a651049 100644 --- a/tests-build/tests/fail/macros_invalid_input.rs +++ b/tests-build/tests/fail/macros_invalid_input.rs @@ -68,4 +68,7 @@ async fn test_has_second_test_attr_rust_2021() {} #[tokio::test] async fn test_has_generated_second_test_attr() {} +#[tokio::test(name = 123)] +async fn test_name_not_string() {} + fn main() {} diff --git a/tests-build/tests/fail/macros_invalid_input.stderr b/tests-build/tests/fail/macros_invalid_input.stderr index cfb2f84448a..8bc26668135 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`, `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`, `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`, `name`. --> tests/fail/macros_invalid_input.rs:20:15 | 20 | #[tokio::test(foo = 123)] @@ -119,3 +119,9 @@ error: second test attribute is supplied, consider removing or changing the orde | ^^^^^^^^^^^^^^ | = note: this error originates in the attribute macro `tokio::test` (in Nightly builds, run with -Z macro-backtrace for more info) + +error: Failed to parse value of `name` as string. + --> tests/fail/macros_invalid_input.rs:71:22 + | +71 | #[tokio::test(name = 123)] + | ^^^ diff --git a/tokio-macros/Cargo.toml b/tokio-macros/Cargo.toml index 601774e26ae..50fb7b2b50f 100644 --- a/tokio-macros/Cargo.toml +++ b/tokio-macros/Cargo.toml @@ -27,7 +27,7 @@ quote = "1" syn = { version = "2.0", features = ["full"] } [dev-dependencies] -tokio = { version = "1.0.0", features = ["full", "test-util"] } +tokio = { version = "1.0.0", path = "../tokio", features = ["full", "test-util"] } [package.metadata.docs.rs] all-features = true diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs index 5f52f29468f..476006703d1 100644 --- a/tokio-macros/src/entry.rs +++ b/tokio-macros/src/entry.rs @@ -53,6 +53,7 @@ impl UnhandledPanic { } struct FinalConfig { + name: Option, flavor: RuntimeFlavor, worker_threads: Option, start_paused: Option, @@ -62,6 +63,7 @@ struct FinalConfig { /// Config used in case of the attribute not being able to build a valid config const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig { + name: None, flavor: RuntimeFlavor::CurrentThread, worker_threads: None, start_paused: None, @@ -70,6 +72,7 @@ const DEFAULT_ERROR_CONFIG: FinalConfig = FinalConfig { }; struct Configuration { + name: Option, rt_multi_thread_available: bool, default_flavor: RuntimeFlavor, flavor: Option, @@ -83,6 +86,7 @@ struct Configuration { impl Configuration { fn new(is_test: bool, rt_multi_thread: bool) -> Self { Configuration { + name: None, rt_multi_thread_available: rt_multi_thread, default_flavor: match is_test { true => RuntimeFlavor::CurrentThread, @@ -97,6 +101,16 @@ impl Configuration { } } + fn set_name(&mut self, name: syn::Lit, span: Span) -> Result<(), syn::Error> { + if self.name.is_some() { + return Err(syn::Error::new(span, "`name` set multiple times.")); + } + + let runtime_name = parse_string(name, span, "name")?; + self.name = Some(runtime_name); + Ok(()) + } + fn set_flavor(&mut self, runtime: syn::Lit, span: Span) -> Result<(), syn::Error> { if self.flavor.is_some() { return Err(syn::Error::new(span, "`flavor` set multiple times.")); @@ -227,6 +241,7 @@ impl Configuration { }; Ok(FinalConfig { + name: self.name.clone(), crate_name: self.crate_name.clone(), flavor, worker_threads, @@ -372,9 +387,12 @@ fn build_config( config .set_unhandled_panic(lit.clone(), syn::spanned::Spanned::span(lit))?; } + "name" => { + config.set_name(lit.clone(), syn::spanned::Spanned::span(lit))?; + } name => { let msg = format!( - "Unknown attribute {name} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`", + "Unknown attribute {name} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `name`.", ); return Err(syn::Error::new_spanned(namevalue, msg)); } @@ -397,11 +415,12 @@ fn build_config( "Set the runtime flavor with #[{macro_name}(flavor = \"current_thread\")]." ) } - "flavor" | "worker_threads" | "start_paused" | "crate" | "unhandled_panic" => { + "flavor" | "worker_threads" | "start_paused" | "crate" | "unhandled_panic" + | "name" => { format!("The `{name}` attribute requires an argument.") } name => { - format!("Unknown attribute {name} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`.") + format!("Unknown attribute {name} is specified; expected one of: `flavor`, `worker_threads`, `start_paused`, `crate`, `unhandled_panic`, `name`.") } }; return Err(syn::Error::new_spanned(path, msg)); @@ -478,6 +497,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.name { + rt = quote_spanned! {last_stmt_start_span=> #rt.name(#v) }; + } let generated_attrs = if is_test { quote! { diff --git a/tokio-macros/src/lib.rs b/tokio-macros/src/lib.rs index cd4cba0a24b..9f999df6f9e 100644 --- a/tokio-macros/src/lib.rs +++ b/tokio-macros/src/lib.rs @@ -91,6 +91,30 @@ use proc_macro::TokenStream; /// /// # Usage /// +/// ## Set the name of the runtime +/// +/// ```rust +/// #[tokio::main(name = "my-runtime")] +/// async fn main() { +/// println!("Hello world"); +/// } +/// ``` +/// +/// Equivalent code not using `#[tokio::main]` +/// +/// ```rust +/// fn main() { +/// tokio::runtime::Builder::new_multi_thread() +/// .enable_all() +/// .name("my-runtime") +/// .build() +/// .unwrap() +/// .block_on(async { +/// println!("Hello world"); +/// }) +/// } +/// ``` +/// /// ## Using the multi-threaded runtime /// /// ```rust @@ -408,6 +432,31 @@ pub fn main_rt(args: TokenStream, item: TokenStream) -> TokenStream { /// /// ## Usage /// +/// ### Set the name of the runtime +/// +/// ```no_run +/// #[tokio::test(name = "my-test-runtime")] +/// async fn my_test() { +/// assert!(true); +/// } +/// ``` +/// +/// Equivalent code not using `#[tokio::test]` +/// +/// ```no_run +/// #[test] +/// fn my_test() { +/// tokio::runtime::Builder::new_current_thread() +/// .enable_all() +/// .name("my-test-runtime") +/// .build() +/// .unwrap() +/// .block_on(async { +/// assert!(true); +/// }) +/// } +/// ``` +/// /// ### Using the multi-thread runtime /// /// ```no_run diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index aea4de7503e..283e48e56f8 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -54,6 +54,9 @@ pub struct Builder { /// Runtime type kind: Kind, + /// Name of the runtime. + name: Option, + /// Whether or not to enable the I/O driver enable_io: bool, nevents: usize, @@ -271,6 +274,9 @@ impl Builder { Builder { kind, + // Default runtime name + name: None, + // I/O defaults to "off" enable_io: false, nevents: 1024, @@ -538,6 +544,28 @@ impl Builder { self } + /// Sets the name of the runtime. + /// + /// # Examples + /// + /// ``` + /// # #[cfg(not(target_family = "wasm"))] + /// # { + /// # use tokio::runtime; + /// + /// # pub fn main() { + /// let rt = runtime::Builder::new_multi_thread() + /// .name("my-runtime") + /// .build(); + /// # } + /// # } + /// ``` + pub fn name(&mut self, val: impl Into) -> &mut Self { + let val = val.into(); + self.name = Some(val); + self + } + /// Sets a function used to generate the name of threads spawned by the `Runtime`'s thread pool. /// /// The default name fn is `|| "tokio-rt-worker".into()`. @@ -1633,6 +1661,7 @@ impl Builder { metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), }, local_tid, + self.name.clone(), ); let handle = Handle { @@ -1814,6 +1843,7 @@ cfg_rt_multi_thread! { metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), }, self.timer_flavor, + self.name.clone(), ); let handle = Handle { inner: scheduler::Handle::MultiThread(handle) }; diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 230e6a0c38e..79fea39da15 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -474,6 +474,27 @@ impl Handle { runtime::Id::new(owned_id) } + /// Returns the name of the current `Runtime`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main(flavor = "current_thread", name = "my-runtime")] + /// async fn main() { + /// println!("Current runtime name: {}", Handle::current().name().unwrap()); + /// } + /// ``` + /// + pub fn name(&self) -> Option<&str> { + match &self.inner { + scheduler::Handle::CurrentThread(handle) => handle.name(), + #[cfg(feature = "rt-multi-thread")] + scheduler::Handle::MultiThread(handle) => handle.name(), + } + } + /// Returns a view that lets you get information about how the runtime /// is performing. pub fn metrics(&self) -> RuntimeMetrics { diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index 68ab17f1402..41b693e35fa 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -34,6 +34,9 @@ pub(crate) struct CurrentThread { /// Handle to the current thread scheduler pub(crate) struct Handle { + /// The name of the runtime + pub(crate) name: Option, + /// Scheduler state shared across threads shared: Shared, @@ -132,6 +135,7 @@ impl CurrentThread { seed_generator: RngSeedGenerator, config: Config, local_tid: Option, + name: Option, ) -> (CurrentThread, Arc) { let worker_metrics = WorkerMetrics::from_config(&config); worker_metrics.set_thread_id(thread::current().id()); @@ -142,6 +146,7 @@ impl CurrentThread { .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL); let handle = Arc::new(Handle { + name, task_hooks: TaskHooks { task_spawn_callback: config.before_spawn.clone(), task_terminate_callback: config.after_termination.clone(), @@ -639,6 +644,10 @@ impl Handle { pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } + + pub(crate) fn name(&self) -> Option<&str> { + self.name.as_deref() + } } impl fmt::Debug for Handle { diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 3d4226c4017..d66f91ed65f 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -23,6 +23,9 @@ use crate::loom::sync::atomic::{AtomicBool, Ordering::SeqCst}; /// Handle to the multi thread scheduler pub(crate) struct Handle { + /// The name of the runtime + pub(crate) name: Option, + /// Task spawner pub(super) shared: worker::Shared, @@ -123,6 +126,10 @@ impl Handle { pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } + + pub(crate) fn name(&self) -> Option<&str> { + self.name.as_deref() + } } impl fmt::Debug for Handle { diff --git a/tokio/src/runtime/scheduler/multi_thread/mod.rs b/tokio/src/runtime/scheduler/multi_thread/mod.rs index 1c5e1a88884..c04a3ef55a6 100644 --- a/tokio/src/runtime/scheduler/multi_thread/mod.rs +++ b/tokio/src/runtime/scheduler/multi_thread/mod.rs @@ -54,6 +54,7 @@ pub(crate) struct MultiThread; // ===== impl MultiThread ===== impl MultiThread { + #[allow(clippy::too_many_arguments)] pub(crate) fn new( size: usize, driver: Driver, @@ -62,6 +63,7 @@ impl MultiThread { seed_generator: RngSeedGenerator, config: Config, timer_flavor: TimerFlavor, + name: Option, ) -> (MultiThread, Arc, Launch) { let parker = Parker::new(driver); let (handle, launch) = worker::create( @@ -72,6 +74,7 @@ impl MultiThread { seed_generator, config, timer_flavor, + name, ); (MultiThread, handle, launch) diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index f48e6ba5271..4d530fa2eae 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -254,6 +254,7 @@ type Notified = task::Notified>; /// improvements. const MAX_LIFO_POLLS_PER_TICK: usize = 3; +#[allow(clippy::too_many_arguments)] pub(super) fn create( size: usize, park: Parker, @@ -262,6 +263,7 @@ pub(super) fn create( seed_generator: RngSeedGenerator, config: Config, timer_flavor: TimerFlavor, + name: Option, ) -> (Arc, Launch) { let mut cores = Vec::with_capacity(size); let mut remotes = Vec::with_capacity(size); @@ -301,6 +303,7 @@ pub(super) fn create( let remotes_len = remotes.len(); let handle = Arc::new(Handle { + name, task_hooks: TaskHooks::from_config(&config), shared: Shared { remotes: remotes.into_boxed_slice(), diff --git a/tokio/tests/rt_threaded.rs b/tokio/tests/rt_threaded.rs index 432add50f38..f1730e881a6 100644 --- a/tokio/tests/rt_threaded.rs +++ b/tokio/tests/rt_threaded.rs @@ -860,6 +860,20 @@ fn test_tuning() { flag.store(false, Relaxed); } +#[test] +fn different_runtime_names() { + let rt1 = runtime::Builder::new_multi_thread() + .name("test-runtime-1") + .build() + .unwrap(); + let rt2 = runtime::Builder::new_multi_thread() + .name("test-runtime-2") + .build() + .unwrap(); + + assert_ne!(rt1.handle().name().unwrap(), rt2.handle().name().unwrap()); +} + fn rt() -> runtime::Runtime { runtime::Runtime::new().unwrap() }