1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
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()
    }
}