From 9815d97db161c0226e8eba2ab69505f4961d0fee Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Fri, 4 May 2018 06:59:33 +0200
Subject: [PATCH] Add the ability to create scoped thread pools

---
 rayon-core/src/registry.rs        |  3 +--
 rayon-core/src/thread_pool/mod.rs | 34 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs
index bfa7eee8c..96b514655 100644
--- a/rayon-core/src/registry.rs
+++ b/rayon-core/src/registry.rs
@@ -219,8 +219,7 @@ 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 fn wait_until_stopped(&self) {
+    pub(crate) fn wait_until_stopped(&self) {
         for info in &self.thread_infos {
             info.stopped.wait();
         }
diff --git a/rayon-core/src/thread_pool/mod.rs b/rayon-core/src/thread_pool/mod.rs
index d13f4e5f9..690052a90 100644
--- a/rayon-core/src/thread_pool/mod.rs
+++ b/rayon-core/src/thread_pool/mod.rs
@@ -86,6 +86,40 @@ impl ThreadPool {
         &DEFAULT_THREAD_POOL
     }
 
+    /// Creates a scoped thread pool
+    pub fn scoped_pool<F, R, H>(builder: ThreadPoolBuilder,
+                                main_handler: H,
+                                with_pool: F) -> Result<R, ThreadPoolBuildError>
+    where F: FnOnce(&ThreadPool) -> R,
+          H: Fn(&mut FnMut()) + Send + Sync
+    {
+        struct Handler(*const ());
+        unsafe impl Send for Handler {}
+        unsafe impl Sync for Handler {}
+
+        let handler = Handler(&main_handler as *const _ as *const ());
+
+        let builder = builder.main_handler(move |_, worker| {
+            let handler = unsafe { &*(handler.0 as *const H) };
+            handler(worker);
+        });
+
+        let pool = builder.build()?;
+
+        struct JoinRegistry(Arc<Registry>);
+
+        impl Drop for JoinRegistry {
+            fn drop(&mut self) {
+                self.0.terminate();
+                self.0.wait_until_stopped();
+            }
+        }
+
+        let _join_registry = JoinRegistry(pool.registry.clone());
+
+        Ok(with_pool(&pool))
+    }
+
     /// Executes `op` within the threadpool. Any attempts to use
     /// `join`, `scope`, or parallel iterators will then operate
     /// within that threadpool.