Skip to content

Commit

Permalink
rt: add runtime ID (#5864)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
hds committed Jul 19, 2023
1 parent 0254454 commit 63577cd
Show file tree
Hide file tree
Showing 10 changed files with 172 additions and 3 deletions.
32 changes: 32 additions & 0 deletions tokio/src/runtime/handle.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#[cfg(tokio_unstable)]
use crate::runtime;
use crate::runtime::{context, scheduler, RuntimeFlavor};

/// Handle to the runtime.
Expand Down Expand Up @@ -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! {
Expand Down
46 changes: 46 additions & 0 deletions tokio/src/runtime/id.rs
Original file line number Diff line number Diff line change
@@ -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<NonZeroU64> 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)
}
}
4 changes: 4 additions & 0 deletions tokio/src/runtime/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
10 changes: 10 additions & 0 deletions tokio/src/runtime/scheduler/current_thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
10 changes: 10 additions & 0 deletions tokio/src/runtime/scheduler/multi_thread/handle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
2 changes: 1 addition & 1 deletion tokio/src/runtime/scheduler/multi_thread/worker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ pub(crate) struct Shared {
idle: Idle,

/// Collection of all active tasks spawned onto this executor.
pub(super) owned: OwnedTasks<Arc<Handle>>,
pub(crate) owned: OwnedTasks<Arc<Handle>>,

/// Data synchronized by the scheduler mutex
pub(super) synced: Mutex<Synced>,
Expand Down
4 changes: 2 additions & 2 deletions tokio/src/runtime/task/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,15 +56,15 @@ cfg_not_has_atomic_u64! {

pub(crate) struct OwnedTasks<S: 'static> {
inner: Mutex<CountedOwnedTasksInner<S>>,
id: NonZeroU64,
pub(crate) id: NonZeroU64,
}
struct CountedOwnedTasksInner<S: 'static> {
list: CountedLinkedList<Task<S>, <Task<S> as Link>::Target>,
closed: bool,
}
pub(crate) struct LocalOwnedTasks<S: 'static> {
inner: UnsafeCell<OwnedTasksInner<S>>,
id: NonZeroU64,
pub(crate) id: NonZeroU64,
_not_send_or_sync: PhantomData<*const ()>,
}
struct OwnedTasksInner<S: 'static> {
Expand Down
26 changes: 26 additions & 0 deletions tokio/src/task/local.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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()
}
}
}

Expand Down
23 changes: 23 additions & 0 deletions tokio/tests/rt_handle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
18 changes: 18 additions & 0 deletions tokio/tests/rt_threaded.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
}

0 comments on commit 63577cd

Please sign in to comment.