Add callbacks for when threads start and stop doing work

Author: Zoxc

rayon-core/src/lib.rs

@@ -152,6 +152,12 @@ pub struct ThreadPoolBuilder {
     /// Closure invoked on worker thread start.
     main_handler: Option<Box<MainHandler>>,
 
+    /// 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.
@@ -190,6 +196,14 @@ type ExitHandler = Fn(usize) + Send + Sync;
 /// Note that this same closure may be invoked multiple times in parallel.
 type MainHandler = Fn(usize, &mut FnMut()) + Send + Sync;
 
+/// 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 = 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 = Fn() + Send + Sync;
+
 impl ThreadPoolBuilder {
     /// Creates and returns a valid rayon thread pool builder, but does not initialize it.
     pub fn new() -> ThreadPoolBuilder {
@@ -363,6 +377,32 @@ impl ThreadPoolBuilder {
         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) -> ThreadPoolBuilder
+        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) -> ThreadPoolBuilder
+        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()
@@ -546,10 +586,12 @@ impl fmt::Debug for ThreadPoolBuilder {
             ref get_thread_name,
             ref panic_handler,
             ref stack_size,
-            ref deadlock_handler, 
+            ref deadlock_handler,
             ref start_handler,
             ref main_handler,
             ref exit_handler,
+            ref acquire_thread_handler,
+            ref release_thread_handler,
             ref breadth_first,
         } = *self;
 
@@ -567,6 +609,8 @@ impl fmt::Debug for ThreadPoolBuilder {
         let start_handler = start_handler.as_ref().map(|_| ClosurePlaceholder);
         let exit_handler = exit_handler.as_ref().map(|_| ClosurePlaceholder);
         let main_handler = main_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)
@@ -577,6 +621,8 @@ impl fmt::Debug for ThreadPoolBuilder {
             .field("start_handler", &start_handler)
             .field("exit_handler", &exit_handler)
             .field("main_handler", &main_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

@@ -18,19 +18,21 @@ use std::thread;
 use std::usize;
 use unwind;
 use util::leak;
-use {ErrorKind, ExitHandler, PanicHandler, DeadlockHandler, StartHandler,
-     MainHandler, ThreadPoolBuildError, ThreadPoolBuilder};
+use {ErrorKind, ExitHandler, PanicHandler, DeadlockHandler, StartHandler, AcquireThreadHandler, 
+     MainHandler, ThreadPoolBuildError, ThreadPoolBuilder, ReleaseThreadHandler};
 
 pub struct Registry {
     thread_infos: Vec<ThreadInfo>,
     state: Mutex<RegistryState>,
     sleep: Sleep,
     job_uninjector: Stealer<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>>,
     main_handler: Option<Box<MainHandler>>,
+    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:
@@ -123,6 +125,8 @@ impl Registry {
             start_handler: builder.take_start_handler(),
             main_handler: builder.take_main_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.
@@ -222,9 +226,23 @@ impl Registry {
     /// Waits for the worker threads to stop. This is used for testing
     /// -- so we can check that termination actually works.
     pub(crate) 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();
+        }
     }
 
     /// ////////////////////////////////////////////////////////////////////////
@@ -375,7 +393,9 @@ impl Registry {
             op(&*worker_thread, true)
         }, LockLatch::new());
         self.inject(&[job.as_job_ref()]);
+        self.release_thread();
         job.latch.wait();
+        self.acquire_thread();
         job.into_result()
     }
 
@@ -614,7 +634,7 @@ impl WorkerThread {
                 yields = self.registry.sleep.no_work_found(
                     self.index,
                     yields,
-                    &self.registry.deadlock_handler
+                    &self.registry
                 );
             }
         }
@@ -716,6 +736,8 @@ unsafe fn main_loop(
         worker_thread.wait_until(&registry.terminate_latch);
     };
 
+    registry.acquire_thread();
+
     if let Some(ref handler) = registry.main_handler {
         match unwind::halt_unwinding(|| handler(index, &mut work)) {
             Ok(()) => {
@@ -748,6 +770,8 @@ unsafe fn main_loop(
         }
         // 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

@@ -3,6 +3,7 @@
 
 use DeadlockHandler;
 use log::Event::*;
+use registry::Registry;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Condvar, Mutex};
 use std::thread;
@@ -113,7 +114,12 @@ impl Sleep {
     }
 
     #[inline]
-    pub fn no_work_found(&self, worker_index: usize, yields: usize, deadlock_handler: &Option<Box<DeadlockHandler>>) -> usize {
+    pub fn no_work_found(
+        &self,
+        worker_index: usize,
+        yields: usize,
+        registry: &Registry,
+    ) -> usize {
         log!(DidNotFindWork {
             worker: worker_index,
             yields: yields,
@@ -140,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
         }
     }
@@ -243,7 +249,11 @@ 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
@@ -322,12 +332,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 {