From 6615423c4014680341700934e98144387fbbed65 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Thu, 13 Jul 2023 15:35:44 +0200 Subject: [PATCH 1/4] 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 | 31 ++++++++++++ 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, 139 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..661cd8b5022 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,35 @@ 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 { + 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()); + } } From a9a5e113ad674e2086bb9fad85d7ada48b509058 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Fri, 14 Jul 2023 14:16:02 +0200 Subject: [PATCH 2/4] use the `OwnedTasks` and `LocalOwnedTasks` ids Instead of generating our own ID, use the one that has already been implemented for the owned tasks structs. The u64 value is then returned as a `runtime::Id` to keep the implementation details hidden. --- tokio/src/runtime/handle.rs | 9 ++++--- tokio/src/runtime/id.rs | 18 +++++-------- tokio/src/runtime/scheduler/current_thread.rs | 13 +++++----- .../runtime/scheduler/multi_thread/handle.rs | 9 +++---- .../runtime/scheduler/multi_thread/worker.rs | 4 +-- tokio/src/runtime/task/list.rs | 4 +-- tokio/src/task/local.rs | 26 +++++++++++++++++++ 7 files changed, 51 insertions(+), 32 deletions(-) diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 661cd8b5022..1476ec8e6ca 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -381,11 +381,12 @@ impl Handle { /// [unstable]: crate#unstable-features /// [`Id`]: struct@crate::runtime::Id pub fn id(&self) -> runtime::Id { - match &self.inner { - scheduler::Handle::CurrentThread(handle) => handle.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.runtime_id, - } + scheduler::Handle::MultiThread(handle) => handle.owned_id(), + }; + runtime::Id::from_u64(owned_id) } } } diff --git a/tokio/src/runtime/id.rs b/tokio/src/runtime/id.rs index fa9af639f64..d675dca7657 100644 --- a/tokio/src/runtime/id.rs +++ b/tokio/src/runtime/id.rs @@ -6,7 +6,7 @@ use std::fmt; /// # Notes /// /// - Runtime IDs are unique relative to other *currently running* runtimes. -/// When a task completes, the same ID may be used for another task. +/// 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 @@ -32,18 +32,14 @@ use std::fmt; #[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 from_u64(val: u64) -> Self { + Id(val) } } -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)) +impl fmt::Display for Id { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) } } diff --git a/tokio/src/runtime/scheduler/current_thread.rs b/tokio/src/runtime/scheduler/current_thread.rs index 3d263084508..22c7a1e6955 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -1,8 +1,6 @@ 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}; @@ -43,9 +41,6 @@ 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 @@ -146,8 +141,6 @@ 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 { @@ -548,6 +541,12 @@ cfg_metrics! { } } +impl Handle { + pub(crate) fn owned_id(&self) -> u64 { + 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 3ba3a1497a7..10022c31ff9 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -1,7 +1,5 @@ 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, @@ -32,9 +30,6 @@ 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 { @@ -64,6 +59,10 @@ impl Handle { handle } + + pub(crate) fn owned_id(&self) -> u64 { + self.shared.owned.id + } } impl fmt::Debug for Handle { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 5c65e79c91d..f674385f31c 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, @@ -302,8 +302,6 @@ 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/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index fb7dbdc1d95..b11010d4f6f 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -55,7 +55,7 @@ cfg_not_has_atomic_u64! { pub(crate) struct OwnedTasks { inner: Mutex>, - id: u64, + pub(crate) id: u64, } struct CountedOwnedTasksInner { list: CountedLinkedList, as Link>::Target>, @@ -63,7 +63,7 @@ struct CountedOwnedTasksInner { } pub(crate) struct LocalOwnedTasks { inner: UnsafeCell>, - id: u64, + pub(crate) id: u64, _not_send_or_sync: PhantomData<*const ()>, } struct OwnedTasksInner { diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 734b95587b9..19d40eb2cb1 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 { + runtime::Id::from_u64(self.context.shared.local_state.owned.id) + } } } From 139c7c78d9051c0f44c8516c7f67a2d29403cd31 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Fri, 14 Jul 2023 14:31:22 +0200 Subject: [PATCH 3/4] fix visibility of methods only needed for unstable --- tokio/src/runtime/scheduler/current_thread.rs | 8 +++++--- tokio/src/runtime/scheduler/multi_thread/handle.rs | 8 ++++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/tokio/src/runtime/scheduler/current_thread.rs b/tokio/src/runtime/scheduler/current_thread.rs index 22c7a1e6955..e3b176b2059 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -541,9 +541,11 @@ cfg_metrics! { } } -impl Handle { - pub(crate) fn owned_id(&self) -> u64 { - self.shared.owned.id +cfg_unstable! { + impl Handle { + pub(crate) fn owned_id(&self) -> u64 { + self.shared.owned.id + } } } diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 10022c31ff9..d92a6725486 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -59,9 +59,13 @@ impl Handle { handle } +} - pub(crate) fn owned_id(&self) -> u64 { - self.shared.owned.id +cfg_unstable! { + impl Handle { + pub(crate) fn owned_id(&self) -> u64 { + self.shared.owned.id + } } } From 9a0ee60b6cab37d7f8171c4391446fb94cecfd37 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Tue, 18 Jul 2023 14:55:02 +0200 Subject: [PATCH 4/4] adapt runtime ID to use `NonZeroU64` --- tokio/src/runtime/handle.rs | 2 +- tokio/src/runtime/id.rs | 9 +++++---- tokio/src/runtime/scheduler/current_thread.rs | 4 +++- tokio/src/runtime/scheduler/multi_thread/handle.rs | 4 +++- tokio/src/task/local.rs | 2 +- 5 files changed, 13 insertions(+), 8 deletions(-) diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 1476ec8e6ca..36431df49c0 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -386,7 +386,7 @@ impl Handle { #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))] scheduler::Handle::MultiThread(handle) => handle.owned_id(), }; - runtime::Id::from_u64(owned_id) + owned_id.into() } } } diff --git a/tokio/src/runtime/id.rs b/tokio/src/runtime/id.rs index d675dca7657..58551d49989 100644 --- a/tokio/src/runtime/id.rs +++ b/tokio/src/runtime/id.rs @@ -1,4 +1,5 @@ use std::fmt; +use std::num::NonZeroU64; /// An opaque ID that uniquely identifies a runtime relative to all other currently /// running runtimes. @@ -30,11 +31,11 @@ use std::fmt; /// [unstable]: crate#unstable-features #[cfg_attr(not(tokio_unstable), allow(unreachable_pub))] #[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)] -pub struct Id(u64); +pub struct Id(NonZeroU64); -impl Id { - pub(crate) fn from_u64(val: u64) -> Self { - Id(val) +impl From for Id { + fn from(value: NonZeroU64) -> Self { + Id(value) } } diff --git a/tokio/src/runtime/scheduler/current_thread.rs b/tokio/src/runtime/scheduler/current_thread.rs index e3b176b2059..80943aea87b 100644 --- a/tokio/src/runtime/scheduler/current_thread.rs +++ b/tokio/src/runtime/scheduler/current_thread.rs @@ -542,8 +542,10 @@ cfg_metrics! { } cfg_unstable! { + use std::num::NonZeroU64; + impl Handle { - pub(crate) fn owned_id(&self) -> u64 { + pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } } diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index f95d8cfaaa2..568eb80af8b 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -60,8 +60,10 @@ impl Handle { } cfg_unstable! { + use std::num::NonZeroU64; + impl Handle { - pub(crate) fn owned_id(&self) -> u64 { + pub(crate) fn owned_id(&self) -> NonZeroU64 { self.shared.owned.id } } diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index 19d40eb2cb1..59b68b8d2cf 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -809,7 +809,7 @@ cfg_unstable! { /// [unstable]: crate#unstable-features /// [`Id`]: struct@crate::runtime::Id pub fn id(&self) -> runtime::Id { - runtime::Id::from_u64(self.context.shared.local_state.owned.id) + self.context.shared.local_state.owned.id.into() } } }