subetha_cxc/
shared_once_cell.rs

1//! `SharedOnceCell<T>` - cross-process init-once cell.
2//!
3//! State machine:
4//! - EMPTY (0): no value; first writer to CAS to INITIALIZING wins.
5//! - INITIALIZING (1): a writer is filling the payload; other
6//!   writers spin until state advances.
7//! - INITIALIZED (2): payload is stable; readers may consume.
8//!
9//! The winner of the EMPTY -> INITIALIZING CAS performs the write
10//! and advances to INITIALIZED. Losers see INITIALIZED and read the
11//! winner's bytes.
12//!
13//! This is the cross-process analog of `once_cell::sync::OnceCell`,
14//! with cross-process safety guaranteed by the atomic CAS protocol
15//! over shared memory.
16
17use std::fs::{File, OpenOptions};
18use std::marker::PhantomData;
19use std::mem::{align_of, size_of};
20use std::path::Path;
21use std::sync::atomic::{AtomicU8, Ordering};
22
23use memmap2::{MmapMut, MmapOptions};
24
25pub const ONCE_MAGIC: u32 = 0x4F4E_4346;
26pub const ONCE_PAYLOAD_BYTES: usize = 56;
27
28pub const STATE_EMPTY: u8 = 0;
29pub const STATE_INITIALIZING: u8 = 1;
30pub const STATE_INITIALIZED: u8 = 2;
31
32#[repr(C, align(64))]
33pub struct OnceHeader {
34    pub magic: u32,
35    pub size: u32,
36    pub state: AtomicU8,
37    pub _pad_to_payload: [u8; 7],
38    pub payload: [u8; ONCE_PAYLOAD_BYTES],
39}
40
41pub const ONCE_FILE_SIZE: usize = size_of::<OnceHeader>();
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum SharedOnceError {
45    LayoutMismatch,
46    PayloadTooLarge,
47    IoError(std::io::ErrorKind),
48}
49
50impl From<std::io::Error> for SharedOnceError {
51    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
52}
53
54pub struct SharedOnceCell<T: Copy + 'static> {
55    _file: File,
56    mmap: MmapMut,
57    _phantom: PhantomData<T>,
58    header_sidecar: subetha_core::HandshakeHeader,
59    ring_sidecar: Box<subetha_core::ObservationRing>,
60}
61
62unsafe impl<T: Copy + Send + 'static> Send for SharedOnceCell<T> {}
63unsafe impl<T: Copy + Sync + 'static> Sync for SharedOnceCell<T> {}
64
65impl<T: Copy + Send + Sync + 'static> subetha_sidecar::AdaptiveInstance for SharedOnceCell<T> {
66    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
67    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
68    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
69        Box::new(subetha_sidecar::NoMigrationPolicy)
70    }
71}
72
73impl<T: Copy + 'static> SharedOnceCell<T> {
74    /// Obtain the cell at `path`, initializing an empty one if the path
75    /// does not yet exist and attaching to it if it does. Attaching
76    /// leaves an initialized value in place; a region built for a
77    /// different payload type is a `LayoutMismatch`.
78    /// [`reset`](Self::reset) reinitializes.
79    pub fn create(path: impl AsRef<Path>) -> Result<Self, SharedOnceError> {
80        Self::check_layout()?;
81        let (file, mmap) = crate::mmf_attach::create_or_attach(
82            path.as_ref(),
83            ONCE_FILE_SIZE,
84            |ptr| unsafe { Self::init_region(ptr) },
85            |ptr| unsafe { (*(ptr as *const OnceHeader)).magic == ONCE_MAGIC },
86        )?;
87        Self::from_region(file, mmap)
88    }
89
90    /// Truncate the cell at `path` and initialize an empty one,
91    /// discarding whatever value a live peer holds. For a caller that
92    /// knows it owns the path.
93    pub fn reset(path: impl AsRef<Path>) -> Result<Self, SharedOnceError> {
94        Self::check_layout()?;
95        let (file, mmap) = crate::mmf_attach::reset(path.as_ref(), ONCE_FILE_SIZE, |ptr| unsafe {
96            Self::init_region(ptr)
97        })?;
98        Self::from_region(file, mmap)
99    }
100
101    /// Lay out an empty cell: the zeroed region is already
102    /// `STATE_EMPTY` with a zero payload, so only the size and then the
103    /// magic are written, magic last, because attachers spin on it.
104    ///
105    /// # Safety
106    /// `ptr` addresses at least [`ONCE_FILE_SIZE`] writable zeroed
107    /// bytes.
108    unsafe fn init_region(ptr: *mut u8) {
109        let hdr = ptr as *mut OnceHeader;
110        unsafe {
111            (*hdr).size = size_of::<T>() as u32;
112            std::ptr::write_volatile(&raw mut (*hdr).magic, ONCE_MAGIC);
113        }
114    }
115
116    /// Wrap an initialized region, refusing one built for a different
117    /// payload type.
118    fn from_region(file: File, mmap: MmapMut) -> Result<Self, SharedOnceError> {
119        let header = unsafe { &*(mmap.as_ptr() as *const OnceHeader) };
120        if header.magic != ONCE_MAGIC || header.size as usize != size_of::<T>() {
121            return Err(SharedOnceError::LayoutMismatch);
122        }
123        Ok(Self {
124            _file: file, mmap, _phantom: PhantomData,
125            header_sidecar: subetha_core::HandshakeHeader::new(),
126            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
127        })
128    }
129
130    pub fn open(path: impl AsRef<Path>) -> Result<Self, SharedOnceError> {
131        Self::check_layout()?;
132        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
133        if file.metadata()?.len() < ONCE_FILE_SIZE as u64 {
134            return Err(SharedOnceError::LayoutMismatch);
135        }
136        let mmap = unsafe { MmapOptions::new().len(ONCE_FILE_SIZE).map_mut(&file)? };
137        Self::from_region(file, mmap)
138    }
139
140    fn check_layout() -> Result<(), SharedOnceError> {
141        if size_of::<T>() > ONCE_PAYLOAD_BYTES {
142            return Err(SharedOnceError::PayloadTooLarge);
143        }
144        if align_of::<T>() > 8 {
145            return Err(SharedOnceError::PayloadTooLarge);
146        }
147        Ok(())
148    }
149
150    fn header(&self) -> &OnceHeader {
151        unsafe { &*(self.mmap.as_ptr() as *const OnceHeader) }
152    }
153
154    /// True when the cell has been initialised.
155    pub fn is_initialized(&self) -> bool {
156        self.header().state.load(Ordering::Acquire) == STATE_INITIALIZED
157    }
158
159    /// Get the value if initialised; otherwise return None.
160    /// Non-blocking; never invokes the initialiser.
161    pub fn get(&self) -> Option<T> {
162        let header = self.header();
163        if header.state.load(Ordering::Acquire) != STATE_INITIALIZED {
164            self.ring_sidecar
165                .push_op(crate::sidecar_ops::cell::OP_GET, 2); // empty / uninitialised
166            return None;
167        }
168        let value: T = unsafe {
169            let src = header.payload.as_ptr() as *const T;
170            std::ptr::read_unaligned(src)
171        };
172        self.ring_sidecar
173            .push_op(crate::sidecar_ops::cell::OP_GET, 0);
174        Some(value)
175    }
176
177    /// Try to write the value. Returns `true` if this caller won
178    /// the init race, `false` if the cell was already initialised
179    /// or another init is in progress.
180    pub fn set(&self, value: T) -> bool {
181        let header = self.header();
182        if header.state.compare_exchange(
183            STATE_EMPTY, STATE_INITIALIZING,
184            Ordering::AcqRel, Ordering::Acquire,
185        ).is_err() {
186            self.ring_sidecar
187                .push_op(crate::sidecar_ops::cell::OP_SET, 1); // lost the init race
188            return false;
189        }
190        unsafe {
191            let dst = header.payload.as_ptr() as *mut T;
192            std::ptr::write_unaligned(dst, value);
193        }
194        header.state.store(STATE_INITIALIZED, Ordering::Release);
195        self.ring_sidecar
196            .push_op(crate::sidecar_ops::cell::OP_SET, 0);
197        true
198    }
199
200    /// Get the cached value, or run `init` to produce it. The first
201    /// caller across all processes runs `init`; subsequent callers
202    /// spin until the value is published and return that value.
203    pub fn get_or_init<F: FnOnce() -> T>(&self, init: F) -> T {
204        if let Some(v) = self.get() { return v; }
205        let header = self.header();
206        match header.state.compare_exchange(
207            STATE_EMPTY, STATE_INITIALIZING,
208            Ordering::AcqRel, Ordering::Acquire,
209        ) {
210            Ok(_) => {
211                // We won; produce the value and publish.
212                let v = init();
213                unsafe {
214                    let dst = header.payload.as_ptr() as *mut T;
215                    std::ptr::write_unaligned(dst, v);
216                }
217                header.state.store(STATE_INITIALIZED, Ordering::Release);
218                self.ring_sidecar
219                    .push_op(crate::sidecar_ops::cell::OP_SET, 0);
220                v
221            }
222            Err(_) => {
223                // Spin until the winner publishes.
224                while header.state.load(Ordering::Acquire) != STATE_INITIALIZED {
225                    std::hint::spin_loop();
226                }
227                self.get().expect("INITIALIZED implies value present")
228            }
229        }
230    }
231
232    /// Non-blocking flush: schedules a writeback via the OS.
233    /// Note: Windows is only partially async (sync to page cache,
234    /// not to disk).
235    pub fn flush_async(&self) -> Result<(), SharedOnceError> {
236        self.mmap.flush_async()?;
237        Ok(())
238    }
239
240    pub fn flush(&self) -> Result<(), SharedOnceError> {
241        self.mmap.flush()?;
242        Ok(())
243    }
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249
250    fn tmp(name: &str) -> std::path::PathBuf {
251        let mut p = std::env::temp_dir();
252        let pid = std::process::id();
253        p.push(format!("subetha-once-{name}-{pid}.bin"));
254        p
255    }
256
257    #[test]
258    fn fresh_cell_is_empty() {
259        let p = tmp("fresh");
260        let c: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
261        assert!(!c.is_initialized());
262        assert_eq!(c.get(), None);
263        std::fs::remove_file(&p).ok();
264    }
265
266    #[test]
267    fn first_set_wins_subsequent_sets_lose() {
268        let p = tmp("first-wins");
269        let c: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
270        assert!(c.set(42));
271        assert!(!c.set(99));
272        assert_eq!(c.get(), Some(42));
273        std::fs::remove_file(&p).ok();
274    }
275
276    /// A second create attaches with the initialized value in place -
277    /// truncation here would break the once guarantee; reset is what
278    /// strips it.
279    #[test]
280    fn second_create_attaches_and_keeps_the_value() {
281        let p = tmp("attach");
282        std::fs::remove_file(&p).ok();
283        let c: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
284        assert!(c.set(42));
285
286        let c2: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
287        assert_eq!(c2.get(), Some(42), "attach lost the initialized value");
288        assert!(!c2.set(99), "attach reopened a spent cell");
289
290        // Windows refuses to truncate a mapped file, so every handle goes
291        // before the reset.
292        drop(c);
293        drop(c2);
294        let fresh: SharedOnceCell<u64> = SharedOnceCell::reset(&p).unwrap();
295        assert_eq!(fresh.get(), None, "reset left a value behind");
296        assert!(fresh.set(7));
297        drop(fresh);
298        std::fs::remove_file(&p).ok();
299    }
300
301    /// Attaching with a different payload type is refused.
302    #[test]
303    fn create_refuses_a_mismatched_region() {
304        let p = tmp("mismatch");
305        std::fs::remove_file(&p).ok();
306        let c: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
307        assert!(matches!(
308            SharedOnceCell::<u32>::create(&p),
309            Err(SharedOnceError::LayoutMismatch),
310        ));
311        drop(c);
312        std::fs::remove_file(&p).ok();
313    }
314
315    #[test]
316    fn cross_handle_init_visible_to_other_handle() {
317        let p = tmp("cross-handle");
318        let writer: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
319        let reader: SharedOnceCell<u64> = SharedOnceCell::open(&p).unwrap();
320        assert!(!reader.is_initialized());
321        writer.set(7777);
322        assert!(reader.is_initialized());
323        assert_eq!(reader.get(), Some(7777));
324        std::fs::remove_file(&p).ok();
325    }
326
327    #[test]
328    fn get_or_init_runs_closure_at_most_once() {
329        use std::sync::Arc;
330        use std::sync::atomic::AtomicU32;
331        use std::thread;
332        let p = tmp("get-or-init");
333        let c: Arc<SharedOnceCell<u64>> = Arc::new(SharedOnceCell::create(&p).unwrap());
334        let runs = Arc::new(AtomicU32::new(0));
335        let mut handles = vec![];
336        for _ in 0..8 {
337            let c = c.clone();
338            let runs = runs.clone();
339            handles.push(thread::spawn(move || {
340                c.get_or_init(|| {
341                    runs.fetch_add(1, Ordering::AcqRel);
342                    1234u64
343                })
344            }));
345        }
346        let results: Vec<u64> = handles.into_iter().map(|h| h.join().unwrap()).collect();
347        assert!(results.iter().all(|v| *v == 1234));
348        assert_eq!(runs.load(Ordering::Acquire), 1,
349                   "init closure must run exactly once across 8 threads");
350        std::fs::remove_file(&p).ok();
351    }
352
353    #[test]
354    fn disk_persistence_survives_reopen() {
355        let p = tmp("disk-persist");
356        {
357            let c: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
358            c.set(8888);
359            c.flush().unwrap();
360        }
361        let c2: SharedOnceCell<u64> = SharedOnceCell::open(&p).unwrap();
362        assert_eq!(c2.get(), Some(8888));
363        assert!(c2.is_initialized());
364        // Set must fail on reopen because the cell is already init.
365        assert!(!c2.set(9999));
366        assert_eq!(c2.get(), Some(8888));
367        std::fs::remove_file(&p).ok();
368    }
369
370    #[test]
371    fn payload_too_large_at_create() {
372        #[allow(dead_code)] // size_of<Big> is the test signal, not the field
373        struct Big([u8; ONCE_PAYLOAD_BYTES + 1]);
374        impl Copy for Big {}
375        impl Clone for Big { fn clone(&self) -> Self { *self } }
376        let p = tmp("too-large");
377        match SharedOnceCell::<Big>::create(&p) {
378            Err(SharedOnceError::PayloadTooLarge) => {}
379            other => panic!("expected PayloadTooLarge, got {:?}", other.as_ref().err()),
380        }
381        std::fs::remove_file(&p).ok();
382    }
383
384    #[test]
385    fn open_rejects_wrong_payload_size() {
386        let p = tmp("wrong-size");
387        let _c: SharedOnceCell<u64> = SharedOnceCell::create(&p).unwrap();
388        match SharedOnceCell::<u32>::open(&p) {
389            Err(SharedOnceError::LayoutMismatch) => {}
390            other => panic!("expected LayoutMismatch, got {:?}", other.as_ref().err()),
391        }
392        std::fs::remove_file(&p).ok();
393    }
394}
subetha_cxc/shared_once_cell.rs

subetha_cxc/
shared_once_cell.rs
