Add a WorkerLocal type which allow you to hold a value per Rayon worker thread

Co-authored-by: Ammar Askar <[email protected]>

Author: Zoxc

rayon-core/src/lib.rs

@@ -89,6 +89,7 @@ mod sleep;
 mod spawn;
 mod thread_pool;
 mod unwind;
+mod worker_local;
 
 mod compile_fail;
 mod test;
@@ -105,6 +106,7 @@ pub use self::thread_pool::current_thread_has_pending_tasks;
 pub use self::thread_pool::current_thread_index;
 pub use self::thread_pool::ThreadPool;
 pub use self::thread_pool::{yield_local, yield_now, Yield};
+pub use worker_local::WorkerLocal;
 
 use self::registry::{CustomSpawn, DefaultSpawn, ThreadSpawn};
 

rayon-core/src/registry.rs

@@ -676,12 +676,12 @@ pub(super) struct WorkerThread {
     /// local queue used for `spawn_fifo` indirection
     fifo: JobFifo,
 
-    index: usize,
+    pub(crate) index: usize,
 
     /// A weak random number generator.
     rng: XorShift64Star,
 
-    registry: Arc<Registry>,
+    pub(crate) registry: Arc<Registry>,
 }
 
 // This is a bit sketchy, but basically: the WorkerThread is

rayon-core/src/worker_local.rs

@@ -0,0 +1,78 @@
+use crate::registry::{Registry, WorkerThread};
+use std::fmt;
+use std::ops::Deref;
+use std::sync::Arc;
+
+#[repr(align(64))]
+#[derive(Debug)]
+struct CacheAligned<T>(T);
+
+/// Holds worker-locals values for each thread in a thread pool.
+/// You can only access the worker local value through the Deref impl
+/// on the thread pool it was constructed on. It will panic otherwise
+pub struct WorkerLocal<T> {
+    locals: Vec<CacheAligned<T>>,
+    registry: Arc<Registry>,
+}
+
+/// We prevent concurrent access to the underlying value in the
+/// Deref impl, thus any values safe to send across threads can
+/// be used with WorkerLocal.
+unsafe impl<T: Send> Sync for WorkerLocal<T> {}
+
+impl<T> WorkerLocal<T> {
+    /// Creates a new worker local where the `initial` closure computes the
+    /// value this worker local should take for each thread in the thread pool.
+    #[inline]
+    pub fn new<F: FnMut(usize) -> T>(mut initial: F) -> WorkerLocal<T> {
+        let registry = Registry::current();
+        WorkerLocal {
+            locals: (0..registry.num_threads())
+                .map(|i| CacheAligned(initial(i)))
+                .collect(),
+            registry,
+        }
+    }
+
+    /// Returns the worker-local value for each thread
+    #[inline]
+    pub fn into_inner(self) -> Vec<T> {
+        self.locals.into_iter().map(|c| c.0).collect()
+    }
+
+    fn current(&self) -> &T {
+        unsafe {
+            let worker_thread = WorkerThread::current();
+            if worker_thread.is_null()
+                || &*(*worker_thread).registry as *const _ != &*self.registry as *const _
+            {
+                panic!("WorkerLocal can only be used on the thread pool it was created on")
+            }
+            &self.locals[(*worker_thread).index].0
+        }
+    }
+}
+
+impl<T> WorkerLocal<Vec<T>> {
+    /// Joins the elements of all the worker locals into one Vec
+    pub fn join(self) -> Vec<T> {
+        self.into_inner().into_iter().flat_map(|v| v).collect()
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for WorkerLocal<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("WorkerLocal")
+            .field("registry", &self.registry.id())
+            .finish()
+    }
+}
+
+impl<T> Deref for WorkerLocal<T> {
+    type Target = T;
+
+    #[inline(always)]
+    fn deref(&self) -> &T {
+        self.current()
+    }
+}