Add ThreadPoolBuilder::build_scoped()

Author: cuviper

rayon-core/src/lib.rs

@@ -34,6 +34,7 @@ use std::str::FromStr;
 
 extern crate crossbeam_deque;
 extern crate crossbeam_queue;
+extern crate crossbeam_utils;
 #[cfg(any(debug_assertions, rayon_unstable))]
 #[macro_use]
 extern crate lazy_static;
@@ -186,6 +187,37 @@ impl ThreadPoolBuilder {
         ThreadPool::build(self)
     }
 
+    /// Create a scoped `ThreadPool` initialized using this configuration.
+    ///
+    /// The threads in this pool will start by calling `wrapper`, which should
+    /// do initialization and continue by calling `ThreadBuilder::run()`.
+    pub fn build_scoped<W, F, R>(self, wrapper: W, with_pool: F) -> Result<R, ThreadPoolBuildError>
+    where
+        W: Fn(ThreadBuilder) + Sync, // expected to call `run()`
+        F: FnOnce(&ThreadPool) -> R,
+    {
+        let result = crossbeam_utils::thread::scope(|scope| {
+            let wrapper = &wrapper;
+            let pool = self.spawn(|thread| {
+                let mut builder = scope.builder();
+                if let Some(name) = thread.name() {
+                    builder = builder.name(name.to_string());
+                }
+                if let Some(size) = thread.stack_size() {
+                    builder = builder.stack_size(size);
+                }
+                builder.spawn(move |_| wrapper(thread))?;
+                Ok(())
+            })?;
+            Ok(with_pool(&pool))
+        });
+
+        match result {
+            Ok(result) => result,
+            Err(err) => unwind::resume_unwinding(err),
+        }
+    }
+
     /// Create a new `ThreadPool` initialized using this configuration and a
     /// custom function for spawning threads.
     ///

rayon-core/tests/scoped_threadpool.rs

@@ -13,7 +13,7 @@ struct Local(i32);
 scoped_thread_local!(static LOCAL: Local);
 
 #[test]
-fn scoped_tls_missing() {
+fn missing_scoped_tls() {
     LOCAL.set(&Local(42), || {
         let pool = ThreadPoolBuilder::new()
             .build()
@@ -27,7 +27,7 @@ fn scoped_tls_missing() {
 }
 
 #[test]
-fn scoped_tls_threadpool() {
+fn spawn_scoped_tls_threadpool() {
     LOCAL.set(&Local(42), || {
         LOCAL.with(|x| {
             thread::scope(|scope| {
@@ -66,3 +66,36 @@ fn scoped_tls_threadpool() {
         });
     });
 }
+
+#[test]
+fn build_scoped_tls_threadpool() {
+    LOCAL.set(&Local(42), || {
+        LOCAL.with(|x| {
+            ThreadPoolBuilder::new()
+                .build_scoped(
+                    move |thread| LOCAL.set(x, || thread.run()),
+                    |pool| {
+                        // The pool matches our local value.
+                        pool.install(|| {
+                            assert!(LOCAL.is_set());
+                            LOCAL.with(|y| {
+                                assert_eq!(x, y);
+                            });
+                        });
+
+                        // If we change our local value, the pool is not affected.
+                        LOCAL.set(&Local(-1), || {
+                            pool.install(|| {
+                                assert!(LOCAL.is_set());
+                                LOCAL.with(|y| {
+                                    assert_eq!(x, y);
+                                });
+                            });
+                        });
+                    },
+                )
+                .expect("thread pool created");
+            // Internally, `crossbeam::scope` will wait for the threads to exit before returning.
+        });
+    });
+}