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

add actions

Browse Source
This commit is contained in:
noah
2025-05-04 10:55:03 -05:00
parent 77eb27a9fb
commit 200cd93a9b
4 changed files with 182 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/runtime/scheduler/current_thread/mod.rs
+24 -16
View File
@@ -462,10 +462,12 @@ impl Handle {
parent parent
.map(|parent| { .map(|parent| {
if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| { if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
parent.on_child_spawn(&mut OnChildTaskSpawnContext { parent
id, .on_child_spawn(&mut OnChildTaskSpawnContext {
_phantom: Default::default(), id,
}) _phantom: Default::default(),
})
.hooks
})) { })) {
r r
} else { } else {
@@ -479,10 +481,12 @@ impl Handle {
.or_else(|| { .or_else(|| {
if let Some(hooks) = me.hooks_factory_ref() { if let Some(hooks) = me.hooks_factory_ref() {
if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| { if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
hooks.on_top_level_spawn(&mut OnTopLevelTaskSpawnContext { hooks
id, .on_top_level_spawn(&mut OnTopLevelTaskSpawnContext {
_phantom: Default::default(), id,
}) _phantom: Default::default(),
})
.hooks
})) { })) {
r r
} else { } else {
@@ -533,10 +537,12 @@ impl Handle {
parent parent
.map(|parent| { .map(|parent| {
if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| { if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
parent.on_child_spawn(&mut OnChildTaskSpawnContext { parent
id, .on_child_spawn(&mut OnChildTaskSpawnContext {
_phantom: Default::default(), id,
}) _phantom: Default::default(),
})
.hooks
})) { })) {
r r
} else { } else {
@@ -550,10 +556,12 @@ impl Handle {
.or_else(|| { .or_else(|| {
if let Some(hooks) = me.hooks_factory_ref() { if let Some(hooks) = me.hooks_factory_ref() {
if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| { if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
hooks.on_top_level_spawn(&mut OnTopLevelTaskSpawnContext { hooks
id, .on_top_level_spawn(&mut OnTopLevelTaskSpawnContext {
_phantom: Default::default(), id,
}) _phantom: Default::default(),
})
.hooks
})) { })) {
r r
} else { } else {
tokio/src/runtime/scheduler/multi_thread/handle.rs
+12 -8
View File
@@ -80,10 +80,12 @@ impl Handle {
parent parent
.map(|parent| { .map(|parent| {
if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| { if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
parent.on_child_spawn(&mut OnChildTaskSpawnContext { parent
id, .on_child_spawn(&mut OnChildTaskSpawnContext {
_phantom: Default::default(), id,
}) _phantom: Default::default(),
})
.hooks
})) { })) {
r r
} else { } else {
@@ -97,10 +99,12 @@ impl Handle {
.or_else(|| { .or_else(|| {
if let Some(hooks) = me.hooks_factory_ref() { if let Some(hooks) = me.hooks_factory_ref() {
if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| { if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
hooks.on_top_level_spawn(&mut OnTopLevelTaskSpawnContext { hooks
id, .on_top_level_spawn(&mut OnTopLevelTaskSpawnContext {
_phantom: Default::default(), id,
}) _phantom: Default::default(),
})
.hooks
})) { })) {
r r
} else { } else {
tokio/src/runtime/task_hooks/mod.rs
+70 -13
View File
@@ -8,11 +8,9 @@ use std::sync::Arc;
/// spawned in "detached mode" via [`crate::task::spawn_with_hooks`], or which were spawned from outside the runtime or /// spawned in "detached mode" via [`crate::task::spawn_with_hooks`], or which were spawned from outside the runtime or
/// from another context where no [`TaskHookHarness`] was present. /// from another context where no [`TaskHookHarness`] was present.
pub trait TaskHookHarnessFactory { pub trait TaskHookHarnessFactory {
/// Create a new [`TaskHookHarness`] object which the runtime will attach to a given task. /// Runs a hook which may produce a new [`TaskHookHarness`] object which the runtime will attach to a given task.
fn on_top_level_spawn( fn on_top_level_spawn(&self, ctx: &mut OnTopLevelTaskSpawnContext<'_>)
&self, -> OnTopLevelSpawnAction;
ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>>;
} }
/// Trait for user-provided "harness" objects which are attached to tasks and provide hook /// Trait for user-provided "harness" objects which are attached to tasks and provide hook
@@ -20,24 +18,27 @@ pub trait TaskHookHarnessFactory {
#[allow(unused_variables)] #[allow(unused_variables)]
pub trait TaskHookHarness { pub trait TaskHookHarness {
/// Pre-poll task hook which runs arbitrary user logic. /// Pre-poll task hook which runs arbitrary user logic.
fn before_poll(&mut self, ctx: &mut BeforeTaskPollContext<'_>) {} fn before_poll(&mut self, ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
BeforeTaskPollAction::default()
}
/// Post-poll task hook which runs arbitrary user logic. /// Post-poll task hook which runs arbitrary user logic.
fn after_poll(&mut self, ctx: &mut AfterTaskPollContext<'_>) {} fn after_poll(&mut self, ctx: &mut AfterTaskPollContext<'_>) -> AfterTaskPollAction {
AfterTaskPollAction::default()
}
/// Task hook which runs when this task spawns a child, unless that child is explicitly spawned /// Task hook which runs when this task spawns a child, unless that child is explicitly spawned
/// detached from the parent. /// detached from the parent.
/// ///
/// This hook creates a harness for the child, or detaches the child from any instrumentation. /// This hook creates a harness for the child, or detaches the child from any instrumentation.
fn on_child_spawn( fn on_child_spawn(&mut self, ctx: &mut OnChildTaskSpawnContext<'_>) -> OnChildSpawnAction {
&mut self, OnChildSpawnAction::default()
ctx: &mut OnChildTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> {
None
} }
/// Task hook which runs on task termination. /// Task hook which runs on task termination.
fn on_task_terminate(&mut self, ctx: &mut OnTaskTerminateContext<'_>) {} fn on_task_terminate(&mut self, ctx: &mut OnTaskTerminateContext<'_>) -> OnTaskTerminateAction {
OnTaskTerminateAction::default()
}
} }
pub(crate) type OptionalTaskHooksFactory = pub(crate) type OptionalTaskHooksFactory =
@@ -97,3 +98,59 @@ pub struct BeforeTaskPollContext<'a> {
pub struct AfterTaskPollContext<'a> { pub struct AfterTaskPollContext<'a> {
pub(crate) _phantom: PhantomData<&'a ()>, pub(crate) _phantom: PhantomData<&'a ()>,
} }
#[derive(Default)]
#[allow(missing_debug_implementations, missing_docs)]
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[non_exhaustive]
pub struct OnTopLevelSpawnAction {
pub(crate) hooks: Option<Box<dyn TaskHookHarness + Send + Sync + 'static>>,
}
impl OnTopLevelSpawnAction {
/// Pass in a set of task hooks for the task.
pub fn set_hooks<T>(&mut self, hooks: T) -> &mut Self
where
T: TaskHookHarness + Send + Sync + 'static,
{
self.hooks = Some(Box::new(hooks));
self
}
}
#[derive(Default)]
#[allow(missing_debug_implementations, missing_docs)]
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[non_exhaustive]
pub struct OnChildSpawnAction {
pub(crate) hooks: Option<Box<dyn TaskHookHarness + Send + Sync + 'static>>,
}
impl OnChildSpawnAction {
/// Pass in a set of task hooks for the child task.
pub fn set_hooks<T>(&mut self, hooks: T) -> &mut Self
where
T: TaskHookHarness + Send + Sync + 'static,
{
self.hooks = Some(Box::new(hooks));
self
}
}
#[derive(Default)]
#[allow(missing_debug_implementations, missing_docs)]
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[non_exhaustive]
pub struct OnTaskTerminateAction {}
#[derive(Default)]
#[allow(missing_debug_implementations, missing_docs)]
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[non_exhaustive]
pub struct BeforeTaskPollAction {}
#[derive(Default)]
#[allow(missing_debug_implementations, missing_docs)]
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[non_exhaustive]
pub struct AfterTaskPollAction {}
tokio/tests/task_hooks.rs
+76 -38
View File
@@ -9,8 +9,9 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc; use std::sync::Arc;
use tokio::runtime; use tokio::runtime;
use tokio::runtime::{ use tokio::runtime::{
AfterTaskPollContext, BeforeTaskPollContext, OnChildTaskSpawnContext, OnTaskTerminateContext, AfterTaskPollAction, AfterTaskPollContext, BeforeTaskPollAction, BeforeTaskPollContext,
OnTopLevelTaskSpawnContext, TaskHookHarness, TaskHookHarnessFactory, OnChildSpawnAction, OnChildTaskSpawnContext, OnTaskTerminateAction, OnTaskTerminateContext,
OnTopLevelSpawnAction, OnTopLevelTaskSpawnContext, TaskHookHarness, TaskHookHarnessFactory,
}; };
#[test] #[test]
@@ -83,9 +84,10 @@ fn run_runtime_default_factory(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
self.counter.fetch_add(1, Ordering::SeqCst); self.counter.fetch_add(1, Ordering::SeqCst);
None
Default::default()
} }
} }
@@ -138,25 +140,30 @@ fn run_parent_child_chaining(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
self.parent_spawns.fetch_add(1, Ordering::SeqCst); self.parent_spawns.fetch_add(1, Ordering::SeqCst);
Some(Box::new(TestHooks { let mut a = OnTopLevelSpawnAction::default();
a.set_hooks(TestHooks {
spawns: self.child_spawns.clone(), spawns: self.child_spawns.clone(),
})) });
a
} }
} }
impl TaskHookHarness for TestHooks { impl TaskHookHarness for TestHooks {
fn on_child_spawn( fn on_child_spawn(&mut self, _ctx: &mut OnChildTaskSpawnContext<'_>) -> OnChildSpawnAction {
&mut self,
_ctx: &mut OnChildTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> {
self.spawns.fetch_add(1, Ordering::SeqCst); self.spawns.fetch_add(1, Ordering::SeqCst);
Some(Box::new(Self { let mut a = OnChildSpawnAction::default();
a.set_hooks(Self {
spawns: self.spawns.clone(), spawns: self.spawns.clone(),
})) });
a
} }
} }
@@ -195,16 +202,22 @@ fn run_before_poll(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
Some(Box::new(TestHooks { let mut a = OnTopLevelSpawnAction::default();
a.set_hooks(TestHooks {
polls: self.polls.clone(), polls: self.polls.clone(),
})) });
a
} }
} }
impl TaskHookHarness for TestHooks { impl TaskHookHarness for TestHooks {
fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) { fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
self.polls.fetch_add(1, Ordering::SeqCst); self.polls.fetch_add(1, Ordering::SeqCst);
Default::default()
} }
} }
@@ -240,16 +253,22 @@ fn run_after_poll(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
Some(Box::new(TestHooks { let mut a = OnTopLevelSpawnAction::default();
a.set_hooks(TestHooks {
polls: self.polls.clone(), polls: self.polls.clone(),
})) });
a
} }
} }
impl TaskHookHarness for TestHooks { impl TaskHookHarness for TestHooks {
fn after_poll(&mut self, _ctx: &mut AfterTaskPollContext<'_>) { fn after_poll(&mut self, _ctx: &mut AfterTaskPollContext<'_>) -> AfterTaskPollAction {
self.polls.fetch_add(1, Ordering::SeqCst); self.polls.fetch_add(1, Ordering::SeqCst);
Default::default()
} }
} }
@@ -285,16 +304,25 @@ fn run_terminate(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
Some(Box::new(TestHooks { let mut a = OnTopLevelSpawnAction::default();
a.set_hooks(TestHooks {
terminations: self.terminations.clone(), terminations: self.terminations.clone(),
})) });
a
} }
} }
impl TaskHookHarness for TestHooks { impl TaskHookHarness for TestHooks {
fn on_task_terminate(&mut self, _ctx: &mut OnTaskTerminateContext<'_>) { fn on_task_terminate(
&mut self,
_ctx: &mut OnTaskTerminateContext<'_>,
) -> OnTaskTerminateAction {
self.terminations.fetch_add(1, Ordering::SeqCst); self.terminations.fetch_add(1, Ordering::SeqCst);
Default::default()
} }
} }
@@ -327,17 +355,23 @@ fn run_hook_switching(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
Some(Box::new(TestHooks { let mut a = OnTopLevelSpawnAction::default();
a.set_hooks(TestHooks {
id: self.next_id.fetch_add(1, Ordering::SeqCst), id: self.next_id.fetch_add(1, Ordering::SeqCst),
flag: self.flag.clone(), flag: self.flag.clone(),
})) });
a
} }
} }
impl TaskHookHarness for TestHooks { impl TaskHookHarness for TestHooks {
fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) { fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
self.flag.store(self.id, Ordering::SeqCst); self.flag.store(self.id, Ordering::SeqCst);
Default::default()
} }
} }
@@ -371,17 +405,20 @@ fn run_override(mut builder: runtime::Builder) {
} }
impl TaskHookHarness for TestHooks { impl TaskHookHarness for TestHooks {
fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) { fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
self.counter.fetch_add(1, Ordering::SeqCst); self.counter.fetch_add(1, Ordering::SeqCst);
Default::default()
} }
fn on_child_spawn( fn on_child_spawn(&mut self, _ctx: &mut OnChildTaskSpawnContext<'_>) -> OnChildSpawnAction {
&mut self, let mut a = OnChildSpawnAction::default();
_ctx: &mut OnChildTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { a.set_hooks(Self {
Some(Box::new(Self {
counter: self.counter.clone(), counter: self.counter.clone(),
})) });
a
} }
} }
@@ -389,9 +426,10 @@ fn run_override(mut builder: runtime::Builder) {
fn on_top_level_spawn( fn on_top_level_spawn(
&self, &self,
_ctx: &mut OnTopLevelTaskSpawnContext<'_>, _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
) -> Option<Box<dyn TaskHookHarness + Send + Sync + 'static>> { ) -> OnTopLevelSpawnAction {
self.counter.fetch_add(1, Ordering::SeqCst); self.counter.fetch_add(1, Ordering::SeqCst);
None
Default::default()
} }
} }