Add the ability to customize thread spawning

cuviper · cuviper · commit 6410e7942332 · 2019-05-08T14:57:38.000-07:00
As an alternative to `ThreadPoolBuilder::build()` and `build_global()`,
the new `spawn()` and `spawn_global()` methods take a closure which will
be responsible for spawning the actual threads. This is called with a
`ThreadBuilder` argument that provides the thread index, name, and stack
size, with the expectation to call its `run()` method in the new thread.

The motivating use cases for this are:
- experimental WASM threading, to be externally implemented.
- scoped threads, like the new test using `scoped_tls`.
diff --git a/rayon-core/Cargo.toml b/rayon-core/Cargo.toml
@@ -19,10 +19,12 @@ num_cpus = "1.2"
 lazy_static = "1"
 crossbeam-deque = "0.6.3"
 crossbeam-queue = "0.1.2"
+crossbeam-utils = "0.6.5"
 
 [dev-dependencies]
 rand = "0.6"
 rand_xorshift = "0.1"
+scoped-tls = "1.0"
 
 [target.'cfg(unix)'.dev-dependencies]
 libc = "0.2"
@@ -49,3 +51,7 @@ path = "tests/scope_join.rs"
 [[test]]
 name = "simple_panic"
 path = "tests/simple_panic.rs"
+
+[[test]]
+name = "scoped_threadpool"
+path = "tests/scoped_threadpool.rs"
diff --git a/rayon-core/src/lib.rs b/rayon-core/src/lib.rs
@@ -64,6 +64,7 @@ mod test;
 #[cfg(rayon_unstable)]
 pub mod internal;
 pub use join::{join, join_context};
+pub use registry::ThreadBuilder;
 pub use scope::{scope, Scope};
 pub use scope::{scope_fifo, ScopeFifo};
 pub use spawn::{spawn, spawn_fifo};
@@ -185,6 +186,21 @@ impl ThreadPoolBuilder {
         ThreadPool::build(self)
     }
 
+    /// Create a new `ThreadPool` initialized using this configuration and a
+    /// custom function for spawning threads.
+    ///
+    /// Note that the threads will not exit until after the pool is dropped. It
+    /// is up to the caller to wait for thread termination if that is important
+    /// for any invariants. For instance, threads created in `crossbeam::scope`
+    /// will be joined before that scope returns, and this will block indefinitely
+    /// if the pool is leaked.
+    pub fn spawn(
+        self,
+        spawn: impl FnMut(ThreadBuilder) -> io::Result<()>,
+    ) -> Result<ThreadPool, ThreadPoolBuildError> {
+        ThreadPool::build_spawn(self, spawn)
+    }
+
     /// Initializes the global thread pool. This initialization is
     /// **optional**.  If you do not call this function, the thread pool
     /// will be automatically initialized with the default
@@ -208,6 +224,21 @@ impl ThreadPoolBuilder {
         Ok(())
     }
 
