Skip to content

Commit b37f0de

Browse files
Noah-KennedyNoah-Kennedy
authored
runtime: implement initial set of task hooks (#6742)
1 parent c9fad08 commit b37f0de

File tree

19 files changed

+384
-16
lines changed

19 files changed

+384
-16
lines changed

‎tokio/src/runtime/blocking/schedule.rs

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#[cfg(feature = "test-util")]
22
use crate::runtime::scheduler;
3-
use crate::runtime::task::{self, Task};
3+
use crate::runtime::task::{self, Task, TaskHarnessScheduleHooks};
44
use crate::runtime::Handle;
55

66
/// `task::Schedule` implementation that does nothing (except some bookkeeping
@@ -12,6 +12,7 @@ use crate::runtime::Handle;
1212
pub(crate) struct BlockingSchedule {
1313
#[cfg(feature = "test-util")]
1414
handle: Handle,
15+
hooks: TaskHarnessScheduleHooks,
1516
}
1617

1718
impl BlockingSchedule {
@@ -32,6 +33,9 @@ impl BlockingSchedule {
3233
BlockingSchedule {
3334
#[cfg(feature = "test-util")]
3435
handle: handle.clone(),
36+
hooks: TaskHarnessScheduleHooks {
37+
task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(),
38+
},
3539
}
3640
}
3741
}
@@ -57,4 +61,10 @@ impl task::Schedule for BlockingSchedule {
5761
fn schedule(&self, _task: task::Notified<Self>) {
5862
unreachable!();
5963
}
64+
65+
fn hooks(&self) -> TaskHarnessScheduleHooks {
66+
TaskHarnessScheduleHooks {
67+
task_terminate_callback: self.hooks.task_terminate_callback.clone(),
68+
}
69+
}
6070
}

‎tokio/src/runtime/builder.rs

Lines changed: 105 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
1+
#![cfg_attr(loom, allow(unused_imports))]
2+
13
use crate::runtime::handle::Handle;
2-
use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime};
4+
#[cfg(tokio_unstable)]
5+
use crate::runtime::TaskMeta;
6+
use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime, TaskCallback};
37
use crate::util::rand::{RngSeed, RngSeedGenerator};
48

