metrics: add worker thread id

tokio/src/runtime/metrics/mock.rs

@@ -1,5 +1,7 @@
 //! This file contains mocks of the types in src/runtime/metrics
 
+use std::thread::ThreadId;
+
 pub(crate) struct SchedulerMetrics {}
 
 pub(crate) struct WorkerMetrics {}
@@ -30,6 +32,7 @@ impl WorkerMetrics {
  }
 
  pub(crate) fn set_queue_depth(&self, _len: usize) {}
+ pub(crate) fn set_thread_id(&self, _thread_id: ThreadId) {}
 }
 
 impl MetricsBatch {

tokio/src/runtime/metrics/runtime.rs

@@ -2,6 +2,7 @@ use crate::runtime::Handle;
 
 cfg_unstable_metrics! {
  use std::ops::Range;
+ use std::thread::ThreadId;
  cfg_64bit_metrics! {
  use std::sync::atomic::Ordering::Relaxed;
  }
@@ -127,6 +128,49 @@ impl RuntimeMetrics {
  self.handle.inner.num_idle_blocking_threads()
  }
 
+ /// Returns the thread id of the given worker thread.
+ ///
+ /// The returned value is `None` if the worker thread has not yet finished
+ /// starting up.
+ ///
+ /// If additional information about the thread, such as its native id, are
+ /// required, those can be collected in [`on_thread_start`] and correlated
+ /// using the thread id.
+ ///
+ /// [`on_thread_start`]: crate::runtime::Builder::on_thread_start
+ ///
+ /// # Arguments
+ ///
+ /// `worker` is the index of the worker being queried. The given value must
+ /// be between 0 and `num_workers()`. The index uniquely identifies a single
+ /// worker and will continue to identify the worker throughout the lifetime
+ /// of the runtime instance.
+ ///
+ /// # Panics
+ ///
+ /// The method panics when `worker` represents an invalid worker, i.e. is
+ /// greater than or equal to `num_workers()`.
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// use tokio::runtime::Handle;
+ ///
+ /// #[tokio::main]
+ /// async fn main() {
+ /// let metrics = Handle::current().metrics();
+ ///
+ /// let id = metrics.worker_thread_id(0);
+ /// println!("worker 0 has id {:?}", id);
+ /// }
+ /// ```
+ pub fn worker_thread_id(&self, worker: usize) -> Option<ThreadId> {
+ self.handle
+ .inner
+ .worker_metrics(worker)
+ .thread_id()
+ }
+
  cfg_64bit_metrics! {
  /// Returns the number of tasks spawned in this runtime since it was created.
  ///

tokio/src/runtime/metrics/worker.rs

@@ -2,6 +2,8 @@ use crate::runtime::metrics::Histogram;
 use crate::runtime::Config;
 use crate::util::metric_atomics::{MetricAtomicU64, MetricAtomicUsize};
 use std::sync::atomic::Ordering::Relaxed;
+use std::sync::Mutex;
+use std::thread::ThreadId;
 
 /// Retrieve runtime worker metrics.
 ///
@@ -49,6 +51,9 @@ pub(crate) struct WorkerMetrics {
 
  /// If `Some`, tracks the number of polls by duration range.
  pub(super) poll_count_histogram: Option<Histogram>,
+
+ /// Thread id of worker thread.
+ thread_id: Mutex<Option<ThreadId>>,
 }
 
 impl WorkerMetrics {
@@ -72,4 +77,12 @@ impl WorkerMetrics {
  pub(crate) fn set_queue_depth(&self, len: usize) {
  self.queue_depth.store(len, Relaxed);
  }
+
+ pub(crate) fn thread_id(&self) -> Option<ThreadId> {
+ *self.thread_id.lock().unwrap()
+ }
+
+ pub(crate) fn set_thread_id(&self, thread_id: ThreadId) {
+ *self.thread_id.lock().unwrap() = Some(thread_id);
+ }
 }

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -11,12 +11,12 @@ use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
 
 use std::cell::RefCell;
 use std::collections::VecDeque;
-use std::fmt;
 use std::future::Future;
 use std::sync::atomic::Ordering::{AcqRel, Release};
 use std::task::Poll::{Pending, Ready};
 use std::task::Waker;
 use std::time::Duration;
+use std::{fmt, thread};
 
 /// Executes tasks on the current thread
 pub(crate) struct CurrentThread {
@@ -123,6 +123,7 @@ impl CurrentThread {
  config: Config,
  ) -> (CurrentThread, Arc<Handle>) {
  let worker_metrics = WorkerMetrics::from_config(&config);
+ worker_metrics.set_thread_id(thread::current().id());
 
  // Get the configured global queue interval, or use the default.
  let global_queue_interval = config
@@ -172,6 +173,10 @@ impl CurrentThread {
  // available or the future is complete.
  loop {
  if let Some(core) = self.take_core(handle) {
+ handle
+ .shared
+ .worker_metrics
+ .set_thread_id(thread::current().id());
  return core.block_on(future);
  } else {
  let notified = self.notify.notified();

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

@@ -72,6 +72,7 @@ use crate::util::rand::{FastRand, RngSeedGenerator};
 
 use std::cell::RefCell;
 use std::task::Waker;
+use std::thread;
 use std::time::Duration;
 
 cfg_unstable_metrics! {
@@ -334,6 +335,12 @@ where
  if let Some(cx) = maybe_cx {
  if self.take_core {
  let core = cx.worker.core.take();
+
+ if core.is_some() {
+ cx.worker.handle.shared.worker_metrics[cx.worker.index]
+ .set_thread_id(thread::current().id());
+ }
+
  let mut cx_core = cx.core.borrow_mut();
  assert!(cx_core.is_none());
  *cx_core = core;
@@ -482,6 +489,8 @@ fn run(worker: Arc<Worker>) {
  None => return,
  };
 
+ worker.handle.shared.worker_metrics[worker.index].set_thread_id(thread::current().id());
+
  let handle = scheduler::Handle::MultiThread(worker.handle.clone());
 
  crate::runtime::context::enter_runtime(&handle, true, |_| {

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

@@ -70,9 +70,9 @@ use crate::util::atomic_cell::AtomicCell;
 use crate::util::rand::{FastRand, RngSeedGenerator};
 
 use std::cell::{Cell, RefCell};
-use std::cmp;
 use std::task::Waker;
 use std::time::Duration;
+use std::{cmp, thread};
 
 cfg_unstable_metrics! {
  mod metrics;
@@ -569,6 +569,7 @@ impl Worker {
  }
  };
 
+ cx.shared().worker_metrics[core.index].set_thread_id(thread::current().id());
  core.stats.start_processing_scheduled_tasks(&mut self.stats);
 
  if let Some(task) = maybe_task {

tokio/tests/rt_unstable_metrics.rs

@@ -10,6 +10,7 @@
 use std::future::Future;
 use std::sync::{Arc, Barrier, Mutex};
 use std::task::Poll;
+use std::thread;
 use tokio::macros::support::poll_fn;
 
 use tokio::runtime::Runtime;
@@ -150,6 +151,68 @@ fn remote_schedule_count() {
  assert_eq!(1, rt.metrics().remote_schedule_count());
 }
 
+#[test]
+fn worker_thread_id_current_thread() {
+ let rt = current_thread();
+ let metrics = rt.metrics();
+
+ // Check that runtime is on this thread.
+ rt.block_on(async {});
+ assert_eq!(Some(thread::current().id()), metrics.worker_thread_id(0));
+
+ // Move runtime to another thread.
+ let thread_id = std::thread::scope(|scope| {
+ let join_handle = scope.spawn(|| {
+ rt.block_on(async {});
+ });
+ join_handle.thread().id()
+ });
+ assert_eq!(Some(thread_id), metrics.worker_thread_id(0));
+
+ // Move runtime back to this thread.
+ rt.block_on(async {});
+ assert_eq!(Some(thread::current().id()), metrics.worker_thread_id(0));
+}
+
+#[test]
+fn worker_thread_id_threaded() {
+ let rt = threaded();
+ let metrics = rt.metrics();
+
+ rt.block_on(rt.spawn(async move {
+ // Check that we are running on a worker thread and determine
+ // the index of our worker.
+ let thread_id = std::thread::current().id();
+ let this_worker = (0..2)
+ .position(|w| metrics.worker_thread_id(w) == Some(thread_id))
+ .expect("task not running on any worker thread");
+
+ // Force worker to another thread.
+ let moved_thread_id = tokio::task::block_in_place(|| {
+ assert_eq!(thread_id, std::thread::current().id());
+
+ // Wait for worker to move to another thread.
+ for _ in 0..100 {
+ let new_id = metrics.worker_thread_id(this_worker).unwrap();
+ if thread_id != new_id {
+ return new_id;
+ }
+ std::thread::sleep(Duration::from_millis(100));
+ }
+
+ panic!("worker did not move to new thread");
+ });
+
+ // After blocking task worker either stays on new thread or
+ // is moved back to current thread.
+ assert!(
+ metrics.worker_thread_id(this_worker) == Some(moved_thread_id)
+ || metrics.worker_thread_id(this_worker) == Some(thread_id)
+ );
+ }))
+ .unwrap()
+}
+
 #[test]
 fn worker_park_count() {
  let rt = current_thread();