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.

Author: hds

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! {

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<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)
+    }
+}

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;
     }

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

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

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<Arc<Handle>>,
+    pub(crate) owned: OwnedTasks<Arc<Handle>>,
 
     /// Data synchronized by the scheduler mutex
     pub(super) synced: Mutex<Synced>,

tokio/src/runtime/task/list.rs

@@ -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> {

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

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

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());
+    }
 }