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)]
as_usize(&self) -> usize52     fn as_usize(&self) -> usize {
53         self.id
54     }
55 
56     #[inline(always)]
from_usize(id: usize) -> Self57     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]
current() -> Self68     pub(crate) fn current() -> Self {
69         REGISTRATION
70             .try_with(Registration::current)
71             .unwrap_or_else(|_| Self::poisoned())
72     }
73 
is_current(self) -> bool74     pub(crate) fn is_current(self) -> bool {
75         REGISTRATION
76             .try_with(|r| self == r.current::<C>())
77             .unwrap_or(false)
78     }
79 
80     #[inline(always)]
new(id: usize) -> Self81     pub fn new(id: usize) -> Self {
82         Self::from_usize(id)
83     }
84 }
85 
86 impl<C> Tid<C> {
87     #[cold]
poisoned() -> Self88     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.
is_poisoned(&self) -> bool97     pub(crate) fn is_poisoned(&self) -> bool {
98         self.id == std::usize::MAX
99     }
100 }
101 
102 impl<C> fmt::Debug for Tid<C> {
fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result103     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104         if self.is_poisoned() {
105             f.debug_tuple("Tid")
106                 .field(&format_args!("<poisoned>"))
107                 .finish()
108         } else {
109             f.debug_tuple("Tid")
110                 .field(&format_args!("{}", self.id))
111                 .finish()
112         }
113     }
114 }
115 
116 impl<C> PartialEq for Tid<C> {
eq(&self, other: &Self) -> bool117     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> {
clone(&self) -> Self125     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 {
new() -> Self135     fn new() -> Self {
136         Self(Cell::new(None))
137     }
138 
139     #[inline(always)]
current<C: cfg::Config>(&self) -> Tid<C>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]
register<C: cfg::Config>(&self) -> Tid<C>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 {
drop(&mut self)188     fn drop(&mut self) {
189         if let Some(id) = self.0.get() {
190             let mut free_list = REGISTRY.free.lock().unwrap_or_else(PoisonError::into_inner);
191             free_list.push_back(id);
192         }
193     }
194 }
195