dualcache_ff/componant/
qsbr.rs

1#![allow(clippy::not_unsafe_ptr_arg_deref)]
2#![allow(clippy::missing_safety_doc)]
3use ::core::sync::atomic::{AtomicPtr, AtomicUsize, AtomicBool, Ordering};
4use core::ptr;
5
6static GLOBAL_EPOCH: AtomicUsize = AtomicUsize::new(1);
7static THREAD_STATES: AtomicPtr<ThreadStateNode> = AtomicPtr::new(ptr::null_mut());
8
9#[doc(hidden)]
10pub unsafe fn reset() {
11    GLOBAL_EPOCH.store(1, Ordering::SeqCst);
12    THREAD_STATES.store(ptr::null_mut(), Ordering::SeqCst);
13}
14
15#[derive(Copy, Clone)]
16struct RetiredNode {
17    index: u32,
18    epoch: u64,
19}
20
21const GARBAGE_CAP: usize = 16384;
22pub struct GarbageQueue {
23    items: [RetiredNode; GARBAGE_CAP],
24    head: AtomicUsize,
25    tail: AtomicUsize,
26}
27
28impl GarbageQueue {
29    const fn new() -> Self {
30        Self {
31            items: [RetiredNode { index: 0, epoch: 0 }; GARBAGE_CAP],
32            head: AtomicUsize::new(0),
33            tail: AtomicUsize::new(0),
34        }
35    }
36}
37
38pub struct MpmcQueue<const N: usize> {
39    head: AtomicUsize,
40    tail: AtomicUsize,
41    buffer: [::core::sync::atomic::AtomicU32; N],
42}
43
44impl<const N: usize> Default for MpmcQueue<N> {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl<const N: usize> MpmcQueue<N> {
51    pub const fn new() -> Self {
52        Self {
53            head: AtomicUsize::new(0),
54            tail: AtomicUsize::new(0),
55            buffer: [const { ::core::sync::atomic::AtomicU32::new(u32::MAX) }; N],
56        }
57    }
58    pub fn push(&self, val: u32) -> bool {
59        let mut tail = self.tail.load(Ordering::Relaxed);
60        loop {
61            if tail.wrapping_sub(self.head.load(Ordering::Acquire)) >= N {
62                return false;
63            }
64            if self.tail.compare_exchange_weak(tail, tail + 1, Ordering::SeqCst, Ordering::Relaxed).is_ok() {
65                while self.buffer[tail % N].load(Ordering::Acquire) != u32::MAX {
66                    ::core::hint::spin_loop();
67                }
68                self.buffer[tail % N].store(val, Ordering::Release);
69                return true;
70            }
71            tail = self.tail.load(Ordering::Relaxed);
72        }
73    }
74    pub fn pop(&self) -> Option<u32> {
75        let mut head = self.head.load(Ordering::Relaxed);
76        loop {
77            if head == self.tail.load(Ordering::Acquire) {
78                return None;
79            }
80            let val = self.buffer[head % N].load(Ordering::Acquire);
81            if val != u32::MAX {
82                if self.head.compare_exchange_weak(head, head + 1, Ordering::SeqCst, Ordering::Relaxed).is_ok() {
83                    self.buffer[head % N].store(u32::MAX, Ordering::Release);
84                    return Some(val);
85                }
86            } else {
87                return None;
88            }
89            head = self.head.load(Ordering::Relaxed);
90        }
91    }
92}
93
94
95pub struct LocalFreeQueue {
96    items: [u32; 16384],
97    len: usize,
98}
99
100impl Default for LocalFreeQueue {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106impl LocalFreeQueue {
107    pub const fn new() -> Self {
108        Self { items: [0; 16384], len: 0 }
109    }
110    pub fn pop(&mut self) -> Option<u32> {
111        if self.len > 0 {
112            self.len -= 1;
113            Some(self.items[self.len])
114        } else {
115            None
116        }
117    }
118    #[must_use]
119    pub fn push(&mut self, val: u32) -> bool {
120        if self.len < 16384 {
121            self.items[self.len] = val;
122            self.len += 1;
123            true
124        } else {
125            false
126        }
127    }
128}
129
130pub struct QsbrToken;
131
132#[cfg_attr(not(feature = "std"), no_std_tool::auto_static(capacity = 256, partition = "qsbr"))]
133pub struct ThreadStateNode {
134    pub active: AtomicBool,
135    pub epoch: AtomicUsize,
136    pub next: *mut ThreadStateNode,
137    pub garbage_queue: GarbageQueue,
138    pub local_free: core::cell::UnsafeCell<LocalFreeQueue>,
139}
140
141impl Default for ThreadStateNode {
142    fn default() -> Self {
143        Self::new()
144    }
145}
146
147impl ThreadStateNode {
148    pub const fn new() -> Self {
149        Self {
150            active: AtomicBool::new(false),
151            epoch: AtomicUsize::new(0),
152            next: ptr::null_mut(),
153            garbage_queue: GarbageQueue::new(),
154            local_free: core::cell::UnsafeCell::new(LocalFreeQueue::new()),
155        }
156    }
157}
158
159unsafe impl Send for ThreadStateNode {}
160unsafe impl Sync for ThreadStateNode {}
161
162/// Register a pre-allocated thread state node. The caller should allocate this node
163/// locally or in the static TLS blocks.
164#[cfg(feature = "std")]
165struct QsbrCleanup {
166    node: *mut ThreadStateNode,
167    id: usize,
168    registry_ptr: *const core::ffi::c_void,
169    free_id_fn: Option<unsafe fn(*const core::ffi::c_void, usize)>,
170}
171
172#[cfg(feature = "std")]
173impl Drop for QsbrCleanup {
174    fn drop(&mut self) {
175        unsafe {
176            (*self.node).active.store(false, Ordering::Release);
177            if let (Some(free_fn), false) = (self.free_id_fn, self.registry_ptr.is_null()) {
178                free_fn(self.registry_ptr, self.id);
179            }
180        }
181    }
182}
183
184#[cfg(feature = "std")]
185thread_local! {
186    static QSBR_CLEANUP: std::cell::RefCell<std::vec::Vec<QsbrCleanup>> = const { std::cell::RefCell::new(std::vec::Vec::new()) };
187}
188
189#[cfg(feature = "std")]
190pub fn reregister_cleanup(
191    node: *mut ThreadStateNode,
192    id: usize,
193    registry_ptr: *const core::ffi::c_void,
194    free_id_fn: Option<unsafe fn(*const core::ffi::c_void, usize)>
195) {
196    QSBR_CLEANUP.with(|c| c.borrow_mut().push(QsbrCleanup { node, id, registry_ptr, free_id_fn }));
197}
198
199pub fn register_node(
200    node: *mut ThreadStateNode,
201    id: usize,
202    registry_ptr: *const core::ffi::c_void,
203    free_id_fn: Option<unsafe fn(*const core::ffi::c_void, usize)>
204) {
205    let mut head = THREAD_STATES.load(Ordering::Acquire);
206    loop {
207        unsafe { (*node).next = head };
208        
209        // Yield to encourage a CAS collision for coverage
210        #[cfg(test)]
211        std::thread::yield_now();
212
213        match THREAD_STATES.compare_exchange_weak(
214            head,
215            node,
216            Ordering::Release,
217            Ordering::Relaxed,
218        ) {
219            Ok(_) => {
220                #[cfg(feature = "std")]
221                QSBR_CLEANUP.with(|c| c.borrow_mut().push(QsbrCleanup { node, id, registry_ptr, free_id_fn }));
222                break;
223            }
224            Err(new_head) => {
225                head = new_head;
226                ::core::hint::spin_loop();
227            }
228        }
229    }
230}
231
232pub fn get_global_epoch() -> usize {
233    GLOBAL_EPOCH.load(Ordering::Relaxed)
234}
235
236/// A guard that pins the current thread to an epoch.
237pub struct Guard {
238    node: *mut ThreadStateNode,
239}
240
241impl Guard {
242    /// Create a new Guard using the explicitly provided ThreadStateNode.
243    pub fn new(node: *mut ThreadStateNode) -> Self {
244        let global_epoch = GLOBAL_EPOCH.load(Ordering::Acquire);
245        unsafe {
246            (*node).epoch.store(global_epoch, Ordering::Release);
247            (*node).active.store(true, Ordering::Release);
248        }
249        Self { node }
250    }
251
252    #[inline(always)]
253    pub unsafe fn unpinned(node: *mut ThreadStateNode) -> Self {
254        Self { node }
255    }
256    
257    #[inline(always)]
258    pub unsafe fn dummy() -> Self {
259        Self { node: core::ptr::null_mut() }
260    }
261    
262    /// Get the underlying ThreadStateNode pointer.
263    #[inline(always)]
264    pub fn node(&self) -> *mut ThreadStateNode {
265        self.node
266    }
267}
268
269impl Drop for Guard {
270    #[inline(always)]
271    fn drop(&mut self) {
272        if !self.node.is_null() {
273            let node_ref = unsafe { &*self.node };
274            node_ref.active.store(false, Ordering::Release);
275        }
276    }
277}
278
279pub fn pin_relaxed(node: *mut ThreadStateNode) -> Guard {
280    let node_ref = unsafe { &*node };
281    node_ref.active.store(true, Ordering::Relaxed);
282    Guard { node }
283}
284
285#[inline(always)]
286pub fn pin(node: *mut ThreadStateNode) -> Guard {
287    Guard::new(node)
288}
289
290
291/// Retire a node index into the thread-local garbage queue safely using QSBR.
292/// This prevents ABA by ensuring the index is not freed to the Arena until all threads observing it have advanced.
293pub fn retire<F: FnMut(u32)>(index: usize, node: *mut ThreadStateNode, mut _free_fn: F) {
294    let epoch = GLOBAL_EPOCH.load(Ordering::Acquire) as u64;
295    unsafe {
296        let q = &mut (*node).garbage_queue;
297        
298        // Wait until there's space (Daemon is slow, but we shouldn't drop)
299        let mut head = q.head.load(Ordering::Relaxed);
300        while head.wrapping_sub(q.tail.load(Ordering::Acquire)) >= GARBAGE_CAP {
301            // UPDATE EPOCH WHILE WAITING so daemon can advance min_epoch!
302            let cur_epoch = GLOBAL_EPOCH.load(Ordering::Relaxed);
303            (*node).epoch.store(cur_epoch, Ordering::Release);
304            ::core::hint::spin_loop();
305            head = q.head.load(Ordering::Relaxed);
306        }
307        
308        let idx = head % GARBAGE_CAP;
309        q.items[idx] = RetiredNode {
310            index: index as u32,
311            epoch,
312        };
313        q.head.store(head.wrapping_add(1), Ordering::Release);
314    }
315}
316
317/// The Daemon calls this globally to move safe nodes from GarbageQueues to Arena.
318pub fn daemon_reclaim<F: FnMut(&[u32])>(mut free_batch_fn: F) {
319    GLOBAL_EPOCH.fetch_add(1, Ordering::Relaxed);
320    let min_epoch = get_min_epoch();
321
322    let mut batch = [0u32; 128];
323    let mut batch_len = 0;
324
325    let mut node = THREAD_STATES.load(Ordering::Acquire);
326    while !node.is_null() {
327        unsafe {
328            let q = &mut (*node).garbage_queue;
329            let mut tail = q.tail.load(Ordering::Relaxed);
330            let head = q.head.load(Ordering::Acquire);
331            
332            while tail < head {
333                let idx = tail % GARBAGE_CAP;
334                let retired = q.items[idx];
335                
336                if retired.epoch < min_epoch as u64 {
337                    batch[batch_len] = retired.index;
338                    batch_len += 1;
339                    tail += 1;
340                    
341                    if batch_len == 128 {
342                        free_batch_fn(&batch[..batch_len]);
343                        batch_len = 0;
344                    }
345                } else {
346                    break;
347                }
348            }
349            q.tail.store(tail, Ordering::Release);
350            
351            node = (*node).next;
352        }
353    }
354    if batch_len > 0 {
355        free_batch_fn(&batch[..batch_len]);
356    }
357}
358
359fn get_min_epoch() -> usize {
360    let mut min_epoch = GLOBAL_EPOCH.load(Ordering::Acquire);
361    let mut node = THREAD_STATES.load(Ordering::Acquire);
362    while !node.is_null() {
363        unsafe {
364            if (*node).active.load(Ordering::Acquire) {
365                let e = (*node).epoch.load(Ordering::Acquire);
366                if e < min_epoch { min_epoch = e; }
367            }
368            node = (*node).next;
369        }
370    }
371    min_epoch
372}
373
374
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379
380    #[test]
381    fn test_qsbr_thread_state_node_new() {
382        let node = ThreadStateNode::new();
383        assert_eq!(node.epoch.load(Ordering::Relaxed), 0);
384    }
385}
dualcache_ff/componant/qsbr.rs

dualcache_ff/componant/
qsbr.rs
