|
| 1 | +use crate::cell::Cell; |
| 2 | +use crate::iter; |
| 3 | +use crate::sync::Arc; |
| 4 | +use crate::thread::Thread; |
| 5 | + |
| 6 | +crate::thread_local! { |
| 7 | + /// A thread local linked list of spawn hooks. |
| 8 | + /// |
| 9 | + /// It is a linked list of Arcs, such that it can very cheaply be inhereted by spawned threads. |
| 10 | + /// |
| 11 | + /// (That technically makes it a set of linked lists with shared tails, so a linked tree.) |
| 12 | + static SPAWN_HOOKS: Cell<SpawnHooks> = const { Cell::new(SpawnHooks { first: None }) }; |
| 13 | +} |
| 14 | + |
| 15 | +#[derive(Default, Clone)] |
| 16 | +struct SpawnHooks { |
| 17 | + first: Option<Arc<SpawnHook>>, |
| 18 | +} |
| 19 | + |
| 20 | +// Manually implement drop to prevent deep recursion when dropping linked Arc list. |
| 21 | +impl Drop for SpawnHooks { |
| 22 | + fn drop(&mut self) { |
| 23 | + let mut next = self.first.take(); |
| 24 | + while let Some(SpawnHook { hook, next: n }) = next.and_then(|n| Arc::into_inner(n)) { |
| 25 | + drop(hook); |
| 26 | + next = n; |
| 27 | + } |
| 28 | + } |
| 29 | +} |
| 30 | + |
| 31 | +struct SpawnHook { |
| 32 | + hook: Box<dyn Send + Sync + Fn(&Thread) -> Box<dyn Send + FnOnce()>>, |
| 33 | + next: Option<Arc<SpawnHook>>, |
| 34 | +} |
| 35 | + |
| 36 | +/// Registers a function to run for every newly thread spawned. |
| 37 | +/// |
| 38 | +/// The hook is executed in the parent thread, and returns a function |
| 39 | +/// that will be executed in the new thread. |
| 40 | +/// |
| 41 | +/// The hook is called with the `Thread` handle for the new thread. |
| 42 | +/// |
| 43 | +/// The hook will only be added for the current thread and is inherited by the threads it spawns. |
| 44 | +/// In other words, adding a hook has no effect on already running threads (other than the current |
| 45 | +/// thread) and the threads they might spawn in the future. |
| 46 | +/// |
| 47 | +/// Hooks can only be added, not removed. |
| 48 | +/// |
| 49 | +/// The hooks will run in reverse order, starting with the most recently added. |
| 50 | +/// |
| 51 | +/// # Usage |
| 52 | +/// |
| 53 | +/// ``` |
| 54 | +/// #![feature(thread_spawn_hook)] |
| 55 | +/// |
| 56 | +/// std::thread::add_spawn_hook(|_| { |
| 57 | +/// ..; // This will run in the parent (spawning) thread. |
| 58 | +/// move || { |
| 59 | +/// ..; // This will run it the child (spawned) thread. |
| 60 | +/// } |
| 61 | +/// }); |
| 62 | +/// ``` |
| 63 | +/// |
| 64 | +/// # Example |
| 65 | +/// |
| 66 | +/// A spawn hook can be used to "inherit" a thread local from the parent thread: |
| 67 | +/// |
| 68 | +/// ``` |
| 69 | +/// #![feature(thread_spawn_hook)] |
| 70 | +/// |
| 71 | +/// use std::cell::Cell; |
| 72 | +/// |
| 73 | +/// thread_local! { |
| 74 | +/// static X: Cell<u32> = Cell::new(0); |
| 75 | +/// } |
| 76 | +/// |
| 77 | +/// // This needs to be done once in the main thread before spawning any threads. |
| 78 | +/// std::thread::add_spawn_hook(|_| { |
| 79 | +/// // Get the value of X in the spawning thread. |
| 80 | +/// let value = X.get(); |
| 81 | +/// // Set the value of X in the newly spawned thread. |
| 82 | +/// move || X.set(value) |
| 83 | +/// }); |
| 84 | +/// |
| 85 | +/// X.set(123); |
| 86 | +/// |
| 87 | +/// std::thread::spawn(|| { |
| 88 | +/// assert_eq!(X.get(), 123); |
| 89 | +/// }).join().unwrap(); |
| 90 | +/// ``` |
| 91 | +#[unstable(feature = "thread_spawn_hook", issue = "132951")] |
| 92 | +pub fn add_spawn_hook<F, G>(hook: F) |
| 93 | +where |
| 94 | + F: 'static + Send + Sync + Fn(&Thread) -> G, |
| 95 | + G: 'static + Send + FnOnce(), |
| 96 | +{ |
| 97 | + SPAWN_HOOKS.with(|h| { |
| 98 | + let mut hooks = h.take(); |
| 99 | + let next = hooks.first.take(); |
| 100 | + hooks.first = Some(Arc::new(SpawnHook { |
| 101 | + hook: Box::new(move |thread| Box::new(hook(thread))), |
| 102 | + next, |
| 103 | + })); |
| 104 | + h.set(hooks); |
| 105 | + }); |
| 106 | +} |
| 107 | + |
| 108 | +/// Runs all the spawn hooks. |
| 109 | +/// |
| 110 | +/// Called on the parent thread. |
| 111 | +/// |
| 112 | +/// Returns the functions to be called on the newly spawned thread. |
| 113 | +pub(super) fn run_spawn_hooks(thread: &Thread) -> ChildSpawnHooks { |
| 114 | + // Get a snapshot of the spawn hooks. |
| 115 | + // (Increments the refcount to the first node.) |
| 116 | + let hooks = SPAWN_HOOKS.with(|hooks| { |
| 117 | + let snapshot = hooks.take(); |
| 118 | + hooks.set(snapshot.clone()); |
| 119 | + snapshot |
| 120 | + }); |
| 121 | + // Iterate over the hooks, run them, and collect the results in a vector. |
| 122 | + let to_run: Vec<_> = iter::successors(hooks.first.as_deref(), |hook| hook.next.as_deref()) |
| 123 | + .map(|hook| (hook.hook)(thread)) |
| 124 | + .collect(); |
| 125 | + // Pass on the snapshot of the hooks and the results to the new thread, |
| 126 | + // which will then run SpawnHookResults::run(). |
| 127 | + ChildSpawnHooks { hooks, to_run } |
| 128 | +} |
| 129 | + |
| 130 | +/// The results of running the spawn hooks. |
| 131 | +/// |
| 132 | +/// This struct is sent to the new thread. |
| 133 | +/// It contains the inherited hooks and the closures to be run. |
| 134 | +#[derive(Default)] |
| 135 | +pub(super) struct ChildSpawnHooks { |
| 136 | + hooks: SpawnHooks, |
| 137 | + to_run: Vec<Box<dyn FnOnce() + Send>>, |
| 138 | +} |
| 139 | + |
| 140 | +impl ChildSpawnHooks { |
| 141 | + // This is run on the newly spawned thread, directly at the start. |
| 142 | + pub(super) fn run(self) { |
| 143 | + SPAWN_HOOKS.set(self.hooks); |
| 144 | + for run in self.to_run { |
| 145 | + run(); |
| 146 | + } |
| 147 | + } |
| 148 | +} |
0 commit comments