rt: add runtime name (#7924)

Author: mattiapitossi

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

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)]
+   |                      ^^^

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

tokio-macros/src/entry.rs

@@ -53,6 +53,7 @@ impl UnhandledPanic {
 }
 
 struct FinalConfig {
+    name: Option<String>,
     flavor: RuntimeFlavor,
     worker_threads: Option<usize>,
     start_paused: Option<bool>,
@@ -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<String>,
     rt_multi_thread_available: bool,
     default_flavor: RuntimeFlavor,
     flavor: Option<RuntimeFlavor>,
@@ -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! {

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

tokio/Cargo.toml

@@ -90,7 +90,7 @@ io-uring = ["dep:io-uring", "libc", "mio/os-poll", "mio/os-ext", "dep:slab"]
 taskdump = ["dep:backtrace"]
 
 [dependencies]
-tokio-macros = { version = "~2.6.0", optional = true }
+tokio-macros = { version = "~2.6.0", path = "../tokio-macros", optional = true }
 
 pin-project-lite = "0.2.11"
 

tokio/src/runtime/builder.rs

@@ -54,6 +54,9 @@ pub struct Builder {
     /// Runtime type
     kind: Kind,
 
+    /// Name of the runtime.
+    name: Option<String>,
+
     /// 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,34 @@ 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();
+    /// # }
+    /// # }
+    /// ```
+    /// # Panics
+    ///
+    /// This function will panic if an empty value is passed as an argument.
+    ///
+    #[track_caller]
+    pub fn name(&mut self, val: impl Into<String>) -> &mut Self {
+        let val = val.into();
+        assert!(!val.trim().is_empty(), "runtime name shouldn't be empty");
+        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 +1667,7 @@ impl Builder {
                 metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(),
             },
             local_tid,
+            self.name.clone(),
         );
 
         let handle = Handle {
@@ -1814,6 +1849,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) };
@@ -1829,7 +1865,13 @@ cfg_rt_multi_thread! {
 
 impl fmt::Debug for Builder {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
-        fmt.debug_struct("Builder")
+        let mut debug = fmt.debug_struct("Builder");
+
+        if let Some(name) = &self.name {
+            debug.field("name", name);
+        }
+
+        debug
             .field("worker_threads", &self.worker_threads)
             .field("max_blocking_threads", &self.max_blocking_threads)
             .field(
@@ -1840,7 +1882,12 @@ impl fmt::Debug for Builder {
             .field("after_start", &self.after_start.as_ref().map(|_| "..."))
             .field("before_stop", &self.before_stop.as_ref().map(|_| "..."))
             .field("before_park", &self.before_park.as_ref().map(|_| "..."))
-            .field("after_unpark", &self.after_unpark.as_ref().map(|_| "..."))
-            .finish()
+            .field("after_unpark", &self.after_unpark.as_ref().map(|_| "..."));
+
+        if self.name.is_none() {
+            debug.finish_non_exhaustive()
+        } else {
+            debug.finish()
+        }
     }
 }

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 {

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
+    name: Option<String>,
+
     /// Scheduler state shared across threads
     shared: Shared,
 
@@ -132,6 +135,7 @@ impl CurrentThread {
         seed_generator: RngSeedGenerator,
         config: Config,
         local_tid: Option<ThreadId>,
+        name: Option<String>,
     ) -> (CurrentThread, Arc<Handle>) {
         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 {