+    /// Initializes the global thread pool using a custom function for spawning
+    /// threads.
+    ///
+    /// Note that the global thread pool doesn't terminate until the entire process
+    /// exits! If this is used with something like `crossbeam::scope` that tries to
+    /// join threads, that will block indefinitely.
+    pub fn spawn_global(
+        self,
+        spawn: impl FnMut(ThreadBuilder) -> io::Result<()>,
+    ) -> Result<(), ThreadPoolBuildError> {
+        let registry = registry::spawn_global_registry(self, spawn)?;
+        registry.wait_until_primed();
+        Ok(())
+    }
+
     /// Get the number of threads that will be used for the thread
     /// pool. See `num_threads()` for more information.
     fn get_num_threads(&self) -> usize {
diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs
@@ -11,7 +11,9 @@ use sleep::Sleep;
 use std::any::Any;
 use std::cell::Cell;
 use std::collections::hash_map::DefaultHasher;
+use std::fmt;
 use std::hash::Hasher;
+use std::io;
 use std::mem;
 use std::ptr;
 #[allow(deprecated)]
@@ -24,6 +26,49 @@ use unwind;
 use util::leak;
 use {ErrorKind, ExitHandler, PanicHandler, StartHandler, ThreadPoolBuildError, ThreadPoolBuilder};
 
+/// Thread builder used for customization via `ThreadPoolBuilder::spawn`.
+pub struct ThreadBuilder {
+    name: Option<String>,
+    stack_size: Option<usize>,
+    worker: Worker<JobRef>,
+    registry: Arc<Registry>,
+    index: usize,
+}
+
+impl ThreadBuilder {
+    /// Get the index of this thread in the pool, within `0..num_threads`.
+    pub fn index(&self) -> usize {
+        self.index
+    }
+
+    /// Get the string that was specified by `ThreadPoolBuilder::name()`.
+    pub fn name(&self) -> Option<&str> {
+        self.name.as_ref().map(String::as_str)
+    }
+
+    /// Get the value that was specified by `ThreadPoolBuilder::stack_size()`.
+    pub fn stack_size(&self) -> Option<usize> {
+        self.stack_size
+    }
+
+    /// Execute the main loop for this thread. This will not return until the
+    /// thread pool is dropped.
+    pub fn run(self) {
+        unsafe { main_loop(self.worker, self.registry, self.index) }
+    }
+}
+
+impl fmt::Debug for ThreadBuilder {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("ThreadBuilder")
+            .field("pool", &self.registry.id())
+            .field("index", &self.index)
+            .field("name", &self.name)
+            .field("stack_size", &self.stack_size)
+            .finish()
+    }
+}
+
 pub(super) struct Registry {
     thread_infos: Vec<ThreadInfo>,
     sleep: Sleep,
@@ -58,39 +103,50 @@ static THE_REGISTRY_SET: Once = ONCE_INIT;
 /// initialization has not already occurred, use the default
 /// configuration.
 fn global_registry() -> &'static Arc<Registry> {
-    THE_REGISTRY_SET.call_once(|| unsafe { init_registry(ThreadPoolBuilder::new()).unwrap() });
-    unsafe { THE_REGISTRY.expect("The global thread pool has not been initialized.") }
+    unsafe {
+        THE_REGISTRY.unwrap_or_else(|| {
+            init_global_registry(ThreadPoolBuilder::new())
+                .expect("The global thread pool has not been initialized.")
+        })
+    }
 }
 
 /// Starts the worker threads (if that has not already happened) with
 /// the given builder.
 pub(super) fn init_global_registry(
     builder: ThreadPoolBuilder,
-) -> Result<&'static Registry, ThreadPoolBuildError> {
-    let mut called = false;
-    let mut init_result = Ok(());;
-    THE_REGISTRY_SET.call_once(|| unsafe {
-        init_result = init_registry(builder);
-        called = true;
-    });
-    if called {
-        init_result?;
-        Ok(&**global_registry())
-    } else {
-        Err(ThreadPoolBuildError::new(
-            ErrorKind::GlobalPoolAlreadyInitialized,
-        ))
-    }
+) -> Result<&'static Arc<Registry>, ThreadPoolBuildError> {
+    set_global_registry(|| Registry::new(builder))
+}
+
+/// Starts the worker threads (if that has not already happened) with
+/// the given builder.
+pub(super) fn spawn_global_registry(
+    builder: ThreadPoolBuilder,
+    spawn: impl FnMut(ThreadBuilder) -> io::Result<()>,
+) -> Result<&'static Arc<Registry>, ThreadPoolBuildError> {
+    set_global_registry(|| Registry::spawn(builder, spawn))
 }
 
-/// Initializes the global registry with the given builder.
-/// Meant to be called from within the `THE_REGISTRY_SET` once
-/// function. Declared `unsafe` because it writes to `THE_REGISTRY` in
-/// an unsynchronized fashion.
-unsafe fn init_registry(builder: ThreadPoolBuilder) -> Result<(), ThreadPoolBuildError> {
-    let registry = Registry::new(builder)?;
-    THE_REGISTRY = Some(leak(registry));
-    Ok(())
+/// Starts the worker threads (if that has not already happened)
+/// by creating a registry with the given callback.
+fn set_global_registry<F>(registry: F) -> Result<&'static Arc<Registry>, ThreadPoolBuildError>
+where
+    F: FnOnce() -> Result<Arc<Registry>, ThreadPoolBuildError>,
+{
+    let mut result = Err(ThreadPoolBuildError::new(
+        ErrorKind::GlobalPoolAlreadyInitialized,
+    ));
+    THE_REGISTRY_SET.call_once(|| {
+        result = registry().map(|registry| {
+            let registry = leak(registry);
+            unsafe {
+                THE_REGISTRY = Some(registry);
+            }
+            registry
+        });
+    });
+    result
 }
 
 struct Terminator<'a>(&'a Arc<Registry>);
