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