From b8ac94ed70df22f885bad7ea3c0ff51c536bad4a Mon Sep 17 00:00:00 2001 From: Jason Gin <67525213+GJason88@users.noreply.github.com> Date: Thu, 30 Jan 2025 22:14:00 +0100 Subject: [PATCH] rt: add before and after task poll callbacks (#7120) Add callbacks for poll start and stop, enabling users to instrument these points in the runtime's life cycle. --- tokio/src/runtime/builder.rs | 111 +++++++++++++++ tokio/src/runtime/config.rs | 8 ++ .../runtime/scheduler/current_thread/mod.rs | 13 ++ .../runtime/scheduler/multi_thread/worker.rs | 27 +++- .../scheduler/multi_thread_alt/worker.rs | 5 +- tokio/src/runtime/task/mod.rs | 32 +++-- tokio/src/runtime/task_hooks.rs | 40 ++++++ tokio/tests/rt_poll_callbacks.rs | 128 ++++++++++++++++++ 8 files changed, 347 insertions(+), 17 deletions(-) create mode 100644 tokio/tests/rt_poll_callbacks.rs diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index c9a47c386..11538a098 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -88,6 +88,14 @@ pub struct Builder { /// To run before each task is spawned. pub(super) before_spawn: Option, + /// To run before each poll + #[cfg(tokio_unstable)] + pub(super) before_poll: Option, + + /// To run after each poll + #[cfg(tokio_unstable)] + pub(super) after_poll: Option, + /// To run after each task is terminated. pub(super) after_termination: Option, @@ -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(&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(&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, diff --git a/tokio/src/runtime/config.rs b/tokio/src/runtime/config.rs index eb4bf81aa..43ce5aebd 100644 --- a/tokio/src/runtime/config.rs +++ b/tokio/src/runtime/config.rs @@ -27,6 +27,14 @@ pub(crate) struct Config { /// To run after each task is terminated. pub(crate) after_termination: Option, + /// To run before each poll + #[cfg(tokio_unstable)] + pub(crate) before_poll: Option, + + /// To run after each poll + #[cfg(tokio_unstable)] + pub(crate) after_poll: Option, + /// 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 diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index c66635e7b..37f37a4e9 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -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; diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index ec15106fe..8866ea54b 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -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) -> 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); } }) } diff --git a/tokio/src/runtime/scheduler/multi_thread_alt/worker.rs b/tokio/src/runtime/scheduler/multi_thread_alt/worker.rs index d88eb5e89..206c9855b 100644 --- a/tokio/src/runtime/scheduler/multi_thread_alt/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread_alt/worker.rs @@ -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, diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 15c5a8f4a..7d314c3b1 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -256,6 +256,13 @@ pub(crate) struct LocalNotified { _not_send: PhantomData<*const ()>, } +impl LocalNotified { + #[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 { @@ -386,6 +393,16 @@ impl Task { 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 Task { } } - /// 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 Notified { 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 Notified { diff --git a/tokio/src/runtime/task_hooks.rs b/tokio/src/runtime/task_hooks.rs index 2c884af74..13865ed51 100644 --- a/tokio/src/runtime/task_hooks.rs +++ b/tokio/src/runtime/task_hooks.rs @@ -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, pub(crate) task_terminate_callback: Option, + #[cfg(tokio_unstable)] + pub(crate) before_poll_callback: Option, + #[cfg(tokio_unstable)] + pub(crate) after_poll_callback: Option, } /// Task metadata supplied to user-provided hooks for task events. diff --git a/tokio/tests/rt_poll_callbacks.rs b/tokio/tests/rt_poll_callbacks.rs new file mode 100644 index 000000000..8ccff3857 --- /dev/null +++ b/tokio/tests/rt_poll_callbacks.rs @@ -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>> = + Arc::new(Mutex::new(None)); + let after_task_poll_callback_task_id: Arc>> = + 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>> = + Arc::new(Mutex::new(None)); + let after_task_poll_callback_task_id: Arc>> = + 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); +}