tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-20 00:00:08 +02:00
Code Issues Packages Projects Releases Wiki Activity

rt: add before and after task poll callbacks (#7120)

Browse Source 
Add callbacks for poll start and stop, enabling users to instrument these points
in the runtime's life cycle.
This commit is contained in:
Jason Gin
2025-01-30 21:14:00 +00:00
committed by GitHub
parent 5086e56dcb
commit b8ac94ed70
8 changed files with 347 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/builder.rs
+111
View File
@@ -88,6 +88,14 @@ pub struct Builder {
/// To run before each task is spawned.
pub(super) before_spawn: Option<TaskCallback>,
/// To run before each poll
#[cfg(tokio_unstable)]
pub(super) before_poll: Option<TaskCallback>,
/// To run after each poll
#[cfg(tokio_unstable)]
pub(super) after_poll: Option<TaskCallback>,
/// To run after each task is terminated.
pub(super) after_termination: Option<TaskCallback>,
@@ -306,6 +314,11 @@ impl Builder {
before_spawn: None,
after_termination: None,
#[cfg(tokio_unstable)]
before_poll: None,
#[cfg(tokio_unstable)]
after_poll: None,
keep_alive: None,
// Defaults for these values depend on the scheduler kind, so we get them
@@ -743,6 +756,92 @@ impl Builder {
self
}
/// Executes function `f` just before a task is polled
///
/// `f` is called within the Tokio context, so functions like
/// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being
/// invoked immediately.
///
/// **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.
///
/// [unstable]: crate#unstable-features
///
/// # Examples
///
/// ```
/// # use std::sync::{atomic::AtomicUsize, Arc};
/// # use tokio::task::yield_now;
/// # pub fn main() {
/// let poll_start_counter = Arc::new(AtomicUsize::new(0));
/// let poll_start = poll_start_counter.clone();
/// let rt = tokio::runtime::Builder::new_multi_thread()
/// .enable_all()
/// .on_before_task_poll(move |meta| {
/// println!("task {} is about to be polled", meta.id())
/// })
/// .build()
/// .unwrap();
/// let task = rt.spawn(async {
/// yield_now().await;
/// });
/// let _ = rt.block_on(task);
///
/// # }
/// ```
#[cfg(tokio_unstable)]
pub fn on_before_task_poll<F>(&mut self, f: F) -> &mut Self
where
F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
{
self.before_poll = Some(std::sync::Arc::new(f));
self
}
/// Executes function `f` just after a task is polled
///
/// `f` is called within the Tokio context, so functions like
/// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being
/// invoked immediately.
///
/// **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.
///
/// [unstable]: crate#unstable-features
///
/// # Examples
///
/// ```
/// # use std::sync::{atomic::AtomicUsize, Arc};
/// # use tokio::task::yield_now;
/// # pub fn main() {
/// let poll_stop_counter = Arc::new(AtomicUsize::new(0));
/// let poll_stop = poll_stop_counter.clone();
/// let rt = tokio::runtime::Builder::new_multi_thread()
/// .enable_all()
/// .on_after_task_poll(move |meta| {
/// println!("task {} completed polling", meta.id());
/// })
/// .build()
/// .unwrap();
/// let task = rt.spawn(async {
/// yield_now().await;
/// });
/// let _ = rt.block_on(task);
///
/// # }
/// ```
#[cfg(tokio_unstable)]
pub fn on_after_task_poll<F>(&mut self, f: F) -> &mut Self
where
F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
{
self.after_poll = Some(std::sync::Arc::new(f));
self
}
/// Executes function `f` just after a task is terminated.
///
/// `f` is called within the Tokio context, so functions like
@@ -1410,6 +1509,10 @@ impl Builder {
before_park: self.before_park.clone(),
after_unpark: self.after_unpark.clone(),
before_spawn: self.before_spawn.clone(),
#[cfg(tokio_unstable)]
before_poll: self.before_poll.clone(),
#[cfg(tokio_unstable)]
after_poll: self.after_poll.clone(),
after_termination: self.after_termination.clone(),
global_queue_interval: self.global_queue_interval,
event_interval: self.event_interval,
@@ -1560,6 +1663,10 @@ cfg_rt_multi_thread! {
before_park: self.before_park.clone(),
after_unpark: self.after_unpark.clone(),
before_spawn: self.before_spawn.clone(),
#[cfg(tokio_unstable)]
before_poll: self.before_poll.clone(),
#[cfg(tokio_unstable)]
after_poll: self.after_poll.clone(),
after_termination: self.after_termination.clone(),
global_queue_interval: self.global_queue_interval,
event_interval: self.event_interval,
@@ -1610,6 +1717,10 @@ cfg_rt_multi_thread! {
after_unpark: self.after_unpark.clone(),
before_spawn: self.before_spawn.clone(),
after_termination: self.after_termination.clone(),
#[cfg(tokio_unstable)]
before_poll: self.before_poll.clone(),
#[cfg(tokio_unstable)]
after_poll: self.after_poll.clone(),
global_queue_interval: self.global_queue_interval,
event_interval: self.event_interval,
local_queue_capacity: self.local_queue_capacity,
tokio/src/runtime/config.rs
+8
View File
@@ -27,6 +27,14 @@ pub(crate) struct Config {
/// To run after each task is terminated.
pub(crate) after_termination: Option<TaskCallback>,
/// To run before each poll
#[cfg(tokio_unstable)]
pub(crate) before_poll: Option<TaskCallback>,
/// To run after each poll
#[cfg(tokio_unstable)]
pub(crate) after_poll: Option<TaskCallback>,
/// The multi-threaded scheduler includes a per-worker LIFO slot used to
/// store the last scheduled task. This can improve certain usage patterns,
/// especially message passing between tasks. However, this LIFO slot is not
tokio/src/runtime/scheduler/current_thread/mod.rs
+13
View File
@@ -145,6 +145,10 @@ impl CurrentThread {
task_hooks: TaskHooks {
task_spawn_callback: config.before_spawn.clone(),
task_terminate_callback: config.after_termination.clone(),
#[cfg(tokio_unstable)]
before_poll_callback: config.before_poll.clone(),
#[cfg(tokio_unstable)]
after_poll_callback: config.after_poll.clone(),
},
shared: Shared {
inject: Inject::new(),
@@ -766,8 +770,17 @@ impl CoreGuard<'_> {
let task = context.handle.shared.owned.assert_owner(task);
#[cfg(tokio_unstable)]
let task_id = task.task_id();
let (c, ()) = context.run_task(core, || {
#[cfg(tokio_unstable)]
context.handle.task_hooks.poll_start_callback(task_id);
task.run();
#[cfg(tokio_unstable)]
context.handle.task_hooks.poll_stop_callback(task_id);
});
core = c;
tokio/src/runtime/scheduler/multi_thread/worker.rs
+23 -4
View File
@@ -282,10 +282,7 @@ pub(super) fn create(
let remotes_len = remotes.len();
let handle = Arc::new(Handle {
task_hooks: TaskHooks {
task_spawn_callback: config.before_spawn.clone(),
task_terminate_callback: config.after_termination.clone(),
},
task_hooks: TaskHooks::from_config(&config),
shared: Shared {
remotes: remotes.into_boxed_slice(),
inject,
@@ -574,6 +571,9 @@ impl Context {
}
fn run_task(&self, task: Notified, mut core: Box<Core>) -> RunResult {
#[cfg(tokio_unstable)]
let task_id = task.task_id();
let task = self.worker.handle.shared.owned.assert_owner(task);
// Make sure the worker is not in the **searching** state. This enables
@@ -593,7 +593,16 @@ impl Context {
// Run the task
coop::budget(|| {
// Unlike the poll time above, poll start callback is attached to the task id,
// so it is tightly associated with the actual poll invocation.
#[cfg(tokio_unstable)]
self.worker.handle.task_hooks.poll_start_callback(task_id);
task.run();
#[cfg(tokio_unstable)]
self.worker.handle.task_hooks.poll_stop_callback(task_id);
let mut lifo_polls = 0;
// As long as there is budget remaining and a task exists in the
@@ -656,7 +665,17 @@ impl Context {
// Run the LIFO task, then loop
*self.core.borrow_mut() = Some(core);
let task = self.worker.handle.shared.owned.assert_owner(task);
#[cfg(tokio_unstable)]
let task_id = task.task_id();
#[cfg(tokio_unstable)]
self.worker.handle.task_hooks.poll_start_callback(task_id);
task.run();
#[cfg(tokio_unstable)]
self.worker.handle.task_hooks.poll_stop_callback(task_id);
}
})
}
tokio/src/runtime/scheduler/multi_thread_alt/worker.rs
+1 -4
View File
@@ -303,10 +303,7 @@ pub(super) fn create(
let (inject, inject_synced) = inject::Shared::new();
let handle = Arc::new(Handle {
task_hooks: TaskHooks {
task_spawn_callback: config.before_spawn.clone(),
task_terminate_callback: config.after_termination.clone(),
},
task_hooks: TaskHooks::from_config(&config),
shared: Shared {
remotes: remotes.into_boxed_slice(),
inject,
tokio/src/runtime/task/mod.rs
+23 -9
View File
@@ -256,6 +256,13 @@ pub(crate) struct LocalNotified<S: 'static> {
_not_send: PhantomData<*const ()>,
}
impl<S> LocalNotified<S> {
#[cfg(tokio_unstable)]
pub(crate) fn task_id(&self) -> Id {
self.task.id()
}
}
/// A task that is not owned by any `OwnedTasks`. Used for blocking tasks.
/// This type holds two ref-counts.
pub(crate) struct UnownedTask<S: 'static> {
@@ -386,6 +393,16 @@ impl<S: 'static> Task<S> {
self.raw.header_ptr()
}
/// Returns a [task ID] that uniquely identifies this task relative to other
/// currently spawned tasks.
///
/// [task ID]: crate::task::Id
#[cfg(tokio_unstable)]
pub(crate) fn id(&self) -> crate::task::Id {
// Safety: The header pointer is valid.
unsafe { Header::get_id(self.raw.header_ptr()) }
}
cfg_taskdump! {
/// Notify the task for task dumping.
///
@@ -400,15 +417,6 @@ impl<S: 'static> Task<S> {
}
}
/// Returns a [task ID] that uniquely identifies this task relative to other
/// currently spawned tasks.
///
/// [task ID]: crate::task::Id
#[cfg(tokio_unstable)]
pub(crate) fn id(&self) -> crate::task::Id {
// Safety: The header pointer is valid.
unsafe { Header::get_id(self.raw.header_ptr()) }
}
}
}
@@ -416,6 +424,12 @@ impl<S: 'static> Notified<S> {
fn header(&self) -> &Header {
self.0.header()
}
#[cfg(tokio_unstable)]
#[allow(dead_code)]
pub(crate) fn task_id(&self) -> crate::task::Id {
self.0.id()
}
}
impl<S: 'static> Notified<S> {
tokio/src/runtime/task_hooks.rs
+40
View File
@@ -1,17 +1,57 @@
use std::marker::PhantomData;
use super::Config;
impl TaskHooks {
pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) {
if let Some(f) = self.task_spawn_callback.as_ref() {
f(meta)
}
}
#[allow(dead_code)]
pub(crate) fn from_config(config: &Config) -> Self {
Self {
task_spawn_callback: config.before_spawn.clone(),
task_terminate_callback: config.after_termination.clone(),
#[cfg(tokio_unstable)]
before_poll_callback: config.before_poll.clone(),
#[cfg(tokio_unstable)]
after_poll_callback: config.after_poll.clone(),
}
}
#[cfg(tokio_unstable)]
#[inline]
pub(crate) fn poll_start_callback(&self, id: super::task::Id) {
if let Some(poll_start) = &self.before_poll_callback {
(poll_start)(&TaskMeta {
id,
_phantom: std::marker::PhantomData,
})
}
}
#[cfg(tokio_unstable)]
#[inline]
pub(crate) fn poll_stop_callback(&self, id: super::task::Id) {
if let Some(poll_stop) = &self.after_poll_callback {
(poll_stop)(&TaskMeta {
id,
_phantom: std::marker::PhantomData,
})
}
}
}
#[derive(Clone)]
pub(crate) struct TaskHooks {
pub(crate) task_spawn_callback: Option<TaskCallback>,
pub(crate) task_terminate_callback: Option<TaskCallback>,
#[cfg(tokio_unstable)]
pub(crate) before_poll_callback: Option<TaskCallback>,
#[cfg(tokio_unstable)]
pub(crate) after_poll_callback: Option<TaskCallback>,
}
/// Task metadata supplied to user-provided hooks for task events.
tokio/tests/rt_poll_callbacks.rs
+128
View File
@@ -0,0 +1,128 @@
#![allow(unknown_lints, unexpected_cfgs)]
#![cfg(tokio_unstable)]
use std::sync::{atomic::AtomicUsize, Arc, Mutex};
use tokio::task::yield_now;
#[cfg(not(target_os = "wasi"))]
#[test]
fn callbacks_fire_multi_thread() {
let poll_start_counter = Arc::new(AtomicUsize::new(0));
let poll_stop_counter = Arc::new(AtomicUsize::new(0));
let poll_start = poll_start_counter.clone();
let poll_stop = poll_stop_counter.clone();
let before_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
Arc::new(Mutex::new(None));
let after_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
Arc::new(Mutex::new(None));
let before_task_poll_callback_task_id_ref = Arc::clone(&before_task_poll_callback_task_id);
let after_task_poll_callback_task_id_ref = Arc::clone(&after_task_poll_callback_task_id);
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.on_before_task_poll(move |task_meta| {
before_task_poll_callback_task_id_ref
.lock()
.unwrap()
.replace(task_meta.id());
poll_start_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
})
.on_after_task_poll(move |task_meta| {
after_task_poll_callback_task_id_ref
.lock()
.unwrap()
.replace(task_meta.id());
poll_stop_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
})
.build()
.unwrap();
let task = rt.spawn(async {
yield_now().await;
yield_now().await;
yield_now().await;
});
let spawned_task_id = task.id();
rt.block_on(task).expect("task should succeed");
// We need to drop the runtime to guarantee the workers have exited (and thus called the callback)
drop(rt);
assert_eq!(
before_task_poll_callback_task_id.lock().unwrap().unwrap(),
spawned_task_id
);
assert_eq!(
after_task_poll_callback_task_id.lock().unwrap().unwrap(),
spawned_task_id
);
let actual_count = 4;
assert_eq!(
poll_start.load(std::sync::atomic::Ordering::Relaxed),
actual_count,
"unexpected number of poll starts"
);
assert_eq!(
poll_stop.load(std::sync::atomic::Ordering::Relaxed),
actual_count,
"unexpected number of poll stops"
);
}
#[test]
fn callbacks_fire_current_thread() {
let poll_start_counter = Arc::new(AtomicUsize::new(0));
let poll_stop_counter = Arc::new(AtomicUsize::new(0));
let poll_start = poll_start_counter.clone();
let poll_stop = poll_stop_counter.clone();
let before_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
Arc::new(Mutex::new(None));
let after_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
Arc::new(Mutex::new(None));
let before_task_poll_callback_task_id_ref = Arc::clone(&before_task_poll_callback_task_id);
let after_task_poll_callback_task_id_ref = Arc::clone(&after_task_poll_callback_task_id);
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.on_before_task_poll(move |task_meta| {
before_task_poll_callback_task_id_ref
.lock()
.unwrap()
.replace(task_meta.id());
poll_start_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
})
.on_after_task_poll(move |task_meta| {
after_task_poll_callback_task_id_ref
.lock()
.unwrap()
.replace(task_meta.id());
poll_stop_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
})
.build()
.unwrap();
let task = rt.spawn(async {
yield_now().await;
yield_now().await;
yield_now().await;
});
let spawned_task_id = task.id();
let _ = rt.block_on(task);
drop(rt);
assert_eq!(
before_task_poll_callback_task_id.lock().unwrap().unwrap(),
spawned_task_id
);
assert_eq!(
after_task_poll_callback_task_id.lock().unwrap().unwrap(),
spawned_task_id
);
assert_eq!(poll_start.load(std::sync::atomic::Ordering::Relaxed), 4);
assert_eq!(poll_stop.load(std::sync::atomic::Ordering::Relaxed), 4);
}