diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index cf3cf52af..98e476585 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -9,6 +9,14 @@ use crate::util::RngSeedGenerator; use std::fmt; +cfg_metrics! { + mod metrics; +} + +cfg_taskdump! { + mod taskdump; +} + /// Handle to the multi thread scheduler pub(crate) struct Handle { /// Task spawner @@ -53,73 +61,6 @@ impl Handle { } } -cfg_metrics! { - use crate::runtime::{SchedulerMetrics, WorkerMetrics}; - - impl Handle { - pub(crate) fn num_workers(&self) -> usize { - self.shared.worker_metrics.len() - } - - pub(crate) fn num_blocking_threads(&self) -> usize { - self.blocking_spawner.num_threads() - } - - pub(crate) fn num_idle_blocking_threads(&self) -> usize { - self.blocking_spawner.num_idle_threads() - } - - pub(crate) fn active_tasks_count(&self) -> usize { - self.shared.owned.active_tasks_count() - } - - pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { - &self.shared.scheduler_metrics - } - - pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { - &self.shared.worker_metrics[worker] - } - - pub(crate) fn injection_queue_depth(&self) -> usize { - self.shared.injection_queue_depth() - } - - pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { - self.shared.worker_local_queue_depth(worker) - } - - pub(crate) fn blocking_queue_depth(&self) -> usize { - self.blocking_spawner.queue_depth() - } - } -} - -cfg_taskdump! { - impl Handle { - pub(crate) async fn dump(&self) -> crate::runtime::Dump { - let trace_status = &self.shared.trace_status; - - // If a dump is in progress, block. - trace_status.start_trace_request(&self).await; - - let result = loop { - if let Some(result) = trace_status.take_result() { - break result; - } else { - self.notify_all(); - trace_status.result_ready.notified().await; - } - }; - - // Allow other queued dumps to proceed. - trace_status.end_trace_request(&self).await; - - result - } - } -} - impl fmt::Debug for Handle { fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { fmt.debug_struct("multi_thread::Handle { ... }").finish() diff --git a/tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs b/tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs new file mode 100644 index 000000000..838694fc8 --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/handle/metrics.rs @@ -0,0 +1,41 @@ +use super::Handle; + +use crate::runtime::{SchedulerMetrics, WorkerMetrics}; + +impl Handle { + pub(crate) fn num_workers(&self) -> usize { + self.shared.worker_metrics.len() + } + + pub(crate) fn num_blocking_threads(&self) -> usize { + self.blocking_spawner.num_threads() + } + + pub(crate) fn num_idle_blocking_threads(&self) -> usize { + self.blocking_spawner.num_idle_threads() + } + + pub(crate) fn active_tasks_count(&self) -> usize { + self.shared.owned.active_tasks_count() + } + + pub(crate) fn scheduler_metrics(&self) -> &SchedulerMetrics { + &self.shared.scheduler_metrics + } + + pub(crate) fn worker_metrics(&self, worker: usize) -> &WorkerMetrics { + &self.shared.worker_metrics[worker] + } + + pub(crate) fn injection_queue_depth(&self) -> usize { + self.shared.injection_queue_depth() + } + + pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { + self.shared.worker_local_queue_depth(worker) + } + + pub(crate) fn blocking_queue_depth(&self) -> usize { + self.blocking_spawner.queue_depth() + } +} diff --git a/tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs b/tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs new file mode 100644 index 000000000..477d857d8 --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/handle/taskdump.rs @@ -0,0 +1,26 @@ +use super::Handle; + +use crate::runtime::Dump; + +impl Handle { + pub(crate) async fn dump(&self) -> Dump { + let trace_status = &self.shared.trace_status; + + // If a dump is in progress, block. + trace_status.start_trace_request(&self).await; + + let result = loop { + if let Some(result) = trace_status.take_result() { + break result; + } else { + self.notify_all(); + trace_status.result_ready.notified().await; + } + }; + + // Allow other queued dumps to proceed. + trace_status.end_trace_request(&self).await; + + result + } +} diff --git a/tokio/src/runtime/scheduler/multi_thread/mod.rs b/tokio/src/runtime/scheduler/multi_thread/mod.rs index 29cdbba50..d85a0ae0a 100644 --- a/tokio/src/runtime/scheduler/multi_thread/mod.rs +++ b/tokio/src/runtime/scheduler/multi_thread/mod.rs @@ -24,9 +24,17 @@ mod worker; pub(crate) use worker::{Context, Launch, Shared}; cfg_taskdump! { + mod trace; + use trace::TraceStatus; + pub(crate) use worker::Synced; } +cfg_not_taskdump! { + mod trace_mock; + use trace_mock::TraceStatus; +} + pub(crate) use worker::block_in_place; use crate::loom::sync::Arc; diff --git a/tokio/src/runtime/scheduler/multi_thread/trace.rs b/tokio/src/runtime/scheduler/multi_thread/trace.rs new file mode 100644 index 000000000..7b4aeb5c1 --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/trace.rs @@ -0,0 +1,61 @@ +use crate::loom::sync::atomic::{AtomicBool, Ordering}; +use crate::loom::sync::{Barrier, Mutex}; +use crate::runtime::dump::Dump; +use crate::runtime::scheduler::multi_thread::Handle; +use crate::sync::notify::Notify; + +/// Tracing status of the worker. +pub(super) struct TraceStatus { + pub(super) trace_requested: AtomicBool, + pub(super) trace_start: Barrier, + pub(super) trace_end: Barrier, + pub(super) result_ready: Notify, + pub(super) trace_result: Mutex>, +} + +impl TraceStatus { + pub(super) fn new(remotes_len: usize) -> Self { + Self { + trace_requested: AtomicBool::new(false), + trace_start: Barrier::new(remotes_len), + trace_end: Barrier::new(remotes_len), + result_ready: Notify::new(), + trace_result: Mutex::new(None), + } + } + + pub(super) fn trace_requested(&self) -> bool { + self.trace_requested.load(Ordering::Relaxed) + } + + pub(super) async fn start_trace_request(&self, handle: &Handle) { + while self + .trace_requested + .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) + .is_err() + { + handle.notify_all(); + crate::task::yield_now().await; + } + } + + pub(super) fn stash_result(&self, dump: Dump) { + let _ = self.trace_result.lock().insert(dump); + self.result_ready.notify_one(); + } + + pub(super) fn take_result(&self) -> Option { + self.trace_result.lock().take() + } + + pub(super) async fn end_trace_request(&self, handle: &Handle) { + while self + .trace_requested + .compare_exchange(true, false, Ordering::Acquire, Ordering::Relaxed) + .is_err() + { + handle.notify_all(); + crate::task::yield_now().await; + } + } +} diff --git a/tokio/src/runtime/scheduler/multi_thread/trace_mock.rs b/tokio/src/runtime/scheduler/multi_thread/trace_mock.rs new file mode 100644 index 000000000..2c17a4e38 --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/trace_mock.rs @@ -0,0 +1,11 @@ +pub(super) struct TraceStatus {} + +impl TraceStatus { + pub(super) fn new(_: usize) -> Self { + Self {} + } + + pub(super) fn trace_requested(&self) -> bool { + false + } +} diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 581c13ad4..47ff86a5c 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -60,7 +60,7 @@ use crate::loom::sync::{Arc, Mutex}; use crate::runtime; use crate::runtime::context; use crate::runtime::scheduler::multi_thread::{ - idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, Unparker, + idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker, }; use crate::runtime::scheduler::{inject, Defer, Lock}; use crate::runtime::task::OwnedTasks; @@ -74,8 +74,16 @@ use std::cell::RefCell; use std::task::Waker; use std::time::Duration; +cfg_metrics! { + mod metrics; +} + cfg_taskdump! { - use crate::loom::sync::Barrier; + mod taskdump; +} + +cfg_not_taskdump! { + mod taskdump_mock; } /// A scheduler worker @@ -214,82 +222,6 @@ pub(crate) struct Context { /// Starts the workers pub(crate) struct Launch(Vec>); -cfg_not_taskdump! { - pub(super) struct TraceStatus {} - - impl TraceStatus { - fn new(_: usize) -> Self { - Self {} - } - - fn trace_requested(&self) -> bool { - false - } - } -} - -cfg_taskdump! { - use crate::sync::notify::Notify; - use crate::runtime::dump::Dump; - use crate::loom::sync::atomic::{AtomicBool, Ordering}; - - /// Tracing status of the worker. - pub(super) struct TraceStatus { - pub(super) trace_requested: AtomicBool, - trace_start: Barrier, - trace_end: Barrier, - pub(super) result_ready: Notify, - pub(super) trace_result: Mutex>, - } - - impl TraceStatus { - fn new(remotes_len: usize) -> Self { - Self { - trace_requested: AtomicBool::new(false), - trace_start: Barrier::new(remotes_len), - trace_end: Barrier::new(remotes_len), - result_ready: Notify::new(), - trace_result: Mutex::new(None), - } - } - - fn trace_requested(&self) -> bool { - self.trace_requested.load(Ordering::Relaxed) - } - - pub(super) async fn start_trace_request(&self, handle: &Handle) { - while self.trace_requested.compare_exchange(false, - true, - Ordering::Acquire, - Ordering::Relaxed).is_err() - { - handle.notify_all(); - crate::task::yield_now().await; - } - } - - fn stash_result(&self, dump: Dump) { - let _ = self.trace_result.lock().insert(dump); - self.result_ready.notify_one(); - } - - pub(super) fn take_result(&self) -> Option { - self.trace_result.lock().take() - } - - pub(super) async fn end_trace_request(&self, handle: &Handle) { - while self.trace_requested.compare_exchange(true, - false, - Ordering::Acquire, - Ordering::Relaxed).is_err() - { - handle.notify_all(); - crate::task::yield_now().await; - } - } - } -} - /// Running a task may consume the core. If the core is still available when /// running the task completes, it is returned. Otherwise, the worker will need /// to stop processing. @@ -1200,64 +1132,6 @@ impl Handle { } } - cfg_not_taskdump! { - fn trace_core(&self, core: Box) -> Box { - core - } - } - - cfg_taskdump! { - fn trace_core(&self, mut core: Box) -> Box { - use crate::runtime::dump; - use task::trace::trace_multi_thread; - - core.is_traced = false; - - if core.is_shutdown { - return core; - } - - // wait for other workers, or timeout without tracing - let timeout = Duration::from_millis(250); // a _very_ generous timeout - let barrier = if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) { - barrier - } else { - // don't attempt to trace - return core; - }; - - if !barrier.is_leader() { - // wait for leader to finish tracing - self.shared.trace_status.trace_end.wait(); - return core; - } - - // trace - - let owned = &self.shared.owned; - let mut local = self.shared.steal_all(); - let synced = &self.shared.synced; - let injection = &self.shared.inject; - - // safety: `trace_multi_thread` is invoked with the same `synced` that `injection` - // was created with. - let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) } - .into_iter() - .map(dump::Task::new) - .collect(); - - let result = dump::Dump::new(traces); - - // stash the result - self.shared.trace_status.stash_result(result); - - // allow other workers to proceed - self.shared.trace_status.trace_end.wait(); - - core - } - } - fn ptr_eq(&self, other: &Handle) -> bool { std::ptr::eq(self, other) } @@ -1308,41 +1182,6 @@ fn with_current(f: impl FnOnce(Option<&Context>) -> R) -> R { }) } -cfg_metrics! { - impl Shared { - pub(super) fn injection_queue_depth(&self) -> usize { - self.inject.len() - } - - pub(super) fn worker_local_queue_depth(&self, worker: usize) -> usize { - self.remotes[worker].steal.len() - } - } -} - -cfg_taskdump! { - impl Shared { - /// Steal all tasks from remotes into a single local queue. - pub(super) fn steal_all(&self) -> super::queue::Local> { - let (_steal, mut local) = super::queue::local(); - - let worker_metrics = WorkerMetrics::new(); - let mut stats = Stats::new(&worker_metrics); - - for remote in self.remotes.iter() { - let steal = &remote.steal; - while !steal.is_empty() { - if let Some(task) = steal.steal_into(&mut local, &mut stats) { - local.push_back([task].into_iter()); - } - } - } - - local - } - } -} - // `u32::abs_diff` is not available on Tokio's MSRV. fn abs_diff(a: u32, b: u32) -> u32 { if a > b { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs b/tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs new file mode 100644 index 000000000..a9a5ab3ed --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/worker/metrics.rs @@ -0,0 +1,11 @@ +use super::Shared; + +impl Shared { + pub(crate) fn injection_queue_depth(&self) -> usize { + self.inject.len() + } + + pub(crate) fn worker_local_queue_depth(&self, worker: usize) -> usize { + self.remotes[worker].steal.len() + } +} diff --git a/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs b/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs new file mode 100644 index 000000000..d310d9f6d --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/worker/taskdump.rs @@ -0,0 +1,79 @@ +use super::{Core, Handle, Shared}; + +use crate::loom::sync::Arc; +use crate::runtime::scheduler::multi_thread::Stats; +use crate::runtime::task::trace::trace_multi_thread; +use crate::runtime::{dump, WorkerMetrics}; + +use std::time::Duration; + +impl Handle { + pub(super) fn trace_core(&self, mut core: Box) -> Box { + core.is_traced = false; + + if core.is_shutdown { + return core; + } + + // wait for other workers, or timeout without tracing + let timeout = Duration::from_millis(250); // a _very_ generous timeout + let barrier = + if let Some(barrier) = self.shared.trace_status.trace_start.wait_timeout(timeout) { + barrier + } else { + // don't attempt to trace + return core; + }; + + if !barrier.is_leader() { + // wait for leader to finish tracing + self.shared.trace_status.trace_end.wait(); + return core; + } + + // trace + + let owned = &self.shared.owned; + let mut local = self.shared.steal_all(); + let synced = &self.shared.synced; + let injection = &self.shared.inject; + + // safety: `trace_multi_thread` is invoked with the same `synced` that `injection` + // was created with. + let traces = unsafe { trace_multi_thread(owned, &mut local, synced, injection) } + .into_iter() + .map(dump::Task::new) + .collect(); + + let result = dump::Dump::new(traces); + + // stash the result + self.shared.trace_status.stash_result(result); + + // allow other workers to proceed + self.shared.trace_status.trace_end.wait(); + + core + } +} + +impl Shared { + /// Steal all tasks from remotes into a single local queue. + pub(super) fn steal_all(&self) -> super::queue::Local> { + let (_steal, mut local) = super::queue::local(); + + let worker_metrics = WorkerMetrics::new(); + let mut stats = Stats::new(&worker_metrics); + + for remote in self.remotes.iter() { + let steal = &remote.steal; + while !steal.is_empty() { + if let Some(task) = steal.steal_into(&mut local, &mut stats) { + local.push_back([task].into_iter()); + } + } + } + + local + } +} diff --git a/tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs b/tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs new file mode 100644 index 000000000..24c5600ce --- /dev/null +++ b/tokio/src/runtime/scheduler/multi_thread/worker/taskdump_mock.rs @@ -0,0 +1,7 @@ +use super::{Core, Handle}; + +impl Handle { + pub(super) fn trace_core(&self, core: Box) -> Box { + core + } +}