sharded_slab/
tid.rs

Help
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(false)
78    }
79
80    #[inline(always)]
81    pub fn new(id: usize) -> Self {
82        Self::from_usize(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.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
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) = 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}
sharded_slab/tid.rs

sharded_slab/
tid.rs
