Add std:🧵:add_spawn_hook.

This commit is contained in:
Mara Bos 2024-05-22 14:10:52 +02:00
parent 145f9cf95d
commit 7ac4c04731
2 changed files with 103 additions and 0 deletions

View File

@ -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 = "none")]
pub use spawnhook::add_spawn_hook;
////////////////////////////////////////////////////////////////////////////////
// Thread-local storage
////////////////////////////////////////////////////////////////////////////////
@ -485,6 +490,9 @@ impl Builder {
Some(name) => Thread::new(id, name.into()),
None => Thread::new_unnamed(id),
};
let hooks = spawnhook::run_spawn_hooks(&my_thread)?;
let their_thread = my_thread.clone();
let my_packet: Arc<Packet<'scope, T>> = Arc::new(Packet {
@ -535,6 +543,9 @@ impl Builder {
}
crate::io::set_output_capture(output_capture);
for hook in hooks {
hook();
}
let f = f.into_inner();
let try_result = panic::catch_unwind(panic::AssertUnwindSafe(|| {

View File

@ -0,0 +1,92 @@
use crate::io;
use crate::sync::RwLock;
use crate::thread::Thread;
static SPAWN_HOOKS: RwLock<
Vec<&'static (dyn Fn(&Thread) -> io::Result<Box<dyn FnOnce() + Send>> + Sync)>,
> = RwLock::new(Vec::new());
/// Registers a function to run for every new 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.
///
/// Hooks can only be added, not removed.
///
/// The hooks will run in 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.
/// Ok(move || {
/// ..; // This will run it the child (spawned) thread.
/// })
/// });
/// ```
///
/// # Example
///
/// A spawn hook can be used to initialize thread locals from the parent thread:
///
/// ```
/// #![feature(thread_spawn_hook)]
///
/// use std::cell::Cell;
///
/// thread_local! {
/// static X: Cell<u32> = Cell::new(0);
/// }
///
/// 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);
/// })
/// });
///
/// X.set(123);
///
/// std::thread::spawn(|| {
/// assert_eq!(X.get(), 123);
/// }).join().unwrap();
/// ```
#[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>,
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)
},
)));
}
/// 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()
}