std/thread/
spawnhook.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
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();
        }
    }
}