1 | use crate::{ |
2 | cfg::{self, CfgPrivate}, |
3 | page, |
4 | sync::{ |
5 | atomic::{AtomicUsize, Ordering}, |
6 | lazy_static, thread_local, Mutex, |
7 | }, |
8 | Pack, |
9 | }; |
10 | use std::{ |
11 | cell::{Cell, UnsafeCell}, |
12 | collections::VecDeque, |
13 | fmt, |
14 | marker::PhantomData, |
15 | sync::PoisonError, |
16 | }; |
17 | |
18 | /// Uniquely identifies a thread. |
19 | pub(crate) struct Tid<C> { |
20 | id: usize, |
21 | _not_send: PhantomData<UnsafeCell<()>>, |
22 | _cfg: PhantomData<fn(C)>, |
23 | } |
24 | |
25 | #[derive (Debug)] |
26 | struct Registration(Cell<Option<usize>>); |
27 | |
28 | struct Registry { |
29 | next: AtomicUsize, |
30 | free: Mutex<VecDeque<usize>>, |
31 | } |
32 | |
33 | lazy_static! { |
34 | static ref REGISTRY: Registry = Registry { |
35 | next: AtomicUsize::new(0), |
36 | free: Mutex::new(VecDeque::new()), |
37 | }; |
38 | } |
39 | |
40 | thread_local! { |
41 | static REGISTRATION: Registration = Registration::new(); |
42 | } |
43 | |
44 | // === impl Tid === |
45 | |
46 | impl<C: cfg::Config> Pack<C> for Tid<C> { |
47 | const LEN: usize = C::MAX_SHARDS.trailing_zeros() as usize + 1; |
48 | |
49 | type Prev = page::Addr<C>; |
50 | |
51 | #[inline (always)] |
52 | fn as_usize(&self) -> usize { |
53 | self.id |
54 | } |
55 | |
56 | #[inline (always)] |
57 | fn from_usize(id: usize) -> Self { |
58 | Self { |
59 | id, |
60 | _not_send: PhantomData, |
61 | _cfg: PhantomData, |
62 | } |
63 | } |
64 | } |
65 | |
66 | impl<C: cfg::Config> Tid<C> { |
67 | #[inline ] |
68 | pub(crate) fn current() -> Self { |
69 | REGISTRATION |
70 | .try_with(Registration::current) |
71 | .unwrap_or_else(|_| Self::poisoned()) |
72 | } |
73 | |
74 | pub(crate) fn is_current(self) -> bool { |
75 | REGISTRATION |
76 | .try_with(|r| self == r.current::<C>()) |
77 | .unwrap_or(default:false) |
78 | } |
79 | |
80 | #[inline (always)] |
81 | pub fn new(id: usize) -> Self { |
82 | Self::from_usize(val:id) |
83 | } |
84 | } |
85 | |
86 | impl<C> Tid<C> { |
87 | #[cold ] |
88 | fn poisoned() -> Self { |
89 | Self { |
90 | id: std::usize::MAX, |
91 | _not_send: PhantomData, |
92 | _cfg: PhantomData, |
93 | } |
94 | } |
95 | |
96 | /// Returns true if the local thread ID was accessed while unwinding. |
97 | pub(crate) fn is_poisoned(&self) -> bool { |
98 | self.id == std::usize::MAX |
99 | } |
100 | } |
101 | |
102 | impl<C> fmt::Debug for Tid<C> { |
103 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
104 | if self.is_poisoned() { |
105 | f&mut DebugTuple<'_, '_>.debug_tuple(name:"Tid" ) |
106 | .field(&format_args!("<poisoned>" )) |
107 | .finish() |
108 | } else { |
109 | f&mut DebugTuple<'_, '_>.debug_tuple(name:"Tid" ) |
110 | .field(&format_args!(" {}" , self.id)) |
111 | .finish() |
112 | } |
113 | } |
114 | } |
115 | |
116 | impl<C> PartialEq for Tid<C> { |
117 | fn eq(&self, other: &Self) -> bool { |
118 | self.id == other.id |
119 | } |
120 | } |
121 | |
122 | impl<C> Eq for Tid<C> {} |
123 | |
124 | impl<C: cfg::Config> Clone for Tid<C> { |
125 | fn clone(&self) -> Self { |
126 | Self::new(self.id) |
127 | } |
128 | } |
129 | |
130 | impl<C: cfg::Config> Copy for Tid<C> {} |
131 | |
132 | // === impl Registration === |
133 | |
134 | impl Registration { |
135 | fn new() -> Self { |
136 | Self(Cell::new(None)) |
137 | } |
138 | |
139 | #[inline (always)] |
140 | fn current<C: cfg::Config>(&self) -> Tid<C> { |
141 | if let Some(tid) = self.0.get().map(Tid::new) { |
142 | return tid; |
143 | } |
144 | |
145 | self.register() |
146 | } |
147 | |
148 | #[cold ] |
149 | fn register<C: cfg::Config>(&self) -> Tid<C> { |
150 | let id = REGISTRY |
151 | .free |
152 | .lock() |
153 | .ok() |
154 | .and_then(|mut free| { |
155 | if free.len() > 1 { |
156 | free.pop_front() |
157 | } else { |
158 | None |
159 | } |
160 | }) |
161 | .unwrap_or_else(|| { |
162 | let id = REGISTRY.next.fetch_add(1, Ordering::AcqRel); |
163 | if id > Tid::<C>::BITS { |
164 | panic_in_drop!( |
165 | "creating a new thread ID ( {}) would exceed the \ |
166 | maximum number of thread ID bits specified in {} \ |
167 | ( {})" , |
168 | id, |
169 | std::any::type_name::<C>(), |
170 | Tid::<C>::BITS, |
171 | ); |
172 | } |
173 | id |
174 | }); |
175 | |
176 | self.0.set(Some(id)); |
177 | Tid::new(id) |
178 | } |
179 | } |
180 | |
181 | // Reusing thread IDs doesn't work under loom, since this `Drop` impl results in |
182 | // an access to a `loom` lazy_static while the test is shutting down, which |
183 | // panics. T_T |
184 | // Just skip TID reuse and use loom's lazy_static macro to ensure we have a |
185 | // clean initial TID on every iteration, instead. |
186 | #[cfg (not(all(loom, any(feature = "loom" , test))))] |
187 | impl Drop for Registration { |
188 | fn drop(&mut self) { |
189 | if let Some(id: usize) = self.0.get() { |
190 | let mut free_list: MutexGuard<'_, VecDeque> = REGISTRY.free.lock().unwrap_or_else(op:PoisonError::into_inner); |
191 | free_list.push_back(id); |
192 | } |
193 | } |
194 | } |
195 | |