Add callbacks for when threads start and stop doing work

Author: Zoxc

rayon-core/src/lib.rs

@@ -199,6 +199,12 @@ pub struct ThreadPoolBuilder<S = DefaultSpawn> {
     /// Closure invoked to spawn threads.
     spawn_handler: S,
 
+    /// Closure invoked when starting computations in a thread.
+    acquire_thread_handler: Option<Box<AcquireThreadHandler>>,
+
+    /// Closure invoked when blocking in a thread.
+    release_thread_handler: Option<Box<ReleaseThreadHandler>>,
+
     /// If false, worker threads will execute spawned jobs in a
     /// "depth-first" fashion. If true, they will do a "breadth-first"
     /// fashion. Depth-first is the default.
@@ -242,12 +248,22 @@ impl Default for ThreadPoolBuilder {
             start_handler: None,
             exit_handler: None,
             deadlock_handler: None,
+            acquire_thread_handler: None,
+            release_thread_handler: None,
             spawn_handler: DefaultSpawn,
             breadth_first: false,
         }
     }
 }
 
+/// The type for a closure that gets invoked before starting computations in a thread.
+/// Note that this same closure may be invoked multiple times in parallel.
+type AcquireThreadHandler = dyn Fn() + Send + Sync;
+
+/// The type for a closure that gets invoked before blocking in a thread.
+/// Note that this same closure may be invoked multiple times in parallel.
+type ReleaseThreadHandler = dyn Fn() + Send + Sync;
+
 impl ThreadPoolBuilder {
     /// Creates and returns a valid rayon thread pool builder, but does not initialize it.
     pub fn new() -> Self {
@@ -348,7 +364,12 @@ impl ThreadPoolBuilder {
                     Ok(())
                 })
                 .build()?;
-            Ok(with_pool(&pool))
+            let result = unwind::halt_unwinding(|| with_pool(&pool));
+            pool.wait_until_stopped();
+            match result {
+                Ok(result) => Ok(result),
+                Err(err) => unwind::resume_unwinding(err),
+            }
         });
 
         match result {
@@ -460,6 +481,8 @@ impl<S> ThreadPoolBuilder<S> {
             start_handler: self.start_handler,
             exit_handler: self.exit_handler,
             deadlock_handler: self.deadlock_handler,
+            acquire_thread_handler: self.acquire_thread_handler,
+            release_thread_handler: self.release_thread_handler,
             breadth_first: self.breadth_first,
         }
     }
@@ -618,6 +641,34 @@ impl<S> ThreadPoolBuilder<S> {
         self.breadth_first
     }
 
+    /// Takes the current acquire thread callback, leaving `None`.
+    fn take_acquire_thread_handler(&mut self) -> Option<Box<AcquireThreadHandler>> {
+        self.acquire_thread_handler.take()
+    }
+
+    /// Set a callback to be invoked when starting computations in a thread.
+    pub fn acquire_thread_handler<H>(mut self, acquire_thread_handler: H) -> Self
+    where
+        H: Fn() + Send + Sync + 'static,
+    {
+        self.acquire_thread_handler = Some(Box::new(acquire_thread_handler));
+        self
+    }
+
+    /// Takes the current release thread callback, leaving `None`.
+    fn take_release_thread_handler(&mut self) -> Option<Box<ReleaseThreadHandler>> {
+        self.release_thread_handler.take()
+    }
+
+    /// Set a callback to be invoked when blocking in thread.
+    pub fn release_thread_handler<H>(mut self, release_thread_handler: H) -> Self
+    where
+        H: Fn() + Send + Sync + 'static,
+    {
+        self.release_thread_handler = Some(Box::new(release_thread_handler));
+        self
+    }
+
     /// Takes the current deadlock callback, leaving `None`.
     fn take_deadlock_handler(&mut self) -> Option<Box<DeadlockHandler>> {
         self.deadlock_handler.take()
@@ -801,6 +852,8 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             ref deadlock_handler,
             ref start_handler,
             ref exit_handler,
+            ref acquire_thread_handler,
+            ref release_thread_handler,
             spawn_handler: _,
             ref breadth_first,
         } = *self;
@@ -818,6 +871,8 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
         let deadlock_handler = deadlock_handler.as_ref().map(|_| ClosurePlaceholder);
         let start_handler = start_handler.as_ref().map(|_| ClosurePlaceholder);
         let exit_handler = exit_handler.as_ref().map(|_| ClosurePlaceholder);
