std/thread/
spawnhook.rs

1use super::thread::Thread;
2use crate::cell::Cell;
3use crate::iter;
4use crate::sync::{Arc, UniqueArc};
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 inherited 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 = 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
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/// The hooks will run in reverse order, starting with the most recently added.
48///
49/// Hooks can only be added, not removed.
50/// Note that this does *not* mean that they are guaranteed to run. See [`Builder::no_hooks`].
51/// You cannot rely on a hook running for soundness.
52///
53/// # Usage
54///
55/// ```
56/// #![feature(thread_spawn_hook)]
57///
58/// std::thread::add_spawn_hook(|_| {
59///     ..; // This will run in the parent (spawning) thread.
60///     move || {
61///         ..; // This will run it the child (spawned) thread.
62///     }
63/// });
64/// ```
65///
66/// # Example
67///
68/// A spawn hook can be used to "inherit" a thread local from the parent thread:
69///
70/// ```
71/// #![feature(thread_spawn_hook)]
72///
73/// use std::cell::Cell;
74///
75/// thread_local! {
76///     static X: Cell<u32> = Cell::new(0);
77/// }
78///
79/// // This needs to be done once in the main thread before spawning any threads.
80/// std::thread::add_spawn_hook(|_| {
81///     // Get the value of X in the spawning thread.
82///     let value = X.get();
83///     // Set the value of X in the newly spawned thread.
84///     move || X.set(value)
85/// });
86///
87/// X.set(123);
88///
89/// std::thread::spawn(|| {
90///     assert_eq!(X.get(), 123);
91/// }).join().unwrap();
92/// ```
93///
94/// [`Builder::no_hooks`]: crate::thread::Builder::no_hooks
95#[unstable(feature = "thread_spawn_hook", issue = "132951")]
96pub fn add_spawn_hook<F, G>(hook: F)
97where
98    F: 'static + Send + Sync + Fn(&Thread) -> G,
99    G: 'static + Send + FnOnce(),
100{
101    SPAWN_HOOKS.with(|h| {
102        // Perform all fallible operations before taking the current hooks (see #159923)
103        let mut new_first = UniqueArc::new(SpawnHook {
104            hook: Box::new(move |thread| Box::new(hook(thread))),
105            next: None,
106        });
107        let mut hooks = h.take();
108        new_first.next = hooks.first.take();
109        hooks.first = Some(UniqueArc::into_arc(new_first));
110        h.set(hooks);
111    });
112}
113
114/// Runs all the spawn hooks.
115///
116/// Called on the parent thread.
117///
118/// Returns the functions to be called on the newly spawned thread.
119pub(super) fn run_spawn_hooks(thread: &Thread) -> ChildSpawnHooks {
120    // Get a snapshot of the spawn hooks.
121    // (Increments the refcount to the first node.)
122    if let Ok(hooks) = SPAWN_HOOKS.try_with(|hooks| {
123        let snapshot = hooks.take();
124        hooks.set(snapshot.clone());
125        snapshot
126    }) {
127        // Iterate over the hooks, run them, and collect the results in a vector.
128        let to_run: Vec<_> = iter::successors(hooks.first.as_deref(), |hook| hook.next.as_deref())
129            .map(|hook| (hook.hook)(thread))
130            .collect();
131        // Pass on the snapshot of the hooks and the results to the new thread,
132        // which will then run SpawnHookResults::run().
133        ChildSpawnHooks { hooks, to_run }
134    } else {
135        // TLS has been destroyed. Skip running the hooks.
136        // See https://github.com/rust-lang/rust/issues/138696
137        ChildSpawnHooks::default()
138    }
139}
140
141/// The results of running the spawn hooks.
142///
143/// This struct is sent to the new thread.
144/// It contains the inherited hooks and the closures to be run.
145#[derive(Default)]
146pub(super) struct ChildSpawnHooks {
147    hooks: SpawnHooks,
148    to_run: Vec<Box<dyn FnOnce() + Send>>,
149}
150
151impl ChildSpawnHooks {
152    // This is run on the newly spawned thread, directly at the start.
153    pub(super) fn inherit_and_run(self) {
154        SPAWN_HOOKS.set(self.hooks);
155        for run in self.to_run {
156            run();
157        }
158    }
159}
std/thread/spawnhook.rs

std/thread/
spawnhook.rs
