tokio/runtime/
task_hooks.rs

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