Add a main_handler which is passed an argument to run the proper main loop

Author: Zoxc

rayon-core/src/lib.rs

@@ -156,6 +156,9 @@ pub struct ThreadPoolBuilder<S = DefaultSpawn> {
     /// Closure invoked to spawn threads.
     spawn_handler: S,
 
+    /// Closure invoked on worker thread start.
+    main_handler: Option<Box<MainHandler>>,
+
     /// 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.
@@ -194,12 +197,19 @@ impl Default for ThreadPoolBuilder {
             stack_size: None,
             start_handler: None,
             exit_handler: None,
+            main_handler: None,
             spawn_handler: DefaultSpawn,
             breadth_first: false,
         }
     }
 }
 
+/// The type for a closure that gets invoked with a
+/// function which runs rayon tasks.
+/// The closure is passed the index of the thread on which it is invoked.
+/// Note that this same closure may be invoked multiple times in parallel.
+type MainHandler = Fn(usize, &mut FnMut()) + Send + Sync;
+
 impl ThreadPoolBuilder {
     /// Creates and returns a valid rayon thread pool builder, but does not initialize it.
     pub fn new() -> Self {
@@ -380,6 +390,7 @@ impl<S> ThreadPoolBuilder<S> {
             stack_size: self.stack_size,
             start_handler: self.start_handler,
             exit_handler: self.exit_handler,
+            main_handler: self.main_handler,
             breadth_first: self.breadth_first,
         }
     }
@@ -575,6 +586,24 @@ impl<S> ThreadPoolBuilder<S> {
         self.exit_handler = Some(Box::new(exit_handler));
         self
     }
+
+    /// Takes the current thread main callback, leaving `None`.
+    fn take_main_handler(&mut self) -> Option<Box<MainHandler>> {
+        self.main_handler.take()
+    }
+
+    /// Set a callback to be invoked on thread main.
+    ///
+    /// The closure is passed the index of the thread on which it is invoked.
+    /// Note that this same closure may be invoked multiple times in parallel.
+    /// If this closure panics, the panic will be passed to the panic handler.
+    pub fn main_handler<H>(mut self, main_handler: H) -> Self
+    where
+        H: Fn(usize, &mut FnMut()) + Send + Sync + 'static,
+    {
+        self.main_handler = Some(Box::new(main_handler));
+        self
+    }
 }
 
 #[allow(deprecated)]
@@ -692,6 +721,7 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             ref panic_handler,
             ref stack_size,
             ref start_handler,
+            ref main_handler,
             ref exit_handler,
             spawn_handler: _,
             ref breadth_first,
@@ -709,6 +739,7 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
         let panic_handler = panic_handler.as_ref().map(|_| ClosurePlaceholder);
         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);
 
         f.debug_struct("ThreadPoolBuilder")
             .field("num_threads", num_threads)
@@ -717,6 +748,7 @@ impl<S> fmt::Debug for ThreadPoolBuilder<S> {
             .field("stack_size", &stack_size)
             .field("start_handler", &start_handler)
             .field("exit_handler", &exit_handler)
+            .field("main_handler", &main_handler)
             .field("breadth_first", &breadth_first)
             .finish()
     }

rayon-core/src/registry.rs

@@ -24,7 +24,10 @@ use std::thread;
 use std::usize;
 use unwind;
 use util::leak;
-use {ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder};
+use {
+    ErrorKind, ExitHandler, MainHandler, PanicHandler, StartHandler, ThreadPoolBuildError,
+    ThreadPoolBuilder,
+};
 
 /// Thread builder used for customization via
 /// [`ThreadPoolBuilder::spawn_handler`](struct.ThreadPoolBuilder.html#method.spawn_handler).
@@ -140,6 +143,7 @@ pub(super) struct Registry {
     panic_handler: Option<Box<PanicHandler>>,
     start_handler: Option<Box<StartHandler>>,
     exit_handler: Option<Box<ExitHandler>>,
+    main_handler: Option<Box<MainHandler>>,
 
     // When this latch reaches 0, it means that all work on this
     // registry must be complete. This is ensured in the following ways:
@@ -239,6 +243,7 @@ impl Registry {
             terminate_latch: CountLatch::new(),
             panic_handler: builder.take_panic_handler(),
             start_handler: builder.take_start_handler(),
+            main_handler: builder.take_main_handler(),
             exit_handler: builder.take_exit_handler(),
         });
 
@@ -811,7 +816,20 @@ unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usiz
         }
     }
 
-    worker_thread.wait_until(&registry.terminate_latch);
+    let mut work = || {
+        worker_thread.wait_until(&registry.terminate_latch);
+    };
+
+    if let Some(ref handler) = registry.main_handler {
+        match unwind::halt_unwinding(|| handler(index, &mut work)) {
+            Ok(()) => {}
+            Err(err) => {
+                registry.handle_panic(err);
+            }
+        }
+    } else {
+        work();
+    }
 
     // Should not be any work left in our queue.
     debug_assert!(worker_thread.take_local_job().is_none());