+        let acquire_thread_handler = acquire_thread_handler.as_ref().map(|_| ClosurePlaceholder);
+        let release_thread_handler = release_thread_handler.as_ref().map(|_| ClosurePlaceholder);
 
         f.debug_struct("ThreadPoolBuilder")
             .field("num_threads", num_threads)
@@ -827,6 +882,8 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             .field("deadlock_handler", &deadlock_handler)
             .field("start_handler", &start_handler)
             .field("exit_handler", &exit_handler)
+            .field("acquire_thread_handler", &acquire_thread_handler)
+            .field("release_thread_handler", &release_thread_handler)
             .field("breadth_first", &breadth_first)
             .finish()
     }

rayon-core/src/registry.rs

@@ -6,8 +6,8 @@ use crate::sleep::Sleep;
 use crate::tlv::Tlv;
 use crate::unwind;
 use crate::{
-    DeadlockHandler, ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError,
-    ThreadPoolBuilder, Yield,
+    AcquireThreadHandler, DeadlockHandler, ErrorKind, ExitHandler, PanicHandler,
+    ReleaseThreadHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder, Yield,
 };
 use crossbeam_deque::{Injector, Steal, Stealer, Worker};
 use std::cell::Cell;
@@ -137,9 +137,11 @@ pub struct Registry {
     injected_jobs: Injector<JobRef>,
     broadcasts: Mutex<Vec<Worker<JobRef>>>,
     panic_handler: Option<Box<PanicHandler>>,
-    deadlock_handler: Option<Box<DeadlockHandler>>,
+    pub(crate) deadlock_handler: Option<Box<DeadlockHandler>>,
     start_handler: Option<Box<StartHandler>>,
     exit_handler: Option<Box<ExitHandler>>,
+    pub(crate) acquire_thread_handler: Option<Box<AcquireThreadHandler>>,
+    pub(crate) release_thread_handler: Option<Box<ReleaseThreadHandler>>,
 
     // When this latch reaches 0, it means that all work on this
     // registry must be complete. This is ensured in the following ways:
@@ -294,6 +296,8 @@ impl Registry {
             deadlock_handler: builder.take_deadlock_handler(),
             start_handler: builder.take_start_handler(),
             exit_handler: builder.take_exit_handler(),
+            acquire_thread_handler: builder.take_acquire_thread_handler(),
+            release_thread_handler: builder.take_release_thread_handler(),
         });
 
         // If we return early or panic, make sure to terminate existing threads.
@@ -398,11 +402,24 @@ impl Registry {
 
     /// Waits for the worker threads to stop. This is used for testing
     /// -- so we can check that termination actually works.
-    #[cfg(test)]
     pub(super) fn wait_until_stopped(&self) {
+        self.release_thread();
         for info in &self.thread_infos {
             info.stopped.wait();
         }
+        self.acquire_thread();
+    }
+
+    pub(crate) fn acquire_thread(&self) {
+        if let Some(ref acquire_thread_handler) = self.acquire_thread_handler {
+            acquire_thread_handler();
+        }
+    }
+
+    pub(crate) fn release_thread(&self) {
+        if let Some(ref release_thread_handler) = self.release_thread_handler {
+            release_thread_handler();
+        }
     }
 
     /// ////////////////////////////////////////////////////////////////////////
@@ -448,7 +465,7 @@ impl Registry {
         self.sleep.new_injected_jobs(usize::MAX, 1, queue_was_empty);
     }
 
-    fn has_injected_job(&self) -> bool {
+    pub(crate) fn has_injected_job(&self) -> bool {
         !self.injected_jobs.is_empty()
     }
 
@@ -547,7 +564,9 @@ impl Registry {
                 LatchRef::new(l),
             );
             self.inject(job.as_job_ref());
+            self.release_thread();
             job.latch.wait_and_reset(); // Make sure we can use the same latch again next time.
+            self.acquire_thread();
 
             // flush accumulated logs as we exit the thread
             self.logger.log(|| Flush);
@@ -813,7 +832,7 @@ impl WorkerThread {
         }
     }
 
-    fn has_injected_job(&self) -> bool {
+    pub(super) fn has_injected_job(&self) -> bool {
         !self.stealer.is_empty() || self.registry.has_injected_job()
     }
 
@@ -843,12 +862,9 @@ impl WorkerThread {
                 self.execute(job);
                 idle_state = self.registry.sleep.start_looking(self.index, latch);
             } else {
-                self.registry.sleep.no_work_found(
-                    &mut idle_state,
-                    latch,
-                    || self.has_injected_job(),
-                    &self.registry.deadlock_handler,
-                )
+                self.registry
+                    .sleep
+                    .no_work_found(&mut idle_state, latch, &self)
             }
         }
 
@@ -971,6 +987,7 @@ unsafe fn main_loop(thread: ThreadBuilder) {
         worker: index,
         terminate_addr: my_terminate_latch.as_core_latch().addr(),
     });