@@ -102,7 +158,24 @@ impl<'a> Drop for Terminator<'a> {
 }
 
 impl Registry {
-    pub(super) fn new(mut builder: ThreadPoolBuilder) -> Result<Arc<Self>, ThreadPoolBuildError> {
+    pub(super) fn new(builder: ThreadPoolBuilder) -> Result<Arc<Self>, ThreadPoolBuildError> {
+        Registry::spawn(builder, |thread| {
+            let mut b = thread::Builder::new();
+            if let Some(ref name) = thread.name {
+                b = b.name(name.clone());
+            }
+            if let Some(stack_size) = thread.stack_size {
+                b = b.stack_size(stack_size);
+            }
+            b.spawn(|| thread.run())?;
+            Ok(())
+        })
+    }
+
+    pub(super) fn spawn(
+        mut builder: ThreadPoolBuilder,
+        mut spawn: impl FnMut(ThreadBuilder) -> io::Result<()>,
+    ) -> Result<Arc<Self>, ThreadPoolBuildError> {
         let n_threads = builder.get_num_threads();
         let breadth_first = builder.get_breadth_first();
 
@@ -130,15 +203,14 @@ impl Registry {
         let t1000 = Terminator(&registry);
 
         for (index, worker) in workers.into_iter().enumerate() {
-            let registry = registry.clone();
-            let mut b = thread::Builder::new();
-            if let Some(name) = builder.get_thread_name(index) {
-                b = b.name(name);
-            }
-            if let Some(stack_size) = builder.get_stack_size() {
-                b = b.stack_size(stack_size);
-            }
-            if let Err(e) = b.spawn(move || unsafe { main_loop(worker, registry, index) }) {
+            let thread = ThreadBuilder {
+                name: builder.get_thread_name(index),
+                stack_size: builder.get_stack_size(),
+                registry: registry.clone(),
+                worker,
+                index,
+            };
+            if let Err(e) = spawn(thread) {
                 return Err(ThreadPoolBuildError::new(ErrorKind::IOError(e)));
             }
         }
diff --git a/rayon-core/src/thread_pool/mod.rs b/rayon-core/src/thread_pool/mod.rs
@@ -8,12 +8,13 @@ use registry::{Registry, WorkerThread};
 use spawn;
 use std::error::Error;
 use std::fmt;
+use std::io;
 use std::sync::Arc;
 #[allow(deprecated)]
 use Configuration;
 use {scope, Scope};
 use {scope_fifo, ScopeFifo};
-use {ThreadPoolBuildError, ThreadPoolBuilder};
+use {ThreadBuilder, ThreadPoolBuildError, ThreadPoolBuilder};
 
 mod internal;
 mod test;
@@ -66,6 +67,14 @@ impl ThreadPool {
         Ok(ThreadPool { registry })
     }
 
+    pub(super) fn build_spawn(
+        builder: ThreadPoolBuilder,
+        spawn: impl FnMut(ThreadBuilder) -> io::Result<()>,
+    ) -> Result<ThreadPool, ThreadPoolBuildError> {
+        let registry = Registry::spawn(builder, spawn)?;
+        Ok(ThreadPool { registry })
+    }
+
     /// Returns a handle to the global thread pool. This is the pool
     /// that Rayon will use by default when you perform a `join()` or
     /// `scope()` operation, if no other thread-pool is installed. If
diff --git a/rayon-core/tests/scoped_threadpool.rs b/rayon-core/tests/scoped_threadpool.rs
@@ -0,0 +1,68 @@
+extern crate crossbeam_utils;
+extern crate rayon_core;
+
+#[macro_use]
+extern crate scoped_tls;
+
+use crossbeam_utils::thread;
+use rayon_core::ThreadPoolBuilder;
+
+#[derive(PartialEq, Eq, Debug)]
+struct Local(i32);
+
+scoped_thread_local!(static LOCAL: Local);
+
+#[test]
+fn scoped_tls_missing() {
+    LOCAL.set(&Local(42), || {
+        let pool = ThreadPoolBuilder::new()
+            .build()
+            .expect("thread pool created");
+
+        // `LOCAL` is not set in the pool.
+        pool.install(|| {
+            assert!(!LOCAL.is_set());
+        });
+    });
+}
+
+#[test]
+fn scoped_tls_threadpool() {
+    LOCAL.set(&Local(42), || {
+        LOCAL.with(|x| {
+            thread::scope(|scope| {
+                let pool = ThreadPoolBuilder::new()
+                    .spawn(move |thread| {
+                        scope
+                            .builder()
+                            .spawn(move |_| {
+                                // Borrow the same local value in the thread pool.
+                                LOCAL.set(x, || thread.run())
+                            })
+                            .map(|_| ())
+                    })
+                    .expect("thread pool created");
+
+                // 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("scope threads ok");
+            // `thread::scope` will wait for the threads to exit before returning.
+        });
+    });
+}
