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 sync::PoisonError,
16};
17
18/// Uniquely identifies a thread.
19pub(crate) struct Tid<C> {
20 id: usize,
21 _not_send: PhantomData<UnsafeCell<()>>,
22 _cfg: PhantomData<fn(C)>,
23}
24
25#[derive(Debug)]
26struct Registration(Cell<Option<usize>>);
27
28struct Registry {
29 next: AtomicUsize,
30 free: Mutex<VecDeque<usize>>,
31}
32
33lazy_static! {
34 static ref REGISTRY: Registry = Registry {
35 next: AtomicUsize::new(0),
36 free: Mutex::new(VecDeque::new()),
37 };
38}
39
40thread_local! {
41 static REGISTRATION: Registration = Registration::new();
42}
43
44// === impl Tid ===
45
46impl<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
66impl<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
86impl<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
102impl<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
116impl<C> PartialEq for Tid<C> {
117 fn eq(&self, other: &Self) -> bool {
118 self.id == other.id
119 }
120}
121
122impl<C> Eq for Tid<C> {}
123
124impl<C: cfg::Config> Clone for Tid<C> {
125 fn clone(&self) -> Self {
126 Self::new(self.id)
127 }
128}
129
130impl<C: cfg::Config> Copy for Tid<C> {}
131
132// === impl Registration ===
133
134impl 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))))]
187impl 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