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