From 22bf1995cd767036eac694a06aae0d5b68b44534 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Thu, 13 Jul 2023 15:35:44 +0200 Subject: [PATCH] rt: add runtime Id There are a number of cases in which being able to identify a runtime is useful. When instrumenting an application, this is particularly true. For example, we would like to be able to add traces for runtimes so that tasks can be differentiated (#5792). It would also allow a way to differentiate runtimes which are have their tasks dumped. Outside of instrumentation, it may be useful to check whether 2 runtime handles are pointing to the same runtime. This change adds an opaque `runtime::Id` struct which serves this purpose, initially behind the `tokio_unstable` cfg flag. It follows the same pattern as the `task::Id` struct. The Id can be compared for equality with another `runtime::Id` and implements `Debug` and `Display` so that it can be output as well. Internally the Id is a `u64`, but that is an implementation detail. There is a degree of code duplication, but that is necessary to ensure that the Ids are not used to compare against one another. The Id is added within the scope of working towards closing #5545. --- tokio/src/runtime/handle.rs | 30 ++++++++++++ tokio/src/runtime/id.rs | 49 +++++++++++++++++++ tokio/src/runtime/mod.rs | 4 ++ tokio/src/runtime/scheduler/current_thread.rs | 7 +++ .../runtime/scheduler/multi_thread/handle.rs | 5 ++ .../runtime/scheduler/multi_thread/worker.rs | 2 + tokio/tests/rt_handle.rs | 23 +++++++++ tokio/tests/rt_threaded.rs | 18 +++++++ 8 files changed, 138 insertions(+) create mode 100644 tokio/src/runtime/id.rs diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index be4743d4775..88288680913 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -1,3 +1,5 @@ +#[cfg(tokio_unstable)] +use crate::runtime; use crate::runtime::{context, scheduler, RuntimeFlavor}; /// Handle to the runtime. @@ -357,6 +359,34 @@ impl Handle { scheduler::Handle::MultiThread(_) => RuntimeFlavor::MultiThread, } } + + cfg_unstable! { + /// Returns the [`Id`] of the current `Runtime`. + /// + /// # Examples + /// + /// ``` + /// use tokio::runtime::Handle; + /// + /// #[tokio::main(flavor = "current_thread")] + /// async fn main() { + /// println!("Current runtime id: {}", Handle::current().id()); + /// } + /// ``` + /// + /// **Note**: This is an [unstable API][unstable]. The public API of this type + /// may break in 1.x releases. See [the documentation on unstable + /// features][unstable] for details. + /// + /// [unstable]: crate#unstable-features + pub fn id(&self) -> runtime::Id { + match &self.inner { + scheduler::Handle::CurrentThread(handle) => handle.runtime_id, + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Handle::MultiThread(handle) => handle.runtime_id, + } + } + } } cfg_metrics! { diff --git a/tokio/src/runtime/id.rs b/tokio/src/runtime/id.rs new file mode 100644 index 00000000000..fa9af639f64 --- /dev/null +++ b/tokio/src/runtime/id.rs @@ -0,0 +1,49 @@ +use std::fmt; + +/// An opaque ID that uniquely identifies a runtime relative to all other currently +/// running runtimes. +/// +/// # Notes +/// +/// - Runtime IDs are unique relative to other *currently running* runtimes. +/// When a task completes, the same ID may be used for another task. +/// - Runtime IDs are *not* sequential, and do not indicate the order in which +/// runtimes are started or any other data. +/// - The runtime ID of the currently running task can be obtained from the +/// Handle. +/// +/// # Examples +/// +/// ``` +/// use tokio::runtime::Handle; +/// +/// #[tokio::main(flavor = "multi_thread", worker_threads = 4)] +/// async fn main() { +/// println!("Current runtime id: {}", Handle::current().id()); +/// } +/// ``` +/// +/// **Note**: This is an [unstable API][unstable]. The public API of this type +/// may break in 1.x releases. See [the documentation on unstable +/// features][unstable] for details. +/// +/// [unstable]: crate#unstable-features +#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] +#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)] +pub struct Id(u64); + +impl fmt::Display for Id { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + +impl Id { + pub(crate) fn next() -> Self { + use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64}; + + static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1); + + Self(NEXT_ID.fetch_add(1, Relaxed)) + } +} diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index cb198f51f0d..f8b651745b2 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -226,6 +226,10 @@ cfg_rt! { mod builder; pub use self::builder::Builder; cfg_unstable! { + mod id; + #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] + pub use id::Id; + pub use self::builder::UnhandledPanic; pub use crate::util::rand::RngSeed; } diff --git a/tokio/src/runtime/scheduler/current_thread.rs b/tokio/src/runtime/scheduler/current_thread.rs index ac4a8d6fac1..3d263084508 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -1,6 +1,8 @@ use crate::future::poll_fn; use crate::loom::sync::atomic::AtomicBool; use crate::loom::sync::Arc; +#[cfg(tokio_unstable)] +use crate::runtime; use crate::runtime::driver::{self, Driver}; use crate::runtime::scheduler::{self, Defer, Inject}; use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task}; @@ -41,6 +43,9 @@ pub(crate) struct Handle { /// Current random number generator seed pub(crate) seed_generator: RngSeedGenerator, + + #[cfg(tokio_unstable)] + pub(crate) runtime_id: runtime::Id, } /// Data required for executing the scheduler. The struct is passed around to @@ -141,6 +146,8 @@ impl CurrentThread { driver: driver_handle, blocking_spawner, seed_generator, + #[cfg(tokio_unstable)] + runtime_id: runtime::Id::next(), }); let core = AtomicCell::new(Some(Box::new(Core { diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 98e47658560..3ba3a1497a7 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -1,5 +1,7 @@ use crate::future::Future; use crate::loom::sync::Arc; +#[cfg(tokio_unstable)] +use crate::runtime; use crate::runtime::scheduler::multi_thread::worker; use crate::runtime::{ blocking, driver, @@ -30,6 +32,9 @@ pub(crate) struct Handle { /// Current random number generator seed pub(crate) seed_generator: RngSeedGenerator, + + #[cfg(tokio_unstable)] + pub(crate) runtime_id: runtime::Id, } impl Handle { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 6ae11463373..5c65e79c91d 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -302,6 +302,8 @@ pub(super) fn create( driver: driver_handle, blocking_spawner, seed_generator, + #[cfg(tokio_unstable)] + runtime_id: runtime::Id::next(), }); let mut launch = Launch(vec![]); diff --git a/tokio/tests/rt_handle.rs b/tokio/tests/rt_handle.rs index 34c99cdaead..14d6524f62e 100644 --- a/tokio/tests/rt_handle.rs +++ b/tokio/tests/rt_handle.rs @@ -60,6 +60,29 @@ fn interleave_then_enter() { let _enter = rt3.enter(); } +#[cfg(tokio_unstable)] +mod unstable { + use super::*; + + #[test] + fn runtime_id_is_same() { + let rt = rt(); + + let handle1 = rt.handle(); + let handle2 = rt.handle(); + + assert_eq!(handle1.id(), handle2.id()); + } + + #[test] + fn runtime_ids_different() { + let rt1 = rt(); + let rt2 = rt(); + + assert_ne!(rt1.handle().id(), rt2.handle().id()); + } +} + fn rt() -> Runtime { tokio::runtime::Builder::new_current_thread() .build() diff --git a/tokio/tests/rt_threaded.rs b/tokio/tests/rt_threaded.rs index 69b186947bd..3c77c7e5c50 100644 --- a/tokio/tests/rt_threaded.rs +++ b/tokio/tests/rt_threaded.rs @@ -762,4 +762,22 @@ mod unstable { .unwrap(); }) } + + #[test] + fn runtime_id_is_same() { + let rt = rt(); + + let handle1 = rt.handle(); + let handle2 = rt.handle(); + + assert_eq!(handle1.id(), handle2.id()); + } + + #[test] + fn runtime_ids_different() { + let rt1 = rt(); + let rt2 = rt(); + + assert_ne!(rt1.handle().id(), rt2.handle().id()); + } }