std/thread/
spawnhook.rs

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 = 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/// 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| {
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.
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    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)]
135pub(super) struct ChildSpawnHooks {
136    hooks: SpawnHooks,
137    to_run: Vec<Box<dyn FnOnce() + Send>>,
138}
139
140impl 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}