1use crate::cell::Cell;
2use crate::iter;
3use crate::sync::Arc;
4use crate::thread::Thread;
5
6crate::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)]
16struct SpawnHooks {
17 first: Option<Arc<SpawnHook>>,
18}
19
20// Manually implement drop to prevent deep recursion when dropping linked Arc list.
21impl Drop for SpawnHooks {
22 fn drop(&mut self) {
23 let mut next: Option> = self.first.take();
24 while let Some(SpawnHook { hook: Box Box<…> + Send + Sync + 'static>, next: n: Option> }) = next.and_then(|n: Arc| Arc::into_inner(this:n)) {
25 drop(hook);
26 next = n;
27 }
28 }
29}
30
31struct 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")]
92pub fn add_spawn_hook<F, G>(hook: F)
93where
94 F: 'static + Send + Sync + Fn(&Thread) -> G,
95 G: 'static + Send + FnOnce(),
96{
97 SPAWN_HOOKS.with(|h: &Cell| {
98 let mut hooks: SpawnHooks = h.take();
99 let next: Option> = hooks.first.take();
100 hooks.first = Some(Arc::new(data:SpawnHook {
101 hook: Box::new(move |thread: &Thread| Box::new(hook(thread))),
102 next,
103 }));
104 h.set(val: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.
113pub(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 if let Ok(hooks: SpawnHooks) = SPAWN_HOOKS.try_with(|hooks: &Cell| {
117 let snapshot: SpawnHooks = hooks.take();
118 hooks.set(val:snapshot.clone());
119 snapshot
120 }) {
121 // Iterate over the hooks, run them, and collect the results in a vector.
122 let to_run: Vec<_> = iterimpl Iterator>::successors(first:hooks.first.as_deref(), |hook: &&SpawnHook| hook.next.as_deref())
123 .map(|hook: &SpawnHook| (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 } else {
129 // TLS has been destroyed. Skip running the hooks.
130 // See https://github.com/rust-lang/rust/issues/138696
131 ChildSpawnHooks::default()
132 }
133}
134
135/// The results of running the spawn hooks.
136///
137/// This struct is sent to the new thread.
138/// It contains the inherited hooks and the closures to be run.
139#[derive(Default)]
140pub(super) struct ChildSpawnHooks {
141 hooks: SpawnHooks,
142 to_run: Vec<Box<dyn FnOnce() + Send>>,
143}
144
145impl ChildSpawnHooks {
146 // This is run on the newly spawned thread, directly at the start.
147 pub(super) fn run(self) {
148 SPAWN_HOOKS.set(self.hooks);
149 for run: Box in self.to_run {
150 run();
151 }
152 }
153}
154