Add callbacks for when threads start and stop doing work

Author: Zoxc

rayon-core/src/lib.rs

@@ -144,6 +144,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.
@@ -186,12 +192,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 {
@@ -292,7 +308,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 {
@@ -371,6 +392,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,
         }
     }
@@ -529,6 +552,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()
@@ -699,6 +750,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;
@@ -716,6 +769,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)
@@ -725,6 +780,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

@@ -5,8 +5,8 @@ use crate::sleep::Sleep;
 use crate::unwind;
 use crate::util::leak;
 use crate::{
-    DeadlockHandler, ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError,
-    ThreadPoolBuilder,
+    AcquireThreadHandler, DeadlockHandler, ErrorKind, ExitHandler, PanicHandler,
+    ReleaseThreadHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder,
 };
 use crossbeam_deque::{Steal, Stealer, Worker};
 use crossbeam_queue::SegQueue;
@@ -137,9 +137,11 @@ pub struct Registry {
     sleep: Sleep,
     injected_jobs: SegQueue<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:
@@ -244,6 +246,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.
@@ -346,11 +350,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();
+        }
     }
 
     /// ////////////////////////////////////////////////////////////////////////
@@ -453,7 +470,9 @@ impl Registry {
                 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();
             job.into_result()
         })
     }
@@ -687,11 +706,10 @@ impl WorkerThread {
                 yields = self.registry.sleep.work_found(self.index, yields);
                 self.execute(job);
             } else {
-                yields = self.registry.sleep.no_work_found(
-                    self.index,
-                    yields,
-                    &self.registry.deadlock_handler,
-                );
+                yields = self
+                    .registry
+                    .sleep
+                    .no_work_found(self.index, yields, &self.registry);
             }
         }
 
@@ -783,6 +801,7 @@ unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usiz
         }
     }
 
+    registry.acquire_thread();
     worker_thread.wait_until(&registry.terminate_latch);
 
     // Should not be any work left in our queue.
@@ -805,6 +824,8 @@ unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usiz
         }
         // 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

@@ -2,6 +2,7 @@
 //! for an overview.
 
 use crate::log::Event::*;
+use crate::registry::Registry;
 use crate::DeadlockHandler;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Condvar, Mutex};
@@ -117,7 +118,7 @@ impl Sleep {
         &self,
         worker_index: usize,
         yields: usize,
-        deadlock_handler: &Option<Box<DeadlockHandler>>,
+        registry: &Registry,
     ) -> usize {
         log!(DidNotFindWork {
             worker: worker_index,
@@ -145,7 +146,7 @@ impl Sleep {
             }
         } else {
             debug_assert_eq!(yields, ROUNDS_UNTIL_ASLEEP);
-            self.sleep(worker_index, deadlock_handler);
+            self.sleep(worker_index, registry);
             0
         }
     }
@@ -248,7 +249,7 @@ impl Sleep {
         self.worker_is_sleepy(state, worker_index)
     }
 
-    fn sleep(&self, worker_index: usize, deadlock_handler: &Option<Box<DeadlockHandler>>) {
+    fn sleep(&self, worker_index: usize, registry: &Registry) {
         loop {
             // Acquire here suffices. If we observe that the current worker is still
             // sleepy, then in fact we know that no writes have occurred, and anyhow
@@ -327,12 +328,15 @@ impl Sleep {
 
                     // Decrement the number of active threads and check for a deadlock
                     data.active_threads -= 1;
-                    data.deadlock_check(deadlock_handler);
+                    data.deadlock_check(&registry.deadlock_handler);
+
+                    registry.release_thread();
 
                     let _ = self.tickle.wait(data).unwrap();
                     log!(GotAwoken {
                         worker: worker_index
                     });
+                    registry.acquire_thread();
                     return;
                 }
             } else {

rayon-core/src/thread_pool/mod.rs

@@ -252,6 +252,12 @@ impl ThreadPool {
         // We assert that `self.registry` has not terminated.
         unsafe { spawn::spawn_fifo_in(op, &self.registry) }
     }
+
+    pub(crate) fn wait_until_stopped(self) {
+        let registry = self.registry.clone();
+        drop(self);
+        registry.wait_until_stopped();
+    }
 }
 
 impl Drop for ThreadPool {