Rollup merge of #125405 - m-ou-se:thread-add-spawn-hook, r=WaffleLapkin

Add std:🧵:add_spawn_hook.

Implementation of https://github.com/rust-lang/rfcs/pull/3642

library/std/src/thread/mod.rs

@@ -188,6 +188,11 @@ mod current;
 pub use current::current;
 pub(crate) use current::{current_id, drop_current, set_current, try_current};
 
+mod spawnhook;
+
+#[unstable(feature = "thread_spawn_hook", issue = "132951")]
+pub use spawnhook::add_spawn_hook;
+
 ////////////////////////////////////////////////////////////////////////////////
 // Thread-local storage
 ////////////////////////////////////////////////////////////////////////////////
@@ -259,6 +264,8 @@ pub struct Builder {
     name: Option<String>,
     // The size of the stack for the spawned thread in bytes
     stack_size: Option<usize>,
+    // Skip running and inheriting the thread spawn hooks
+    no_hooks: bool,
 }
 
 impl Builder {
@@ -282,7 +289,7 @@ impl Builder {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn new() -> Builder {
-        Builder { name: None, stack_size: None }
+        Builder { name: None, stack_size: None, no_hooks: false }
     }
 
     /// Names the thread-to-be. Currently the name is used for identification
@@ -338,6 +345,16 @@ impl Builder {
         self
     }
 
+    /// Disables running and inheriting [spawn hooks](add_spawn_hook).
+    ///
+    /// Use this if the parent thread is in no way relevant for the child thread.
+    /// For example, when lazily spawning threads for a thread pool.
+    #[unstable(feature = "thread_spawn_hook", issue = "132951")]
+    pub fn no_hooks(mut self) -> Builder {
+        self.no_hooks = true;
+        self
+    }
+
     /// Spawns a new thread by taking ownership of the `Builder`, and returns an
     /// [`io::Result`] to its [`JoinHandle`].
     ///
@@ -460,7 +477,7 @@ impl Builder {
         F: Send,
         T: Send,
     {
-        let Builder { name, stack_size } = self;
+        let Builder { name, stack_size, no_hooks } = self;
 
         let stack_size = stack_size.unwrap_or_else(|| {
             static MIN: AtomicUsize = AtomicUsize::new(0);
@@ -485,6 +502,13 @@ impl Builder {
             Some(name) => Thread::new(id, name.into()),
             None => Thread::new_unnamed(id),
         };
+
+        let hooks = if no_hooks {
+            spawnhook::ChildSpawnHooks::default()
+        } else {
+            spawnhook::run_spawn_hooks(&my_thread)
+        };
+
         let their_thread = my_thread.clone();
 
         let my_packet: Arc<Packet<'scope, T>> = Arc::new(Packet {
@@ -494,9 +518,6 @@ impl Builder {
         });
         let their_packet = my_packet.clone();
 
-        let output_capture = crate::io::set_output_capture(None);
-        crate::io::set_output_capture(output_capture.clone());
-
         // Pass `f` in `MaybeUninit` because actually that closure might *run longer than the lifetime of `F`*.
         // See <https://github.com/rust-lang/rust/issues/101983> for more details.
         // To prevent leaks we use a wrapper that drops its contents.
@@ -534,10 +555,9 @@ impl Builder {
                 imp::Thread::set_name(name);
             }
 
-            crate::io::set_output_capture(output_capture);
-
             let f = f.into_inner();
             let try_result = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                crate::sys::backtrace::__rust_begin_short_backtrace(|| hooks.run());
                 crate::sys::backtrace::__rust_begin_short_backtrace(f)
             }));
             // SAFETY: `their_packet` as been built just above and moved by the

library/std/src/thread/spawnhook.rs

@@ -0,0 +1,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();
+        }
+    }
+}

library/test/src/lib.rs

@@ -24,6 +24,7 @@
 #![feature(process_exitcode_internals)]
 #![feature(panic_can_unwind)]
 #![feature(test)]
+#![feature(thread_spawn_hook)]
 #![allow(internal_features)]
 #![warn(rustdoc::unescaped_backticks)]
 
@@ -134,6 +135,16 @@ pub fn test_main(args: &[String], tests: Vec<TestDescAndFn>, options: Option<Opt
                 }
             });
             panic::set_hook(hook);
+            // Use a thread spawning hook to make new threads inherit output capturing.
+            std::thread::add_spawn_hook(|_| {
+                // Get and clone the output capture of the current thread.
+                let output_capture = io::set_output_capture(None);
+                io::set_output_capture(output_capture.clone());
+                // Set the output capture of the new thread.
+                || {
+                    io::set_output_capture(output_capture);
+                }
+            });
         }
         let res = console::run_tests_console(&opts, tests);
         // Prevent Valgrind from reporting reachable blocks in users' unit tests.