rustc_thread_pool/
worker_local.rs

1use std::fmt;
2use std::ops::Deref;
3use std::sync::Arc;
4
5use crate::registry::{Registry, WorkerThread};
6
7#[repr(align(64))]
8#[derive(Debug)]
9struct CacheAligned<T>(T);
10
11/// Holds worker-locals values for each thread in a thread pool.
12/// You can only access the worker local value through the Deref impl
13/// on the thread pool it was constructed on. It will panic otherwise
14pub struct WorkerLocal<T> {
15    locals: Vec<CacheAligned<T>>,
16    registry: Arc<Registry>,
17}
18
19/// We prevent concurrent access to the underlying value in the
20/// Deref impl, thus any values safe to send across threads can
21/// be used with WorkerLocal.
22unsafe impl<T: Send> Sync for WorkerLocal<T> {}
23
24impl<T> WorkerLocal<T> {
25    /// Creates a new worker local where the `initial` closure computes the
26    /// value this worker local should take for each thread in the thread pool.
27    #[inline]
28    pub fn new<F: FnMut(usize) -> T>(mut initial: F) -> WorkerLocal<T> {
29        let registry = Registry::current();
30        WorkerLocal {
31            locals: (0..registry.num_threads()).map(|i| CacheAligned(initial(i))).collect(),
32            registry,
33        }
34    }
35
36    /// Returns the worker-local value for each thread
37    #[inline]
38    pub fn into_inner(self) -> Vec<T> {
39        self.locals.into_iter().map(|c| c.0).collect()
40    }
41
42    fn current(&self) -> &T {
43        unsafe {
44            let worker_thread = WorkerThread::current();
45            if worker_thread.is_null()
46                || &*(*worker_thread).registry as *const _ != &*self.registry as *const _
47            {
48                panic!("WorkerLocal can only be used on the thread pool it was created on")
49            }
50            &self.locals[(*worker_thread).index].0
51        }
52    }
53}
54
55impl<T> WorkerLocal<Vec<T>> {
56    /// Joins the elements of all the worker locals into one Vec
57    pub fn join(self) -> Vec<T> {
58        self.into_inner().into_iter().flat_map(|v| v).collect()
59    }
60}
61
62impl<T: fmt::Debug> fmt::Debug for WorkerLocal<T> {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        f.debug_struct("WorkerLocal").field("registry", &self.registry.id()).finish()
65    }
66}
67
68impl<T> Deref for WorkerLocal<T> {
69    type Target = T;
70
71    #[inline(always)]
72    fn deref(&self) -> &T {
73        self.current()
74    }
75}
rustc_thread_pool/worker_local.rs

rustc_thread_pool/
worker_local.rs
