tokio/runtime/
task_hooks.rs

1use super::Config;
2use std::marker::PhantomData;
3
4impl TaskHooks {
5    pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) {
6        if let Some(f) = self.task_spawn_callback.as_ref() {
7            f(meta)
8        }
9    }
10
11    #[allow(dead_code)]
12    pub(crate) fn from_config(config: &Config) -> Self {
13        Self {
14            task_spawn_callback: config.before_spawn.clone(),
15            task_terminate_callback: config.after_termination.clone(),
16            #[cfg(tokio_unstable)]
17            before_poll_callback: config.before_poll.clone(),
18            #[cfg(tokio_unstable)]
19            after_poll_callback: config.after_poll.clone(),
20        }
21    }
22
23    #[cfg(tokio_unstable)]
24    #[inline]
25    pub(crate) fn poll_start_callback(&self, meta: &TaskMeta<'_>) {
26        if let Some(poll_start) = &self.before_poll_callback {
27            (poll_start)(meta);
28        }
29    }
30
31    #[cfg(tokio_unstable)]
32    #[inline]
33    pub(crate) fn poll_stop_callback(&self, meta: &TaskMeta<'_>) {
34        if let Some(poll_stop) = &self.after_poll_callback {
35            (poll_stop)(meta);
36        }
37    }
38}
39
40#[derive(Clone)]
41pub(crate) struct TaskHooks {
42    pub(crate) task_spawn_callback: Option<TaskCallback>,
43    pub(crate) task_terminate_callback: Option<TaskCallback>,
44    #[cfg(tokio_unstable)]
45    pub(crate) before_poll_callback: Option<TaskCallback>,
46    #[cfg(tokio_unstable)]
47    pub(crate) after_poll_callback: Option<TaskCallback>,
48}
49
50/// Task metadata supplied to user-provided hooks for task events.
51///
52/// **Note**: This is an [unstable API][unstable]. The public API of this type
53/// may break in 1.x releases. See [the documentation on unstable
54/// features][unstable] for details.
55///
56/// [unstable]: crate#unstable-features
57#[allow(missing_debug_implementations)]
58#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
59pub struct TaskMeta<'a> {
60    /// The opaque ID of the task.
61    pub(crate) id: super::task::Id,
62    /// The location where the task was spawned.
63    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
64    pub(crate) spawned_at: crate::runtime::task::SpawnLocation,
65    pub(crate) _phantom: PhantomData<&'a ()>,
66}
67
68impl<'a> TaskMeta<'a> {
69    /// Return the opaque ID of the task.
70    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
71    pub fn id(&self) -> super::task::Id {
72        self.id
73    }
74
75    /// Return the source code location where the task was spawned.
76    #[cfg(tokio_unstable)]
77    pub fn spawned_at(&self) -> &'static std::panic::Location<'static> {
78        self.spawned_at.0
79    }
80}
81
82/// Runs on specific task-related events
83pub(crate) type TaskCallback = std::sync::Arc<dyn Fn(&TaskMeta<'_>) + Send + Sync>;
tokio/runtime/task_hooks.rs

tokio/runtime/
task_hooks.rs
