runtime: expose schedule latency in task hooks (#8282)

Add an explicit tracking knob and expose the sampled task schedule
latency through TaskMeta. Reuse the histogram poll timestamp where
possible and preserve grouped histogram accounting for LIFO polls.

Document activation and interval semantics and cover current-thread,
multi-thread, LIFO, disabled, and non-poll callback behavior.
This commit is contained in:
Russell Cohen
2026-08-09 14:38:39 +00:00
committed by GitHub
parent ecd621dd2c
commit 6b62ac48ed
15 changed files with 473 additions and 83 deletions
+58
View File
@@ -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,
+4
View File
@@ -48,7 +48,11 @@ pub(crate) struct Config {
/// How to build poll time histograms
pub(crate) metrics_poll_count_histogram: Option<crate::runtime::HistogramBuilder>,
/// 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<crate::runtime::HistogramBuilder>,
#[cfg(tokio_unstable)]
+78 -15
View File
@@ -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<ScheduleLatencyContext>) {}
pub(crate) fn start_poll(
&mut self,
_schedule_latency_context: Option<ScheduleLatencyContext>,
) -> Option<Duration> {
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<ScheduleLatencyContext>) {
/// `schedule_latency_context` is used to calculate task schedule latency.
pub(crate) fn start_poll(
&mut self,
schedule_latency_context: Option<ScheduleLatencyContext>,
) -> Option<Duration> {
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<ScheduleLatencyContext>,
) -> Option<Duration> {
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<ScheduleLatencyContext>,
) -> Option<Duration> {
#[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<ScheduleLatencyContext>,
poll_started_at: Option<Instant>,
) -> Option<Duration> {
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
+2 -2
View File
@@ -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};
}
@@ -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 {
@@ -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")
}
}
@@ -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<Instant>,
schedule_latency_start: Option<Instant>,
}
/// 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<Arc<Handle>>, mut core: Box<Core>) -> Box<Core> {
#[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<Handle> {
fn schedule(&self, task: task::Notified<Self>) {
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 {
@@ -93,6 +93,8 @@ impl Handle {
me.task_hooks.spawn(&TaskMeta {
id,
spawned_at,
#[cfg(feature = "schedule-latency")]
schedule_latency: None,
_phantom: Default::default(),
});
@@ -114,10 +114,21 @@ impl Stats {
}
}
pub(crate) fn start_poll(&mut self, task_scheduled_at: Option<ScheduleLatencyContext>) {
self.batch.start_poll(task_scheduled_at);
pub(crate) fn start_poll(
&mut self,
schedule_latency_context: Option<ScheduleLatencyContext>,
) -> Option<Duration> {
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<ScheduleLatencyContext>,
) -> Option<Duration> {
self.batch.record_schedule_latency(schedule_latency_context)
}
pub(crate) fn end_poll(&mut self) {
@@ -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<Instant>,
schedule_latency_start: Option<Instant>,
/// 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<Core>) -> 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| {
+2
View File
@@ -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(),
})
}));
+13 -9
View File
@@ -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<S> Sync for Task<S> {}
pub(crate) struct Notified<S: 'static>(Task<S>);
impl<S> Notified<S> {
#[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<S: 'static> {
impl<S> LocalNotified<S> {
#[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<Duration>,
) -> crate::runtime::TaskMeta<'meta> {
self.task.task_meta(schedule_latency)
}
pub(crate) fn get_scheduled_at(&self) -> ScheduleLatencyInstant {
@@ -448,10 +447,15 @@ impl<S: 'static> Task<S> {
// 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<Duration>,
) -> 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,
}
}
+31
View File
@@ -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<Duration>,
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<Duration> {
self.schedule_latency
}
}
/// Runs on specific task-related events
+51
View File
@@ -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::<Vec<_>>();
assert_eq!(bucket_counts.iter().sum::<u64>(), 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::<u64>();
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();
+172
View File
@@ -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<Mutex<Option<tokio::task::Id>>>,
latencies: &Arc<Mutex<Vec<Duration>>>,
) -> 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<AtomicUsize>,