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    /// Obtain the stack at `path`, initializing an empty one if the
122    /// path does not yet exist and attaching to it if it does.
123    /// Attaching leaves pushed entries and the free list in place; a
124    /// region built with a different capacity or payload type is a
125    /// `LayoutMismatch`. [`reset`](Self::reset) reinitializes.
126    pub fn create(
127        path: impl AsRef<Path>, capacity: usize,
128    ) -> Result<Self, StackError> {
129        assert!(capacity >= 1);
130        assert!(capacity < STACK_NIL as usize, "capacity must be < u32::MAX");
131        let (file, mmap) = crate::mmf_attach::create_or_attach(
132            path.as_ref(),
133            stack_file_size(capacity, size_of::<T>()),
134            |ptr| unsafe { Self::init_region(ptr, capacity) },
135            |ptr| unsafe { (*(ptr as *const StackHeader)).magic == STACK_MAGIC },
136        )?;
137        Self::from_region(file, mmap, capacity)
138    }
139
140    /// Truncate the stack at `path` and initialize an empty one,
141    /// discarding every entry live peers share. For a caller that
142    /// knows it owns the path.
143    pub fn reset(
144        path: impl AsRef<Path>, capacity: usize,
145    ) -> Result<Self, StackError> {
146        assert!(capacity >= 1);
147        assert!(capacity < STACK_NIL as usize, "capacity must be < u32::MAX");
148        let (file, mmap) = crate::mmf_attach::reset(
149            path.as_ref(),
150            stack_file_size(capacity, size_of::<T>()),
151            |ptr| unsafe { Self::init_region(ptr, capacity) },
152        )?;
153        Self::from_region(file, mmap, capacity)
154    }
155
156    /// Lay out an empty stack: config and the NIL head, free head and
157    /// bump cursor first, magic last, because attachers spin on it.
158    ///
159    /// # Safety
160    /// `ptr` addresses at least `stack_file_size(capacity,
161    /// size_of::<T>())` writable zeroed bytes.
162    unsafe fn init_region(ptr: *mut u8, capacity: usize) {
163        let hdr = ptr as *mut StackHeader;
164        unsafe {
165            (*hdr).capacity = capacity as u32;
166            (*hdr).slot_size = size_of::<T>() as u32;
167            std::ptr::write(&raw mut (*hdr).head, AtomicU64::new(pack(0, STACK_NIL)));
168            std::ptr::write(&raw mut (*hdr).free_head, AtomicU64::new(pack(0, STACK_NIL)));
169            std::ptr::write_volatile(&raw mut (*hdr).magic, STACK_MAGIC);
170        }
171    }
172
173    /// Wrap an initialized region, refusing one built with a different
174    /// capacity or payload type.
175    fn from_region(
176        file: File,
177        mmap: MmapMut,
178        capacity: usize,
179    ) -> Result<Self, StackError> {
180        let hdr = unsafe { &*(mmap.as_ptr() as *const StackHeader) };
181        if hdr.magic != STACK_MAGIC
182            || hdr.capacity != capacity as u32
183            || hdr.slot_size != size_of::<T>() as u32
184        {
185            return Err(StackError::LayoutMismatch);
186        }
187        let next_offset = size_of::<StackHeader>();
188        let slots_offset = next_offset + capacity * size_of::<AtomicU32>();
189        Ok(Self {
190            _file: file, mmap, capacity, next_offset, slots_offset,
191            _phantom: PhantomData,
192            header_sidecar: subetha_core::HandshakeHeader::new(),
193            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
194        })
195    }
196
197    pub fn open(
198        path: impl AsRef<Path>, expected_capacity: usize,
199    ) -> Result<Self, StackError> {
200        let total = stack_file_size(expected_capacity, size_of::<T>());
201        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
202        if file.metadata()?.len() < total as u64 {
203            return Err(StackError::LayoutMismatch);
204        }
205        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
206        Self::from_region(file, mmap, expected_capacity)
207    }
208
209    #[inline]
210    pub fn capacity(&self) -> usize { self.capacity }
211
212    fn header(&self) -> &StackHeader {
213        unsafe { &*(self.mmap.as_ptr() as *const StackHeader) }
214    }
215
216    fn next_link(&self, idx: usize) -> &AtomicU32 {
217        let base = unsafe { self.mmap.as_ptr().add(self.next_offset) };
218        unsafe { &*(base.add(idx * size_of::<AtomicU32>()) as *const AtomicU32) }
219    }
220
221    fn slot_ptr(&self, idx: usize) -> *mut T {
222        let base = unsafe { self.mmap.as_ptr().add(self.slots_offset) };
223        unsafe { base.add(idx * size_of::<T>()) as *mut T }
224    }
225
226    /// Acquire a slot index via the free list, falling back to bump
227    /// alloc. Returns `Err(Full)` when capacity is exhausted.
228    fn acquire_slot(&self) -> Result<u32, StackError> {
229        // Try free-list pop.
230        loop {
231            let head = self.header().free_head.load(Ordering::Acquire);
232            let (counter, idx) = unpack(head);
233            if idx == STACK_NIL { break; }
234            let next_idx = self.next_link(idx as usize).load(Ordering::Acquire);
235            let new_head = pack(counter.wrapping_add(1), next_idx);
236            if self.header().free_head.compare_exchange(
237                head, new_head, Ordering::AcqRel, Ordering::Acquire,
238            ).is_ok() {
239                return Ok(idx);
240            }
241        }
242        // Bump allocation.
243        let idx = self.header().bump_next.fetch_add(1, Ordering::AcqRel);
244        if (idx as usize) >= self.capacity {
245            self.header().bump_next.fetch_sub(1, Ordering::AcqRel);
246            return Err(StackError::Full);
247        }
248        Ok(idx)
249    }
250
251    /// Return a slot to the free list (Treiber push onto free_head).
252    fn release_slot(&self, idx: u32) {
253        loop {
254            let head = self.header().free_head.load(Ordering::Acquire);
255            let (counter, old_top) = unpack(head);
256            self.next_link(idx as usize).store(old_top, Ordering::Release);
257            let new_head = pack(counter.wrapping_add(1), idx);
258            if self.header().free_head.compare_exchange(
259                head, new_head, Ordering::AcqRel, Ordering::Acquire,
260            ).is_ok() {
261                return;
262            }
263        }
264    }
265
266    /// Push a value onto the stack.
267    pub fn push(&self, value: T) -> Result<(), StackError> {
268        let idx = match self.acquire_slot() {
269            Ok(i) => i,
270            Err(e) => {
271                self.ring_sidecar
272                    .push_op(crate::sidecar_ops::ordered::OP_INSERT, 1); // full
273                return Err(e);
274            }
275        };
276        unsafe { std::ptr::write(self.slot_ptr(idx as usize), value); }
277        // Treiber push: CAS head from (c, old_top) to (c+1, idx),
278        // with next_link[idx] = old_top.
279        loop {
280            let head = self.header().head.load(Ordering::Acquire);
281            let (counter, old_top) = unpack(head);
282            self.next_link(idx as usize).store(old_top, Ordering::Release);
283            let new_head = pack(counter.wrapping_add(1), idx);
284            if self.header().head.compare_exchange(
285                head, new_head, Ordering::AcqRel, Ordering::Acquire,
286            ).is_ok() {
287                self.ring_sidecar
288                    .push_op(crate::sidecar_ops::ordered::OP_INSERT, 0);
289                return Ok(());
290            }
291        }
292    }
293
294    /// Pop a value off the stack. Returns `None` if empty.
295    pub fn pop(&self) -> Option<T> {
296        loop {
297            let head = self.header().head.load(Ordering::Acquire);
298            let (counter, top) = unpack(head);
299            if top == STACK_NIL {
300                self.ring_sidecar
301                    .push_op(crate::sidecar_ops::ordered::OP_POP, 2); // empty
302                return None;
303            }
304            let next_top = self.next_link(top as usize).load(Ordering::Acquire);
305            let new_head = pack(counter.wrapping_add(1), next_top);
306            if self.header().head.compare_exchange(
307                head, new_head, Ordering::AcqRel, Ordering::Acquire,
308            ).is_ok() {
309                let value = unsafe { std::ptr::read(self.slot_ptr(top as usize)) };
310                self.release_slot(top);
311                self.ring_sidecar
312                    .push_op(crate::sidecar_ops::ordered::OP_POP, 0);
313                return Some(value);
314            }
315        }
316    }
317
318    /// Peek at the top without popping. Returns `None` if empty.
319    pub fn peek(&self) -> Option<T> {
320        let head = self.header().head.load(Ordering::Acquire);
321        let (_, top) = unpack(head);
322        if top == STACK_NIL {
323            self.ring_sidecar
324                .push_op(crate::sidecar_ops::ordered::OP_GET, 2); // empty
325            return None;
326        }
327        let v = unsafe { std::ptr::read(self.slot_ptr(top as usize)) };
328        self.ring_sidecar
329            .push_op(crate::sidecar_ops::ordered::OP_GET, 0);
330        Some(v)
331    }
332
333    /// True when the stack is empty.
334    pub fn is_empty(&self) -> bool {
335        let head = self.header().head.load(Ordering::Acquire);
336        unpack(head).1 == STACK_NIL
337    }
338
339    /// Approximate len: walks the linked list from head. O(N).
340    /// Subject to race with concurrent push/pop.
341    pub fn approx_len(&self) -> usize {
342        let head = self.header().head.load(Ordering::Acquire);
343        let (_, mut idx) = unpack(head);
344        let mut count = 0usize;
345        let mut visited = 0;
346        while idx != STACK_NIL && visited < self.capacity {
347            count += 1;
348            visited += 1;
349            idx = self.next_link(idx as usize).load(Ordering::Acquire);
350        }
351        count
352    }
353
354    pub fn flush(&self) -> Result<(), StackError> {
355        self.mmap.flush()?;
356        Ok(())
357    }
358    pub fn flush_async(&self) -> Result<(), StackError> {
359        self.mmap.flush_async()?;
360        Ok(())
361    }
362}
363
364#[cfg(test)]
365mod tests {
366    use super::*;
367    use std::sync::Arc;
368    use std::thread;
369
370    fn tmp(name: &str) -> std::path::PathBuf {
371        let mut p = std::env::temp_dir();
372        let pid = std::process::id();
373        p.push(format!("subetha-stack-{name}-{pid}.bin"));
374        p
375    }
376
377    #[test]
378    fn create_initial_state_is_empty() {
379        let p = tmp("init");
380        let s: SharedTreiberStack<u64> = SharedTreiberStack::create(&p, 16).unwrap();
381        assert!(s.is_empty());
382        assert_eq!(s.pop(), None);
383        assert_eq!(s.peek(), None);
384        std::fs::remove_file(&p).ok();
385    }
386
387    /// A second create attaches with pushed entries in place; reset is
388    /// what strips them.
389    #[test]
390    fn second_create_attaches_and_keeps_entries() {
391        let p = tmp("attach");
392        std::fs::remove_file(&p).ok();
393        let s: SharedTreiberStack<u64> = SharedTreiberStack::create(&p, 16).unwrap();
394        s.push(777).unwrap();
395
396        let s2: SharedTreiberStack<u64> = SharedTreiberStack::create(&p, 16).unwrap();
397        assert_eq!(s2.peek(), Some(777), "attach lost a pushed entry");
398        assert!(matches!(
399            SharedTreiberStack::<u64>::create(&p, 8),
400            Err(StackError::LayoutMismatch),
401        ));
402
403        // Windows refuses to truncate a mapped file, so every handle goes
404        // before the reset.
405        drop(s);
406        drop(s2);
407        let fresh: SharedTreiberStack<u64> = SharedTreiberStack::reset(&p, 16).unwrap();
408        assert!(fresh.is_empty(), "reset kept an entry");
409        drop(fresh);
410        std::fs::remove_file(&p).ok();
411    }
412
413    #[test]
414    fn push_pop_lifo_order() {
415        let p = tmp("lifo");
416        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 16).unwrap();
417        s.push(10).unwrap();
418        s.push(20).unwrap();
419        s.push(30).unwrap();
420        assert_eq!(s.pop(), Some(30));
421        assert_eq!(s.pop(), Some(20));
422        assert_eq!(s.pop(), Some(10));
423        assert_eq!(s.pop(), None);
424        std::fs::remove_file(&p).ok();
425    }
426
427    #[test]
428    fn peek_does_not_remove() {
429        let p = tmp("peek");
430        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 8).unwrap();
431        s.push(42).unwrap();
432        assert_eq!(s.peek(), Some(42));
433        assert_eq!(s.peek(), Some(42));
434        assert_eq!(s.pop(), Some(42));
435        assert_eq!(s.peek(), None);
436        std::fs::remove_file(&p).ok();
437    }
438
439    #[test]
440    fn full_capacity_returns_error() {
441        let p = tmp("full");
442        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 4).unwrap();
443        for i in 0..4 { s.push(i).unwrap(); }
444        assert_eq!(s.push(99).err(), Some(StackError::Full));
445        // After popping, can push again.
446        s.pop();
447        s.push(99).unwrap();
448        std::fs::remove_file(&p).ok();
449    }
450
451    #[test]
452    fn free_list_reuse_after_pop() {
453        let p = tmp("reuse");
454        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 4).unwrap();
455        for i in 0..4 { s.push(i).unwrap(); }
456        for _ in 0..4 { s.pop(); }
457        // After full drain, push 4 more should succeed (slots reused).
458        for i in 100..104 { s.push(i).unwrap(); }
459        assert_eq!(s.pop(), Some(103));
460        assert_eq!(s.pop(), Some(102));
461        std::fs::remove_file(&p).ok();
462    }
463
464    #[test]
465    fn approx_len_tracks_size() {
466        let p = tmp("len");
467        let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 16).unwrap();
468        assert_eq!(s.approx_len(), 0);
469        s.push(1).unwrap();
470        s.push(2).unwrap();
471        s.push(3).unwrap();
472        assert_eq!(s.approx_len(), 3);
473        s.pop();
474        assert_eq!(s.approx_len(), 2);
475        std::fs::remove_file(&p).ok();
476    }
477
478    #[test]
479    fn cross_handle_visibility() {
480        let p = tmp("cross-handle");
481        let w: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 8).unwrap();
482        let r: SharedTreiberStack<u32> = SharedTreiberStack::open(&p, 8).unwrap();
483        w.push(42).unwrap();
484        w.push(7).unwrap();
485        assert_eq!(r.peek(), Some(7));
486        assert_eq!(r.pop(), Some(7));
487        assert_eq!(w.pop(), Some(42));
488        assert!(r.is_empty());
489        std::fs::remove_file(&p).ok();
490    }
491
492    #[test]
493    fn struct_payload_round_trip() {
494        #[derive(Clone, Copy, Debug, PartialEq)]
495        #[repr(C)]
496        struct Frame { pc: u64, sp: u64 }
497        let p = tmp("struct");
498        let s: SharedTreiberStack<Frame> = SharedTreiberStack::create(&p, 8).unwrap();
499        s.push(Frame { pc: 0x1000, sp: 0xFF00 }).unwrap();
500        s.push(Frame { pc: 0x2000, sp: 0xFE00 }).unwrap();
501        assert_eq!(s.pop(), Some(Frame { pc: 0x2000, sp: 0xFE00 }));
502        assert_eq!(s.pop(), Some(Frame { pc: 0x1000, sp: 0xFF00 }));
503        std::fs::remove_file(&p).ok();
504    }
505
506    #[test]
507    fn concurrent_pushers_all_succeed() {
508        let p = tmp("concurrent-push");
509        let s: Arc<SharedTreiberStack<u32>>
510            = Arc::new(SharedTreiberStack::create(&p, 1024).unwrap());
511        let n_threads = 4;
512        let per_thread = 100;
513        let mut handles = vec![];
514        for t in 0..n_threads as u32 {
515            let s = s.clone();
516            handles.push(thread::spawn(move || {
517                for i in 0..per_thread as u32 {
518                    s.push(t * 1000 + i).unwrap();
519                }
520            }));
521        }
522        for h in handles { h.join().unwrap(); }
523        assert_eq!(s.approx_len(), n_threads * per_thread);
524        // Drain and collect.
525        let mut all = Vec::new();
526        while let Some(v) = s.pop() { all.push(v); }
527        all.sort();
528        // Expect 4 threads * 100 values: t=0..4, i=0..100.
529        let mut expected: Vec<u32> = (0..n_threads as u32)
530            .flat_map(|t| (0..per_thread as u32).map(move |i| t * 1000 + i))
531            .collect();
532        expected.sort();
533        assert_eq!(all, expected);
534        std::fs::remove_file(&p).ok();
535    }
536
537    #[test]
538    fn concurrent_push_pop_no_corruption() {
539        // Producers push, consumers pop. After joining, no items
540        // lost or duplicated.
541        let p = tmp("concurrent-pp");
542        let s: Arc<SharedTreiberStack<u32>>
543            = Arc::new(SharedTreiberStack::create(&p, 1024).unwrap());
544        // Pre-fill with 500 known values.
545        for i in 0..500u32 { s.push(i).unwrap(); }
546        // 4 consumer threads pop everything they can, collecting locally.
547        let mut handles = vec![];
548        for _ in 0..4 {
549            let s = s.clone();
550            handles.push(thread::spawn(move || {
551                let mut got = Vec::new();
552                while let Some(v) = s.pop() { got.push(v); }
553                got
554            }));
555        }
556        let mut total: Vec<u32> = handles.into_iter()
557            .flat_map(|h| h.join().unwrap()).collect();
558        total.sort();
559        let expected: Vec<u32> = (0..500u32).collect();
560        assert_eq!(total, expected, "no items should be lost or duplicated");
561        std::fs::remove_file(&p).ok();
562    }
563
564    #[test]
565    fn disk_persistence_survives_reopen() {
566        let p = tmp("disk");
567        {
568            let s: SharedTreiberStack<u32> = SharedTreiberStack::create(&p, 8).unwrap();
569            s.push(1).unwrap();
570            s.push(2).unwrap();
571            s.push(3).unwrap();
572            s.flush().unwrap();
573        }
574        let s2: SharedTreiberStack<u32> = SharedTreiberStack::open(&p, 8).unwrap();
575        assert_eq!(s2.pop(), Some(3));
576        assert_eq!(s2.pop(), Some(2));
577        assert_eq!(s2.pop(), Some(1));
578        std::fs::remove_file(&p).ok();
579    }
580}
subetha_cxc/shared_treiber_stack.rs

subetha_cxc/
shared_treiber_stack.rs
