Skip to main content

subetha_cxc/
shared_atomic.rs

1//! `SharedAtomic<T>` - cross-process atomic counter / flag.
2//!
3//! Backed by an MMF cell whose payload is interpreted directly as
4//! `AtomicU8 / AtomicU16 / AtomicU32 / AtomicU64`. The atomic ops
5//! are cross-process safe on every modern CPU because hardware
6//! cache coherence guarantees the atomic semantics across address
7//! spaces; the only requirement is that both processes map the
8//! same physical page (which the OS guarantees when they open the
9//! same MMF file).
10//!
11//! Three concrete types:
12//! - `SharedAtomicU32`
13//! - `SharedAtomicU64`
14//! - `SharedAtomicBool` (one byte, but enforced bool semantics)
15//!
16//! Type-erased layout: header + payload region the size of the
17//! native atomic, aligned naturally.
18
19use std::fs::{File, OpenOptions};
20use std::path::Path;
21use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
22
23use memmap2::{MmapMut, MmapOptions};
24
25pub const ATOMIC_MAGIC: u32 = 0x4150_5443;
26
27#[repr(C, align(64))]
28struct AtomicHeader {
29    magic: u32,
30    width: u32,  // 1, 4, or 8 bytes
31    payload_u64: AtomicU64,  // also covers u32, u8 via punning
32}
33
34const ATOMIC_FILE_SIZE: usize = std::mem::size_of::<AtomicHeader>();
35
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum SharedAtomicError {
38    LayoutMismatch,
39    IoError(std::io::ErrorKind),
40}
41
42impl From<std::io::Error> for SharedAtomicError {
43    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
44}
45
46macro_rules! shared_atomic_impl {
47    ($name:ident, $atomic:ty, $native:ty, $width:expr) => {
48        pub struct $name {
49            _file: File,
50            mmap: MmapMut,
51            header_sidecar: subetha_core::HandshakeHeader,
52            ring_sidecar: Box<subetha_core::ObservationRing>,
53        }
54
55        unsafe impl Send for $name {}
56        unsafe impl Sync for $name {}
57
58        impl subetha_sidecar::AdaptiveInstance for $name {
59            fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
60            fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
61            fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
62                Box::new(subetha_sidecar::NoMigrationPolicy)
63            }
64        }
65
66        impl $name {
67            /// Obtain the atomic at `path`, initializing it to `init` if
68            /// the path does not yet exist and attaching to it if it
69            /// does. Attaching leaves the live value in place; `init` is
70            /// then unused. A region built for a different width is a
71            /// `LayoutMismatch`. [`reset`](Self::reset) reinitializes.
72            pub fn create(path: impl AsRef<Path>, init: $native) -> Result<Self, SharedAtomicError> {
73                let (file, mmap) = crate::mmf_attach::create_or_attach(
74                    path.as_ref(),
75                    ATOMIC_FILE_SIZE,
76                    |ptr| unsafe { Self::init_region(ptr, init) },
77                    |ptr| unsafe { (*(ptr as *const AtomicHeader)).magic == ATOMIC_MAGIC },
78                )?;
79                Self::from_region(file, mmap)
80            }
81
82            /// Truncate the atomic at `path` and initialize it to
83            /// `init`, discarding the value a live peer holds. For a
84            /// caller that knows it owns the path.
85            pub fn reset(path: impl AsRef<Path>, init: $native) -> Result<Self, SharedAtomicError> {
86                let (file, mmap) = crate::mmf_attach::reset(
87                    path.as_ref(),
88                    ATOMIC_FILE_SIZE,
89                    |ptr| unsafe { Self::init_region(ptr, init) },
90                )?;
91                Self::from_region(file, mmap)
92            }
93
94            /// Lay out a fresh atomic: width and value first, magic
95            /// last, because attachers spin on it.
96            ///
97            /// # Safety
98            /// `ptr` addresses at least `ATOMIC_FILE_SIZE` writable
99            /// zeroed bytes.
100            unsafe fn init_region(ptr: *mut u8, init: $native) {
101                let hdr = ptr as *mut AtomicHeader;
102                unsafe {
103                    (*hdr).width = $width;
104                    let payload = (&raw mut (*hdr).payload_u64) as *mut $atomic;
105                    std::ptr::write(payload, <$atomic>::new(init));
106                    std::ptr::write_volatile(&raw mut (*hdr).magic, ATOMIC_MAGIC);
107                }
108            }
109
110            /// Wrap an initialized region, refusing one built for a
111            /// different width.
112            fn from_region(file: File, mmap: MmapMut) -> Result<Self, SharedAtomicError> {
113                let hdr = unsafe { &*(mmap.as_ptr() as *const AtomicHeader) };
114                if hdr.magic != ATOMIC_MAGIC || hdr.width != $width {
115                    return Err(SharedAtomicError::LayoutMismatch);
116                }
117                Ok(Self {
118                    _file: file, mmap,
119                    header_sidecar: subetha_core::HandshakeHeader::new(),
120                    ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
121                })
122            }
123
124            pub fn open(path: impl AsRef<Path>) -> Result<Self, SharedAtomicError> {
125                let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
126                if file.metadata()?.len() < ATOMIC_FILE_SIZE as u64 {
127                    return Err(SharedAtomicError::LayoutMismatch);
128                }
129                let mmap = unsafe { MmapOptions::new().len(ATOMIC_FILE_SIZE).map_mut(&file)? };
130                Self::from_region(file, mmap)
131            }
132
133            #[inline]
134            fn atomic(&self) -> &$atomic {
135                let base = unsafe {
136                    self.mmap.as_ptr()
137                        .add(std::mem::offset_of!(AtomicHeader, payload_u64))
138                };
139                unsafe { &*(base as *const $atomic) }
140            }
141
142            #[inline]
143            pub fn load(&self, ord: Ordering) -> $native {
144                let v = self.atomic().load(ord);
145                self.ring_sidecar
146                    .push_op(crate::sidecar_ops::atomic::OP_LOAD, 0);
147                v
148            }
149
150            #[inline]
151            pub fn store(&self, v: $native, ord: Ordering) {
152                self.atomic().store(v, ord);
153                self.ring_sidecar
154                    .push_op(crate::sidecar_ops::atomic::OP_STORE, 0);
155            }
156
157            #[inline]
158            pub fn fetch_add(&self, v: $native, ord: Ordering) -> $native {
159                let prev = self.atomic().fetch_add(v, ord);
160                self.ring_sidecar
161                    .push_op(crate::sidecar_ops::atomic::OP_FETCH_ADD, 0);
162                prev
163            }
164
165            #[inline]
166            pub fn fetch_sub(&self, v: $native, ord: Ordering) -> $native {
167                let prev = self.atomic().fetch_sub(v, ord);
168                self.ring_sidecar
169                    .push_op(crate::sidecar_ops::atomic::OP_FETCH_ADD, 0);
170                prev
171            }
172
173            #[inline]
174            pub fn fetch_or(&self, v: $native, ord: Ordering) -> $native {
175                let prev = self.atomic().fetch_or(v, ord);
176                self.ring_sidecar
177                    .push_op(crate::sidecar_ops::atomic::OP_FETCH_ADD, 0);
178                prev
179            }
180
181            #[inline]
182            pub fn fetch_and(&self, v: $native, ord: Ordering) -> $native {
183                let prev = self.atomic().fetch_and(v, ord);
184                self.ring_sidecar
185                    .push_op(crate::sidecar_ops::atomic::OP_FETCH_ADD, 0);
186                prev
187            }
188
189            #[inline]
190            pub fn fetch_xor(&self, v: $native, ord: Ordering) -> $native {
191                let prev = self.atomic().fetch_xor(v, ord);
192                self.ring_sidecar
193                    .push_op(crate::sidecar_ops::atomic::OP_FETCH_ADD, 0);
194                prev
195            }
196
197            #[inline]
198            pub fn swap(&self, v: $native, ord: Ordering) -> $native {
199                let prev = self.atomic().swap(v, ord);
200                self.ring_sidecar
201                    .push_op(crate::sidecar_ops::atomic::OP_CAS, 0);
202                prev
203            }
204
205            #[inline]
206            pub fn compare_exchange(
207                &self, current: $native, new: $native,
208                success: Ordering, failure: Ordering,
209            ) -> Result<$native, $native> {
210                let r = self.atomic().compare_exchange(current, new, success, failure);
211                self.ring_sidecar
212                    .push_op(crate::sidecar_ops::atomic::OP_CAS, if r.is_err() { 1 } else { 0 });
213                r
214            }
215
216            pub fn flush(&self) -> Result<(), SharedAtomicError> {
217                self.mmap.flush()?;
218                Ok(())
219            }
220
221            /// Non-blocking flush: schedules a writeback via the OS
222            /// (msync(MS_ASYNC) on Linux; FlushViewOfFile without
223            /// FlushFileBuffers on Windows). Note: Windows is only
224            /// partially async (sync to page cache, not to disk).
225            pub fn flush_async(&self) -> Result<(), SharedAtomicError> {
226                self.mmap.flush_async()?;
227                Ok(())
228            }
229        }
230    };
231}
232
233shared_atomic_impl!(SharedAtomicU32, AtomicU32, u32, 4);
234shared_atomic_impl!(SharedAtomicU64, AtomicU64, u64, 8);
235
236pub struct SharedAtomicBool {
237    _file: File,
238    mmap: MmapMut,
239    header_sidecar: subetha_core::HandshakeHeader,
240    ring_sidecar: Box<subetha_core::ObservationRing>,
241}
242
243unsafe impl Send for SharedAtomicBool {}
244unsafe impl Sync for SharedAtomicBool {}
245
246impl subetha_sidecar::AdaptiveInstance for SharedAtomicBool {
247    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
248    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
249    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
250        Box::new(subetha_sidecar::NoMigrationPolicy)
251    }
252}
253
254impl SharedAtomicBool {
255    /// Obtain the flag at `path`, initializing it to `init` if the path
256    /// does not yet exist and attaching to it if it does. Attaching
257    /// leaves the live value in place, so a racing peer never resets a
258    /// flag another process already set.
259    /// [`reset`](Self::reset) reinitializes.
260    pub fn create(path: impl AsRef<Path>, init: bool) -> Result<Self, SharedAtomicError> {
261        let (file, mmap) = crate::mmf_attach::create_or_attach(
262            path.as_ref(),
263            ATOMIC_FILE_SIZE,
264            |ptr| unsafe { Self::init_region(ptr, init) },
265            |ptr| unsafe { (*(ptr as *const AtomicHeader)).magic == ATOMIC_MAGIC },
266        )?;
267        Self::from_region(file, mmap)
268    }
269
270    /// Truncate the flag at `path` and initialize it to `init`,
271    /// discarding the value live peers share. For a caller that knows
272    /// it owns the path.
273    pub fn reset(path: impl AsRef<Path>, init: bool) -> Result<Self, SharedAtomicError> {
274        let (file, mmap) = crate::mmf_attach::reset(
275            path.as_ref(),
276            ATOMIC_FILE_SIZE,
277            |ptr| unsafe { Self::init_region(ptr, init) },
278        )?;
279        Self::from_region(file, mmap)
280    }
281
282    /// Lay out the flag: width and payload first, magic last, because
283    /// attachers spin on it.
284    ///
285    /// # Safety
286    /// `ptr` addresses at least `ATOMIC_FILE_SIZE` writable zeroed bytes.
287    unsafe fn init_region(ptr: *mut u8, init: bool) {
288        let hdr = ptr as *mut AtomicHeader;
289        unsafe {
290            (*hdr).width = 1;
291            std::ptr::write(&raw mut (*hdr).payload_u64, AtomicU64::new(u64::from(init)));
292            std::ptr::write_volatile(&raw mut (*hdr).magic, ATOMIC_MAGIC);
293        }
294    }
295
296    /// Wrap an initialized region, refusing one laid out for a
297    /// different width.
298    fn from_region(file: File, mmap: MmapMut) -> Result<Self, SharedAtomicError> {
299        let hdr = unsafe { &*(mmap.as_ptr() as *const AtomicHeader) };
300        if hdr.magic != ATOMIC_MAGIC || hdr.width != 1 {
301            return Err(SharedAtomicError::LayoutMismatch);
302        }
303        Ok(Self {
304            _file: file, mmap,
305            header_sidecar: subetha_core::HandshakeHeader::new(),
306            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
307        })
308    }
309
310    pub fn open(path: impl AsRef<Path>) -> Result<Self, SharedAtomicError> {
311        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
312        if file.metadata()?.len() < ATOMIC_FILE_SIZE as u64 {
313            return Err(SharedAtomicError::LayoutMismatch);
314        }
315        let mmap = unsafe { MmapOptions::new().len(ATOMIC_FILE_SIZE).map_mut(&file)? };
316        let hdr = unsafe { &*(mmap.as_ptr() as *const AtomicHeader) };
317        if hdr.magic != ATOMIC_MAGIC || hdr.width != 1 {
318            return Err(SharedAtomicError::LayoutMismatch);
319        }
320        Ok(Self {
321            _file: file, mmap,
322            header_sidecar: subetha_core::HandshakeHeader::new(),
323            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
324        })
325    }
326
327    fn atomic(&self) -> &AtomicBool {
328        let base = unsafe {
329            self.mmap.as_ptr().add(std::mem::offset_of!(AtomicHeader, payload_u64))
330        };
331        unsafe { &*(base as *const AtomicBool) }
332    }
333
334    pub fn load(&self, ord: Ordering) -> bool {
335        let v = self.atomic().load(ord);
336        self.ring_sidecar
337            .push_op(crate::sidecar_ops::atomic::OP_LOAD, 0);
338        v
339    }
340    pub fn store(&self, v: bool, ord: Ordering) {
341        self.atomic().store(v, ord);
342        self.ring_sidecar
343            .push_op(crate::sidecar_ops::atomic::OP_STORE, 0);
344    }
345    pub fn swap(&self, v: bool, ord: Ordering) -> bool {
346        let prev = self.atomic().swap(v, ord);
347        self.ring_sidecar
348            .push_op(crate::sidecar_ops::atomic::OP_CAS, 0);
349        prev
350    }
351
352    pub fn flush(&self) -> Result<(), SharedAtomicError> {
353        self.mmap.flush()?;
354        Ok(())
355    }
356
357    /// Non-blocking flush: schedules a writeback via the OS.
358    /// Note: Windows is only partially async (sync to page cache,
359    /// not to disk).
360    pub fn flush_async(&self) -> Result<(), SharedAtomicError> {
361        self.mmap.flush_async()?;
362        Ok(())
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
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-atomic-{name}-{pid}.bin"));
374        p
375    }
376
377    #[test]
378    fn u64_load_store_round_trip() {
379        let p = tmp("u64-rt");
380        let a = SharedAtomicU64::create(&p, 42).unwrap();
381        assert_eq!(a.load(Ordering::Acquire), 42);
382        a.store(99, Ordering::Release);
383        assert_eq!(a.load(Ordering::Acquire), 99);
384        std::fs::remove_file(&p).ok();
385    }
386
387    #[test]
388    fn u32_fetch_add_increments() {
389        let p = tmp("u32-add");
390        let a = SharedAtomicU32::create(&p, 0).unwrap();
391        for _ in 0..100 { a.fetch_add(1, Ordering::AcqRel); }
392        assert_eq!(a.load(Ordering::Acquire), 100);
393        std::fs::remove_file(&p).ok();
394    }
395
396    /// A second create attaches with the live value in place, ignoring
397    /// its init argument; reset is what re-seeds.
398    #[test]
399    fn second_create_attaches_and_keeps_the_value() {
400        let p = tmp("attach");
401        std::fs::remove_file(&p).ok();
402        let a = SharedAtomicU64::create(&p, 42).unwrap();
403        a.store(777, Ordering::Release);
404
405        let b = SharedAtomicU64::create(&p, 0).unwrap();
406        assert_eq!(b.load(Ordering::Acquire), 777, "attach clobbered the value");
407
408        // Windows refuses to truncate a mapped file, so every handle goes
409        // before the reset.
410        drop(a);
411        drop(b);
412        let fresh = SharedAtomicU64::reset(&p, 5).unwrap();
413        assert_eq!(fresh.load(Ordering::Acquire), 5, "reset did not re-seed");
414        drop(fresh);
415        std::fs::remove_file(&p).ok();
416    }
417
418    /// Attaching with a different width is refused.
419    #[test]
420    fn create_refuses_a_mismatched_width() {
421        let p = tmp("mismatch");
422        std::fs::remove_file(&p).ok();
423        let a = SharedAtomicU64::create(&p, 1).unwrap();
424        assert!(matches!(
425            SharedAtomicU32::create(&p, 1),
426            Err(SharedAtomicError::LayoutMismatch),
427        ));
428        drop(a);
429        std::fs::remove_file(&p).ok();
430    }
431
432    /// A second create attaches to the live flag rather than resetting
433    /// it to its init value; reset is what discards it. The other widths
434    /// obtain the same way, so the bool must not be the odd one out.
435    #[test]
436    fn bool_second_create_attaches_and_keeps_the_value() {
437        let p = tmp("bool-attach");
438        std::fs::remove_file(&p).ok();
439
440        let first = SharedAtomicBool::create(&p, false).unwrap();
441        first.store(true, Ordering::Release);
442
443        let second = SharedAtomicBool::create(&p, false).unwrap();
444        assert!(
445            second.load(Ordering::Acquire),
446            "attach reset a flag another handle had already set",
447        );
448
449        // Windows refuses to truncate a mapped file, so every handle goes
450        // before the reset.
451        drop(first);
452        drop(second);
453        let fresh = SharedAtomicBool::reset(&p, false).unwrap();
454        assert!(!fresh.load(Ordering::Acquire), "reset kept the old value");
455        drop(fresh);
456        std::fs::remove_file(&p).ok();
457    }
458
459    #[test]
460    fn cross_handle_visibility() {
461        let p = tmp("cross-handle");
462        let writer = SharedAtomicU64::create(&p, 0).unwrap();
463        let reader = SharedAtomicU64::open(&p).unwrap();
464        writer.store(7777, Ordering::Release);
465        assert_eq!(reader.load(Ordering::Acquire), 7777);
466        std::fs::remove_file(&p).ok();
467    }
468
469    #[test]
470    fn concurrent_fetch_add_sums_correctly() {
471        use std::sync::Arc;
472        use std::thread;
473        let p = tmp("concurrent");
474        let a = Arc::new(SharedAtomicU64::create(&p, 0).unwrap());
475        let mut handles = vec![];
476        for _ in 0..8 {
477            let a = a.clone();
478            handles.push(thread::spawn(move || {
479                for _ in 0..1000 { a.fetch_add(1, Ordering::AcqRel); }
480            }));
481        }
482        for h in handles { h.join().unwrap(); }
483        assert_eq!(a.load(Ordering::Acquire), 8000);
484        std::fs::remove_file(&p).ok();
485    }
486
487    #[test]
488    fn compare_exchange_wins_once() {
489        let p = tmp("cas");
490        let a = SharedAtomicU64::create(&p, 5).unwrap();
491        let r1 = a.compare_exchange(5, 10, Ordering::AcqRel, Ordering::Acquire);
492        let r2 = a.compare_exchange(5, 20, Ordering::AcqRel, Ordering::Acquire);
493        assert_eq!(r1, Ok(5));
494        assert_eq!(r2, Err(10));
495        assert_eq!(a.load(Ordering::Acquire), 10);
496        std::fs::remove_file(&p).ok();
497    }
498
499    #[test]
500    fn bool_load_store_swap() {
501        let p = tmp("bool");
502        let b = SharedAtomicBool::create(&p, false).unwrap();
503        assert!(!b.load(Ordering::Acquire));
504        b.store(true, Ordering::Release);
505        assert!(b.load(Ordering::Acquire));
506        let prev = b.swap(false, Ordering::AcqRel);
507        assert!(prev);
508        assert!(!b.load(Ordering::Acquire));
509        std::fs::remove_file(&p).ok();
510    }
511
512    #[test]
513    fn disk_persistence_survives_reopen() {
514        let p = tmp("disk-persist");
515        {
516            let a = SharedAtomicU64::create(&p, 12345).unwrap();
517            a.flush().unwrap();
518        }
519        let a2 = SharedAtomicU64::open(&p).unwrap();
520        assert_eq!(a2.load(Ordering::Acquire), 12345);
521        std::fs::remove_file(&p).ok();
522    }
523
524    #[test]
525    fn open_rejects_wrong_width() {
526        let p = tmp("wrong-width");
527        let _a = SharedAtomicU64::create(&p, 0).unwrap();
528        match SharedAtomicU32::open(&p) {
529            Err(SharedAtomicError::LayoutMismatch) => {}
530            other => panic!("expected LayoutMismatch, got {:?}", other.as_ref().err()),
531        }
532        std::fs::remove_file(&p).ok();
533    }
534}