+    registry.acquire_thread();
     worker_thread.wait_until(my_terminate_latch);
 
     // Should not be any work left in our queue.
@@ -989,6 +1006,8 @@ unsafe fn main_loop(thread: ThreadBuilder) {
         registry.catch_unwind(|| handler(index));
         // We're already exiting the thread, there's nothing else to do.
     }
+
+    registry.release_thread();
 }
 
 /// If already in a worker-thread, just execute `op`.  Otherwise,

rayon-core/src/sleep/mod.rs

@@ -4,6 +4,7 @@
 use crate::latch::CoreLatch;
 use crate::log::Event::*;
 use crate::log::Logger;
+use crate::registry::WorkerThread;
 use crate::DeadlockHandler;
 use crossbeam_utils::CachePadded;
 use std::sync::atomic::Ordering;
@@ -160,8 +161,7 @@ impl Sleep {
         &self,
         idle_state: &mut IdleState,
         latch: &CoreLatch,
-        has_injected_jobs: impl FnOnce() -> bool,
-        deadlock_handler: &Option<Box<DeadlockHandler>>,
+        thread: &WorkerThread,
     ) {
         if idle_state.rounds < ROUNDS_UNTIL_SLEEPY {
             thread::yield_now();
@@ -175,7 +175,7 @@ impl Sleep {
             thread::yield_now();
         } else {
             debug_assert_eq!(idle_state.rounds, ROUNDS_UNTIL_SLEEPING);
-            self.sleep(idle_state, latch, has_injected_jobs, deadlock_handler);
+            self.sleep(idle_state, latch, thread);
         }
     }
 
@@ -193,13 +193,7 @@ impl Sleep {
     }
 
     #[cold]
-    fn sleep(
-        &self,
-        idle_state: &mut IdleState,
-        latch: &CoreLatch,
-        has_injected_jobs: impl FnOnce() -> bool,
-        deadlock_handler: &Option<Box<DeadlockHandler>>,
-    ) {
+    fn sleep(&self, idle_state: &mut IdleState, latch: &CoreLatch, thread: &WorkerThread) {
         let worker_index = idle_state.worker_index;
 
         if !latch.get_sleepy() {
@@ -266,7 +260,7 @@ impl Sleep {
         // - that job triggers the rollover over the JEC such that we don't see it
         // - we are the last active worker thread
         std::sync::atomic::fence(Ordering::SeqCst);
-        if has_injected_jobs() {
+        if thread.has_injected_job() {
             // If we see an externally injected job, then we have to 'wake
             // ourselves up'. (Ordinarily, `sub_sleeping_thread` is invoked by
             // the one that wakes us.)
@@ -276,7 +270,7 @@ impl Sleep {
                 // Decrement the number of active threads and check for a deadlock
                 let mut data = self.data.lock().unwrap();
                 data.active_threads -= 1;
-                data.deadlock_check(deadlock_handler);
+                data.deadlock_check(&thread.registry.deadlock_handler);
             }
 
             // If we don't see an injected job (the normal case), then flag
@@ -287,10 +281,16 @@ impl Sleep {
             // that whomever is coming to wake us will have to wait until we
             // release the mutex in the call to `wait`, so they will see this
             // boolean as true.)
+            thread.registry.release_thread();
             *is_blocked = true;
             while *is_blocked {
                 is_blocked = sleep_state.condvar.wait(is_blocked).unwrap();
             }
+
+            // Drop `is_blocked` now in case `acquire_thread` blocks
+            drop(is_blocked);
+
+            thread.registry.acquire_thread();
         }
 
         // Update other state:

rayon-core/src/thread_pool/mod.rs

@@ -363,6 +363,12 @@ impl ThreadPool {
         let curr = self.registry.current_thread()?;
         Some(curr.yield_local())
     }
+
+    pub(crate) fn wait_until_stopped(self) {
+        let registry = self.registry.clone();
+        drop(self);
+        registry.wait_until_stopped();
+    }
 }
 
 impl Drop for ThreadPool {