Add callbacks for when threads start and stop doing work

Author: Zoxc

rayon-core/src/lib.rs

@@ -162,6 +162,12 @@ pub struct ThreadPoolBuilder<S = DefaultSpawn> {
     /// 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.
@@ -205,6 +211,8 @@ impl Default for ThreadPoolBuilder {
             exit_handler: None,
             main_handler: None,
             deadlock_handler: None,
+            acquire_thread_handler: None,
+            release_thread_handler: None,
             spawn_handler: DefaultSpawn,
             breadth_first: false,
         }
@@ -217,6 +225,14 @@ impl Default for ThreadPoolBuilder {
 /// 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() -> Self {
@@ -399,6 +415,8 @@ impl<S> ThreadPoolBuilder<S> {
             exit_handler: self.exit_handler,
             main_handler: self.main_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,
         }
     }
@@ -557,6 +575,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()
@@ -746,6 +792,8 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             ref start_handler,
             ref main_handler,
             ref exit_handler,
+            ref acquire_thread_handler,
+            ref release_thread_handler,
             spawn_handler: _,
             ref breadth_first,
         } = *self;
@@ -764,6 +812,8 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
         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)
@@ -774,6 +824,8 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             .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

@@ -25,8 +25,8 @@ use std::usize;
 use unwind;
 use util::leak;
 use {
-    DeadlockHandler, ErrorKind, ExitHandler, MainHandler, PanicHandler, StartHandler,
-    ThreadPoolBuildError, ThreadPoolBuilder,
+    AcquireThreadHandler, DeadlockHandler, ErrorKind, ExitHandler, MainHandler, PanicHandler,
+    ReleaseThreadHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder,
 };
 
 /// Thread builder used for customization via
@@ -141,10 +141,12 @@ 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>>,
     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:
@@ -247,6 +249,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.
@@ -355,9 +359,23 @@ impl Registry {
     /// Waits for the worker threads to stop. This is used for testing
     /// -- so we can check that termination actually works.
     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();
+        }
     }
 
     /// ////////////////////////////////////////////////////////////////////////
@@ -504,7 +522,9 @@ impl Registry {
             LockLatch::new(),
         );
         self.inject(&[job.as_job_ref()]);
+        self.release_thread();
         job.latch.wait();
+        self.acquire_thread();
         job.into_result()
     }
 
@@ -743,11 +763,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);
             }
         }
 
@@ -843,6 +862,8 @@ unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usiz
         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(()) => {}
@@ -874,6 +895,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 log::Event::*;
+use registry::Registry;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Condvar, Mutex};
 use std::thread;
@@ -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 {