- Sponsor
-
Notifications
You must be signed in to change notification settings - Fork 2.6k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Loading status checks…
rt: implement initial set of task hooks
This change implements two hooks for per-task actions, one which is invoked on task spawn, and one which is invoked during task termination. These hooks initially are only supplied with the task ID (on unstable only), but more information can be added in the future, as the struct used to supply parameters is opaque. Fixes #3181.
1 parent
1077b0b
commit de28721
Showing
18 changed files
with
358 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,35 @@ | ||
| use std::marker::PhantomData; | ||
|
|
||
| #[derive(Clone)] | ||
| pub(crate) struct TaskHooks { | ||
| pub(crate) task_spawn_callback: Option<TaskCallback>, | ||
| pub(crate) task_terminate_callback: Option<TaskCallback>, | ||
| } | ||
|
|
||
| impl TaskHooks { | ||
| pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) { | ||
| if let Some(f) = self.task_spawn_callback.as_ref() { | ||
| f(meta) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// Task metadata supplied to user-provided hooks for task events. | ||
| #[allow(missing_debug_implementations)] | ||
| pub struct TaskMeta<'a> { | ||
| /// The opaque ID of the task. | ||
| #[cfg(tokio_unstable)] | ||
| pub(crate) id: super::task::Id, | ||
| pub(crate) _phantom: PhantomData<&'a ()>, | ||
| } | ||
|
|
||
| impl<'a> TaskMeta<'a> { | ||
| /// Return the opaque ID of the task. | ||
| #[cfg(tokio_unstable)] | ||
| pub fn id(&self) -> super::task::Id { | ||
| self.id | ||
| } | ||
| } | ||
|
|
||
| /// Runs on specific task-related events | ||
| pub(crate) type TaskCallback = std::sync::Arc<dyn Fn(&TaskMeta<'_>) + Send + Sync>; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,76 @@ | ||
| #![allow(unknown_lints, unexpected_cfgs)] | ||
| #![warn(rust_2018_idioms)] | ||
| #![cfg(all(feature = "full", tokio_unstable, target_has_atomic = "64"))] | ||
|
|
||
| use std::collections::HashSet; | ||
| use std::sync::atomic::{AtomicUsize, Ordering}; | ||
| use std::sync::{Arc, Mutex}; | ||
|
|
||
| use tokio::runtime::Builder; | ||
|
|
||
| const TASKS: usize = 8; | ||
| const ITERATIONS: usize = 64; | ||
| /// Assert that the spawn task hook always fires when set. | ||
| #[test] | ||
| fn spawn_task_hook_fires() { | ||
| let count = Arc::new(AtomicUsize::new(0)); | ||
| let count2 = Arc::clone(&count); | ||
|
|
||
| let ids = Arc::new(Mutex::new(HashSet::new())); | ||
| let ids2 = Arc::clone(&ids); | ||
|
|
||
| let runtime = Builder::new_current_thread() | ||
| .on_task_spawn(move |data| { | ||
| ids2.lock().unwrap().insert(data.id()); | ||
|
|
||
| count2.fetch_add(1, Ordering::SeqCst); | ||
| }) | ||
| .build() | ||
| .unwrap(); | ||
|
|
||
| for _ in 0..TASKS { | ||
| runtime.spawn(std::future::pending::<()>()); | ||
| } | ||
|
|
||
| let count_realized = count.load(Ordering::SeqCst); | ||
| assert_eq!( | ||
| TASKS, count_realized, | ||
| "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}", | ||
| count_realized | ||
| ); | ||
|
|
||
| let count_ids_realized = ids.lock().unwrap().len(); | ||
|
|
||
| assert_eq!( | ||
| TASKS, count_ids_realized, | ||
| "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}", | ||
| count_realized | ||
| ); | ||
| } | ||
|
|
||
| /// Assert that the terminate task hook always fires when set. | ||
| #[test] | ||
| fn terminate_task_hook_fires() { | ||
| let count = Arc::new(AtomicUsize::new(0)); | ||
| let count2 = Arc::clone(&count); | ||
|
|
||
| let runtime = Builder::new_current_thread() | ||
| .on_task_terminate(move |_data| { | ||
| count2.fetch_add(1, Ordering::SeqCst); | ||
| }) | ||
| .build() | ||
| .unwrap(); | ||
|
|
||
| for _ in 0..TASKS { | ||
| runtime.spawn(std::future::ready(())); | ||
| } | ||
|
|
||
| runtime.block_on(async { | ||
| // tick the runtime a bunch to close out tasks | ||
| for _ in 0..ITERATIONS { | ||
| tokio::task::yield_now().await; | ||
| } | ||
| }); | ||
|
|
||
| assert_eq!(TASKS, count.load(Ordering::SeqCst)); | ||
| } |