tokio/runtime/
task_hooks.rs

1use super::Config;
2use std::marker::PhantomData;
3#[cfg(all(tokio_unstable, feature = "schedule-latency"))]
4use std::time::Duration;
5
6impl TaskHooks {
7    pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) {
8        if let Some(f) = self.task_spawn_callback.as_ref() {
9            f(meta)
10        }
11    }
12
13    #[allow(dead_code)]
14    pub(crate) fn from_config(config: &Config) -> Self {
15        Self {
16            task_spawn_callback: config.before_spawn.clone(),
17            task_terminate_callback: config.after_termination.clone(),
18            #[cfg(tokio_unstable)]
19            before_poll_callback: config.before_poll.clone(),
20            #[cfg(tokio_unstable)]
21            after_poll_callback: config.after_poll.clone(),
22        }
23    }
24
25    #[cfg(tokio_unstable)]
26    #[inline]
27    pub(crate) fn poll_start_callback(&self, meta: &TaskMeta<'_>) {
28        if let Some(poll_start) = &self.before_poll_callback {
29            (poll_start)(meta);
30        }
31    }
32
33    #[cfg(tokio_unstable)]
34    #[inline]
35    pub(crate) fn poll_stop_callback(&self, meta: &TaskMeta<'_>) {
36        if let Some(poll_stop) = &self.after_poll_callback {
37            (poll_stop)(meta);
38        }
39    }
40}
41
42#[derive(Clone)]
43pub(crate) struct TaskHooks {
44    pub(crate) task_spawn_callback: Option<TaskCallback>,
45    pub(crate) task_terminate_callback: Option<TaskCallback>,
46    #[cfg(tokio_unstable)]
47    pub(crate) before_poll_callback: Option<TaskCallback>,
48    #[cfg(tokio_unstable)]
49    pub(crate) after_poll_callback: Option<TaskCallback>,
50}
51
52/// Task metadata supplied to user-provided hooks for task events.
53///
54/// **Note**: This is an [unstable API][unstable]. The public API of this type
55/// may break in 1.x releases. See [the documentation on unstable
56/// features][unstable] for details.
57///
58/// [unstable]: crate#unstable-features
59#[allow(missing_debug_implementations)]
60#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
61pub struct TaskMeta<'a> {
62    /// The opaque ID of the task.
63    pub(crate) id: super::task::Id,
64    /// The location where the task was spawned.
65    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
66    pub(crate) spawned_at: crate::runtime::task::SpawnLocation,
67    /// The latency between scheduling the task and starting its current poll.
68    #[cfg(all(tokio_unstable, feature = "schedule-latency"))]
69    pub(crate) schedule_latency: Option<Duration>,
70    pub(crate) _phantom: PhantomData<&'a ()>,
71}
72
73impl<'a> TaskMeta<'a> {
74    /// Return the opaque ID of the task.
75    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
76    pub fn id(&self) -> super::task::Id {
77        self.id
78    }
79
80    /// Return the source code location where the task was spawned.
81    #[cfg(tokio_unstable)]
82    pub fn spawned_at(&self) -> &'static std::panic::Location<'static> {
83        self.spawned_at.0
84    }
85
86    /// Returns the latency between scheduling the task and starting its
87    /// current poll.
88    ///
89    /// Task schedule latency is measured from the task's most recent transition
90    /// to the scheduled state until immediately before the
91    /// [`on_before_task_poll`] callback. Waking a task that is already scheduled
92    /// does not reset the measurement.
93    ///
94    /// This returns `Some` from [`on_before_task_poll`] and
95    /// [`on_after_task_poll`] callbacks when tracking is enabled with
96    /// [`track_task_schedule_latency`] or
97    /// [`enable_metrics_schedule_latency_histogram`]. Both callbacks receive
98    /// the same value, so time spent polling the task is not included. It
99    /// returns `None` when tracking is disabled and from task spawn and
100    /// termination callbacks.
101    ///
102    /// [`track_task_schedule_latency`]: crate::runtime::Builder::track_task_schedule_latency
103    /// [`enable_metrics_schedule_latency_histogram`]: crate::runtime::Builder::enable_metrics_schedule_latency_histogram
104    /// [`on_before_task_poll`]: crate::runtime::Builder::on_before_task_poll
105    /// [`on_after_task_poll`]: crate::runtime::Builder::on_after_task_poll
106    #[cfg(all(tokio_unstable, feature = "schedule-latency"))]
107    #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "schedule-latency"))))]
108    pub fn schedule_latency(&self) -> Option<Duration> {
109        self.schedule_latency
110    }
111}
112
113/// Runs on specific task-related events
114pub(crate) type TaskCallback = std::sync::Arc<dyn Fn(&TaskMeta<'_>) + Send + Sync>;
tokio/runtime/task_hooks.rs

tokio/runtime/
task_hooks.rs
