Update thread spawn hooks.

1. Make the effect thread local.
2. Don't return a io::Result from hooks.
This commit is contained in:
Mara Bos
2024-10-24 10:52:47 +02:00
parent ef9055f3ee
commit 947354fbec
3 changed files with 86 additions and 37 deletions

View File

@@ -491,7 +491,7 @@ impl Builder {
None => Thread::new_unnamed(id), 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(); let their_thread = my_thread.clone();
@@ -539,12 +539,9 @@ impl Builder {
imp::Thread::set_name(name); imp::Thread::set_name(name);
} }
for hook in hooks {
hook();
}
let f = f.into_inner(); let f = f.into_inner();
let try_result = panic::catch_unwind(panic::AssertUnwindSafe(|| { 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) crate::sys::backtrace::__rust_begin_short_backtrace(f)
})); }));
// SAFETY: `their_packet` as been built just above and moved by the // SAFETY: `their_packet` as been built just above and moved by the

View File

@@ -1,21 +1,43 @@
use crate::io; use crate::cell::Cell;
use crate::sync::RwLock; use crate::sync::Arc;
use crate::thread::Thread; use crate::thread::Thread;
static SPAWN_HOOKS: RwLock< // A thread local linked list of spawn hooks.
Vec<&'static (dyn Fn(&Thread) -> io::Result<Box<dyn FnOnce() + Send>> + Sync)>, crate::thread_local! {
> = RwLock::new(Vec::new()); static SPAWN_HOOKS: Cell<SpawnHooks> = const { Cell::new(SpawnHooks { first: None }) };
}
/// Registers a function to run for every new thread spawned. #[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 newly thread spawned.
/// ///
/// The hook is executed in the parent thread, and returns a function /// The hook is executed in the parent thread, and returns a function
/// that will be executed in the new thread. /// that will be executed in the new thread.
/// ///
/// The hook is called with the `Thread` handle for 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 /// The hook will only be added for the current thread and is inherited by the threads it spawns.
/// function used to spawn the thread (e.g. `std::thread::spawn`) will return /// In other words, adding a hook has no effect on already running threads (other than the current
/// the error returned by the hook. /// thread) and the threads they might spawn in the future.
/// ///
/// Hooks can only be added, not removed. /// Hooks can only be added, not removed.
/// ///
@@ -28,15 +50,15 @@ static SPAWN_HOOKS: RwLock<
/// ///
/// std::thread::add_spawn_hook(|_| { /// std::thread::add_spawn_hook(|_| {
/// ..; // This will run in the parent (spawning) thread. /// ..; // This will run in the parent (spawning) thread.
/// Ok(move || { /// move || {
/// ..; // This will run it the child (spawned) thread. /// ..; // This will run it the child (spawned) thread.
/// }) /// }
/// }); /// });
/// ``` /// ```
/// ///
/// # Example /// # 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)] /// #![feature(thread_spawn_hook)]
@@ -47,13 +69,12 @@ static SPAWN_HOOKS: RwLock<
/// static X: Cell<u32> = Cell::new(0); /// 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(|_| { /// std::thread::add_spawn_hook(|_| {
/// // Get the value of X in the spawning thread. /// // Get the value of X in the spawning thread.
/// let value = X.get(); /// let value = X.get();
/// // Set the value of X in the newly spawned thread. /// // Set the value of X in the newly spawned thread.
/// Ok(move || { /// move || X.set(value)
/// X.set(value);
/// })
/// }); /// });
/// ///
/// X.set(123); /// X.set(123);
@@ -65,15 +86,17 @@ static SPAWN_HOOKS: RwLock<
#[unstable(feature = "thread_spawn_hook", issue = "none")] #[unstable(feature = "thread_spawn_hook", issue = "none")]
pub fn add_spawn_hook<F, G>(hook: F) pub fn add_spawn_hook<F, G>(hook: F)
where where
F: 'static + Sync + Fn(&Thread) -> io::Result<G>, F: 'static + Sync + Fn(&Thread) -> G,
G: 'static + Send + FnOnce(), G: 'static + Send + FnOnce(),
{ {
SPAWN_HOOKS.write().unwrap_or_else(|e| e.into_inner()).push(Box::leak(Box::new( SPAWN_HOOKS.with(|h| {
move |thread: &Thread| -> io::Result<_> { let mut hooks = h.take();
let f: Box<dyn FnOnce() + Send> = Box::new(hook(thread)?); hooks.first = Some(Arc::new(SpawnHook {
Ok(f) hook: Box::new(move |thread| Box::new(hook(thread))),
}, next: hooks.first.take(),
))); }));
h.set(hooks);
});
} }
/// Runs all the spawn hooks. /// Runs all the spawn hooks.
@@ -81,12 +104,41 @@ where
/// Called on the parent thread. /// Called on the parent thread.
/// ///
/// Returns the functions to be called on the newly spawned 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>>> { pub(super) fn run_spawn_hooks(thread: &Thread) -> SpawnHookResults {
SPAWN_HOOKS // Get a snapshot of the spawn hooks.
.read() // (Increments the refcount to the first node.)
.unwrap_or_else(|e| e.into_inner()) let hooks = SPAWN_HOOKS.with(|hooks| {
.iter() let snapshot = hooks.take();
.rev() hooks.set(snapshot.clone());
.map(|hook| hook(thread)) snapshot
.collect() });
// 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();
}
}
} }

View File

@@ -141,9 +141,9 @@ pub fn test_main(args: &[String], tests: Vec<TestDescAndFn>, options: Option<Opt
let output_capture = io::set_output_capture(None); let output_capture = io::set_output_capture(None);
io::set_output_capture(output_capture.clone()); io::set_output_capture(output_capture.clone());
// Set the output capture of the new thread. // Set the output capture of the new thread.
Ok(|| { || {
io::set_output_capture(output_capture); io::set_output_capture(output_capture);
}) }
}); });
} }
let res = console::run_tests_console(&opts, tests); let res = console::run_tests_console(&opts, tests);