spawnhook.rs - source
Description: Source of the Rust file `library/std/src/thread/spawnhook.rs`.

#
std/thread/
spawnhook.rs

    1use super::thread::Thread;
    2use crate::cell::Cell;
    3use crate::iter;
    4use crate::sync::Arc;
    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/// 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    if let Ok(hooks) = SPAWN_HOOKS.try_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    } 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 in self.to_run {
    150            run();
    151        }
    152    }
    153}