runtime: add tokio::runtime::worker_index() (#7921)

tokio/src/runtime/context.rs

@@ -159,6 +159,11 @@ cfg_rt! {
         CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None)
     }
 
+    #[cfg(tokio_unstable)]
+    pub(crate) fn worker_index() -> Option<usize> {
+        with_scheduler(|ctx| ctx.and_then(|c| c.worker_index()))
+    }
+
     #[track_caller]
     pub(crate) fn defer(waker: &Waker) {
         with_scheduler(|maybe_scheduler| {

tokio/src/runtime/mod.rs

@@ -569,6 +569,43 @@ cfg_rt! {
 
         mod local_runtime;
         pub use local_runtime::{LocalRuntime, LocalOptions};
+
+        /// Returns the index of the current worker thread, if called from a
+        /// runtime worker thread.
+        ///
+        /// The returned value is a 0-based index matching the worker indices
+        /// used by [`RuntimeMetrics`] methods such as
+        /// [`worker_total_busy_duration`](RuntimeMetrics::worker_total_busy_duration).
+        ///
+        /// Returns `None` when called from outside a runtime worker thread
+        /// (for example, from a blocking thread or a non-Tokio thread). On the
+        /// multi-thread runtime, the thread that calls [`Runtime::block_on`] is
+        /// not a worker thread, so this also returns `None` there.
+        ///
+        /// For the current-thread runtime and [`LocalRuntime`], this always
+        /// returns `Some(0)` (including inside `block_on`, since the calling
+        /// thread *is* the worker thread).
+        ///
+        /// Note that the result may change across `.await` points, as the
+        /// task may be moved to a different worker thread by the scheduler.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// # #[cfg(not(target_family = "wasm"))]
+        /// # {
+        /// #[tokio::main(flavor = "multi_thread", worker_threads = 4)]
+        /// async fn main() {
+        ///     let index = tokio::spawn(async {
+        ///         tokio::runtime::worker_index()
+        ///     }).await.unwrap();
+        ///     println!("Task ran on worker {:?}", index);
+        /// }
+        /// # }
+        /// ```
+        pub fn worker_index() -> Option<usize> {
+            context::worker_index()
+        }
     }
 
     cfg_taskdump! {

tokio/src/runtime/scheduler/mod.rs

@@ -293,6 +293,15 @@ cfg_rt! {
             match_flavor!(self, Context(context) => context.defer(waker));
         }
 
+        #[cfg(tokio_unstable)]
+        pub(crate) fn worker_index(&self) -> Option<usize> {
+            match self {
+                Context::CurrentThread(_) => Some(0),
+                #[cfg(feature = "rt-multi-thread")]
+                Context::MultiThread(context) => Some(context.worker_index()),
+            }
+        }
+
         #[cfg(all(tokio_unstable, feature = "time", feature = "rt-multi-thread"))]
         pub(crate) fn with_time_temp_local_context<F, R>(&self, f: F) -> R
         where

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -1006,6 +1006,11 @@ impl Context {
             None => f(None),
         })
     }
+
+    #[cfg(tokio_unstable)]
+    pub(crate) fn worker_index(&self) -> usize {
+        self.worker.index
+    }
 }
 
 impl Core {

tokio/tests/rt_worker_index.rs

@@ -0,0 +1,83 @@
+#![warn(rust_2018_idioms)]
+#![cfg(all(feature = "full", tokio_unstable, not(target_os = "wasi"),))]
+
+use tokio::runtime::{self, Runtime};
+
+#[test]
+fn worker_index_current_thread() {
+    let rt = runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    rt.block_on(async {
+        let index = runtime::worker_index();
+        assert_eq!(index, Some(0));
+    });
+}
+
+#[test]
+fn worker_index_local_runtime() {
+    let rt = runtime::LocalRuntime::new().unwrap();
+    rt.block_on(async {
+        let index = runtime::worker_index();
+        assert_eq!(index, Some(0));
+    });
+}
+
+#[test]
+fn worker_index_outside_runtime() {
+    assert_eq!(runtime::worker_index(), None);
+}
+
+#[cfg(target_has_atomic = "64")]
+#[test]
+fn worker_index_matches_metrics_worker_thread_id() {
+    let rt = runtime::Builder::new_multi_thread()
+        .worker_threads(4)
+        .enable_all()
+        .build()
+        .unwrap();
+    let metrics = rt.metrics();
+
+    rt.block_on(async {
+        // Spawn a task and verify the worker_index matches the metrics index
+        tokio::task::spawn(async move {
+            let index = runtime::worker_index().expect("should be on worker thread");
+            let current_thread = std::thread::current().id();
+            let metrics_thread = metrics.worker_thread_id(index);
+            assert_eq!(
+                metrics_thread,
+                Some(current_thread),
+                "worker_index() returned {index} but metrics.worker_thread_id({index}) \
+                 does not match current thread"
+            );
+        })
+        .await
+        .unwrap();
+    });
+}
+
+#[test]
+fn worker_index_from_spawn_blocking() {
+    let rt = Runtime::new().unwrap();
+    rt.block_on(async {
+        let index = tokio::task::spawn_blocking(runtime::worker_index)
+            .await
+            .unwrap();
+        assert_eq!(
+            index, None,
+            "spawn_blocking should not be on a worker thread"
+        );
+    });
+}
+
+#[test]
+fn worker_index_block_on_multi_thread() {
+    let rt = Runtime::new().unwrap();
+    // block_on runs on the calling thread, not a worker thread
+    let index = rt.block_on(async { runtime::worker_index() });
+    assert_eq!(
+        index, None,
+        "block_on thread is not a worker thread on multi-thread runtime"
+    );
+}