Update thread spawn hooks.

m-ou-se · m-ou-se · commit 49ac15b192b8 · 2024-11-19T18:54:20.000+01:00
1. Make the effect thread local.
2. Don't return a io::Result from hooks.
diff --git a/std/src/thread/mod.rs b/std/src/thread/mod.rs
@@ -491,7 +491,7 @@ impl Builder {
             None => Thread::new_unnamed(id),
         };
 
-        let hooks = spawnhook::run_spawn_hooks(&my_thread)?;
+        let hooks = spawnhook::run_spawn_hooks(&my_thread);
 
         let their_thread = my_thread.clone();
 
@@ -539,12 +539,9 @@ impl Builder {
                 imp::Thread::set_name(name);
             }
 
-            for hook in hooks {
-                hook();
-            }
-
             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
diff --git a/std/src/thread/spawnhook.rs b/std/src/thread/spawnhook.rs
@@ -1,21 +1,43 @@
-use crate::io;
-use crate::sync::RwLock;
+use crate::cell::Cell;
+use crate::sync::Arc;
 use crate::thread::Thread;
 
-static SPAWN_HOOKS: RwLock<
-    Vec<&'static (dyn Fn(&Thread) -> io::Result<Box<dyn FnOnce() + Send>> + Sync)>,
-> = RwLock::new(Vec::new());
+// A thread local linked list of spawn hooks.
+crate::thread_local! {
+    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 Sync + Fn(&Thread) -> Box<dyn Send + FnOnce()>>,
+    next: Option<Arc<SpawnHook>>,
+}
 
-/// Registers a function to run for every new thread spawned.
+/// 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.
 ///
-/// If the hook returns an `Err`, thread spawning is aborted. In that case, the
-/// function used to spawn the thread (e.g. `std::thread::spawn`) will return
-/// the error returned by the hook.
+/// 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.
 ///
@@ -28,15 +50,15 @@ static SPAWN_HOOKS: RwLock<
 ///
 /// std::thread::add_spawn_hook(|_| {
 ///     ..; // This will run in the parent (spawning) thread.
-///     Ok(move || {
+///     move || {
 ///         ..; // This will run it the child (spawned) thread.
-///     })
+///     }
 /// });
 /// ```
 ///
 /// # Example
 ///
-/// A spawn hook can be used to initialize thread locals from the parent thread:
+/// A spawn hook can be used to "inherit" a thread local from the parent thread:
 ///
 /// ```
 /// #![feature(thread_spawn_hook)]
@@ -47,13 +69,12 @@ static SPAWN_HOOKS: RwLock<
 ///     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.
-///     Ok(move || {
-///         X.set(value);
-///     })
+///     move || X.set(value)
 /// });
 ///
 /// X.set(123);
@@ -65,28 +86,59 @@ static SPAWN_HOOKS: RwLock<
 #[unstable(feature = "thread_spawn_hook", issue = "none")]
 pub fn add_spawn_hook<F, G>(hook: F)
 where
-    F: 'static + Sync + Fn(&Thread) -> io::Result<G>,
+    F: 'static + Sync + Fn(&Thread) -> G,
     G: 'static + Send + FnOnce(),
 {
-    SPAWN_HOOKS.write().unwrap_or_else(|e| e.into_inner()).push(Box::leak(Box::new(
-        move |thread: &Thread| -> io::Result<_> {
-            let f: Box<dyn FnOnce() + Send> = Box::new(hook(thread)?);
-            Ok(f)
-        },
-    )));
+    SPAWN_HOOKS.with(|h| {
+        let mut hooks = h.take();
+        hooks.first = Some(Arc::new(SpawnHook {
+            hook: Box::new(move |thread| Box::new(hook(thread))),
+            next: hooks.first.take(),
+        }));
+        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) -> io::Result<Vec<Box<dyn FnOnce() + Send>>> {
-    SPAWN_HOOKS
-        .read()
-        .unwrap_or_else(|e| e.into_inner())
-        .iter()
-        .rev()
-        .map(|hook| hook(thread))
-        .collect()
+pub(super) fn run_spawn_hooks(thread: &Thread) -> SpawnHookResults {
+    // 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 mut next: &Option<Arc<SpawnHook>> = &hooks.first;
+    let mut to_run = Vec::new();
+    while let Some(hook) = next {
+        to_run.push((hook.hook)(thread));
+        next = &hook.next;
+    }
+    // Pass on the snapshot of the hooks and the results to the new thread,
+    // which will then run SpawnHookResults::run().
+    SpawnHookResults { 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.
+pub(super) struct SpawnHookResults {
+    hooks: SpawnHooks,
+    to_run: Vec<Box<dyn FnOnce() + Send>>,
+}
+
+impl SpawnHookResults {
+    // 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();
+        }
+    }
 }
diff --git a/test/src/lib.rs b/test/src/lib.rs
@@ -141,9 +141,9 @@ pub fn test_main(args: &[String], tests: Vec<TestDescAndFn>, options: Option<Opt
                 let output_capture = io::set_output_capture(None);
                 io::set_output_capture(output_capture.clone());
                 // Set the output capture of the new thread.
-                Ok(|| {
+                || {
                     io::set_output_capture(output_capture);
-                })
+                }
             });
         }
         let res = console::run_tests_console(&opts, tests);
