Skip to main content

subetha_cxc/
shared_string_arena.rs

1//! `SharedStringArena` - append-only position-independent string
2//! pool backed by an MMF.
3//!
4//! # Why this exists
5//!
6//! Variable-length strings can't be stored inline in fixed-size
7//! slots (SharedHashMap, SharedVec, etc.) without padding waste or
8//! truncation. The natural cross-process solution is a shared byte
9//! arena: every process maps the same file at (potentially)
10//! different base addresses, and string references are
11//! position-independent `(offset, len)` pairs. Adding
12//! `mmap_base + offset` in any process resolves to the same bytes.
13//!
14//! # Layout
15//!
16//! ```text
17//! +---------------------------+
18//! | ArenaHeader (64B)         |
19//! |   magic, capacity_bytes   |
20//! |   used_bytes: AtomicU64   |
21//! +---------------------------+
22//! | bytes[0 .. capacity]      |
23//! +---------------------------+
24//! ```
25//!
26//! # Protocol
27//!
28//! `intern(s)`:
29//! 1. `offset = used_bytes.fetch_add(len)`.
30//! 2. If `offset + len > capacity`, rollback with
31//!    `fetch_sub(len)` and return `Full`. (Note: the
32//!    rollback is best-effort; if two threads race-overflow
33//!    simultaneously, both fetch_subs leave the counter
34//!    deterministic without "losing" bytes.)
35//! 3. Memcpy `s.bytes()` into `arena[offset..offset+len]`.
36//! 4. Return `StringRef { offset, len }`.
37//!
38//! `get(r)`:
39//! - Bounds-check `r.offset + r.len <= used_bytes` (sanity), then
40//!   return `&arena[r.offset..r.offset+r.len]` as a `&str`.
41//!
42//! # Concurrency
43//!
44//! Concurrent interners get distinct slices via fetch_add. Once
45//! the bytes are written, they are never moved (append-only). A
46//! reader holding a StringRef can always resolve it correctly,
47//! provided their `get` happens AFTER the interner returned the
48//! ref (which is the natural happens-before edge: the interner
49//! does the write, then makes the ref visible to the reader).
50//!
51//! # Deduplication
52//!
53//! Not provided here. For dedup, layer a `SharedHashMap<u64 hash,
54//! StringRef>` over the arena and consult it before each intern.
55//!
56//! # No deletion
57//!
58//! Append-only. The whole arena is reclaimed via `clear` (callers
59//! must ensure no concurrent readers); fine-grained deletion
60//! requires a free-list / compaction protocol that defeats the
61//! point of an arena.
62
63use std::fs::{File, OpenOptions};
64use std::mem::size_of;
65use std::path::Path;
66use std::sync::atomic::{AtomicU64, Ordering};
67
68use memmap2::{MmapMut, MmapOptions};
69
70pub const ARENA_MAGIC: u64 = 0x4150_5341_524E_4131;
71
72#[repr(C, align(64))]
73pub struct ArenaHeader {
74    pub magic: u64,
75    pub capacity_bytes: u64,
76    pub used_bytes: AtomicU64,
77    _pad: [u8; 40],
78}
79
80const _: () = {
81    assert!(size_of::<ArenaHeader>() == 64);
82};
83
84pub const fn arena_file_size(capacity_bytes: usize) -> usize {
85    size_of::<ArenaHeader>() + capacity_bytes
86}
87
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum ArenaError {
90    Full,
91    InvalidRef,
92    InvalidUtf8,
93    LayoutMismatch,
94    IoError(std::io::ErrorKind),
95}
96
97impl From<std::io::Error> for ArenaError {
98    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
99}
100
101/// Position-independent reference to a string in a SharedStringArena.
102/// Encoded as a `u64` (offset:u32, len:u32) for stable cross-process
103/// passing (the same u64 resolves to the same bytes in every process
104/// that maps the arena).
105#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
106pub struct StringRef {
107    pub offset: u32,
108    pub len: u32,
109}
110
111impl StringRef {
112    #[inline]
113    pub fn to_u64(self) -> u64 {
114        ((self.offset as u64) << 32) | (self.len as u64)
115    }
116    #[inline]
117    pub fn from_u64(v: u64) -> Self {
118        Self {
119            offset: (v >> 32) as u32,
120            len: v as u32,
121        }
122    }
123}
124
125pub struct SharedStringArena {
126    _file: File,
127    mmap: MmapMut,
128    capacity_bytes: usize,
129    header_sidecar: subetha_core::HandshakeHeader,
130    ring_sidecar: Box<subetha_core::ObservationRing>,
131}
132
133unsafe impl Send for SharedStringArena {}
134unsafe impl Sync for SharedStringArena {}
135
136impl subetha_sidecar::AdaptiveInstance for SharedStringArena {
137    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
138    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
139    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
140        Box::new(subetha_sidecar::NoMigrationPolicy)
141    }
142}
143
144impl SharedStringArena {
145    pub fn create(
146        path: impl AsRef<Path>, capacity_bytes: usize,
147    ) -> Result<Self, ArenaError> {
148        assert!(capacity_bytes >= 1);
149        assert!(capacity_bytes <= u32::MAX as usize,
150            "capacity_bytes must fit in u32 for StringRef offset");
151        let total = arena_file_size(capacity_bytes);
152        let file = OpenOptions::new()
153            .read(true).write(true).create(true).truncate(true)
154            .open(path.as_ref())?;
155        file.set_len(total as u64)?;
156        let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
157        let hdr = mmap.as_mut_ptr() as *mut ArenaHeader;
158        unsafe {
159            std::ptr::write(hdr, ArenaHeader {
160                magic: ARENA_MAGIC,
161                capacity_bytes: capacity_bytes as u64,
162                used_bytes: AtomicU64::new(0),
163                _pad: [0; 40],
164            });
165        }
166        Ok(Self {
167            _file: file, mmap, capacity_bytes,
168            header_sidecar: subetha_core::HandshakeHeader::new(),
169            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
170        })
171    }
172
173    pub fn open(
174        path: impl AsRef<Path>, expected_capacity_bytes: usize,
175    ) -> Result<Self, ArenaError> {
176        let total = arena_file_size(expected_capacity_bytes);
177        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
178        if file.metadata()?.len() < total as u64 {
179            return Err(ArenaError::LayoutMismatch);
180        }
181        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
182        let hdr = unsafe { &*(mmap.as_ptr() as *const ArenaHeader) };
183        if hdr.magic != ARENA_MAGIC || hdr.capacity_bytes != expected_capacity_bytes as u64 {
184            return Err(ArenaError::LayoutMismatch);
185        }
186        Ok(Self {
187            _file: file, mmap, capacity_bytes: expected_capacity_bytes,
188            header_sidecar: subetha_core::HandshakeHeader::new(),
189            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
190        })
191    }
192
193    #[inline]
194    pub fn capacity_bytes(&self) -> usize { self.capacity_bytes }
195
196    #[inline]
197    pub fn used_bytes(&self) -> usize {
198        self.header().used_bytes.load(Ordering::Acquire) as usize
199    }
200
201    #[inline]
202    pub fn remaining_bytes(&self) -> usize {
203        self.capacity_bytes.saturating_sub(self.used_bytes())
204    }
205
206    fn header(&self) -> &ArenaHeader {
207        unsafe { &*(self.mmap.as_ptr() as *const ArenaHeader) }
208    }
209
210    /// Append a string to the arena. Returns a StringRef that
211    /// resolves to the bytes in any mapping of the same file.
212    ///
213    /// Returns `Err(Full)` when the arena has no room. The empty
214    /// string `""` interns at the current offset with `len = 0`.
215    pub fn intern(&self, s: &str) -> Result<StringRef, ArenaError> {
216        self.intern_bytes(s.as_bytes())
217    }
218
219    /// Append arbitrary bytes (not necessarily UTF-8) to the arena.
220    /// Useful for storing binary blobs alongside strings. Retrieve
221    /// with `get_bytes`; `get` will reject non-UTF-8 with
222    /// `InvalidUtf8`.
223    pub fn intern_bytes(&self, bytes: &[u8]) -> Result<StringRef, ArenaError> {
224        let len = bytes.len() as u64;
225        if len > self.capacity_bytes as u64 {
226            self.ring_sidecar
227                .push_op(crate::sidecar_ops::string_arena::OP_INTERN, 1);
228            return Err(ArenaError::Full);
229        }
230        let offset = self.header().used_bytes.fetch_add(len, Ordering::AcqRel);
231        if offset.saturating_add(len) > self.capacity_bytes as u64 {
232            self.header().used_bytes.fetch_sub(len, Ordering::AcqRel);
233            self.ring_sidecar
234                .push_op(crate::sidecar_ops::string_arena::OP_INTERN, 1);
235            return Err(ArenaError::Full);
236        }
237        let dst = unsafe {
238            self.mmap.as_ptr()
239                .add(size_of::<ArenaHeader>())
240                .add(offset as usize)
241                as *mut u8
242        };
243        unsafe {
244            std::ptr::copy_nonoverlapping(bytes.as_ptr(), dst, bytes.len());
245        }
246        self.ring_sidecar
247            .push_op(crate::sidecar_ops::string_arena::OP_INTERN, 0);
248        Ok(StringRef { offset: offset as u32, len: len as u32 })
249    }
250
251    /// Resolve a StringRef to its `&[u8]`. Returns `Err(InvalidRef)`
252    /// when the ref doesn't fall inside the arena's used region.
253    pub fn get_bytes(&self, r: StringRef) -> Result<&[u8], ArenaError> {
254        let end = (r.offset as u64).saturating_add(r.len as u64);
255        if end > self.header().used_bytes.load(Ordering::Acquire) {
256            self.ring_sidecar
257                .push_op(crate::sidecar_ops::string_arena::OP_GET_BYTES, 1);
258            return Err(ArenaError::InvalidRef);
259        }
260        if end > self.capacity_bytes as u64 {
261            self.ring_sidecar
262                .push_op(crate::sidecar_ops::string_arena::OP_GET_BYTES, 1);
263            return Err(ArenaError::InvalidRef);
264        }
265        self.ring_sidecar
266            .push_op(crate::sidecar_ops::string_arena::OP_GET_BYTES, 0);
267        let base = unsafe {
268            self.mmap.as_ptr()
269                .add(size_of::<ArenaHeader>())
270                .add(r.offset as usize)
271        };
272        Ok(unsafe { std::slice::from_raw_parts(base, r.len as usize) })
273    }
274
275    /// Resolve a StringRef to a `&str`. Returns `Err(InvalidUtf8)`
276    /// when the bytes aren't valid UTF-8 (the arena doesn't enforce
277    /// validity per-segment; it's checked on read).
278    pub fn get(&self, r: StringRef) -> Result<&str, ArenaError> {
279        let bytes = self.get_bytes(r)?;
280        std::str::from_utf8(bytes).map_err(|_| ArenaError::InvalidUtf8)
281    }
282
283    /// Convenience: intern AND return a `&str` view into the
284    /// just-written bytes plus the ref.
285    pub fn intern_and_get(&self, s: &str) -> Result<(StringRef, &str), ArenaError> {
286        let r = self.intern(s)?;
287        let got = self.get(r)?;
288        Ok((r, got))
289    }
290
291    /// Reset the arena to empty. NOT concurrency-safe; callers must
292    /// ensure no other threads/processes are interning or reading.
293    /// Existing StringRefs become invalid (their bytes may be
294    /// overwritten by subsequent interns).
295    pub fn clear(&self) {
296        self.header().used_bytes.store(0, Ordering::Release);
297        self.ring_sidecar
298            .push_op(crate::sidecar_ops::string_arena::OP_CLEAR, 0);
299    }
300
301    pub fn flush(&self) -> Result<(), ArenaError> {
302        self.mmap.flush()?;
303        Ok(())
304    }
305
306    /// Non-blocking flush: schedules a writeback via the OS.
307    /// Note: Windows is only partially async (sync to page cache,
308    /// not to disk).
309    pub fn flush_async(&self) -> Result<(), ArenaError> {
310        self.mmap.flush_async()?;
311        Ok(())
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318    use std::sync::Arc;
319    use std::thread;
320
321    fn tmp(name: &str) -> std::path::PathBuf {
322        let mut p = std::env::temp_dir();
323        let pid = std::process::id();
324        p.push(format!("subetha-arena-{name}-{pid}.bin"));
325        p
326    }
327
328    #[test]
329    fn create_initial_state_is_empty() {
330        let p = tmp("init");
331        let a = SharedStringArena::create(&p, 1024).unwrap();
332        assert_eq!(a.capacity_bytes(), 1024);
333        assert_eq!(a.used_bytes(), 0);
334        assert_eq!(a.remaining_bytes(), 1024);
335        std::fs::remove_file(&p).ok();
336    }
337
338    #[test]
339    fn intern_and_get_round_trip() {
340        let p = tmp("rt");
341        let a = SharedStringArena::create(&p, 1024).unwrap();
342        let r1 = a.intern("hello").unwrap();
343        let r2 = a.intern("world").unwrap();
344        assert_eq!(a.get(r1).unwrap(), "hello");
345        assert_eq!(a.get(r2).unwrap(), "world");
346        assert_eq!(a.used_bytes(), 10);
347        std::fs::remove_file(&p).ok();
348    }
349
350    #[test]
351    fn empty_string_interns_with_zero_len() {
352        let p = tmp("empty");
353        let a = SharedStringArena::create(&p, 16).unwrap();
354        let r = a.intern("").unwrap();
355        assert_eq!(r.len, 0);
356        assert_eq!(a.get(r).unwrap(), "");
357        assert_eq!(a.used_bytes(), 0);
358        std::fs::remove_file(&p).ok();
359    }
360
361    #[test]
362    fn full_arena_returns_error() {
363        let p = tmp("full");
364        let a = SharedStringArena::create(&p, 10).unwrap();
365        a.intern("hello").unwrap();
366        a.intern("world").unwrap();
367        assert_eq!(a.intern("more").err(), Some(ArenaError::Full));
368        // Used bytes rolled back, not 14.
369        assert_eq!(a.used_bytes(), 10);
370        std::fs::remove_file(&p).ok();
371    }
372
373    #[test]
374    fn string_too_large_returns_full() {
375        let p = tmp("too-large");
376        let a = SharedStringArena::create(&p, 8).unwrap();
377        let big = "x".repeat(100);
378        assert_eq!(a.intern(&big).err(), Some(ArenaError::Full));
379        assert_eq!(a.used_bytes(), 0);
380        std::fs::remove_file(&p).ok();
381    }
382
383    #[test]
384    fn string_ref_packs_and_unpacks() {
385        let r = StringRef { offset: 0x1234_5678, len: 42 };
386        let packed = r.to_u64();
387        let unpacked = StringRef::from_u64(packed);
388        assert_eq!(unpacked, r);
389    }
390
391    #[test]
392    fn cross_handle_visibility() {
393        let p = tmp("cross-handle");
394        let writer = SharedStringArena::create(&p, 1024).unwrap();
395        let reader = SharedStringArena::open(&p, 1024).unwrap();
396        let r = writer.intern("cross-process").unwrap();
397        assert_eq!(reader.get(r).unwrap(), "cross-process");
398        std::fs::remove_file(&p).ok();
399    }
400
401    #[test]
402    fn invalid_ref_beyond_used_rejected() {
403        let p = tmp("invalid");
404        let a = SharedStringArena::create(&p, 1024).unwrap();
405        a.intern("hi").unwrap();  // used = 2
406        let bad = StringRef { offset: 100, len: 5 };
407        assert_eq!(a.get(bad).err(), Some(ArenaError::InvalidRef));
408        std::fs::remove_file(&p).ok();
409    }
410
411    #[test]
412    fn concurrent_interners_get_distinct_refs() {
413        let p = tmp("concurrent");
414        let a: Arc<SharedStringArena> = Arc::new(SharedStringArena::create(&p, 4096).unwrap());
415        let n_threads = 4;
416        let per_thread = 20;
417        let mut handles = vec![];
418        for t in 0..n_threads {
419            let a = a.clone();
420            handles.push(thread::spawn(move || {
421                let mut refs = vec![];
422                for i in 0..per_thread {
423                    let s = format!("thread-{t}-msg-{i:03}");
424                    let r = a.intern(&s).unwrap();
425                    refs.push((s, r));
426                }
427                refs
428            }));
429        }
430        let all: Vec<(String, StringRef)> = handles.into_iter()
431            .flat_map(|h| h.join().unwrap())
432            .collect();
433        // Every interned string must read back to its original value.
434        for (expected, r) in &all {
435            let got = a.get(*r).unwrap();
436            assert_eq!(got, expected,
437                "ref offset={} len={} should resolve to {expected}",
438                r.offset, r.len);
439        }
440        // No two refs overlap.
441        let mut refs: Vec<StringRef> = all.iter().map(|(_, r)| *r).collect();
442        refs.sort_by_key(|r| r.offset);
443        for w in refs.windows(2) {
444            let r1_end = w[0].offset + w[0].len;
445            assert!(r1_end <= w[1].offset,
446                "ref {:?} overlaps with ref {:?}", w[0], w[1]);
447        }
448        std::fs::remove_file(&p).ok();
449    }
450
451    #[test]
452    fn intern_and_get_helper_returns_both() {
453        let p = tmp("intern-and-get");
454        let a = SharedStringArena::create(&p, 1024).unwrap();
455        let (r, s) = a.intern_and_get("composite").unwrap();
456        assert_eq!(s, "composite");
457        assert_eq!(a.get(r).unwrap(), "composite");
458        std::fs::remove_file(&p).ok();
459    }
460
461    #[test]
462    fn utf8_validation_on_get() {
463        let p = tmp("utf8");
464        let a = SharedStringArena::create(&p, 128).unwrap();
465        // Intern valid UTF-8.
466        let r = a.intern("hello").unwrap();
467        assert!(a.get(r).is_ok());
468        // intern_bytes accepts arbitrary bytes; get() then rejects
469        // non-UTF-8 with InvalidUtf8 while get_bytes returns the raw
470        // bytes without validation.
471        let r2 = a.intern_bytes(&[0xFF, 0xFE, 0xFD]).unwrap();
472        assert_eq!(a.get(r2).err(), Some(ArenaError::InvalidUtf8));
473        assert_eq!(a.get_bytes(r2).unwrap(), &[0xFF, 0xFE, 0xFD]);
474        std::fs::remove_file(&p).ok();
475    }
476
477    #[test]
478    fn clear_resets_used_bytes() {
479        let p = tmp("clear");
480        let a = SharedStringArena::create(&p, 128).unwrap();
481        a.intern("first").unwrap();
482        a.intern("second").unwrap();
483        assert!(a.used_bytes() > 0);
484        a.clear();
485        assert_eq!(a.used_bytes(), 0);
486        // Fresh interns work.
487        let r = a.intern("after-clear").unwrap();
488        assert_eq!(a.get(r).unwrap(), "after-clear");
489        assert_eq!(r.offset, 0);
490        std::fs::remove_file(&p).ok();
491    }
492
493    #[test]
494    fn disk_persistence_survives_reopen() {
495        let p = tmp("disk");
496        let r_persist;
497        {
498            let a = SharedStringArena::create(&p, 1024).unwrap();
499            r_persist = a.intern("persisted-string").unwrap();
500            a.flush().unwrap();
501        }
502        let a2 = SharedStringArena::open(&p, 1024).unwrap();
503        assert_eq!(a2.get(r_persist).unwrap(), "persisted-string");
504        // And it can keep interning.
505        let r2 = a2.intern("more-after-reopen").unwrap();
506        assert_eq!(a2.get(r2).unwrap(), "more-after-reopen");
507        std::fs::remove_file(&p).ok();
508    }
509
510    #[test]
511    fn deduplication_via_hashmap_composition() {
512        // Demonstrate the dedup pattern: layer SharedHashMap<u64, u64>
513        // (hash -> StringRef.to_u64) over the arena.
514        use crate::SharedHashMap;
515        use crate::shared_hash_map::fnv1a_64;
516
517        let p_arena = tmp("dedup-arena");
518        let p_index = tmp("dedup-index");
519        let arena = SharedStringArena::create(&p_arena, 256).unwrap();
520        let index: SharedHashMap<u64, u64> = SharedHashMap::create(&p_index, 32).unwrap();
521
522        let s = "deduplicate-me";
523        let h = fnv1a_64(s.as_bytes());
524
525        // First intern: check index, miss, intern + insert into index.
526        let r = if let Some(packed) = index.get(&h) {
527            StringRef::from_u64(packed)
528        } else {
529            let r = arena.intern(s).unwrap();
530            index.insert(h, r.to_u64()).unwrap();
531            r
532        };
533        let used_after_first = arena.used_bytes();
534
535        // Second intern of same string: hit in index, no arena append.
536        let r2 = if let Some(packed) = index.get(&h) {
537            StringRef::from_u64(packed)
538        } else {
539            let r = arena.intern(s).unwrap();
540            index.insert(h, r.to_u64()).unwrap();
541            r
542        };
543        assert_eq!(r, r2, "dedup should return the same ref");
544        assert_eq!(arena.used_bytes(), used_after_first,
545            "second intern should not consume more bytes");
546
547        std::fs::remove_file(&p_arena).ok();
548        std::fs::remove_file(&p_index).ok();
549    }
550}