use crate::{ cfg::{self, CfgPrivate}, page, sync::{ atomic::{AtomicUsize, Ordering}, lazy_static, thread_local, Mutex, }, Pack, }; use std::{ cell::{Cell, UnsafeCell}, collections::VecDeque, fmt, marker::PhantomData, sync::PoisonError, }; /// Uniquely identifies a thread. pub(crate) struct Tid { id: usize, _not_send: PhantomData>, _cfg: PhantomData, } #[derive(Debug)] struct Registration(Cell>); struct Registry { next: AtomicUsize, free: Mutex>, } lazy_static! { static ref REGISTRY: Registry = Registry { next: AtomicUsize::new(0), free: Mutex::new(VecDeque::new()), }; } thread_local! { static REGISTRATION: Registration = Registration::new(); } // === impl Tid === impl Pack for Tid { const LEN: usize = C::MAX_SHARDS.trailing_zeros() as usize + 1; type Prev = page::Addr; #[inline(always)] fn as_usize(&self) -> usize { self.id } #[inline(always)] fn from_usize(id: usize) -> Self { Self { id, _not_send: PhantomData, _cfg: PhantomData, } } } impl Tid { #[inline] pub(crate) fn current() -> Self { REGISTRATION .try_with(Registration::current) .unwrap_or_else(|_| Self::poisoned()) } pub(crate) fn is_current(self) -> bool { REGISTRATION .try_with(|r| self == r.current::()) .unwrap_or(false) } #[inline(always)] pub fn new(id: usize) -> Self { Self::from_usize(id) } } impl Tid { #[cold] fn poisoned() -> Self { Self { id: std::usize::MAX, _not_send: PhantomData, _cfg: PhantomData, } } /// Returns true if the local thread ID was accessed while unwinding. pub(crate) fn is_poisoned(&self) -> bool { self.id == std::usize::MAX } } impl fmt::Debug for Tid { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { if self.is_poisoned() { f.debug_tuple("Tid") .field(&format_args!("")) .finish() } else { f.debug_tuple("Tid") .field(&format_args!("{}", self.id)) .finish() } } } impl PartialEq for Tid { fn eq(&self, other: &Self) -> bool { self.id == other.id } } impl Eq for Tid {} impl Clone for Tid { fn clone(&self) -> Self { Self::new(self.id) } } impl Copy for Tid {} // === impl Registration === impl Registration { fn new() -> Self { Self(Cell::new(None)) } #[inline(always)] fn current(&self) -> Tid { if let Some(tid) = self.0.get().map(Tid::new) { return tid; } self.register() } #[cold] fn register(&self) -> Tid { let id = REGISTRY .free .lock() .ok() .and_then(|mut free| { if free.len() > 1 { free.pop_front() } else { None } }) .unwrap_or_else(|| { let id = REGISTRY.next.fetch_add(1, Ordering::AcqRel); if id > Tid::::BITS { panic_in_drop!( "creating a new thread ID ({}) would exceed the \ maximum number of thread ID bits specified in {} \ ({})", id, std::any::type_name::(), Tid::::BITS, ); } id }); self.0.set(Some(id)); Tid::new(id) } } // Reusing thread IDs doesn't work under loom, since this `Drop` impl results in // an access to a `loom` lazy_static while the test is shutting down, which // panics. T_T // Just skip TID reuse and use loom's lazy_static macro to ensure we have a // clean initial TID on every iteration, instead. #[cfg(not(all(loom, any(feature = "loom", test))))] impl Drop for Registration { fn drop(&mut self) { if let Some(id) = self.0.get() { let mut free_list = REGISTRY.free.lock().unwrap_or_else(PoisonError::into_inner); free_list.push_back(id); } } }