From 947354fbec27766f3a99c13c90413cc22739108c Mon Sep 17 00:00:00 2001 From: Mara Bos Date: Thu, 24 Oct 2024 10:52:47 +0200 Subject: [PATCH] Update thread spawn hooks. 1. Make the effect thread local. 2. Don't return a io::Result from hooks. --- library/std/src/thread/mod.rs | 7 +- library/std/src/thread/spawnhook.rs | 112 ++++++++++++++++++++-------- library/test/src/lib.rs | 4 +- 3 files changed, 86 insertions(+), 37 deletions(-) diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs index c6b1e0de84c..79621e7fa6b 100644 --- a/library/std/src/thread/mod.rs +++ b/library/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/library/std/src/thread/spawnhook.rs b/library/std/src/thread/spawnhook.rs index c64aea4262b..9c4e9eb6aa1 100644 --- a/library/std/src/thread/spawnhook.rs +++ b/library/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> + Sync)>, -> = RwLock::new(Vec::new()); +// A thread local linked list of spawn hooks. +crate::thread_local! { + static SPAWN_HOOKS: Cell = const { Cell::new(SpawnHooks { first: None }) }; +} -/// Registers a function to run for every new thread spawned. +#[derive(Default, Clone)] +struct SpawnHooks { + first: Option>, +} + +// 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 Box>, + next: Option>, +} + +/// 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 = 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,15 +86,17 @@ static SPAWN_HOOKS: RwLock< #[unstable(feature = "thread_spawn_hook", issue = "none")] pub fn add_spawn_hook(hook: F) where - F: 'static + Sync + Fn(&Thread) -> io::Result, + 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 = 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. @@ -81,12 +104,41 @@ where /// 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>> { - 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> = &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>, +} + +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/library/test/src/lib.rs b/library/test/src/lib.rs index 00a0b55c6a9..47407df909b 100644 --- a/library/test/src/lib.rs +++ b/library/test/src/lib.rs @@ -141,9 +141,9 @@ pub fn test_main(args: &[String], tests: Vec, options: Option