Skip to content

Commit

Permalink
Add std::thread::add_spawn_hook.
Browse files Browse the repository at this point in the history
  • Loading branch information
m-ou-se committed Nov 19, 2024
1 parent 36cfa4e commit 24a0765
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 0 deletions.
11 changes: 11 additions & 0 deletions std/src/thread/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
////////////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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(|| {
Expand Down
92 changes: 92 additions & 0 deletions std/src/thread/spawnhook.rs
Original file line number Diff line number Diff line change
@@ -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()
}

0 comments on commit 24a0765

Please sign in to comment.