Skip to main content

subetha_cxc/
shared_rw_lock.rs

1//! `SharedRWLock` - cross-process reader-writer lock with writer
2//! priority.
3//!
4//! Multiple concurrent readers OR exactly one writer. When a writer
5//! is waiting, new readers block to prevent writer starvation.
6//!
7//! # State encoding
8//!
9//! ONE AtomicU64 packed:
10//! - bit 63: writer active (1 if a writer holds the lock)
11//! - bits 32-62: writers waiting count (31 bits)
12//! - bits 0-31: reader count (32 bits)
13//!
14//! All transitions are single CAS so observers never see torn state.
15
16use std::fs::{File, OpenOptions};
17use std::mem::size_of;
18use std::path::Path;
19use std::sync::atomic::{AtomicU64, Ordering};
20
21use memmap2::{MmapMut, MmapOptions};
22
23pub const RWLOCK_MAGIC: u64 = 0x4150_5257_4C4F_434B;
24
25const WRITER_BIT: u64 = 1u64 << 63;
26const WAITING_SHIFT: u64 = 32;
27const WAITING_MASK: u64 = 0x7FFF_FFFF << WAITING_SHIFT;
28const READERS_MASK: u64 = 0xFFFF_FFFF;
29
30#[repr(C, align(64))]
31pub struct RWLockHeader {
32    pub magic: u64,
33    pub state: AtomicU64,
34    _pad: [u8; 48],
35}
36
37const _: () = {
38    assert!(size_of::<RWLockHeader>() == 64);
39};
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum RWLockError {
43    WouldBlock,
44    LayoutMismatch,
45    IoError(std::io::ErrorKind),
46}
47
48impl From<std::io::Error> for RWLockError {
49    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
50}
51
52pub struct SharedRWLock {
53    _file: File,
54    mmap: MmapMut,
55    header_sidecar: subetha_core::HandshakeHeader,
56    ring_sidecar: Box<subetha_core::ObservationRing>,
57}
58
59unsafe impl Send for SharedRWLock {}
60unsafe impl Sync for SharedRWLock {}
61
62impl subetha_sidecar::AdaptiveInstance for SharedRWLock {
63    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
64    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
65    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
66        Box::new(subetha_sidecar::NoMigrationPolicy)
67    }
68}
69
70impl SharedRWLock {
71    pub fn create(path: impl AsRef<Path>) -> Result<Self, RWLockError> {
72        let total = size_of::<RWLockHeader>();
73        let file = OpenOptions::new()
74            .read(true).write(true).create(true).truncate(true)
75            .open(path.as_ref())?;
76        file.set_len(total as u64)?;
77        let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
78        let hdr = mmap.as_mut_ptr() as *mut RWLockHeader;
79        unsafe {
80            std::ptr::write_bytes(hdr as *mut u8, 0, total);
81            (*hdr).magic = RWLOCK_MAGIC;
82        }
83        Ok(Self {
84            _file: file, mmap,
85            header_sidecar: subetha_core::HandshakeHeader::new(),
86            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
87        })
88    }
89
90    pub fn open(path: impl AsRef<Path>) -> Result<Self, RWLockError> {
91        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
92        if file.metadata()?.len() < size_of::<RWLockHeader>() as u64 {
93            return Err(RWLockError::LayoutMismatch);
94        }
95        let mmap = unsafe {
96            MmapOptions::new().len(size_of::<RWLockHeader>()).map_mut(&file)?
97        };
98        let hdr = unsafe { &*(mmap.as_ptr() as *const RWLockHeader) };
99        if hdr.magic != RWLOCK_MAGIC {
100            return Err(RWLockError::LayoutMismatch);
101        }
102        Ok(Self {
103            _file: file, mmap,
104            header_sidecar: subetha_core::HandshakeHeader::new(),
105            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
106        })
107    }
108
109    fn state(&self) -> &AtomicU64 {
110        unsafe { &(*(self.mmap.as_ptr() as *const RWLockHeader)).state }
111    }
112
113    /// Try to acquire a read lock without blocking.
114    pub fn try_read_lock(&self) -> Result<ReadGuard<'_>, RWLockError> {
115        let r = self.try_read_lock_inner();
116        self.ring_sidecar.push_op(
117            crate::sidecar_ops::rw_lock::OP_TRY_READ,
118            if r.is_err() { 1 } else { 0 },
119        );
120        r
121    }
122
123    fn try_read_lock_inner(&self) -> Result<ReadGuard<'_>, RWLockError> {
124        loop {
125            let s = self.state().load(Ordering::Acquire);
126            let writer_active = (s & WRITER_BIT) != 0;
127            let writers_waiting = (s & WAITING_MASK) >> WAITING_SHIFT;
128            if writer_active || writers_waiting > 0 {
129                return Err(RWLockError::WouldBlock);
130            }
131            let readers = s & READERS_MASK;
132            let new = (s & !READERS_MASK) | (readers + 1);
133            if self.state().compare_exchange(
134                s, new, Ordering::AcqRel, Ordering::Acquire,
135            ).is_ok() {
136                return Ok(ReadGuard { lock: self });
137            }
138        }
139    }
140
141    /// Acquire a read lock, blocking with backoff until available.
142    /// Writer-priority: blocks if any writer is active OR waiting.
143    pub fn read_lock(&self) -> ReadGuard<'_> {
144        let mut spins = 0u32;
145        loop {
146            if let Ok(g) = self.try_read_lock_inner() {
147                self.ring_sidecar.push_op(
148                    crate::sidecar_ops::rw_lock::OP_READ,
149                    if spins > 0 { 1 } else { 0 }, // contention
150                );
151                return g;
152            }
153            spins += 1;
154            if spins < 32 {
155                std::hint::spin_loop();
156            } else if spins < 256 {
157                std::thread::yield_now();
158            } else {
159                std::thread::sleep(std::time::Duration::from_micros(50));
160            }
161        }
162    }
163
164    /// Try to acquire a write lock without blocking.
165    pub fn try_write_lock(&self) -> Result<WriteGuard<'_>, RWLockError> {
166        let r = self.try_write_lock_inner();
167        self.ring_sidecar.push_op(
168            crate::sidecar_ops::rw_lock::OP_TRY_WRITE,
169            if r.is_err() { 1 } else { 0 },
170        );
171        r
172    }
173
174    fn try_write_lock_inner(&self) -> Result<WriteGuard<'_>, RWLockError> {
175        loop {
176            let s = self.state().load(Ordering::Acquire);
177            let writer_active = (s & WRITER_BIT) != 0;
178            let readers = s & READERS_MASK;
179            if writer_active || readers > 0 {
180                return Err(RWLockError::WouldBlock);
181            }
182            let new = (s & !WRITER_BIT) | WRITER_BIT;
183            if self.state().compare_exchange(
184                s, new, Ordering::AcqRel, Ordering::Acquire,
185            ).is_ok() {
186                return Ok(WriteGuard { lock: self });
187            }
188        }
189    }
190
191    /// Acquire a write lock, blocking until available. Registers
192    /// as "waiting" so new readers will block.
193    pub fn write_lock(&self) -> WriteGuard<'_> {
194        // Register as waiting.
195        self.state().fetch_add(1u64 << WAITING_SHIFT, Ordering::AcqRel);
196        let mut spins = 0u32;
197        loop {
198            let s = self.state().load(Ordering::Acquire);
199            let writer_active = (s & WRITER_BIT) != 0;
200            let readers = s & READERS_MASK;
201            if !writer_active && readers == 0 {
202                // Try to claim: set writer bit + decrement waiting.
203                let new = (s & READERS_MASK) | WRITER_BIT
204                    | ((((s & WAITING_MASK) >> WAITING_SHIFT) - 1) << WAITING_SHIFT);
205                if self.state().compare_exchange(
206                    s, new, Ordering::AcqRel, Ordering::Acquire,
207                ).is_ok() {
208                    self.ring_sidecar.push_op(
209                        crate::sidecar_ops::rw_lock::OP_WRITE,
210                        if spins > 0 { 1 } else { 0 }, // contention
211                    );
212                    return WriteGuard { lock: self };
213                }
214            }
215            spins += 1;
216            if spins < 32 {
217                std::hint::spin_loop();
218            } else if spins < 256 {
219                std::thread::yield_now();
220            } else {
221                std::thread::sleep(std::time::Duration::from_micros(50));
222            }
223        }
224    }
225
226    /// Number of active readers (observational; may race).
227    pub fn reader_count(&self) -> u32 {
228        (self.state().load(Ordering::Acquire) & READERS_MASK) as u32
229    }
230
231    /// True if a writer currently holds the lock.
232    pub fn has_writer(&self) -> bool {
233        (self.state().load(Ordering::Acquire) & WRITER_BIT) != 0
234    }
235
236    /// Number of writers currently waiting for the lock.
237    pub fn waiting_writers(&self) -> u32 {
238        ((self.state().load(Ordering::Acquire) & WAITING_MASK) >> WAITING_SHIFT) as u32
239    }
240
241    /// Release one reader. Internal; called by ReadGuard::drop.
242    /// Defensive: checks that the reader count is positive before
243    /// decrementing. In debug builds this panics on protocol
244    /// violation (release without acquire); in release builds it
245    /// silently no-ops to avoid underflow corruption.
246    fn release_read(&self) {
247        loop {
248            let s = self.state().load(Ordering::Acquire);
249            let readers = s & READERS_MASK;
250            debug_assert!(
251                readers > 0,
252                "SharedRWLock::release_read called when reader count is 0 - \
253                 indicates a protocol violation (double-release or release \
254                 without acquire). The lock counter will not be decremented.",
255            );
256            if readers == 0 { return; }
257            let new = (s & !READERS_MASK) | (readers - 1);
258            if self.state().compare_exchange(
259                s, new, Ordering::AcqRel, Ordering::Acquire,
260            ).is_ok() {
261                return;
262            }
263        }
264    }
265
266    /// Release the writer. Internal; called by WriteGuard::drop.
267    /// Defensive: checks that a writer is actually active before
268    /// clearing. In debug builds this panics on protocol violation.
269    fn release_write(&self) {
270        let prev = self.state().fetch_and(!WRITER_BIT, Ordering::AcqRel);
271        debug_assert!(
272            (prev & WRITER_BIT) != 0,
273            "SharedRWLock::release_write called when no writer holds the lock - \
274             indicates a protocol violation (double-release or release without \
275             acquire).",
276        );
277    }
278
279    /// Public hook for the `BlockingRWLock` wrapper to mirror the
280    /// inner `ReadGuard::drop` semantics after the wrapper's own
281    /// guard runs (the wrapper `mem::forget`s the inner guard so it
282    /// can interleave a wake call between the state release and the
283    /// guard's destructor).
284    pub fn release_read_for_blocking(&self) { self.release_read(); }
285
286    /// Public hook for the `BlockingRWLock` wrapper; mirror of
287    /// `WriteGuard::drop`.
288    pub fn release_write_for_blocking(&self) { self.release_write(); }
289
290    pub fn flush(&self) -> Result<(), RWLockError> {
291        self.mmap.flush()?;
292        Ok(())
293    }
294    pub fn flush_async(&self) -> Result<(), RWLockError> {
295        self.mmap.flush_async()?;
296        Ok(())
297    }
298}
299
300pub struct ReadGuard<'a> { lock: &'a SharedRWLock }
301impl Drop for ReadGuard<'_> {
302    fn drop(&mut self) { self.lock.release_read(); }
303}
304
305pub struct WriteGuard<'a> { lock: &'a SharedRWLock }
306impl Drop for WriteGuard<'_> {
307    fn drop(&mut self) { self.lock.release_write(); }
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313    use std::sync::atomic::{AtomicU32, Ordering as O};
314    use std::sync::Arc;
315    use std::thread;
316
317    fn tmp(name: &str) -> std::path::PathBuf {
318        let mut p = std::env::temp_dir();
319        let pid = std::process::id();
320        p.push(format!("subetha-rwlock-{name}-{pid}.bin"));
321        p
322    }
323
324    #[test]
325    fn create_initial_state_is_idle() {
326        let p = tmp("init");
327        let l = SharedRWLock::create(&p).unwrap();
328        assert_eq!(l.reader_count(), 0);
329        assert!(!l.has_writer());
330        assert_eq!(l.waiting_writers(), 0);
331        std::fs::remove_file(&p).ok();
332    }
333
334    #[test]
335    fn try_read_succeeds_when_idle() {
336        let p = tmp("try-read");
337        let l = SharedRWLock::create(&p).unwrap();
338        let g = l.try_read_lock().unwrap();
339        assert_eq!(l.reader_count(), 1);
340        drop(g);
341        assert_eq!(l.reader_count(), 0);
342        std::fs::remove_file(&p).ok();
343    }
344
345    #[test]
346    fn try_write_succeeds_when_idle() {
347        let p = tmp("try-write");
348        let l = SharedRWLock::create(&p).unwrap();
349        let g = l.try_write_lock().unwrap();
350        assert!(l.has_writer());
351        drop(g);
352        assert!(!l.has_writer());
353        std::fs::remove_file(&p).ok();
354    }
355
356    #[test]
357    fn multiple_readers_coexist() {
358        let p = tmp("multi-read");
359        let l = SharedRWLock::create(&p).unwrap();
360        let g1 = l.try_read_lock().unwrap();
361        let g2 = l.try_read_lock().unwrap();
362        let g3 = l.try_read_lock().unwrap();
363        assert_eq!(l.reader_count(), 3);
364        drop(g1); drop(g2); drop(g3);
365        assert_eq!(l.reader_count(), 0);
366        std::fs::remove_file(&p).ok();
367    }
368
369    #[test]
370    fn writer_excludes_readers() {
371        let p = tmp("w-excl-r");
372        let l = SharedRWLock::create(&p).unwrap();
373        let _w = l.try_write_lock().unwrap();
374        assert_eq!(l.try_read_lock().err(), Some(RWLockError::WouldBlock));
375        std::fs::remove_file(&p).ok();
376    }
377
378    #[test]
379    fn reader_excludes_writer() {
380        let p = tmp("r-excl-w");
381        let l = SharedRWLock::create(&p).unwrap();
382        let _r = l.try_read_lock().unwrap();
383        assert_eq!(l.try_write_lock().err(), Some(RWLockError::WouldBlock));
384        std::fs::remove_file(&p).ok();
385    }
386
387    #[test]
388    fn writer_excludes_writer() {
389        let p = tmp("w-excl-w");
390        let l = SharedRWLock::create(&p).unwrap();
391        let _w = l.try_write_lock().unwrap();
392        assert_eq!(l.try_write_lock().err(), Some(RWLockError::WouldBlock));
393        std::fs::remove_file(&p).ok();
394    }
395
396    #[test]
397    fn write_lock_blocks_until_readers_drop() {
398        // Clean pattern: spawn a reader thread that holds its guard
399        // for a known duration. Main thread spawns a writer that
400        // must block until the reader's guard drops. No unsafe
401        // ptr::read; the reader's guard lifetime is tied to its
402        // thread's scope.
403        let p = tmp("w-blocks");
404        let l = Arc::new(SharedRWLock::create(&p).unwrap());
405        let l_reader = l.clone();
406        let reader_done = Arc::new(AtomicU32::new(0));
407        let reader_done_clone = reader_done.clone();
408        let reader = thread::spawn(move || {
409            let _g = l_reader.read_lock();
410            std::thread::sleep(std::time::Duration::from_millis(30));
411            reader_done_clone.store(1, O::Release);
412            // Guard drops here, releasing the lock.
413        });
414        // Wait (bounded) for the reader thread to acquire; a fixed
415        // sleep races the scheduler under full-suite load.
416        let acquire_deadline = std::time::Instant::now()
417            + std::time::Duration::from_secs(5);
418        while l.reader_count() != 1
419            && std::time::Instant::now() < acquire_deadline
420        {
421            std::thread::yield_now();
422        }
423        assert_eq!(l.reader_count(), 1);
424
425        let l_writer = l.clone();
426        let writer_started = std::time::Instant::now();
427        let writer = thread::spawn(move || {
428            let _g = l_writer.write_lock();
429            writer_started.elapsed()
430        });
431
432        let elapsed = writer.join().unwrap();
433        reader.join().unwrap();
434        // Writer should have blocked at least until reader finished
435        // (which is ~30ms - 5ms from when writer was spawned = ~25ms).
436        assert!(
437            elapsed >= std::time::Duration::from_millis(15),
438            "writer should have blocked for the reader's hold time, got {elapsed:?}",
439        );
440        assert_eq!(reader_done.load(O::Acquire), 1);
441        std::fs::remove_file(&p).ok();
442    }
443
444    #[test]
445    fn writer_priority_blocks_new_readers() {
446        // When a writer is waiting, new try_read should fail
447        // (writer priority).
448        let p = tmp("w-priority");
449        let l = SharedRWLock::create(&p).unwrap();
450        // Simulate a waiting writer by bumping the waiting count
451        // directly (real writers do this in write_lock).
452        l.state().fetch_add(1u64 << WAITING_SHIFT, Ordering::AcqRel);
453        assert_eq!(l.try_read_lock().err(), Some(RWLockError::WouldBlock));
454        // Clean up the bumped count for the file teardown.
455        l.state().fetch_sub(1u64 << WAITING_SHIFT, Ordering::AcqRel);
456        std::fs::remove_file(&p).ok();
457    }
458
459    #[test]
460    fn concurrent_readers_all_succeed() {
461        let p = tmp("concurrent-r");
462        let l = Arc::new(SharedRWLock::create(&p).unwrap());
463        let n = 8;
464        let count = Arc::new(AtomicU32::new(0));
465        let mut handles = vec![];
466        for _ in 0..n {
467            let l = l.clone();
468            let count = count.clone();
469            handles.push(thread::spawn(move || {
470                let _g = l.read_lock();
471                count.fetch_add(1, O::AcqRel);
472                std::thread::sleep(std::time::Duration::from_millis(5));
473            }));
474        }
475        for h in handles { h.join().unwrap(); }
476        assert_eq!(count.load(O::Acquire), n);
477        std::fs::remove_file(&p).ok();
478    }
479
480    #[test]
481    fn cross_handle_visibility() {
482        let p = tmp("cross-handle");
483        let w = SharedRWLock::create(&p).unwrap();
484        let r = SharedRWLock::open(&p).unwrap();
485        let _g = w.try_read_lock().unwrap();
486        // Reader handle sees the same state.
487        assert_eq!(r.reader_count(), 1);
488        assert_eq!(r.try_write_lock().err(), Some(RWLockError::WouldBlock));
489        std::fs::remove_file(&p).ok();
490    }
491
492    #[test]
493    fn writer_then_reader_serialized() {
494        let p = tmp("w-then-r");
495        let l = SharedRWLock::create(&p).unwrap();
496        {
497            let _w = l.try_write_lock().unwrap();
498        }
499        // After writer drops, reader can acquire.
500        let _r = l.try_read_lock().unwrap();
501        std::fs::remove_file(&p).ok();
502    }
503}