Add the ability to create scoped thread pools

Author: Zoxc

rayon-core/src/registry.rs

@@ -352,7 +352,6 @@ 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) {
         for info in &self.thread_infos {
             info.stopped.wait();

rayon-core/src/thread_pool/mod.rs

@@ -92,6 +92,43 @@ 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.