|
| 1 | +use crate::sync::Lock; |
| 2 | +use std::cell::Cell; |
| 3 | +use std::cell::OnceCell; |
| 4 | +use std::ops::Deref; |
| 5 | +use std::ptr; |
| 6 | +use std::sync::Arc; |
| 7 | + |
| 8 | +#[cfg(parallel_compiler)] |
| 9 | +use {crate::cold_path, crate::sync::CacheAligned}; |
| 10 | + |
| 11 | +/// A pointer to the `RegistryData` which uniquely identifies a registry. |
| 12 | +/// This identifier can be reused if the registry gets freed. |
| 13 | +#[derive(Clone, Copy, PartialEq)] |
| 14 | +struct RegistryId(*const RegistryData); |
| 15 | + |
| 16 | +impl RegistryId { |
| 17 | + #[inline(always)] |
| 18 | + /// Verifies that the current thread is associated with the registry and returns its unique |
| 19 | + /// index within the registry. This panics if the current thread is not associated with this |
| 20 | + /// registry. |
| 21 | + /// |
| 22 | + /// Note that there's a race possible where the identifer in `THREAD_DATA` could be reused |
| 23 | + /// so this can succeed from a different registry. |
| 24 | + #[cfg(parallel_compiler)] |
| 25 | + fn verify(self) -> usize { |
| 26 | + let (id, index) = THREAD_DATA.with(|data| (data.registry_id.get(), data.index.get())); |
| 27 | + |
| 28 | + if id == self { |
| 29 | + index |
| 30 | + } else { |
| 31 | + cold_path(|| panic!("Unable to verify registry association")) |
| 32 | + } |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +struct RegistryData { |
| 37 | + thread_limit: usize, |
| 38 | + threads: Lock<usize>, |
| 39 | +} |
| 40 | + |
| 41 | +/// Represents a list of threads which can access worker locals. |
| 42 | +#[derive(Clone)] |
| 43 | +pub struct Registry(Arc<RegistryData>); |
| 44 | + |
| 45 | +thread_local! { |
| 46 | + /// The registry associated with the thread. |
| 47 | + /// This allows the `WorkerLocal` type to clone the registry in its constructor. |
| 48 | + static REGISTRY: OnceCell<Registry> = OnceCell::new(); |
| 49 | +} |
| 50 | + |
| 51 | +struct ThreadData { |
| 52 | + registry_id: Cell<RegistryId>, |
| 53 | + index: Cell<usize>, |
| 54 | +} |
| 55 | + |
| 56 | +thread_local! { |
| 57 | + /// A thread local which contains the identifer of `REGISTRY` but allows for faster access. |
| 58 | + /// It also holds the index of the current thread. |
| 59 | + static THREAD_DATA: ThreadData = const { ThreadData { |
| 60 | + registry_id: Cell::new(RegistryId(ptr::null())), |
| 61 | + index: Cell::new(0), |
| 62 | + }}; |
| 63 | +} |
| 64 | + |
| 65 | +impl Registry { |
| 66 | + /// Creates a registry which can hold up to `thread_limit` threads. |
| 67 | + pub fn new(thread_limit: usize) -> Self { |
| 68 | + Registry(Arc::new(RegistryData { thread_limit, threads: Lock::new(0) })) |
| 69 | + } |
| 70 | + |
| 71 | + /// Gets the registry associated with the current thread. Panics if there's no such registry. |
| 72 | + pub fn current() -> Self { |
| 73 | + REGISTRY.with(|registry| registry.get().cloned().expect("No assocated registry")) |
| 74 | + } |
| 75 | + |
| 76 | + /// Registers the current thread with the registry so worker locals can be used on it. |
| 77 | + /// Panics if the thread limit is hit or if the thread already has an associated registry. |
| 78 | + pub fn register(&self) { |
| 79 | + let mut threads = self.0.threads.lock(); |
| 80 | + if *threads < self.0.thread_limit { |
| 81 | + REGISTRY.with(|registry| { |
| 82 | + if registry.get().is_some() { |
| 83 | + drop(threads); |
| 84 | + panic!("Thread already has a registry"); |
| 85 | + } |
| 86 | + registry.set(self.clone()).ok(); |
| 87 | + THREAD_DATA.with(|data| { |
| 88 | + data.registry_id.set(self.id()); |
| 89 | + data.index.set(*threads); |
| 90 | + }); |
| 91 | + *threads += 1; |
| 92 | + }); |
| 93 | + } else { |
| 94 | + drop(threads); |
| 95 | + panic!("Thread limit reached"); |
| 96 | + } |
| 97 | + } |
| 98 | + |
| 99 | + /// Gets the identifer of this registry. |
| 100 | + fn id(&self) -> RegistryId { |
| 101 | + RegistryId(&*self.0) |
| 102 | + } |
| 103 | +} |
| 104 | + |
| 105 | +/// Holds worker local values for each possible thread in a registry. You can only access the |
| 106 | +/// worker local value through the `Deref` impl on the registry associated with the thread it was |
| 107 | +/// created on. It will panic otherwise. |
| 108 | +pub struct WorkerLocal<T> { |
| 109 | + #[cfg(not(parallel_compiler))] |
| 110 | + local: T, |
| 111 | + #[cfg(parallel_compiler)] |
| 112 | + locals: Box<[CacheAligned<T>]>, |
| 113 | + #[cfg(parallel_compiler)] |
| 114 | + registry: Registry, |
| 115 | +} |
| 116 | + |
| 117 | +// This is safe because the `deref` call will return a reference to a `T` unique to each thread |
| 118 | +// or it will panic for threads without an associated local. So there isn't a need for `T` to do |
| 119 | +// it's own synchronization. The `verify` method on `RegistryId` has an issue where the the id |
| 120 | +// can be reused, but `WorkerLocal` has a reference to `Registry` which will prevent any reuse. |
| 121 | +#[cfg(parallel_compiler)] |
| 122 | +unsafe impl<T: Send> Sync for WorkerLocal<T> {} |
| 123 | + |
| 124 | +impl<T> WorkerLocal<T> { |
| 125 | + /// Creates a new worker local where the `initial` closure computes the |
| 126 | + /// value this worker local should take for each thread in the registry. |
| 127 | + #[inline] |
| 128 | + pub fn new<F: FnMut(usize) -> T>(mut initial: F) -> WorkerLocal<T> { |
| 129 | + #[cfg(parallel_compiler)] |
| 130 | + { |
| 131 | + let registry = Registry::current(); |
| 132 | + WorkerLocal { |
| 133 | + locals: (0..registry.0.thread_limit).map(|i| CacheAligned(initial(i))).collect(), |
| 134 | + registry, |
| 135 | + } |
| 136 | + } |
| 137 | + #[cfg(not(parallel_compiler))] |
| 138 | + { |
| 139 | + WorkerLocal { local: initial(0) } |
| 140 | + } |
| 141 | + } |
| 142 | + |
| 143 | + /// Returns the worker-local values for each thread |
| 144 | + #[inline] |
| 145 | + pub fn into_inner(self) -> impl Iterator<Item = T> { |
| 146 | + #[cfg(parallel_compiler)] |
| 147 | + { |
| 148 | + self.locals.into_vec().into_iter().map(|local| local.0) |
| 149 | + } |
| 150 | + #[cfg(not(parallel_compiler))] |
| 151 | + { |
| 152 | + std::iter::once(self.local) |
| 153 | + } |
| 154 | + } |
| 155 | +} |
| 156 | + |
| 157 | +impl<T> WorkerLocal<Vec<T>> { |
| 158 | + /// Joins the elements of all the worker locals into one Vec |
| 159 | + pub fn join(self) -> Vec<T> { |
| 160 | + self.into_inner().into_iter().flat_map(|v| v).collect() |
| 161 | + } |
| 162 | +} |
| 163 | + |
| 164 | +impl<T> Deref for WorkerLocal<T> { |
| 165 | + type Target = T; |
| 166 | + |
| 167 | + #[inline(always)] |
| 168 | + #[cfg(not(parallel_compiler))] |
| 169 | + fn deref(&self) -> &T { |
| 170 | + &self.local |
| 171 | + } |
| 172 | + |
| 173 | + #[inline(always)] |
| 174 | + #[cfg(parallel_compiler)] |
| 175 | + fn deref(&self) -> &T { |
| 176 | + // This is safe because `verify` will only return values less than |
| 177 | + // `self.registry.thread_limit` which is the size of the `self.locals` array. |
| 178 | + unsafe { &self.locals.get_unchecked(self.registry.id().verify()).0 } |
| 179 | + } |
| 180 | +} |
0 commit comments