forked from rust-lang/rust
-
Notifications
You must be signed in to change notification settings - Fork 6
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Rollup merge of rust-lang#125405 - m-ou-se:thread-add-spawn-hook, r=WaffleLapkin Add std::thread::add_spawn_hook. Implementation of rust-lang/rfcs#3642
- Loading branch information
Showing
3 changed files
with
186 additions
and
7 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,148 @@ | ||
| use crate::cell::Cell; | ||
| use crate::iter; | ||
| use crate::sync::Arc; | ||
| use crate::thread::Thread; | ||
|
|
||
| crate::thread_local! { | ||
| /// A thread local linked list of spawn hooks. | ||
| /// | ||
| /// It is a linked list of Arcs, such that it can very cheaply be inhereted by spawned threads. | ||
| /// | ||
| /// (That technically makes it a set of linked lists with shared tails, so a linked tree.) | ||
| static SPAWN_HOOKS: Cell<SpawnHooks> = const { Cell::new(SpawnHooks { first: None }) }; | ||
| } | ||
|
|
||
| #[derive(Default, Clone)] | ||
| struct SpawnHooks { | ||
| first: Option<Arc<SpawnHook>>, | ||
| } | ||
|
|
||
| // Manually implement drop to prevent deep recursion when dropping linked Arc list. | ||
| impl Drop for SpawnHooks { | ||
| fn drop(&mut self) { | ||
| let mut next = self.first.take(); | ||
| while let Some(SpawnHook { hook, next: n }) = next.and_then(|n| Arc::into_inner(n)) { | ||
| drop(hook); | ||
| next = n; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| struct SpawnHook { | ||
| hook: Box<dyn Send + Sync + Fn(&Thread) -> Box<dyn Send + FnOnce()>>, | ||
| next: Option<Arc<SpawnHook>>, | ||
| } | ||
|
|
||
| /// Registers a function to run for every newly thread spawned. | ||
| /// | ||
| /// The hook is executed in the parent thread, and returns a function | ||
| /// that will be executed in the new thread. | ||
| /// | ||
| /// The hook is called with the `Thread` handle for the new thread. | ||
| /// | ||
| /// The hook will only be added for the current thread and is inherited by the threads it spawns. | ||
| /// In other words, adding a hook has no effect on already running threads (other than the current | ||
| /// thread) and the threads they might spawn in the future. | ||
| /// | ||
| /// Hooks can only be added, not removed. | ||
| /// | ||
| /// The hooks will run in reverse order, starting with the most recently added. | ||
| /// | ||
| /// # Usage | ||
| /// | ||
| /// ``` | ||
| /// #![feature(thread_spawn_hook)] | ||
| /// | ||
| /// std::thread::add_spawn_hook(|_| { | ||
| /// ..; // This will run in the parent (spawning) thread. | ||
| /// move || { | ||
| /// ..; // This will run it the child (spawned) thread. | ||
| /// } | ||
| /// }); | ||
| /// ``` | ||
| /// | ||
| /// # Example | ||
| /// | ||
| /// A spawn hook can be used to "inherit" a thread local from the parent thread: | ||
| /// | ||
| /// ``` | ||
| /// #![feature(thread_spawn_hook)] | ||
| /// | ||
| /// use std::cell::Cell; | ||
| /// | ||
| /// thread_local! { | ||
| /// static X: Cell<u32> = Cell::new(0); | ||
| /// } | ||
| /// | ||
| /// // This needs to be done once in the main thread before spawning any threads. | ||
| /// std::thread::add_spawn_hook(|_| { | ||
| /// // Get the value of X in the spawning thread. | ||
| /// let value = X.get(); | ||
| /// // Set the value of X in the newly spawned thread. | ||
| /// move || X.set(value) | ||
| /// }); | ||
| /// | ||
| /// X.set(123); | ||
| /// | ||
| /// std::thread::spawn(|| { | ||
| /// assert_eq!(X.get(), 123); | ||
| /// }).join().unwrap(); | ||
| /// ``` | ||
| #[unstable(feature = "thread_spawn_hook", issue = "132951")] | ||
| pub fn add_spawn_hook<F, G>(hook: F) | ||
| where | ||
| F: 'static + Send + Sync + Fn(&Thread) -> G, | ||
| G: 'static + Send + FnOnce(), | ||
| { | ||
| SPAWN_HOOKS.with(|h| { | ||
| let mut hooks = h.take(); | ||
| let next = hooks.first.take(); | ||
| hooks.first = Some(Arc::new(SpawnHook { | ||
| hook: Box::new(move |thread| Box::new(hook(thread))), | ||
| next, | ||
| })); | ||
| h.set(hooks); | ||
| }); | ||
| } | ||
|
|
||
| /// Runs all the spawn hooks. | ||
| /// | ||
| /// Called on the parent thread. | ||
| /// | ||
| /// Returns the functions to be called on the newly spawned thread. | ||
| pub(super) fn run_spawn_hooks(thread: &Thread) -> ChildSpawnHooks { | ||
| // Get a snapshot of the spawn hooks. | ||
| // (Increments the refcount to the first node.) | ||
| let hooks = SPAWN_HOOKS.with(|hooks| { | ||
| let snapshot = hooks.take(); | ||
| hooks.set(snapshot.clone()); | ||
| snapshot | ||
| }); | ||
| // Iterate over the hooks, run them, and collect the results in a vector. | ||
| let to_run: Vec<_> = iter::successors(hooks.first.as_deref(), |hook| hook.next.as_deref()) | ||
| .map(|hook| (hook.hook)(thread)) | ||
| .collect(); | ||
| // Pass on the snapshot of the hooks and the results to the new thread, | ||
| // which will then run SpawnHookResults::run(). | ||
| ChildSpawnHooks { hooks, to_run } | ||
| } | ||
|
|
||
| /// The results of running the spawn hooks. | ||
| /// | ||
| /// This struct is sent to the new thread. | ||
| /// It contains the inherited hooks and the closures to be run. | ||
| #[derive(Default)] | ||
| pub(super) struct ChildSpawnHooks { | ||
| hooks: SpawnHooks, | ||
| to_run: Vec<Box<dyn FnOnce() + Send>>, | ||
| } | ||
|
|
||
| impl ChildSpawnHooks { | ||
| // This is run on the newly spawned thread, directly at the start. | ||
| pub(super) fn run(self) { | ||
| SPAWN_HOOKS.set(self.hooks); | ||
| for run in self.to_run { | ||
| run(); | ||
| } | ||
| } | ||
| } |
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