diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index 30dd02b73..2966a91b4 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -139,6 +139,9 @@ pub struct Builder { pub(super) metrics_poll_count_histogram: HistogramBuilder, /// When true, enables task schedule latency instrumentation. + pub(super) track_task_schedule_latency: bool, + + /// When true, enables the task schedule latency histogram. pub(super) metrics_schedule_latency_histogram_enabled: bool, /// Configures the task schedule latency histogram. @@ -341,6 +344,8 @@ impl Builder { metrics_poll_count_histogram: HistogramBuilder::default(), + track_task_schedule_latency: false, + metrics_schedule_latency_histogram_enabled: false, metrics_schedule_latency_histogram: HistogramBuilder::default(), @@ -914,6 +919,9 @@ impl Builder { /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being /// invoked immediately. /// + /// When task schedule latency tracking is enabled, the latency is available + /// from `TaskMeta::schedule_latency`. + /// /// **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. @@ -961,6 +969,9 @@ impl Builder { /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being /// invoked immediately. /// + /// When task schedule latency tracking is enabled, the latency is available + /// from `TaskMeta::schedule_latency`. + /// /// **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. @@ -1728,6 +1739,7 @@ impl Builder { enable_eager_driver_handoff: false, seed_generator: seed_generator_1, metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), + track_task_schedule_latency: self.track_task_schedule_latency, metrics_schedule_latency_histogram: self .metrics_schedule_latency_histogram_builder(), }, @@ -1880,6 +1892,46 @@ cfg_test_util! { cfg_schedule_latency! { impl Builder { + /// Enables tracking task schedule latency. + /// + /// Task schedule latency is measured from a task's most recent + /// transition to the scheduled state until immediately before it is + /// polled. Waking a task that is already scheduled does not reset the + /// measurement. Once enabled, the latency is available to task poll + /// hooks through [`TaskMeta::schedule_latency`]. + /// + /// Task schedule latencies are not tracked by default as doing so + /// requires calling [`Instant::now()`] when a task is scheduled and + /// when it is polled, which could add measurable overhead. + /// + /// The [`enable_metrics_schedule_latency_histogram`] method also + /// enables tracking and records the latencies in a histogram. + /// + /// **This feature is only supported on 64-bit targets.** + /// + /// # Examples + /// + /// ``` + /// # use tokio::runtime; + /// let runtime = runtime::Builder::new_current_thread() + /// .track_task_schedule_latency() + /// .on_before_task_poll(|meta| { + /// if let Some(latency) = meta.schedule_latency() { + /// println!("task schedule latency: {latency:?}"); + /// } + /// }) + /// .build() + /// .unwrap(); + /// ``` + /// + /// [`TaskMeta::schedule_latency`]: crate::runtime::TaskMeta::schedule_latency + /// [`Instant::now()`]: std::time::Instant::now + /// [`enable_metrics_schedule_latency_histogram`]: Builder::enable_metrics_schedule_latency_histogram + pub fn track_task_schedule_latency(&mut self) -> &mut Self { + self.track_task_schedule_latency = true; + self + } + /// Enables tracking the distribution of task schedule latencies. Task /// schedule latency is the time between when a task is scheduled for /// execution and when it is polled. @@ -1896,6 +1948,10 @@ cfg_schedule_latency! { /// better granularity with low memory usage, use [`metrics_schedule_latency_histogram_configuration()`] /// to select [`LogHistogram`] instead. /// + /// On the multi-thread runtime, each task polled from the LIFO slot is + /// recorded as a separate schedule-latency sample. Task poll hooks + /// receive the same per-task latency through [`TaskMeta::schedule_latency`]. + /// /// # Examples /// /// ``` @@ -1921,6 +1977,7 @@ cfg_schedule_latency! { /// [`LogHistogram`]: crate::runtime::LogHistogram /// [`metrics_schedule_latency_histogram_configuration()`]: Builder::metrics_schedule_latency_histogram_configuration pub fn enable_metrics_schedule_latency_histogram(&mut self) -> &mut Self { + self.track_task_schedule_latency = true; self.metrics_schedule_latency_histogram_enabled = true; self } @@ -2050,6 +2107,7 @@ cfg_rt_multi_thread! { enable_eager_driver_handoff: self.enable_eager_driver_handoff, seed_generator: seed_generator_1, metrics_poll_count_histogram: self.metrics_poll_count_histogram_builder(), + track_task_schedule_latency: self.track_task_schedule_latency, metrics_schedule_latency_histogram: self.metrics_schedule_latency_histogram_builder(), }, self.timer_flavor, diff --git a/tokio/src/runtime/config.rs b/tokio/src/runtime/config.rs index 30d3a5266..8c16b1ca5 100644 --- a/tokio/src/runtime/config.rs +++ b/tokio/src/runtime/config.rs @@ -48,7 +48,11 @@ pub(crate) struct Config { /// How to build poll time histograms pub(crate) metrics_poll_count_histogram: Option, + /// Whether to track task schedule latency. + pub(crate) track_task_schedule_latency: bool, + /// How to build schedule latency histograms + #[cfg_attr(not(feature = "schedule-latency"), allow(dead_code))] pub(crate) metrics_schedule_latency_histogram: Option, #[cfg(tokio_unstable)] diff --git a/tokio/src/runtime/metrics/batch.rs b/tokio/src/runtime/metrics/batch.rs index 91af8d82d..8234c9cf3 100644 --- a/tokio/src/runtime/metrics/batch.rs +++ b/tokio/src/runtime/metrics/batch.rs @@ -226,34 +226,97 @@ impl MetricsBatch { cfg_metrics_variant! { stable: { /// Start polling an individual task - pub(crate) fn start_poll(&mut self, _task_scheduled_at: Option) {} + pub(crate) fn start_poll( + &mut self, + _schedule_latency_context: Option, + ) -> Option { + None + } }, unstable: { /// Start polling an individual task /// /// # Arguments /// - /// `task_scheduled_at` is used to calculate task schedule latency. - /// A `ScheduleLatencyContext` can be obtained by calling `prepare` on a task's - /// `ScheduleLatencyInstant`. - pub(crate) fn start_poll(&mut self, _task_scheduled_at: Option) { + /// `schedule_latency_context` is used to calculate task schedule latency. + pub(crate) fn start_poll( + &mut self, + schedule_latency_context: Option, + ) -> Option { self.poll_count += 1; - if let Some(poll_timer) = &mut self.poll_timer { - poll_timer.poll_started_at = Instant::now(); - } + let poll_started_at = self.poll_timer.as_mut().map(|poll_timer| { + let now = Instant::now(); + poll_timer.poll_started_at = now; + now + }); + #[cfg(feature = "schedule-latency")] - if let Some(task_scheduled_at) = _task_scheduled_at { - if let Some(schedule_latencies) = &mut self.schedule_latencies { - if let Some(now) = self.poll_timer.as_ref().map(|p| p.poll_started_at).or_else(now) { - let elapsed = task_scheduled_at.elapsed_nanos(now); - schedule_latencies.measure(elapsed, 1); - } - } + { + self.record_schedule_latency_at(schedule_latency_context, poll_started_at) + } + + #[cfg(not(feature = "schedule-latency"))] + { + let _ = (poll_started_at, schedule_latency_context); + None } } } } + cfg_metrics_variant! { + stable: { + /// Record the schedule latency of an additional task polled as part + /// of the current poll operation. + #[cfg(feature = "rt-multi-thread")] + pub(crate) fn record_schedule_latency( + &mut self, + _schedule_latency_context: Option, + ) -> Option { + None + } + }, + unstable: { + /// Record the schedule latency of an additional task polled as part + /// of the current poll operation. + #[cfg(feature = "rt-multi-thread")] + pub(crate) fn record_schedule_latency( + &mut self, + schedule_latency_context: Option, + ) -> Option { + #[cfg(feature = "schedule-latency")] + { + self.record_schedule_latency_at(schedule_latency_context, None) + } + + #[cfg(not(feature = "schedule-latency"))] + { + let _ = schedule_latency_context; + None + } + } + } + } + + #[cfg(all(tokio_unstable, feature = "schedule-latency"))] + fn record_schedule_latency_at( + &mut self, + schedule_latency_context: Option, + poll_started_at: Option, + ) -> Option { + let task_schedule_latency = schedule_latency_context.and_then(|schedule_latency_context| { + poll_started_at + .or_else(now) + .map(|now| Duration::from_nanos(schedule_latency_context.elapsed_nanos(now))) + }); + if let (Some(task_schedule_latency), Some(schedule_latencies)) = + (task_schedule_latency, &mut self.schedule_latencies) + { + schedule_latencies.measure(duration_as_u64(task_schedule_latency), 1); + } + task_schedule_latency + } + cfg_metrics_variant! { stable: { /// Stop polling an individual task diff --git a/tokio/src/runtime/metrics/mod.rs b/tokio/src/runtime/metrics/mod.rs index 2227fa508..6e0c5f45b 100644 --- a/tokio/src/runtime/metrics/mod.rs +++ b/tokio/src/runtime/metrics/mod.rs @@ -41,10 +41,10 @@ cfg_not_unstable_metrics! { cfg_schedule_latency! { mod schedule_latency; - pub(crate) use schedule_latency::{ScheduleLatencyInstant, ScheduleLatencyContext}; + pub(crate) use schedule_latency::{ScheduleLatencyContext, ScheduleLatencyInstant}; } cfg_not_schedule_latency! { mod schedule_latency_mock; - pub(crate) use schedule_latency_mock::{ScheduleLatencyInstant, ScheduleLatencyContext}; + pub(crate) use schedule_latency_mock::{ScheduleLatencyContext, ScheduleLatencyInstant}; } diff --git a/tokio/src/runtime/metrics/schedule_latency.rs b/tokio/src/runtime/metrics/schedule_latency.rs index 856726648..9346fe798 100644 --- a/tokio/src/runtime/metrics/schedule_latency.rs +++ b/tokio/src/runtime/metrics/schedule_latency.rs @@ -41,7 +41,7 @@ impl ScheduleLatencyInstant { /// `ScheduleLatencyContext` contains all the data required to calculate the time elapsed /// since a task was scheduled. /// -/// `ScheduleLatencyInstant` on its own in insufficient because it only contains a delta. +/// `ScheduleLatencyInstant` on its own is insufficient because it only contains a delta. /// The scheduler startup time is required to convert the delta back into an actual time /// but is omitted from `ScheduleLatencyInstant` to keep its memory size minimal. pub(crate) struct ScheduleLatencyContext { diff --git a/tokio/src/runtime/metrics/schedule_latency_mock.rs b/tokio/src/runtime/metrics/schedule_latency_mock.rs index 932fd7988..d00ea09a3 100644 --- a/tokio/src/runtime/metrics/schedule_latency_mock.rs +++ b/tokio/src/runtime/metrics/schedule_latency_mock.rs @@ -20,12 +20,3 @@ impl ScheduleLatencyInstant { pub(crate) struct ScheduleLatencyContext { _private: (), } - -impl ScheduleLatencyContext { - // This method is referenced (but never called) when the `schedule-latency` - // feature is disabled and `tokio_unstable` is enabled. - #[allow(dead_code)] - pub(crate) fn elapsed_nanos(&self, _now: Instant) -> u64 { - unimplemented!("This should never be called because prepare() always returns None") - } -} diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index 7902b2094..7e03f6b39 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -106,7 +106,7 @@ struct Shared { /// Startup time of this scheduler. /// /// This instant is used as the basis of task `scheduled_at` measurements. - started_at: Option, + schedule_latency_start: Option, } /// Thread-local context. @@ -152,10 +152,7 @@ impl CurrentThread { .global_queue_interval .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL); - let started_at = config - .metrics_schedule_latency_histogram - .as_ref() - .map(|_| Instant::now()); + let schedule_latency_start = config.track_task_schedule_latency.then(Instant::now); let handle = Arc::new(Handle { name, @@ -174,7 +171,7 @@ impl CurrentThread { config, scheduler_metrics: SchedulerMetrics::new(), worker_metrics, - started_at, + schedule_latency_start, }, driver: driver_handle, blocking_spawner, @@ -382,13 +379,13 @@ impl Context { /// Execute the closure with the given scheduler core stored in the /// thread-local context. fn run_task(&self, task: LocalNotified>, mut core: Box) -> Box { - #[cfg(tokio_unstable)] - let task_meta = task.task_meta(); + let schedule_latency_context = task + .get_scheduled_at() + .prepare(self.handle.shared.schedule_latency_start); + let _task_schedule_latency = core.metrics.start_poll(schedule_latency_context); - core.metrics.start_poll( - task.get_scheduled_at() - .prepare(self.handle.shared.started_at), - ); + #[cfg(tokio_unstable)] + let task_meta = task.task_meta(_task_schedule_latency); let (mut c, ()) = self.enter(core, || { crate::task::coop::budget(|| { @@ -517,6 +514,8 @@ impl Handle { me.task_hooks.spawn(&TaskMeta { id, spawned_at, + #[cfg(feature = "schedule-latency")] + schedule_latency: None, _phantom: Default::default(), }); @@ -555,6 +554,8 @@ impl Handle { me.task_hooks.spawn(&TaskMeta { id, spawned_at, + #[cfg(feature = "schedule-latency")] + schedule_latency: None, _phantom: Default::default(), }); @@ -709,13 +710,10 @@ impl Schedule for Arc { fn schedule(&self, task: task::Notified) { use scheduler::Context::CurrentThread; - if self - .shared - .config - .metrics_schedule_latency_histogram - .is_some() - { - task.set_scheduled_at(ScheduleLatencyInstant::new(self.shared.started_at)); + if self.shared.schedule_latency_start.is_some() { + task.set_scheduled_at(ScheduleLatencyInstant::new( + self.shared.schedule_latency_start, + )); } context::with_scheduler(|maybe_cx| match maybe_cx { diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 657e3fe34..407ac6da6 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -93,6 +93,8 @@ impl Handle { me.task_hooks.spawn(&TaskMeta { id, spawned_at, + #[cfg(feature = "schedule-latency")] + schedule_latency: None, _phantom: Default::default(), }); diff --git a/tokio/src/runtime/scheduler/multi_thread/stats.rs b/tokio/src/runtime/scheduler/multi_thread/stats.rs index a1f1f919f..db152e08f 100644 --- a/tokio/src/runtime/scheduler/multi_thread/stats.rs +++ b/tokio/src/runtime/scheduler/multi_thread/stats.rs @@ -114,10 +114,21 @@ impl Stats { } } - pub(crate) fn start_poll(&mut self, task_scheduled_at: Option) { - self.batch.start_poll(task_scheduled_at); + pub(crate) fn start_poll( + &mut self, + schedule_latency_context: Option, + ) -> Option { + let task_schedule_latency = self.batch.start_poll(schedule_latency_context); self.tasks_polled_in_batch += 1; + task_schedule_latency + } + + pub(crate) fn record_schedule_latency( + &mut self, + schedule_latency_context: Option, + ) -> Option { + self.batch.record_schedule_latency(schedule_latency_context) } pub(crate) fn end_poll(&mut self) { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 423f1f420..97a285c61 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -206,7 +206,7 @@ pub(crate) struct Shared { /// Startup time of this scheduler. /// /// This instant is used as the basis of task `scheduled_at` measurements. - started_at: Option, + schedule_latency_start: Option, /// Only held to trigger some code on drop. This is used to get internal /// runtime metrics that can be useful when doing performance @@ -317,10 +317,7 @@ pub(super) fn create( let (idle, idle_synced) = Idle::new(size); let (inject, inject_synced) = inject::Shared::new(); - let started_at = config - .metrics_schedule_latency_histogram - .as_ref() - .map(|_| Instant::now()); + let schedule_latency_start = config.track_task_schedule_latency.then(Instant::now); let remotes_len = remotes.len(); let handle = Arc::new(Handle { @@ -342,7 +339,7 @@ pub(super) fn create( config, scheduler_metrics: SchedulerMetrics::new(), worker_metrics: worker_metrics.into_boxed_slice(), - started_at, + schedule_latency_start, _counters: Counters, }, driver: driver_handle, @@ -645,9 +642,6 @@ impl Context { } fn run_task(&self, task: Notified, mut core: Box) -> RunResult { - #[cfg(tokio_unstable)] - let task_meta = task.task_meta(); - let task = self.worker.handle.shared.owned.assert_owner(task); // Make sure the worker is not in the **searching** state. This enables @@ -680,13 +674,16 @@ impl Context { self.assert_lifo_enabled_is_correct(&core); // Measure the poll start time. Note that we may end up polling other - // tasks under this measurement. In this case, the tasks came from the - // LIFO slot and are considered part of the current task for scheduling - // purposes. These tasks inherent the "parent"'s limits. - core.stats.start_poll( - task.get_scheduled_at() - .prepare(self.worker.handle.shared.started_at), - ); + // tasks under this poll-time measurement. Tasks from the LIFO slot + // inherit the "parent"'s limits, but their schedule latency is recorded + // separately when each task is polled. + let schedule_latency_context = task + .get_scheduled_at() + .prepare(self.worker.handle.shared.schedule_latency_start); + let _task_schedule_latency = core.stats.start_poll(schedule_latency_context); + + #[cfg(tokio_unstable)] + let task_meta = task.task_meta(_task_schedule_latency); // Make the core available to the runtime context *self.core.borrow_mut() = Some(core); @@ -765,12 +762,21 @@ impl Context { super::counters::inc_lifo_capped(); } - // Run the LIFO task, then loop - *self.core.borrow_mut() = Some(core); let task = self.worker.handle.shared.owned.assert_owner(task); + // Record the LIFO task's schedule latency independently from + // the outer task's poll-time sample. The returned value is the + // same sample that is passed to the histogram recorder. + let schedule_latency_context = task + .get_scheduled_at() + .prepare(self.worker.handle.shared.schedule_latency_start); + let _task_schedule_latency = + core.stats.record_schedule_latency(schedule_latency_context); + + *self.core.borrow_mut() = Some(core); + #[cfg(tokio_unstable)] - let task_meta = task.task_meta(); + let task_meta = task.task_meta(_task_schedule_latency); #[cfg(tokio_unstable)] self.worker @@ -1345,13 +1351,10 @@ impl Worker { impl Handle { pub(super) fn schedule_task(&self, task: Notified, is_yield: bool) { - if self - .shared - .config - .metrics_schedule_latency_histogram - .is_some() - { - task.set_scheduled_at(ScheduleLatencyInstant::new(self.shared.started_at)); + if self.shared.schedule_latency_start.is_some() { + task.set_scheduled_at(ScheduleLatencyInstant::new( + self.shared.schedule_latency_start, + )); } with_current(|maybe_cx| { diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 6f20d66ef..306a638e9 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -374,6 +374,8 @@ where f(&TaskMeta { id: self.core().task_id, spawned_at: self.core().spawned_at.into(), + #[cfg(feature = "schedule-latency")] + schedule_latency: None, _phantom: Default::default(), }) })); diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 422c7f383..f4da653cc 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -226,6 +226,8 @@ use crate::runtime::TaskCallback; use std::marker::PhantomData; use std::panic::Location; use std::ptr::NonNull; +#[cfg(tokio_unstable)] +use std::time::Duration; use std::{fmt, mem}; /// An owned handle to the task, tracked by ref count. @@ -243,12 +245,6 @@ unsafe impl Sync for Task {} pub(crate) struct Notified(Task); impl Notified { - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - #[inline] - pub(crate) fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { - self.0.task_meta() - } - pub(crate) fn set_scheduled_at(&self, scheduled_at: ScheduleLatencyInstant) { // SAFETY: There are no concurrent writes because there is only ever one `Notified` // reference per task. There are no concurrent reads because this field is only read @@ -275,8 +271,11 @@ pub(crate) struct LocalNotified { impl LocalNotified { #[cfg(tokio_unstable)] #[inline] - pub(crate) fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { - self.task.task_meta() + pub(crate) fn task_meta<'meta>( + &self, + schedule_latency: Option, + ) -> crate::runtime::TaskMeta<'meta> { + self.task.task_meta(schedule_latency) } pub(crate) fn get_scheduled_at(&self) -> ScheduleLatencyInstant { @@ -448,10 +447,15 @@ impl Task { // the compiler infers the lifetimes to be the same, and considers the task // to be borrowed for the lifetime of the returned `TaskMeta`. #[cfg(tokio_unstable)] - pub(crate) fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { + pub(crate) fn task_meta<'meta>( + &self, + _schedule_latency: Option, + ) -> crate::runtime::TaskMeta<'meta> { crate::runtime::TaskMeta { id: self.id(), spawned_at: self.spawned_at().into(), + #[cfg(feature = "schedule-latency")] + schedule_latency: _schedule_latency, _phantom: PhantomData, } } diff --git a/tokio/src/runtime/task_hooks.rs b/tokio/src/runtime/task_hooks.rs index 6df3837b5..11a70f49f 100644 --- a/tokio/src/runtime/task_hooks.rs +++ b/tokio/src/runtime/task_hooks.rs @@ -1,5 +1,7 @@ use super::Config; use std::marker::PhantomData; +#[cfg(feature = "schedule-latency")] +use std::time::Duration; impl TaskHooks { pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) { @@ -62,6 +64,9 @@ pub struct TaskMeta<'a> { /// The location where the task was spawned. #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))] pub(crate) spawned_at: crate::runtime::task::SpawnLocation, + /// The latency between scheduling the task and starting its current poll. + #[cfg(feature = "schedule-latency")] + pub(crate) schedule_latency: Option, pub(crate) _phantom: PhantomData<&'a ()>, } @@ -77,6 +82,32 @@ impl<'a> TaskMeta<'a> { pub fn spawned_at(&self) -> &'static std::panic::Location<'static> { self.spawned_at.0 } + + /// Returns the latency between scheduling the task and starting its + /// current poll. + /// + /// Task schedule latency is measured from the task's most recent transition + /// to the scheduled state until immediately before the + /// [`on_before_task_poll`] callback. Waking a task that is already scheduled + /// does not reset the measurement. + /// + /// This returns `Some` from [`on_before_task_poll`] and + /// [`on_after_task_poll`] callbacks when tracking is enabled with + /// [`track_task_schedule_latency`] or + /// [`enable_metrics_schedule_latency_histogram`]. Both callbacks receive + /// the same value, so time spent polling the task is not included. It + /// returns `None` when tracking is disabled and from task spawn and + /// termination callbacks. + /// + /// [`track_task_schedule_latency`]: crate::runtime::Builder::track_task_schedule_latency + /// [`enable_metrics_schedule_latency_histogram`]: crate::runtime::Builder::enable_metrics_schedule_latency_histogram + /// [`on_before_task_poll`]: crate::runtime::Builder::on_before_task_poll + /// [`on_after_task_poll`]: crate::runtime::Builder::on_after_task_poll + #[cfg(feature = "schedule-latency")] + #[cfg_attr(docsrs, doc(cfg(feature = "schedule-latency")))] + pub fn schedule_latency(&self) -> Option { + self.schedule_latency + } } /// Runs on specific task-related events diff --git a/tokio/tests/rt_unstable_metrics.rs b/tokio/tests/rt_unstable_metrics.rs index f4cde15bb..99f877a4e 100644 --- a/tokio/tests/rt_unstable_metrics.rs +++ b/tokio/tests/rt_unstable_metrics.rs @@ -849,6 +849,57 @@ fn schedule_latency_counts() { } } +#[cfg(feature = "schedule-latency")] +#[test] +fn schedule_latency_lifo_polls_are_recorded_individually() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_metrics_schedule_latency_histogram() + .enable_metrics_poll_time_histogram() + .metrics_schedule_latency_histogram_configuration(HistogramConfiguration::linear( + Duration::from_millis(10), + 3, + )) + .build() + .unwrap(); + let metrics = rt.metrics(); + + rt.block_on(async { + tokio::spawn(async { + drop(tokio::spawn(async {})); + wait_for_elapsed(Duration::from_millis(50)); + }) + .await + .unwrap(); + }); + drop(rt); + + let bucket_counts = (0..metrics.schedule_latency_histogram_num_buckets()) + .map(|bucket| metrics.schedule_latency_histogram_bucket_count(0, bucket)) + .collect::>(); + + assert_eq!(bucket_counts.iter().sum::(), 2); + assert!(bucket_counts.last().is_some_and(|count| *count >= 1)); + + assert_eq!(metrics.worker_poll_count(0), 1); + let poll_time_samples = (0..metrics.poll_time_histogram_num_buckets()) + .map(|bucket| metrics.poll_time_histogram_bucket_count(0, bucket)) + .sum::(); + assert_eq!(poll_time_samples, 1); +} + +#[cfg(feature = "schedule-latency")] +fn wait_for_elapsed(duration: Duration) { + let start = std::time::Instant::now(); + loop { + let elapsed = start.elapsed(); + if elapsed >= duration { + return; + } + std::thread::sleep(duration - elapsed); + } +} + async fn try_spawn_stealable_task() -> Result<(), mpsc::RecvTimeoutError> { // We use a blocking channel to synchronize the tasks. let (tx, rx) = mpsc::channel(); diff --git a/tokio/tests/task_hooks.rs b/tokio/tests/task_hooks.rs index 42bb3fd94..5f484f6ce 100644 --- a/tokio/tests/task_hooks.rs +++ b/tokio/tests/task_hooks.rs @@ -4,6 +4,8 @@ use std::collections::HashSet; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; +#[cfg(feature = "schedule-latency")] +use std::time::{Duration, Instant}; use tokio::runtime::Builder; @@ -20,6 +22,8 @@ fn spawn_task_hook_fires() { let runtime = Builder::new_current_thread() .on_task_spawn(move |data| { + #[cfg(feature = "schedule-latency")] + assert_eq!(data.schedule_latency(), None); ids2.lock().unwrap().insert(data.id()); count2.fetch_add(1, Ordering::SeqCst); @@ -53,6 +57,8 @@ fn terminate_task_hook_fires() { let runtime = Builder::new_current_thread() .on_task_terminate(move |_data| { + #[cfg(feature = "schedule-latency")] + assert_eq!(_data.schedule_latency(), None); count2.fetch_add(1, Ordering::SeqCst); }) .build() @@ -174,6 +180,172 @@ fn task_hook_spawn_location_multi_thread() { assert_eq!(poll_starts, poll_ends.fetch_add(0, Ordering::SeqCst)); } +#[cfg(feature = "schedule-latency")] +#[test] +fn task_hook_schedule_latency_non_poll_callbacks() { + let spawn_count = Arc::new(AtomicUsize::new(0)); + let spawn_count2 = Arc::clone(&spawn_count); + let terminate_count = Arc::new(AtomicUsize::new(0)); + let terminate_count2 = Arc::clone(&terminate_count); + + let runtime = Builder::new_current_thread() + .track_task_schedule_latency() + .on_task_spawn(move |data| { + assert_eq!(data.schedule_latency(), None); + spawn_count2.fetch_add(1, Ordering::SeqCst); + }) + .on_task_terminate(move |data| { + assert_eq!(data.schedule_latency(), None); + terminate_count2.fetch_add(1, Ordering::SeqCst); + }) + .build() + .unwrap(); + + runtime.block_on(runtime.spawn(async {})).unwrap(); + assert_eq!(spawn_count.load(Ordering::SeqCst), 1); + assert_eq!(terminate_count.load(Ordering::SeqCst), 1); +} + +#[cfg(feature = "schedule-latency")] +#[test] +fn task_hook_schedule_latency_disabled() { + let runtime = Builder::new_current_thread() + .on_before_task_poll(|data| assert_eq!(data.schedule_latency(), None)) + .build() + .unwrap(); + + runtime.block_on(runtime.spawn(async {})).unwrap(); +} + +#[cfg(feature = "schedule-latency")] +#[test] +fn task_hook_schedule_latency_current_thread() { + let target = Arc::new(Mutex::new(None)); + let latencies = Arc::new(Mutex::new(Vec::new())); + + let after_latencies = Arc::new(Mutex::new(Vec::new())); + + let runtime = Builder::new_current_thread() + .enable_metrics_schedule_latency_histogram() + .on_before_task_poll(schedule_latency_hook(&target, &latencies)) + .on_after_task_poll(schedule_latency_hook(&target, &after_latencies)) + .build() + .unwrap(); + + let task = runtime.spawn(async {}); + *target.lock().unwrap() = Some(task.id()); + + // A current-thread runtime has no background worker, so the spawned task + // cannot be polled until `block_on` starts driving the runtime below. + wait_for_elapsed(Duration::from_millis(50)); + runtime.block_on(task).unwrap(); + + let latencies = latencies.lock().unwrap(); + assert_eq!(latencies.len(), 1); + assert!(latencies[0] >= Duration::from_millis(25)); + assert_eq!(*latencies, *after_latencies.lock().unwrap()); +} + +#[cfg(feature = "schedule-latency")] +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn task_hook_schedule_latency_multi_thread() { + let target = Arc::new(Mutex::new(None)); + let latencies = Arc::new(Mutex::new(Vec::new())); + + let runtime = Builder::new_multi_thread() + .worker_threads(1) + .track_task_schedule_latency() + .on_before_task_poll(schedule_latency_hook(&target, &latencies)) + .build() + .unwrap(); + + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let blocker = runtime.spawn(async move { + started_tx.send(()).unwrap(); + release_rx.recv().unwrap(); + }); + started_rx.recv().unwrap(); + + let task = runtime.spawn(async {}); + *target.lock().unwrap() = Some(task.id()); + + wait_for_elapsed(Duration::from_millis(50)); + release_tx.send(()).unwrap(); + runtime.block_on(async { + blocker.await.unwrap(); + task.await.unwrap(); + }); + + let latencies = latencies.lock().unwrap(); + assert_eq!(latencies.len(), 1); + assert!(latencies[0] >= Duration::from_millis(25)); +} + +#[cfg(feature = "schedule-latency")] +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn task_hook_schedule_latency_multi_thread_lifo() { + let target = Arc::new(Mutex::new(None)); + let latencies = Arc::new(Mutex::new(Vec::new())); + + let runtime = Builder::new_multi_thread() + .worker_threads(1) + .track_task_schedule_latency() + .on_before_task_poll(schedule_latency_hook(&target, &latencies)) + .build() + .unwrap(); + + let target2 = Arc::clone(&target); + let parent = runtime.spawn(async move { + let task = tokio::spawn(async {}); + *target2.lock().unwrap() = Some(task.id()); + wait_for_elapsed(Duration::from_millis(50)); + task.await.unwrap(); + }); + runtime.block_on(parent).unwrap(); + + let latencies = latencies.lock().unwrap(); + assert_eq!(latencies.len(), 1); + assert!(latencies[0] >= Duration::from_millis(25)); +} + +#[cfg(feature = "schedule-latency")] +fn wait_for_elapsed(duration: Duration) { + let start = Instant::now(); + loop { + let elapsed = start.elapsed(); + if elapsed >= duration { + return; + } + std::thread::sleep(duration - elapsed); + } +} + +#[cfg(feature = "schedule-latency")] +fn schedule_latency_hook( + target: &Arc>>, + latencies: &Arc>>, +) -> impl Fn(&tokio::runtime::TaskMeta<'_>) { + let target = Arc::clone(target); + let latencies = Arc::clone(latencies); + move |data| { + if Some(data.id()) == *target.lock().unwrap() { + latencies + .lock() + .unwrap() + .push(data.schedule_latency().unwrap()); + } + } +} + fn mk_spawn_location_hook( event: &'static str, count: &Arc,