subetha_cxc/
shared_treiber_stack.rs

1//! `SharedTreiberStack<T>` - cross-process lock-free LIFO stack.
2//!
3//! Classic Treiber-stack pattern: a packed `(counter, head_index)`
4//! atomic head, ABA-safe via the counter wrap. Push and pop are
5//! CAS loops bounded only by contention rate (not by logical
6//! waiting conditions).
7//!
8//! # Companion to other queue primitives
9//!
10//! - [`SharedRing`](crate::SharedRing): MPMC FIFO with fixed slot ordering
11//! - [`SharedBroadcastRing`](crate::SharedBroadcastRing): 1P+NC pub/sub
12//! - [`SharedTreiberStack`]: MPMC LIFO (this one)
13//!
14//! # Safety properties
15//!
16//! - **Bounded capacity** at create time; push returns `Err(Full)`
17//!   when capacity is exhausted.
18//! - **ABA-safe** via 32-bit counter packed with index in the head
19//!   atomic; same proven design as [`SharedRegion`](crate::SharedRegion)'s
20//!   free list.
21//! - **CAS loops are contention-bounded**, not logical-condition
22//!   bounded. Each retry happens because another writer won the
23//!   race; eventually contention resolves.
24//! - **No RAII guards** with Drop semantics that risk being
25//!   aliased or double-released. Push and pop return owned values.
26//! - **No underflow**: pop returns `None` on empty rather than
27//!   wrapping a counter.
28//!
29//! # Layout
30//!
31//! ```text
32//! +---------------------------+
33//! | StackHeader (64B)         |
34//! |   magic, capacity         |
35//! |   head: AtomicU64         |  // (counter << 32) | top_index, NIL when empty
36//! |   free_head: AtomicU64    |  // free-list of returned slots
37//! |   bump_next: AtomicU32    |
38//! +---------------------------+
39//! | next[capacity: AtomicU32] |  // chain pointers (overlap usage as free + occupied)
40//! +---------------------------+
41//! | slots[capacity * size_of<T>] |
42//! +---------------------------+
43//! ```
44
45use std::fs::{File, OpenOptions};
46use std::marker::PhantomData;
47use std::mem::size_of;
48use std::path::Path;
49use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
50
51use memmap2::{MmapMut, MmapOptions};
52
53pub const STACK_MAGIC: u32 = 0x4150_5354;
54pub const STACK_NIL: u32 = u32::MAX;
55
56#[repr(C, align(64))]
57pub struct StackHeader {
58    pub magic: u32,
59    pub capacity: u32,
60    pub slot_size: u32,
61    _pad1: u32,
62    pub head: AtomicU64,       // (counter << 32) | top_index (NIL when empty)
63    pub free_head: AtomicU64,  // free-list of returned slots
64    pub bump_next: AtomicU32,
65    _pad2: [u8; 28],
66}
67
68const _: () = {
69    assert!(size_of::<StackHeader>() == 64);
70};
71
72pub fn stack_file_size(capacity: usize, slot_size: usize) -> usize {
73    size_of::<StackHeader>()
74        + capacity * size_of::<AtomicU32>()
75        + capacity * slot_size
76}
77
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum StackError {
80    Full,
81    LayoutMismatch,
82    IoError(std::io::ErrorKind),
83}
84
85impl From<std::io::Error> for StackError {
86    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
87}
88
89#[inline]
90fn pack(counter: u32, index: u32) -> u64 {
91    ((counter as u64) << 32) | (index as u64)
92}
93#[inline]
94fn unpack(v: u64) -> (u32, u32) {
95    ((v >> 32) as u32, v as u32)
96}
97
98pub struct SharedTreiberStack<T: Copy + 'static> {
99    _file: File,
100    mmap: MmapMut,
101    capacity: usize,
102    next_offset: usize,
103    slots_offset: usize,
104    _phantom: PhantomData<T>,
105    header_sidecar: subetha_core::HandshakeHeader,
106    ring_sidecar: Box<subetha_core::ObservationRing>,
107}
108
109unsafe impl<T: Copy + Send + 'static> Send for SharedTreiberStack<T> {}
110unsafe impl<T: Copy + Sync + 'static> Sync for SharedTreiberStack<T> {}
111
112impl<T: Copy + Send + Sync + 'static> subetha_sidecar::AdaptiveInstance for SharedTreiberStack<T> {
113    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
114    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
115    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
116        Box::new(subetha_sidecar::NoMigrationPolicy)
117    }
118}
119
120impl<T: Copy + 'static> SharedTreiberStack<T> {
121    pub fn create(
122        path: impl AsRef<Path>, capacity: usize,
123    ) -> Result<Self, StackError> {
124        assert!(capacity >= 1);
125        assert!(capacity < STACK_NIL as usize, "capacity must be < u32::MAX");
126        let slot_size = size_of::<T>();
127        let total = stack_file_size(capacity, slot_size);
128        let file = OpenOptions::new()
129            .read(true).write(true).create(true).truncate(true)
130            .open(path.as_ref())?;
131        file.set_len(total as u64)?;
132        let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
133        let hdr = mmap.as_mut_ptr() as *mut StackHeader;
134        unsafe {
135            std::ptr::write_bytes(hdr as *mut u8, 0, size_of::<StackHeader>());
136            (*hdr).magic = STACK_MAGIC;
137            (*hdr).capacity = capacity as u32;
138            (*hdr).slot_size = slot_size as u32;
139            (*hdr).head.store(pack(0, STACK_NIL), Ordering::Release);
140            (*hdr).free_head.store(pack(0, STACK_NIL), Ordering::Release);
141            (*hdr).bump_next.store(0, Ordering::Release);
142        }
143        let next_offset = size_of::<StackHeader>();
144        let slots_offset = next_offset + capacity * size_of::<AtomicU32>();
145        Ok(Self {
146            _file: file, mmap, capacity, next_offset, slots_offset,
147            _phantom: PhantomData,
148            header_sidecar: subetha_core::HandshakeHeader::new(),
149            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
150        })
151    }
152
153    pub fn open(
154        path: impl AsRef<Path>, expected_capacity: usize,
155    ) -> Result<Self, StackError> {
156        let slot_size = size_of::<T>();
157        let total = stack_file_size(expected_capacity, slot_size);
158        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
159        if file.metadata()?.len() < total as u64 {
160            return Err(StackError::LayoutMismatch);
161        }
162        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
163        let hdr = unsafe { &*(mmap.as_ptr() as *const StackHeader) };
164        if hdr.magic != STACK_MAGIC
165            || hdr.capacity != expected_capacity as u32
166            || hdr.slot_size != slot_size as u32
167        {
168            return Err(StackError::LayoutMismatch);
169        }
170        let next_offset = size_of::<StackHeader>();
171        let slots_offset = next_offset + expected_capacity * size_of::<AtomicU32>();
172        Ok(Self {
173            _file: file, mmap, capacity: expected_capacity,
174            next_offset, slots_offset,
175            _phantom: PhantomData,
176            header_sidecar: subetha_core::HandshakeHeader::new(),
177            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
178        })
179    }
180
181    #[inline]
182    pub fn capacity(&self) -> usize { self.capacity }
183
184    fn header(&self) -> &StackHeader {
185        unsafe { &*(self.mmap.as_ptr() as *const StackHeader) }
186    }
187
188    fn next_link(&self, idx: usize) -> &AtomicU32 {
189        let base = unsafe { self.mmap.as_ptr().add(self.next_offset) };
190        unsafe { &*(base.add(idx * size_of::<AtomicU32>()) as *const AtomicU32) }
191    }
192
193    fn slot_ptr(&self, idx: usize) -> *mut T {
194        let base = unsafe { self.mmap.as_ptr().add(self.slots_offset) };
195        unsafe { base.add(idx * size_of::<T>()) as *mut T }
196    }
197
198    /// Acquire a slot index via the free list, falling back to bump
199    /// alloc. Returns `Err(Full)` when capacity is exhausted.
200    fn acquire_slot(&self) -> Result<u32, StackError> {
201        // Try free-list pop.
202        loop {
203            let head = self.header().free_head.load(Ordering::Acquire);
204            let (counter, idx) = unpack(head);
205            if idx == STACK_NIL { break; }
206            let next_idx = self.next_link(idx as usize).load(Ordering::Acquire);
207            let new_head = pack(counter.wrapping_add(1), next_idx);
208            if self.header().free_head.compare_exchange(
209                head, new_head, Ordering::AcqRel, Ordering::Acquire,
210            ).is_ok() {
211                return Ok(idx);
212            }
213        }
214        // Bump allocation.
215        let idx = self.header().bump_next.fetch_add(1, Ordering::AcqRel);
216        if (idx as usize) >= self.capacity {
217            self.header().bump_next.fetch_sub(1, Ordering::AcqRel);
218            return Err(StackError::Full);
219        }
220        Ok(idx)
221    }
222
223    /// Return a slot to the free list (Treiber push onto free_head).
224    fn release_slot(&self, idx: u32) {
225        loop {
226            let head = self.header().free_head.load(Ordering::Acquire);
227            let (counter, old_top) = unpack(head);
228            self.next_link(idx as usize).store(old_top, Ordering::Release);
229            let new_head = pack(counter.wrapping_add(1), idx);
230            if self.header().free_head.compare_exchange(
231                head, new_head, Ordering::AcqRel, Ordering::Acquire,
232            ).is_ok() {
233                return;
234            }
235        }
236    }
237
238    /// Push a value onto the stack.
239    pub fn push(&self, value: T) -> Result<(), StackError> {
240        let idx = match self.acquire_slot() {
241            Ok(i) => i,
242            Err(e) => {
243                self.ring_sidecar
244                    .push_op(crate::sidecar_ops::ordered::OP_INSERT, 1); // full
245                return Err(e);
246            }
247        };
248        unsafe { std::ptr::write(self.slot_ptr(idx as usize), value); }
249        // Treiber push: CAS head from (c, old_top) to (c+1, idx),
250        // with next_link[idx] = old_top.
251        loop {
252            let head = self.header().head.load(Ordering::Acquire);
253            let (counter, old_top) = unpack(head);
254            self.next_link(idx as usize).store(old_top, Ordering::Release);
255            let new_head = pack(counter.wrapping_add(1), idx);
256            if self.header().head.compare_exchange(
257                head, new_head, Ordering::AcqRel, Ordering::Acquire,
258            ).is_ok() {
259                self.ring_sidecar
260                    .push_op(crate::sidecar_ops::ordered::OP_INSERT, 0);
261                return Ok(());
262            }
263        }
264    }
265
266    /// Pop a value off the stack. Returns `None` if empty.
267    pub fn pop(&self) -> Option<T> {
268        loop {
269            let head = self.header().head.load(Ordering::Acquire);
270            let (counter, top) = unpack(head);
271            if top == STACK_NIL {
272                self.ring_sidecar
273                    .push_op(crate::sidecar_ops::ordered::OP_POP, 2); // empty
274                return None;
275            }
276            let next_top = self.next_link(top as usize).load(Ordering::Acquire);
277            let new_head = pack(counter.wrapping_add(1), next_top);
278            if self.header().head.compare_exchange(
279                head, new_head, Ordering::AcqRel, Ordering::Acquire,
280            ).is_ok() {
281                let value = unsafe { std::ptr::read(self.slot_ptr(top as usize)) };
282                self.release_slot(top);
283                self.ring_sidecar
284                    .push_op(crate::sidecar_ops::ordered::OP_POP, 0);
285                return Some(value);
286            }
287        }
288    }
289
290    /// Peek at the top without popping. Returns `None` if empty.
291    pub fn peek(&self) -> Option<T> {
292        let head = self.header().head.load(Ordering::Acquire);
293        let (_, top) = unpack(head);
294        if top == STACK_NIL {
295            self.ring_sidecar
296                .push_op(crate::sidecar_ops::ordered::OP_GET, 2); // empty
297            return None;
298        }
299        let v = unsafe { std::ptr::read(self.slot_ptr(top as usize)) };
300        self.ring_sidecar
301            .push_op(crate::sidecar_ops::ordered::OP_GET, 0);
302        Some(v)
303    }
304
305    /// True when the stack is empty.
306    pub fn is_empty(&self) -> bool {
307        let head = self.header().head.load(Ordering::Acquire);
308        unpack(head).1 == STACK_NIL
309    }
310
311    /// Approximate len: walks the linked list from head. O(N).
312    /// Subject to race with concurrent push/pop.
313    pub fn approx_len(&self) -> usize {
314        let head = self.header().head.load(Ordering::Acquire);
315        let (_, mut idx) = unpack(head);
316        let mut count = 0usize;
317        let mut visited = 0;
318        while idx != STACK_NIL && visited < self.capacity {
319            count += 1;
320            visited += 1;
321            idx = self.next_link(idx as usize).load(Ordering::Acquire);
322        }
323        count
324    }
325
326    pub fn flush(&self) -> Result<(), StackError> {
327        self.mmap.flush()?;
328        Ok(())
329    }
330    pub fn flush_async(&self) -> Result<(), StackError> {
331        self.mmap.flush_async()?;
332        Ok(())
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use super::*;
339    use std::sync::Arc;
340    use std::thread;
341
342    fn tmp(name: &str) -> std::path::PathBuf {
343        let mut p = std::env::temp_dir();
344        let pid = std::process::id();
345        p.push(format!("subetha-stack-{name}-{pid}.bin"));
346        p
347    }
348
349    #[test]
350    fn create_initial_state_is_empty() {
351        let p = tmp("init");
352        let s: SharedTreiberStack<u64> = SharedTreiberStack::create(&p, 16).unwrap();
353        assert!(s.is_empty());
354        assert_eq!(s.pop(), None);
355        assert_eq!(s.peek(), None);
356        std::fs::remove_file(&p).ok();
357    }
358
359    #[test]
360    fn push_pop_lifo_order() {
361        let p = tmp("lifo");
362        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 16).unwrap();
363        s.push(10).unwrap();
364        s.push(20).unwrap();
365        s.push(30).unwrap();
366        assert_eq!(s.pop(), Some(30));
367        assert_eq!(s.pop(), Some(20));
368        assert_eq!(s.pop(), Some(10));
369        assert_eq!(s.pop(), None);
370        std::fs::remove_file(&p).ok();
371    }
372
373    #[test]
374    fn peek_does_not_remove() {
375        let p = tmp("peek");
376        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 8).unwrap();
377        s.push(42).unwrap();
378        assert_eq!(s.peek(), Some(42));
379        assert_eq!(s.peek(), Some(42));
380        assert_eq!(s.pop(), Some(42));
381        assert_eq!(s.peek(), None);
382        std::fs::remove_file(&p).ok();
383    }
384
385    #[test]
386    fn full_capacity_returns_error() {
387        let p = tmp("full");
388        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 4).unwrap();
389        for i in 0..4 { s.push(i).unwrap(); }
390        assert_eq!(s.push(99).err(), Some(StackError::Full));
391        // After popping, can push again.
392        s.pop();
393        s.push(99).unwrap();
394        std::fs::remove_file(&p).ok();
395    }
396
397    #[test]
398    fn free_list_reuse_after_pop() {
399        let p = tmp("reuse");
400        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 4).unwrap();
401        for i in 0..4 { s.push(i).unwrap(); }
402        for _ in 0..4 { s.pop(); }
403        // After full drain, push 4 more should succeed (slots reused).
404        for i in 100..104 { s.push(i).unwrap(); }
405        assert_eq!(s.pop(), Some(103));
406        assert_eq!(s.pop(), Some(102));
407        std::fs::remove_file(&p).ok();
408    }
409
410    #[test]
411    fn approx_len_tracks_size() {
412        let p = tmp("len");
413        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 16).unwrap();
414        assert_eq!(s.approx_len(), 0);
415        s.push(1).unwrap();
416        s.push(2).unwrap();
417        s.push(3).unwrap();
418        assert_eq!(s.approx_len(), 3);
419        s.pop();
420        assert_eq!(s.approx_len(), 2);
421        std::fs::remove_file(&p).ok();
422    }
423
424    #[test]
425    fn cross_handle_visibility() {
426        let p = tmp("cross-handle");
427        let w: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 8).unwrap();
428        let r: SharedTreiberStack<u32> = SharedTreiberStack::open(&p, 8).unwrap();
429        w.push(42).unwrap();
430        w.push(7).unwrap();
431        assert_eq!(r.peek(), Some(7));
432        assert_eq!(r.pop(), Some(7));
433        assert_eq!(w.pop(), Some(42));
434        assert!(r.is_empty());
435        std::fs::remove_file(&p).ok();
436    }
437
438    #[test]
439    fn struct_payload_round_trip() {
440        #[derive(Clone, Copy, Debug, PartialEq)]
441        #[repr(C)]
442        struct Frame { pc: u64, sp: u64 }
443        let p = tmp("struct");
444        let s: SharedTreiberStack<Frame> = SharedTreiberStack::create(&p, 8).unwrap();
445        s.push(Frame { pc: 0x1000, sp: 0xFF00 }).unwrap();
446        s.push(Frame { pc: 0x2000, sp: 0xFE00 }).unwrap();
447        assert_eq!(s.pop(), Some(Frame { pc: 0x2000, sp: 0xFE00 }));
448        assert_eq!(s.pop(), Some(Frame { pc: 0x1000, sp: 0xFF00 }));
449        std::fs::remove_file(&p).ok();
450    }
451
452    #[test]
453    fn concurrent_pushers_all_succeed() {
454        let p = tmp("concurrent-push");
455        let s: Arc<SharedTreiberStack<u32>>
456            = Arc::new(SharedTreiberStack::create(&p, 1024).unwrap());
457        let n_threads = 4;
458        let per_thread = 100;
459        let mut handles = vec![];
460        for t in 0..n_threads as u32 {
461            let s = s.clone();
462            handles.push(thread::spawn(move || {
463                for i in 0..per_thread as u32 {
464                    s.push(t * 1000 + i).unwrap();
465                }
466            }));
467        }
468        for h in handles { h.join().unwrap(); }
469        assert_eq!(s.approx_len(), n_threads * per_thread);
470        // Drain and collect.
471        let mut all = Vec::new();
472        while let Some(v) = s.pop() { all.push(v); }
473        all.sort();
474        // Expect 4 threads * 100 values: t=0..4, i=0..100.
475        let mut expected: Vec<u32> = (0..n_threads as u32)
476            .flat_map(|t| (0..per_thread as u32).map(move |i| t * 1000 + i))
477            .collect();
478        expected.sort();
479        assert_eq!(all, expected);
480        std::fs::remove_file(&p).ok();
481    }
482
483    #[test]
484    fn concurrent_push_pop_no_corruption() {
485        // Producers push, consumers pop. After joining, no items
486        // lost or duplicated.
487        let p = tmp("concurrent-pp");
488        let s: Arc<SharedTreiberStack<u32>>
489            = Arc::new(SharedTreiberStack::create(&p, 1024).unwrap());
490        // Pre-fill with 500 known values.
491        for i in 0..500u32 { s.push(i).unwrap(); }
492        // 4 consumer threads pop everything they can, collecting locally.
493        let mut handles = vec![];
494        for _ in 0..4 {
495            let s = s.clone();
496            handles.push(thread::spawn(move || {
497                let mut got = Vec::new();
498                while let Some(v) = s.pop() { got.push(v); }
499                got
500            }));
501        }
502        let mut total: Vec<u32> = handles.into_iter()
503            .flat_map(|h| h.join().unwrap()).collect();
504        total.sort();
505        let expected: Vec<u32> = (0..500u32).collect();
506        assert_eq!(total, expected, "no items should be lost or duplicated");
507        std::fs::remove_file(&p).ok();
508    }
509
510    #[test]
511    fn disk_persistence_survives_reopen() {
512        let p = tmp("disk");
513        {
514            let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 8).unwrap();
515            s.push(1).unwrap();
516            s.push(2).unwrap();
517            s.push(3).unwrap();
518            s.flush().unwrap();
519        }
520        let s2: SharedTreiberStack<u32> = SharedTreiberStack::open(&p, 8).unwrap();
521        assert_eq!(s2.pop(), Some(3));
522        assert_eq!(s2.pop(), Some(2));
523        assert_eq!(s2.pop(), Some(1));
524        std::fs::remove_file(&p).ok();
525    }
526}
subetha_cxc/shared_treiber_stack.rs

subetha_cxc/
shared_treiber_stack.rs