59
use std::fmt;
@@ -78,6 +82,12 @@ pub struct Builder {
7882
/// To run after each thread is unparked.
7983
pub(super) after_unpark: Option<Callback>,
8084

85+
/// To run before each task is spawned.
86+
pub(super) before_spawn: Option<TaskCallback>,
87+
88+
/// To run after each task is terminated.
89+
pub(super) after_termination: Option<TaskCallback>,
90+
8191
/// Customizable keep alive timeout for `BlockingPool`
8292
pub(super) keep_alive: Option<Duration>,
8393

@@ -290,6 +300,9 @@ impl Builder {
290300
before_park: None,
291301
after_unpark: None,
292302

303+
before_spawn: None,
304+
after_termination: None,
305+
293306
keep_alive: None,
294307

295308
// Defaults for these values depend on the scheduler kind, so we get them
@@ -677,6 +690,91 @@ impl Builder {
677690
self
678691
}
679692

693+
/// Executes function `f` just before a task is spawned.
694+
///
695+
/// `f` is called within the Tokio context, so functions like
696+
/// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being
697+
/// invoked immediately.
698+
///
699+
/// This can be used for bookkeeping or monitoring purposes.
700+
///
701+
/// Note: There can only be one spawn callback for a runtime; calling this function more
702+
/// than once replaces the last callback defined, rather than adding to it.
703+
///
704+
/// This *does not* support [`LocalSet`](crate::task::LocalSet) at this time.
705+
///
706+
/// # Examples
707+
///
708+
/// ```
709+
/// # use tokio::runtime;
710+
/// # pub fn main() {
711+
/// let runtime = runtime::Builder::new_current_thread()
712+
/// .on_task_spawn(|_| {
713+
/// println!("spawning task");
714+
/// })
715+
/// .build()
716+
/// .unwrap();
717+
///
718+
/// runtime.block_on(async {
719+
/// tokio::task::spawn(std::future::ready(()));
720+
///
721+
/// for _ in 0..64 {
722+
/// tokio::task::yield_now().await;
723+
/// }
724+
/// })
725+
/// # }
726+
/// ```
727+
#[cfg(all(not(loom), tokio_unstable))]
728+
pub fn on_task_spawn<F>(&mut self, f: F) -> &mut Self
729+
where
730+
F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
731+
{
732+
self.before_spawn = Some(std::sync::Arc::new(f));
733+
self
734+
}
735+
736+
/// Executes function `f` just after a task is terminated.
737+
///
738+
/// `f` is called within the Tokio context, so functions like
739+
/// [`tokio::spawn`](crate::spawn) can be called.
740+
///
741+
/// This can be used for bookkeeping or monitoring purposes.
742+
///
743+
/// Note: There can only be one task termination callback for a runtime; calling this
744+
/// function more than once replaces the last callback defined, rather than adding to it.
745+
///
746+
/// This *does not* support [`LocalSet`](crate::task::LocalSet) at this time.
747+
///
748+
/// # Examples
749+
///
750+
/// ```
751+
/// # use tokio::runtime;
752+
/// # pub fn main() {
753+
/// let runtime = runtime::Builder::new_current_thread()
754+
/// .on_task_terminate(|_| {
755+
/// println!("killing task");
756+
/// })
757+
/// .build()
758+
/// .unwrap();
759+
///
760+
/// runtime.block_on(async {
761+
/// tokio::task::spawn(std::future::ready(()));
762+
///
763+
/// for _ in 0..64 {
764+
/// tokio::task::yield_now().await;
765+
/// }
766+
/// })
767+
/// # }
768+
/// ```
769+
#[cfg(all(not(loom), tokio_unstable))]
770+
pub fn on_task_terminate<F>(&mut self, f: F) -> &mut Self
771+
where
772+
F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
773+
{
774+
self.after_termination = Some(std::sync::Arc::new(f));
775+
self
776+
}
777+
680778
/// Creates the configured `Runtime`.
681779
///
682780
/// The returned `Runtime` instance is ready to spawn tasks.
@@ -1118,6 +1216,8 @@ impl Builder {
11181216
Config {
11191217
before_park: self.before_park.clone(),
11201218
after_unpark: self.after_unpark.clone(),
1219+
before_spawn: self.before_spawn.clone(),
1220+
after_termination: self.after_termination.clone(),
11211221
global_queue_interval: self.global_queue_interval,
11221222
event_interval: self.event_interval,
11231223
local_queue_capacity: self.local_queue_capacity,
@@ -1269,6 +1369,8 @@ cfg_rt_multi_thread! {
12691369
Config {
12701370
before_park: self.before_park.clone(),
12711371
after_unpark: self.after_unpark.clone(),
1372+
before_spawn: self.before_spawn.clone(),
1373+
after_termination: self.after_termination.clone(),
12721374
global_queue_interval: self.global_queue_interval,
12731375
event_interval: self.event_interval,
12741376
local_queue_capacity: self.local_queue_capacity,
@@ -1316,6 +1418,8 @@ cfg_rt_multi_thread! {
13161418
Config {
13171419
before_park: self.before_park.clone(),
13181420
after_unpark: self.after_unpark.clone(),
1421+
before_spawn: self.before_spawn.clone(),
1422+
after_termination: self.after_termination.clone(),
13191423
global_queue_interval: self.global_queue_interval,
13201424
event_interval: self.event_interval,
13211425
local_queue_capacity: self.local_queue_capacity,

‎tokio/src/runtime/config.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"),
33
allow(dead_code)
44
)]
5-
use crate::runtime::Callback;
5+
use crate::runtime::{Callback, TaskCallback};
66
use crate::util::RngSeedGenerator;
77

88
pub(crate) struct Config {
@@ -21,6 +21,12 @@ pub(crate) struct Config {
2121
/// Callback for a worker unparking itself
2222
pub(crate) after_unpark: Option<Callback>,
2323

24+
/// To run before each task is spawned.
25+
pub(crate) before_spawn: Option<TaskCallback>,
26+
27+
/// To run after each task is terminated.
28+
pub(crate) after_termination: Option<TaskCallback>,
29+
2430
/// The multi-threaded scheduler includes a per-worker LIFO slot used to
2531
/// store the last scheduled task. This can improve certain usage patterns,
2632
/// especially message passing between tasks. However, this LIFO slot is not

‎tokio/src/runtime/mod.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,13 @@ cfg_rt! {
379379
pub use dump::Dump;
380380
}
381381

382+
mod task_hooks;
383+
pub(crate) use task_hooks::{TaskHooks, TaskCallback};
384+
#[cfg(tokio_unstable)]
385+
pub use task_hooks::TaskMeta;
386+
#[cfg(not(tokio_unstable))]
387+
pub(crate) use task_hooks::TaskMeta;
388+
382389
mod handle;
383390
pub use handle::{EnterGuard, Handle, TryCurrentError};
384391

‎tokio/src/runtime/scheduler/current_thread/mod.rs

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@ use crate::loom::sync::atomic::AtomicBool;
33
use crate::loom::sync::Arc;
44
use crate::runtime::driver::{self, Driver};
55
use crate::runtime::scheduler::{self, Defer, Inject};
6-
use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task};
7-
use crate::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics};
6+
use crate::runtime::task::{
7+
self, JoinHandle, OwnedTasks, Schedule, Task, TaskHarnessScheduleHooks,
8+
};
9+
use crate::runtime::{
10+
blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics,
11+
};
812
use crate::sync::notify::Notify;
913
use crate::util::atomic_cell::AtomicCell;
1014
use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
@@ -41,6 +45,9 @@ pub(crate) struct Handle {
4145

4246
/// Current random number generator seed
4347
pub(crate) seed_generator: RngSeedGenerator,
48+
49+
/// User-supplied hooks to invoke for things
50+
pub(crate) task_hooks: TaskHooks,
4451
}
4552

4653
/// Data required for executing the scheduler. The struct is passed around to
@@ -131,6 +138,10 @@ impl CurrentThread {
131138
.unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL);
132139

133140
let handle = Arc::new(Handle {
141+
task_hooks: TaskHooks {
142+
task_spawn_callback: config.before_spawn.clone(),
143+
task_terminate_callback: config.after_termination.clone(),
144+
},
134145
shared: Shared {
135146
inject: Inject::new(),
136147
owned: OwnedTasks::new(1),
@@ -436,6 +447,12 @@ impl Handle {
436447
{
437448
let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
438449

450+
me.task_hooks.spawn(&TaskMeta {
451+
#[cfg(tokio_unstable)]
452+
id,
453+
_phantom: Default::default(),
454+
});
455+
439456
if let Some(notified) = notified {
440457
me.schedule(notified);
441458
}
@@ -600,6 +617,12 @@ impl Schedule for Arc<Handle> {
600617
});
601618
}
602619

620+
fn hooks(&self) -> TaskHarnessScheduleHooks {
621+
TaskHarnessScheduleHooks {
622+
task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
623+
}
624+
}
625+
603626
cfg_unstable! {
604627
fn unhandled_panic(&self) {
605628
use crate::runtime::UnhandledPanic;

‎tokio/src/runtime/scheduler/mod.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ cfg_rt! {
77

88
pub(crate) mod inject;
99
pub(crate) use inject::Inject;
10+
11+
use crate::runtime::TaskHooks;
1012
}
1113

1214
cfg_rt_multi_thread! {
@@ -151,6 +153,16 @@ cfg_rt! {
151153
}
152154
}
153155

156+
pub(crate) fn hooks(&self) -> &TaskHooks {
157+
match self {
158+
Handle::CurrentThread(h) => &h.task_hooks,
159+
#[cfg(feature = "rt-multi-thread")]
160+
Handle::MultiThread(h) => &h.task_hooks,
161+
#[cfg(all(tokio_unstable, feature = "rt-multi-thread"))]
162+
Handle::MultiThreadAlt(h) => &h.task_hooks,
163+
}
164+
}
165+
154166
cfg_rt_multi_thread! {
155167
cfg_unstable! {
156168
pub(crate) fn expect_multi_thread_alt(&self) -> &Arc<multi_thread_alt::Handle> {

‎tokio/src/runtime/scheduler/multi_thread/handle.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ use crate::runtime::scheduler::multi_thread::worker;
44
use crate::runtime::{
55
blocking, driver,
66
task::{self, JoinHandle},
7+
TaskHooks, TaskMeta,
78
};
89
use crate::util::RngSeedGenerator;
910

@@ -28,6 +29,9 @@ pub(crate) struct Handle {
2829

2930
/// Current random number generator seed
3031
pub(crate) seed_generator: RngSeedGenerator,
32+
33+
/// User-supplied hooks to invoke for things
34+
pub(crate) task_hooks: TaskHooks,
3135
}
3236

3337
impl Handle {
@@ -51,6 +55,12 @@ impl Handle {
5155
{
5256
let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
5357

58+
me.task_hooks.spawn(&TaskMeta {
59+
#[cfg(tokio_unstable)]
60+
id,
61+
_phantom: Default::default(),
62+
});
63+
5464
me.schedule_option_task_without_yield(notified);
5565

5666
handle

‎tokio/src/runtime/scheduler/multi_thread/worker.rs

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,15 +58,15 @@
5858
5959
use crate::loom::sync::{Arc, Mutex};
6060
use crate::runtime;
61-
use crate::runtime::context;
6261
use crate::runtime::scheduler::multi_thread::{
6362
idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
6463
};
6564
use crate::runtime::scheduler::{inject, Defer, Lock};
66-
use crate::runtime::task::OwnedTasks;
65+
use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks};
6766
use crate::runtime::{
6867
blocking, coop, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics,
6968
};
69+
use crate::runtime::{context, TaskHooks};
7070
use crate::util::atomic_cell::AtomicCell;
7171
use crate::util::rand::{FastRand, RngSeedGenerator};
7272

@@ -284,6 +284,10 @@ pub(super) fn create(
284284

285285
let remotes_len = remotes.len();
286286
let handle = Arc::new(Handle {
287+
task_hooks: TaskHooks {
288+
task_spawn_callback: config.before_spawn.clone(),
289+
task_terminate_callback: config.after_termination.clone(),
290+
},
287291
shared: Shared {
288292
remotes: remotes.into_boxed_slice(),
289293
inject,
@@ -1037,6 +1041,12 @@ impl task::Schedule for Arc<Handle> {
10371041
self.schedule_task(task, false);
10381042
}
10391043

1044+
fn hooks(&self) -> TaskHarnessScheduleHooks {
1045+
TaskHarnessScheduleHooks {
1046+
task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
1047+
}
1048+
}
1049+
10401050
fn yield_now(&self, task: Notified) {
10411051
self.schedule_task(task, true);
10421052
}

0 commit comments

Comments
 (0)