From 63577cd8d3f2b1a1c787870125ab808eaa1eaa99 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Wed, 19 Jul 2023 12:53:40 +0200 Subject: [PATCH] rt: add runtime ID (#5864) 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. The inner value of the ID is taken from the `OwnedTasks` or `LocalOwnedTasks` struct which every runtime and local set already has. This will mean that any use of the ID will align with the task dump feature. The ID is added within the scope of working towards closing #5545. --- tokio/src/runtime/handle.rs | 32 +++++++++++++ tokio/src/runtime/id.rs | 46 +++++++++++++++++++ tokio/src/runtime/mod.rs | 4 ++ tokio/src/runtime/scheduler/current_thread.rs | 10 ++++ .../runtime/scheduler/multi_thread/handle.rs | 10 ++++ .../runtime/scheduler/multi_thread/worker.rs | 2 +- tokio/src/runtime/task/list.rs | 4 +- tokio/src/task/local.rs | 26 +++++++++++ tokio/tests/rt_handle.rs | 23 ++++++++++ tokio/tests/rt_threaded.rs | 18 ++++++++ 10 files changed, 172 insertions(+), 3 deletions(-) create mode 100644 tokio/src/runtime/id.rs diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index be4743d4775..36431df49c0 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,36 @@ 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 + /// [`Id`]: struct@crate::runtime::Id + pub fn id(&self) -> runtime::Id { + let owned_id = match &self.inner { + scheduler::Handle::CurrentThread(handle) => handle.owned_id(), + #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] + scheduler::Handle::MultiThread(handle) => handle.owned_id(), + }; + owned_id.into() + } + } } cfg_metrics! { diff --git a/tokio/src/runtime/id.rs b/tokio/src/runtime/id.rs new file mode 100644 index 00000000000..58551d49989 --- /dev/null +++ b/tokio/src/runtime/id.rs @@ -0,0 +1,46 @@ +use std::fmt; +use std::num::NonZeroU64; + +/// 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 runtime completes, the same ID may be used for another runtime. +/// - 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(NonZeroU64); + +impl From for Id { + fn from(value: NonZeroU64) -> Self { + Id(value) + } +} + +impl fmt::Display for Id { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} 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..80943aea87b 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -541,6 +541,16 @@ cfg_metrics! { } } +cfg_unstable! { + use std::num::NonZeroU64; + + impl Handle { + pub(crate) fn owned_id(&self) -> NonZeroU64 { + self.shared.owned.id + } + } +} + impl fmt::Debug for Handle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("current_thread::Handle { ... }").finish() diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index dfa186dda4f..568eb80af8b 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -59,6 +59,16 @@ impl Handle { } } +cfg_unstable! { + use std::num::NonZeroU64; + + impl Handle { + pub(crate) fn owned_id(&self) -> NonZeroU64 { + self.shared.owned.id + } + } +} + impl fmt::Debug for Handle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("multi_thread::Handle { ... }").finish() diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 577c0cc3d84..7db01f3c587 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -158,7 +158,7 @@ pub(crate) struct Shared { idle: Idle, /// Collection of all active tasks spawned onto this executor. - pub(super) owned: OwnedTasks>, + pub(crate) owned: OwnedTasks>, /// Data synchronized by the scheduler mutex pub(super) synced: Mutex, diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index 1c32a1ef361..7f376affda2 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -56,7 +56,7 @@ cfg_not_has_atomic_u64! { pub(crate) struct OwnedTasks { inner: Mutex>, - id: NonZeroU64, + pub(crate) id: NonZeroU64, } struct CountedOwnedTasksInner { list: CountedLinkedList, as Link>::Target>, @@ -64,7 +64,7 @@ struct CountedOwnedTasksInner { } pub(crate) struct LocalOwnedTasks { inner: UnsafeCell>, - id: NonZeroU64, + pub(crate) id: NonZeroU64, _not_send_or_sync: PhantomData<*const ()>, } struct OwnedTasksInner { diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 734b95587b9..59b68b8d2cf 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -1,6 +1,8 @@ //! Runs `!Send` futures on the current thread. use crate::loom::cell::UnsafeCell; use crate::loom::sync::{Arc, Mutex}; +#[cfg(tokio_unstable)] +use crate::runtime; use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task}; use crate::runtime::{context, ThreadId}; use crate::sync::AtomicWaker; @@ -785,6 +787,30 @@ cfg_unstable! { .unhandled_panic = behavior; self } + + /// Returns the [`Id`] of the current `LocalSet` runtime. + /// + /// # Examples + /// + /// ```rust + /// use tokio::task; + /// + /// #[tokio::main] + /// async fn main() { + /// let local_set = task::LocalSet::new(); + /// println!("Local set id: {}", local_set.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 + /// [`Id`]: struct@crate::runtime::Id + pub fn id(&self) -> runtime::Id { + self.context.shared.local_state.owned.id.into() + } } } 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()); + } }