Update thread spawn hooks.
1. Make the effect thread local. 2. Don't return a io::Result from hooks.
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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();
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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);
|
||||||
|
|||||||
Reference in New Issue
Block a user