1- use crate :: io ;
2- use crate :: sync:: RwLock ;
1+ use crate :: cell :: Cell ;
2+ use crate :: sync:: Arc ;
33use crate :: thread:: Thread ;
44
5- static SPAWN_HOOKS : RwLock <
6- Vec < & ' static ( dyn Fn ( & Thread ) -> io:: Result < Box < dyn FnOnce ( ) + Send > > + Sync ) > ,
7- > = RwLock :: new ( Vec :: new ( ) ) ;
5+ // A thread local linked list of spawn hooks.
6+ crate :: thread_local! {
7+ static SPAWN_HOOKS : Cell <SpawnHooks > = const { Cell :: new( SpawnHooks { first: None } ) } ;
8+ }
9+
10+ #[ derive( Default , Clone ) ]
11+ struct SpawnHooks {
12+ first : Option < Arc < SpawnHook > > ,
13+ }
14+
15+ // Manually implement drop to prevent deep recursion when dropping linked Arc list.
16+ impl Drop for SpawnHooks {
17+ fn drop ( & mut self ) {
18+ let mut next = self . first . take ( ) ;
19+ while let Some ( SpawnHook { hook, next : n } ) = next. and_then ( |n| Arc :: into_inner ( n) ) {
20+ drop ( hook) ;
21+ next = n;
22+ }
23+ }
24+ }
25+
26+ struct SpawnHook {
27+ hook : Box < dyn Sync + Fn ( & Thread ) -> Box < dyn Send + FnOnce ( ) > > ,
28+ next : Option < Arc < SpawnHook > > ,
29+ }
830
9- /// Registers a function to run for every new thread spawned.
31+ /// Registers a function to run for every newly thread spawned.
1032///
1133/// The hook is executed in the parent thread, and returns a function
1234/// that will be executed in the new thread.
1335///
1436/// The hook is called with the `Thread` handle for the new thread.
1537///
16- /// If the hook returns an `Err`, thread spawning is aborted. In that case, the
17- /// function used to spawn the thread (e.g. `std::thread::spawn`) will return
18- /// the error returned by the hook .
38+ /// The hook will only be added for the current thread and is inherited by the threads it spawns.
39+ /// In other words, adding a hook has no effect on already running threads (other than the current
40+ /// thread) and the threads they might spawn in the future .
1941///
2042/// Hooks can only be added, not removed.
2143///
@@ -28,15 +50,15 @@ static SPAWN_HOOKS: RwLock<
2850///
2951/// std::thread::add_spawn_hook(|_| {
3052/// ..; // This will run in the parent (spawning) thread.
31- /// Ok( move || {
53+ /// move || {
3254/// ..; // This will run it the child (spawned) thread.
33- /// })
55+ /// }
3456/// });
3557/// ```
3658///
3759/// # Example
3860///
39- /// A spawn hook can be used to initialize thread locals from the parent thread:
61+ /// A spawn hook can be used to "inherit" a thread local from the parent thread:
4062///
4163/// ```
4264/// #![feature(thread_spawn_hook)]
@@ -47,13 +69,12 @@ static SPAWN_HOOKS: RwLock<
4769/// static X: Cell<u32> = Cell::new(0);
4870/// }
4971///
72+ /// // This needs to be done once in the main thread before spawning any threads.
5073/// std::thread::add_spawn_hook(|_| {
5174/// // Get the value of X in the spawning thread.
5275/// let value = X.get();
5376/// // Set the value of X in the newly spawned thread.
54- /// Ok(move || {
55- /// X.set(value);
56- /// })
77+ /// move || X.set(value)
5778/// });
5879///
5980/// X.set(123);
@@ -65,28 +86,59 @@ static SPAWN_HOOKS: RwLock<
6586#[ unstable( feature = "thread_spawn_hook" , issue = "none" ) ]
6687pub fn add_spawn_hook < F , G > ( hook : F )
6788where
68- F : ' static + Sync + Fn ( & Thread ) -> io :: Result < G > ,
89+ F : ' static + Sync + Fn ( & Thread ) -> G ,
6990 G : ' static + Send + FnOnce ( ) ,
7091{
71- SPAWN_HOOKS . write ( ) . unwrap_or_else ( |e| e. into_inner ( ) ) . push ( Box :: leak ( Box :: new (
72- move |thread : & Thread | -> io:: Result < _ > {
73- let f: Box < dyn FnOnce ( ) + Send > = Box :: new ( hook ( thread) ?) ;
74- Ok ( f)
75- } ,
76- ) ) ) ;
92+ SPAWN_HOOKS . with ( |h| {
93+ let mut hooks = h. take ( ) ;
94+ hooks. first = Some ( Arc :: new ( SpawnHook {
95+ hook : Box :: new ( move |thread| Box :: new ( hook ( thread) ) ) ,
96+ next : hooks. first . take ( ) ,
97+ } ) ) ;
98+ h. set ( hooks) ;
99+ } ) ;
77100}
78101
79102/// Runs all the spawn hooks.
80103///
81104/// Called on the parent thread.
82105///
83106/// Returns the functions to be called on the newly spawned thread.
84- pub ( super ) fn run_spawn_hooks ( thread : & Thread ) -> io:: Result < Vec < Box < dyn FnOnce ( ) + Send > > > {
85- SPAWN_HOOKS
86- . read ( )
87- . unwrap_or_else ( |e| e. into_inner ( ) )
88- . iter ( )
89- . rev ( )
90- . map ( |hook| hook ( thread) )
91- . collect ( )
107+ pub ( super ) fn run_spawn_hooks ( thread : & Thread ) -> SpawnHookResults {
108+ // Get a snapshot of the spawn hooks.
109+ // (Increments the refcount to the first node.)
110+ let hooks = SPAWN_HOOKS . with ( |hooks| {
111+ let snapshot = hooks. take ( ) ;
112+ hooks. set ( snapshot. clone ( ) ) ;
113+ snapshot
114+ } ) ;
115+ // Iterate over the hooks, run them, and collect the results in a vector.
116+ let mut next: & Option < Arc < SpawnHook > > = & hooks. first ;
117+ let mut to_run = Vec :: new ( ) ;
118+ while let Some ( hook) = next {
119+ to_run. push ( ( hook. hook ) ( thread) ) ;
120+ next = & hook. next ;
121+ }
122+ // Pass on the snapshot of the hooks and the results to the new thread,
123+ // which will then run SpawnHookResults::run().
124+ SpawnHookResults { hooks, to_run }
125+ }
126+
127+ /// The results of running the spawn hooks.
128+ ///
129+ /// This struct is sent to the new thread.
130+ /// It contains the inherited hooks and the closures to be run.
131+ pub ( super ) struct SpawnHookResults {
132+ hooks : SpawnHooks ,
133+ to_run : Vec < Box < dyn FnOnce ( ) + Send > > ,
134+ }
135+
136+ impl SpawnHookResults {
137+ // This is run on the newly spawned thread, directly at the start.
138+ pub ( super ) fn run ( self ) {
139+ SPAWN_HOOKS . set ( self . hooks ) ;
140+ for run in self . to_run {
141+ run ( ) ;
142+ }
143+ }
92144}
0 commit comments