sharded_slab/
tid.rs

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(false)
77    }
78
79    #[inline(always)]
80    pub fn new(id: usize) -> Self {
81        Self::from_usize(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.debug_tuple("Tid")
105                .field(&format_args!("<poisoned>"))
106                .finish()
107        } else {
108            f.debug_tuple("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) = self.0.get() {
191            let mut free_list = REGISTRY.free.lock().unwrap_or_else(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}
sharded_slab/tid.rs

sharded_slab/
tid.rs
