Skip to main content

orbit_cell/
lib.rs

1//! Typed atomic cells in one fleet-shared SHM table, addressed by id.
2//!
3//! A cell is a place, not a name. [`Cells::allocate`] takes a slot, stamps
4//! its type and hands back an [`Orbital`] handle whose [`CellId`] is the
5//! only thing that ever needs to travel: any process in the fleet that
6//! [`Cells::open`]s the id reads and updates the same 64 bits, atomically,
7//! without a lock and without publishing a frame. `orbit-counter` is the
8//! keyed sibling; this is memory rather than a dictionary.
9//!
10//! Every id carries the slot's generation. [`Cells::release`] bumps it, so a
11//! handle kept past a release answers [`Error::Stale`] instead of touching
12//! whoever took the slot next. Cells have no lifetime of their own: what
13//! allocates one is expected to release it.
14
15use std::collections::HashMap;
16use std::fmt;
17use std::marker::PhantomData;
18use std::str::FromStr;
19use std::sync::atomic::{AtomicU8, AtomicU32, AtomicU64, Ordering};
20use std::sync::{Arc, LazyLock, Mutex, MutexGuard, Weak};
21
22use orbit_core::Fleet;
23#[cfg(unix)]
24use orbit_core::shm::{ShmRegion, ring_segment_name};
25
26mod text;
27
28pub use text::{
29    CELL_TEXT_CAPACITY, CELL_TEXT_KIND, CELL_TEXT_MAX, Text, TextId,
30};
31
32/// Reserved Orbit SHM kind for the cell table.
33pub const CELL_STATE_KIND: u8 = 233;
34/// Number of cells one fleet generation can hold at once.
35///
36/// Compile-time geometry: `ORBIT_CELL_CAPACITY` in the application's
37/// `.cargo/config.toml` overrides the default. Must be a power of two. Peers
38/// built with a different value are refused when they open the table.
39pub const CELL_CAPACITY: usize = orbit_core::compile::usize_from_env(
40    option_env!("ORBIT_CELL_CAPACITY"),
41    4_096,
42);
43
44pub(crate) const SLOT_EMPTY: u8 = 0;
45pub(crate) const SLOT_OCCUPIED: u8 = 1;
46
47/// Whether this build can park on a shared word. A table refuses to open
48/// where it cannot: a sleep loop wearing the shape of a wait is worse
49/// than a clear no, and nothing above here should have to ask again.
50#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
51pub(crate) fn waits_supported() -> bool {
52    orbit_core::sync::supported()
53}
54
55#[cfg(not(any(
56    target_os = "linux",
57    target_os = "freebsd",
58    target_os = "macos"
59)))]
60pub(crate) fn waits_supported() -> bool {
61    false
62}
63
64/// Park until `word` no longer holds `expected`, through the platform's
65/// shared address wait. There is no polling fallback: a table refuses to
66/// open where the platform cannot wait, so by here it can.
67#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
68pub(crate) fn wait_on(
69    word: &AtomicU32,
70    expected: u32,
71) -> Result<()> {
72    orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
73}
74
75#[cfg(not(any(
76    target_os = "linux",
77    target_os = "freebsd",
78    target_os = "macos"
79)))]
80pub(crate) fn wait_on(
81    _word: &AtomicU32,
82    _expected: u32,
83) -> Result<()> {
84    Err(Error::Io(std::io::Error::new(
85        std::io::ErrorKind::Unsupported,
86        "orbit-cell needs a platform that can wait on a shared word",
87    )))
88}
89
90/// Wake everyone parked on `word`; nothing to do where nobody can park.
91pub(crate) fn wake_on(word: &AtomicU32) {
92    #[cfg(any(
93        target_os = "linux",
94        target_os = "freebsd",
95        target_os = "macos"
96    ))]
97    let _ = orbit_core::sync::wake_word(word);
98    #[cfg(not(any(
99        target_os = "linux",
100        target_os = "freebsd",
101        target_os = "macos"
102    )))]
103    let _ = word;
104}
105
106#[cfg(unix)]
107const STATE_MAGIC: u32 = 0x43_45_4c_4c; // "CELL"
108#[cfg(unix)]
109const STATE_VERSION: u16 = 2;
110
111pub type Result<T, E = Error> = std::result::Result<T, E>;
112
113/// Additions are expected: a cause discovered later lands here rather
114/// than in a new major version, so a caller matches what it handles and
115/// leaves the rest to a catch-all.
116#[non_exhaustive]
117#[derive(Debug)]
118pub enum Error {
119    /// The id names a slot nothing occupies, or a generation that has been
120    /// released since the id was minted.
121    Stale(CellId),
122    /// The same, for a text cell.
123    StaleText(TextId),
124    /// The text would not fit the cell; nothing was written.
125    TooLong {
126        len: usize,
127        max: usize,
128    },
129    /// A text cell held bytes that are not UTF-8: something wrote past the
130    /// contract.
131    Corrupt(String),
132    /// The id names a live cell of another type.
133    TypeMismatch {
134        id: CellId,
135        expected: &'static str,
136        found: &'static str,
137    },
138    /// The update would leave the integer range.
139    Overflow,
140    /// Every slot is occupied.
141    Full {
142        capacity: usize,
143    },
144    /// The text is not a cell id.
145    Malformed(String),
146    Io(std::io::Error),
147}
148
149impl fmt::Display for Error {
150    fn fmt(
151        &self,
152        f: &mut fmt::Formatter<'_>,
153    ) -> fmt::Result {
154        match self {
155            Self::Stale(id) => write!(f, "cell {id} has been released"),
156            Self::StaleText(id) => {
157                write!(f, "text cell {id} has been released")
158            }
159            Self::TooLong { len, max } => {
160                write!(f, "text does not fit the cell: len={len} max={max}")
161            }
162            Self::Corrupt(id) => {
163                write!(f, "text cell {id} holds bytes that are not UTF-8")
164            }
165            Self::TypeMismatch { id, expected, found } => {
166                write!(f, "cell {id} holds {found}, not {expected}")
167            }
168            Self::Overflow => {
169                f.write_str("cell value is outside the integer range")
170            }
171            Self::Full { capacity } => {
172                write!(f, "cell table is full: capacity={capacity}")
173            }
174            Self::Malformed(text) => write!(f, "not a cell id: {text:?}"),
175            Self::Io(error) => write!(f, "Orbit cell io error: {error}"),
176        }
177    }
178}
179
180impl std::error::Error for Error {
181    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
182        match self {
183            Self::Io(error) => Some(error),
184            _ => None,
185        }
186    }
187}
188
189impl From<std::io::Error> for Error {
190    fn from(value: std::io::Error) -> Self {
191        Self::Io(value)
192    }
193}
194
195/// The address of one cell: its slot and the generation it was allocated in.
196///
197/// Prints as `cell:<slot>:<generation>` and parses back, which is the form
198/// meant to cross a serialization boundary.
199#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
200pub struct CellId {
201    index: u32,
202    generation: u32,
203}
204
205impl CellId {
206    pub const fn index(self) -> u32 {
207        self.index
208    }
209
210    pub const fn generation(self) -> u32 {
211        self.generation
212    }
213
214    /// The two halves in one integer, slot high, generation low.
215    pub const fn to_bits(self) -> u64 {
216        ((self.index as u64) << 32) | self.generation as u64
217    }
218
219    pub const fn from_bits(bits: u64) -> Self {
220        Self { index: (bits >> 32) as u32, generation: bits as u32 }
221    }
222}
223
224impl fmt::Display for CellId {
225    fn fmt(
226        &self,
227        f: &mut fmt::Formatter<'_>,
228    ) -> fmt::Result {
229        write!(f, "cell:{}:{}", self.index, self.generation)
230    }
231}
232
233impl FromStr for CellId {
234    type Err = Error;
235
236    fn from_str(text: &str) -> Result<Self> {
237        let malformed = || Error::Malformed(text.to_owned());
238        let rest = text.strip_prefix("cell:").ok_or_else(malformed)?;
239        let (index, generation) = rest.split_once(':').ok_or_else(malformed)?;
240        Ok(Self {
241            index: index.parse().map_err(|_| malformed())?,
242            generation: generation.parse().map_err(|_| malformed())?,
243        })
244    }
245}
246
247/// A value that fits one cell: 64 bits, and a tag that says which type they
248/// spell.
249pub trait CellType: Copy + Send + Sync + 'static {
250    const TAG: u8;
251    const NAME: &'static str;
252
253    fn to_bits(self) -> u64;
254    fn from_bits(bits: u64) -> Self;
255}
256
257impl CellType for i64 {
258    const TAG: u8 = 1;
259    const NAME: &'static str = "int";
260
261    fn to_bits(self) -> u64 {
262        self as u64
263    }
264
265    fn from_bits(bits: u64) -> Self {
266        bits as i64
267    }
268}
269
270impl CellType for u64 {
271    const TAG: u8 = 2;
272    const NAME: &'static str = "unsigned";
273
274    fn to_bits(self) -> u64 {
275        self
276    }
277
278    fn from_bits(bits: u64) -> Self {
279        bits
280    }
281}
282
283impl CellType for f64 {
284    const TAG: u8 = 3;
285    const NAME: &'static str = "float";
286
287    fn to_bits(self) -> u64 {
288        self.to_bits()
289    }
290
291    fn from_bits(bits: u64) -> Self {
292        f64::from_bits(bits)
293    }
294}
295
296impl CellType for bool {
297    const TAG: u8 = 4;
298    const NAME: &'static str = "bool";
299
300    fn to_bits(self) -> u64 {
301        u64::from(self)
302    }
303
304    fn from_bits(bits: u64) -> Self {
305        bits != 0
306    }
307}
308
309fn type_name(tag: u8) -> &'static str {
310    match tag {
311        1 => i64::NAME,
312        2 => u64::NAME,
313        3 => f64::NAME,
314        4 => bool::NAME,
315        _ => "unknown",
316    }
317}
318
319/// The fleet's cell tables, scalar and text. Cheap to clone.
320#[derive(Clone)]
321pub struct Cells {
322    backend: CellBackend,
323    text: text::TextBackend,
324}
325
326#[derive(Clone)]
327enum CellBackend {
328    InMemory(Arc<MemoryCellTable>),
329    #[cfg(unix)]
330    Shm(Arc<ShmCellTable>),
331}
332
333impl Cells {
334    pub fn new(fleet: Arc<Fleet>) -> Result<Self> {
335        if !crate::waits_supported() {
336            return Err(Error::Io(std::io::Error::new(
337                std::io::ErrorKind::Unsupported,
338                "orbit-cell needs a platform that can wait on a shared word: Linux, FreeBSD, or macOS 14.4 or later",
339            )));
340        }
341        let backend = if fleet.is_shm() {
342            #[cfg(unix)]
343            {
344                CellBackend::Shm(Arc::new(ShmCellTable::open_or_create(
345                    &ring_segment_name(fleet.name(), CELL_STATE_KIND),
346                )?))
347            }
348            #[cfg(not(unix))]
349            unreachable!("non-Unix fleets cannot use POSIX SHM")
350        } else {
351            CellBackend::InMemory(memory_table(&fleet))
352        };
353        let text = text::TextBackend::new(&fleet)?;
354        Ok(Self { backend, text })
355    }
356
357    /// Take a text cell holding `initial`; up to [`CELL_TEXT_MAX`] bytes.
358    pub fn allocate_text(
359        &self,
360        initial: &str,
361    ) -> Result<Text> {
362        self.text.allocate(initial)
363    }
364
365    /// A handle to a live text cell.
366    pub fn open_text(
367        &self,
368        id: TextId,
369    ) -> Result<Text> {
370        self.text.open(id)
371    }
372
373    pub fn release_text(
374        &self,
375        id: TextId,
376    ) -> Result<()> {
377        self.text.release(id)
378    }
379
380    pub fn is_text_live(
381        &self,
382        id: TextId,
383    ) -> bool {
384        self.text.is_live(id)
385    }
386
387    /// Take a free slot, stamp it with `T` and `initial`, and return the
388    /// handle. The id inside it is what other processes open.
389    pub fn allocate<T: CellType>(
390        &self,
391        initial: T,
392    ) -> Result<Orbital<T>> {
393        let id = self.with_structure(|slots, hint| {
394            let capacity = slots.len();
395            let start = hint.load(Ordering::Relaxed) as usize;
396            for offset in 0..capacity {
397                let index = (start + offset) & (capacity - 1);
398                let slot = &slots[index];
399                if slot.state.load(Ordering::Acquire) == SLOT_EMPTY {
400                    let generation = slot.install(T::TAG, initial.to_bits());
401                    hint.store(
402                        ((index + 1) & (capacity - 1)) as u32,
403                        Ordering::Relaxed,
404                    );
405                    return Ok(CellId { index: index as u32, generation });
406                }
407            }
408            Err(Error::Full { capacity })
409        })?;
410        Ok(Orbital { cells: self.clone(), id, _type: PhantomData })
411    }
412
413    /// A handle to a live cell of type `T`.
414    pub fn open<T: CellType>(
415        &self,
416        id: CellId,
417    ) -> Result<Orbital<T>> {
418        let slot = self.slot(id)?;
419        let found = slot.tag.load(Ordering::Relaxed);
420        if found != T::TAG {
421            return Err(Error::TypeMismatch {
422                id,
423                expected: T::NAME,
424                found: type_name(found),
425            });
426        }
427        Ok(Orbital { cells: self.clone(), id, _type: PhantomData })
428    }
429
430    /// Give the slot back. Every handle to this generation goes stale.
431    pub fn release(
432        &self,
433        id: CellId,
434    ) -> Result<()> {
435        self.with_structure(|slots, _| {
436            let slot = slots
437                .get(id.index as usize)
438                .filter(|slot| slot.is(id))
439                .ok_or(Error::Stale(id))?;
440            slot.state.store(SLOT_EMPTY, Ordering::Release);
441            // Whoever is parked on it finds the slot gone and answers Stale.
442            slot.changed();
443            Ok(())
444        })
445    }
446
447    /// Whether `id` names a live cell right now.
448    pub fn is_live(
449        &self,
450        id: CellId,
451    ) -> bool {
452        self.slot(id).is_ok()
453    }
454
455    /// Clear both tables during quiescent owner boot.
456    pub fn reset_all(&self) -> Result<()> {
457        self.with_structure(|slots, hint| {
458            for slot in slots {
459                slot.state.store(SLOT_EMPTY, Ordering::Release);
460                slot.changed();
461            }
462            hint.store(0, Ordering::Relaxed);
463            Ok(())
464        })?;
465        self.text.reset_all()
466    }
467
468    /// Remove the cell SHM object. Existing mappings remain valid until their
469    /// processes release them.
470    #[cfg(unix)]
471    pub fn unlink(&self) -> Result<()> {
472        match &self.backend {
473            CellBackend::InMemory(_) => self.reset_all()?,
474            CellBackend::Shm(table) => {
475                table.region.unlink().map_err(Error::Io)?
476            }
477        }
478        self.text.unlink()
479    }
480
481    fn slots(&self) -> &[CellSlot] {
482        match &self.backend {
483            CellBackend::InMemory(table) => &table.slots,
484            #[cfg(unix)]
485            CellBackend::Shm(table) => table.slots(),
486        }
487    }
488
489    fn slot(
490        &self,
491        id: CellId,
492    ) -> Result<&CellSlot> {
493        self.slots()
494            .get(id.index as usize)
495            .filter(|slot| slot.is(id))
496            .ok_or(Error::Stale(id))
497    }
498
499    fn with_structure<T>(
500        &self,
501        operation: impl FnOnce(&[CellSlot], &AtomicU32) -> Result<T>,
502    ) -> Result<T> {
503        match &self.backend {
504            CellBackend::InMemory(table) => {
505                let _local = lock_unpoisoned(&table.structural_lock);
506                operation(&table.slots, &table.hint)
507            }
508            #[cfg(unix)]
509            CellBackend::Shm(table) => table.with_structure(operation),
510        }
511    }
512}
513
514/// What [`Orbital::add_until`] did. A bounded add either takes some of what
515/// is left or finds nothing left; there is no third answer, and no value
516/// past the limit is ever written.
517#[derive(Clone, Copy, Debug, Eq, PartialEq)]
518pub enum Claim {
519    /// `taken` can be less than asked for: the last caller gets the
520    /// remainder rather than being refused, so no unit goes unclaimed.
521    Took { taken: i64, after: i64 },
522    /// The limit was already reached. Nothing was added, nothing is left.
523    Exhausted,
524}
525
526/// One cell, typed. Cheap to clone; every clone is the same place.
527#[derive(Clone)]
528pub struct Orbital<T: CellType> {
529    cells: Cells,
530    id: CellId,
531    _type: PhantomData<T>,
532}
533
534impl<T: CellType> Orbital<T> {
535    pub fn id(&self) -> CellId {
536        self.id
537    }
538
539    pub fn load(&self) -> Result<T> {
540        Ok(T::from_bits(self.slot()?.value.load(Ordering::Acquire)))
541    }
542
543    pub fn store(
544        &self,
545        value: T,
546    ) -> Result<()> {
547        let slot = self.slot()?;
548        slot.value.store(value.to_bits(), Ordering::Release);
549        slot.changed();
550        Ok(())
551    }
552
553    /// Replace the value and return what it was.
554    pub fn swap(
555        &self,
556        value: T,
557    ) -> Result<T> {
558        let slot = self.slot()?;
559        let previous = slot.value.swap(value.to_bits(), Ordering::AcqRel);
560        slot.changed();
561        Ok(T::from_bits(previous))
562    }
563
564    /// The write count, the thing [`Self::wait_changed`] waits on. Take it
565    /// before reading the value, and the wait misses nothing in between.
566    pub fn version(&self) -> Result<u32> {
567        Ok(self.slot()?.changes.load(Ordering::Acquire))
568    }
569
570    /// Park until the cell has been written since `since`, then return the
571    /// count now. Coalescing: ten writes while parked wake the caller once,
572    /// and [`Self::load`] gives the latest. A released cell wakes every waiter
573    /// with [`Error::Stale`]. Blocking; an async runtime wraps it.
574    pub fn wait_changed(
575        &self,
576        since: u32,
577    ) -> Result<u32> {
578        loop {
579            let slot = self.slot()?;
580            let now = slot.changes.load(Ordering::SeqCst);
581            if now != since {
582                return Ok(now);
583            }
584            slot.waiters.fetch_add(1, Ordering::SeqCst);
585            let outcome = if slot.changes.load(Ordering::SeqCst) == since {
586                wait_on(&slot.changes, since)
587            } else {
588                Ok(())
589            };
590            slot.waiters.fetch_sub(1, Ordering::SeqCst);
591            outcome?;
592        }
593    }
594
595    /// Store `new` only if the cell still holds `current`. `Ok(Ok(previous))`
596    /// on success, `Ok(Err(actual))` when it held something else.
597    pub fn compare_exchange(
598        &self,
599        current: T,
600        new: T,
601    ) -> Result<std::result::Result<T, T>> {
602        let slot = self.slot()?;
603        let exchanged = slot.value.compare_exchange(
604            current.to_bits(),
605            new.to_bits(),
606            Ordering::AcqRel,
607            Ordering::Acquire,
608        );
609        if exchanged.is_ok() {
610            slot.changed();
611        }
612        Ok(exchanged.map(T::from_bits).map_err(T::from_bits))
613    }
614
615    /// [`Self::wait_changed`], but a second word can end it. The waiter
616    /// still parks until something wakes it; `interrupt` is what tells it
617    /// which kind of wake arrived. `Ok(Some(count))` is a write,
618    /// `Ok(None)` is the interrupt.
619    ///
620    /// This is the shape a caller needs when the wait may outlive its
621    /// reason — a process that has to shut down while a script waits on a
622    /// cell nothing will write again. It is not a timeout: nothing here
623    /// measures time, and a wait nobody interrupts is
624    /// [`Self::wait_changed`] exactly.
625    ///
626    /// The canceller sets `interrupt` and then calls
627    /// [`Self::wake_waiters`], in that order. Set-then-wake is what closes
628    /// the window: a waiter that had already parked is released by the
629    /// wake, and one that had not yet parked sees the flag instead.
630    pub fn wait_changed_until(
631        &self,
632        since: u32,
633        interrupt: &core::sync::atomic::AtomicBool,
634    ) -> Result<Option<u32>> {
635        loop {
636            if interrupt.load(Ordering::SeqCst) {
637                return Ok(None);
638            }
639
640            let slot = self.slot()?;
641            let now = slot.changes.load(Ordering::SeqCst);
642            if now != since {
643                return Ok(Some(now));
644            }
645
646            slot.waiters.fetch_add(1, Ordering::SeqCst);
647            let outcome = if slot.changes.load(Ordering::SeqCst) == since
648                && !interrupt.load(Ordering::SeqCst)
649            {
650                wait_on(&slot.changes, since)
651            } else {
652                Ok(())
653            };
654            slot.waiters.fetch_sub(1, Ordering::SeqCst);
655            outcome?;
656        }
657    }
658
659    /// Wake everyone parked on this cell without writing to it. On its own
660    /// this changes nothing — a waiter re-checks the count, finds it where
661    /// it was and parks again — so it is only useful after the word the
662    /// waiter also watches has been set. See [`Self::wait_changed_until`].
663    pub fn wake_waiters(&self) -> Result<()> {
664        wake_on(&self.slot()?.changes);
665        Ok(())
666    }
667
668    /// Give the cell back; this and every other handle to it go stale.
669    pub fn release(self) -> Result<()> {
670        self.cells.release(self.id)
671    }
672
673    fn slot(&self) -> Result<&CellSlot> {
674        self.cells.slot(self.id)
675    }
676}
677
678impl Orbital<i64> {
679    /// Add `by` (negative to subtract) and return the value after.
680    pub fn fetch_add(
681        &self,
682        by: i64,
683    ) -> Result<i64> {
684        let slot = self.slot()?;
685        slot.value
686            .try_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
687                (bits as i64).checked_add(by).map(|value| value as u64)
688            })
689            .map(|previous| {
690                slot.changed();
691                (previous as i64) + by
692            })
693            .map_err(|_| Error::Overflow)
694    }
695
696    pub fn fetch_sub(
697        &self,
698        by: i64,
699    ) -> Result<i64> {
700        self.fetch_add(by.checked_neg().ok_or(Error::Overflow)?)
701    }
702
703    /// Add up to `by` while the value stays at or under `limit`, in one
704    /// compare-and-swap, and say what was taken.
705    ///
706    /// This is the operation a fleet needs to divide a fixed amount of work
707    /// between processes that cannot see each other. `fetch_add` followed by
708    /// a test cannot do it: the add lands before the caller learns it went
709    /// too far, so N callers overshoot by up to N-1 between them. Here the
710    /// bound is decided inside the swap, so the value **never** passes
711    /// `limit` however many callers race, and the one that finds less than
712    /// `by` left takes the remainder instead of being turned away.
713    ///
714    /// Taking more than one at a time is how the cost comes down: the
715    /// contended word is touched once per claim rather than once per unit.
716    /// `limit` is the ceiling for the total, never for a single call.
717    ///
718    /// Overflow needs no separate guard: a result that cannot pass `limit`
719    /// cannot pass `i64::MAX` either.
720    ///
721    /// **A claim is not a lease.** What this hands back is gone from the
722    /// total the moment the swap lands, and nothing gives it back: a caller
723    /// that takes a hundred, does forty and then dies leaves sixty that
724    /// nobody will ever do, while the count says all hundred were taken. So
725    /// `by` is not only a speed dial, it is the exposure — the most one
726    /// death can lose. It is the same trade a database sequence makes with
727    /// its cache size, and it leaves the same kind of gap. Work that must
728    /// not be lost wants a lease whose owner notices a death, which is
729    /// `orbit-pool`, not a counter.
730    pub fn add_until(
731        &self,
732        by: i64,
733        limit: i64,
734    ) -> Result<Claim> {
735        if by <= 0 {
736            return Err(Error::Overflow);
737        }
738        let slot = self.slot()?;
739        loop {
740            let current = slot.value.load(Ordering::Acquire) as i64;
741            if current >= limit {
742                return Ok(Claim::Exhausted);
743            }
744            let taken = by.min(limit - current);
745            let after = current + taken;
746            if slot
747                .value
748                .compare_exchange(
749                    current as u64,
750                    after as u64,
751                    Ordering::AcqRel,
752                    Ordering::Acquire,
753                )
754                .is_ok()
755            {
756                slot.changed();
757                return Ok(Claim::Took { taken, after });
758            }
759            // Lost the race; somebody else moved the value. Read it again
760            // rather than retry blind: what is left may now be less, or
761            // nothing.
762        }
763    }
764}
765
766impl Orbital<u64> {
767    pub fn fetch_add(
768        &self,
769        by: u64,
770    ) -> Result<u64> {
771        let slot = self.slot()?;
772        slot.value
773            .try_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
774                bits.checked_add(by)
775            })
776            .map(|previous| {
777                slot.changed();
778                previous + by
779            })
780            .map_err(|_| Error::Overflow)
781    }
782
783    pub fn fetch_sub(
784        &self,
785        by: u64,
786    ) -> Result<u64> {
787        let slot = self.slot()?;
788        slot.value
789            .try_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
790                bits.checked_sub(by)
791            })
792            .map(|previous| {
793                slot.changed();
794                previous - by
795            })
796            .map_err(|_| Error::Overflow)
797    }
798}
799
800impl Orbital<f64> {
801    /// Add `by` and return the value after. A compare-and-swap loop: floats
802    /// have no fetch-add of their own.
803    pub fn fetch_add(
804        &self,
805        by: f64,
806    ) -> Result<f64> {
807        let slot = self.slot()?;
808        let previous = slot
809            .value
810            .try_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
811                Some((f64::from_bits(bits) + by).to_bits())
812            })
813            .unwrap_or_else(|bits| bits);
814        slot.changed();
815        Ok(f64::from_bits(previous) + by)
816    }
817}
818
819impl Orbital<bool> {
820    /// Set the flag and return whether it was already set.
821    pub fn set(&self) -> Result<bool> {
822        self.swap(true)
823    }
824
825    /// Clear the flag and return whether it was set.
826    pub fn clear(&self) -> Result<bool> {
827        self.swap(false)
828    }
829}
830
831#[repr(C, align(64))]
832struct CellStateHeader {
833    #[cfg(unix)]
834    magic: u32,
835    #[cfg(not(unix))]
836    _magic: u32,
837    #[cfg(unix)]
838    version: u16,
839    #[cfg(not(unix))]
840    _version: u16,
841    header_size: u16,
842    capacity: u32,
843    slot_size: u32,
844    /// Where the next allocation starts looking; a hint, never a promise.
845    hint: AtomicU32,
846    _reserved: [u8; 44],
847}
848
849impl CellStateHeader {
850    fn new() -> Self {
851        Self {
852            #[cfg(unix)]
853            magic: STATE_MAGIC,
854            #[cfg(not(unix))]
855            _magic: 0,
856            #[cfg(unix)]
857            version: STATE_VERSION,
858            #[cfg(not(unix))]
859            _version: 0,
860            header_size: size_of::<Self>() as u16,
861            capacity: CELL_CAPACITY as u32,
862            slot_size: size_of::<CellSlot>() as u32,
863            hint: AtomicU32::new(0),
864            _reserved: [0; 44],
865        }
866    }
867}
868
869/// One cell, laid out for shared memory: a cache line of its own, so two hot
870/// cells never contend, and every field atomic, so the mapping is never
871/// borrowed `&mut` while another process reads it.
872#[repr(C, align(64))]
873struct CellSlot {
874    state: AtomicU8,
875    tag: AtomicU8,
876    _reserved: [u8; 2],
877    generation: AtomicU32,
878    value: AtomicU64,
879    /// Bumped by every write and by release: the word a waiter parks on.
880    /// Linux futex wants 32 bits, so the value itself cannot be that word.
881    changes: AtomicU32,
882    /// Waiters parked on `changes`; a writer wakes only when this is nonzero,
883    /// so the write path costs one extra load when nobody is listening.
884    waiters: AtomicU32,
885    _padding: [u8; 40],
886}
887
888impl CellSlot {
889    fn empty() -> Self {
890        Self {
891            state: AtomicU8::new(SLOT_EMPTY),
892            tag: AtomicU8::new(0),
893            _reserved: [0; 2],
894            generation: AtomicU32::new(0),
895            value: AtomicU64::new(0),
896            changes: AtomicU32::new(0),
897            waiters: AtomicU32::new(0),
898            _padding: [0; 40],
899        }
900    }
901
902    /// After a write: count it, and wake whoever is waiting for one. `SeqCst`
903    /// on both sides of the exchange with [`Orbital::wait_changed`], so a
904    /// waiter that announced itself just before this write is either seen
905    /// here or sees the new count itself; never neither.
906    fn changed(&self) {
907        self.changes.fetch_add(1, Ordering::SeqCst);
908        if self.waiters.load(Ordering::SeqCst) > 0 {
909            wake_on(&self.changes);
910        }
911    }
912
913    fn is(
914        &self,
915        id: CellId,
916    ) -> bool {
917        // `Acquire` on `state` publishes the generation written by `install`.
918        self.state.load(Ordering::Acquire) == SLOT_OCCUPIED
919            && self.generation.load(Ordering::Relaxed) == id.generation
920    }
921
922    /// Under the structural lock. Returns the generation the new cell lives in;
923    /// it is bumped on every install, so no released id ever matches again.
924    fn install(
925        &self,
926        tag: u8,
927        value: u64,
928    ) -> u32 {
929        let generation =
930            self.generation.load(Ordering::Relaxed).wrapping_add(1).max(1);
931        self.tag.store(tag, Ordering::Relaxed);
932        self.value.store(value, Ordering::Relaxed);
933        self.changes.store(0, Ordering::Relaxed);
934        self.waiters.store(0, Ordering::Relaxed);
935        self.generation.store(generation, Ordering::Relaxed);
936        self.state.store(SLOT_OCCUPIED, Ordering::Release);
937        generation
938    }
939}
940
941pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
942    mutex.lock().unwrap_or_else(|error| error.into_inner())
943}
944
945struct MemoryCellTable {
946    /// Held so the fleet's address, which keys the registry, cannot be
947    /// reused by another fleet while this table lives.
948    _fleet: Arc<Fleet>,
949    slots: Vec<CellSlot>,
950    hint: AtomicU32,
951    structural_lock: Mutex<()>,
952}
953
954impl MemoryCellTable {
955    fn new(fleet: Arc<Fleet>) -> Self {
956        Self {
957            _fleet: fleet,
958            slots: (0..CELL_CAPACITY).map(|_| CellSlot::empty()).collect(),
959            hint: AtomicU32::new(0),
960            structural_lock: Mutex::new(()),
961        }
962    }
963}
964
965type MemoryRegistry = HashMap<usize, Weak<MemoryCellTable>>;
966
967static MEMORY_TABLES: LazyLock<Mutex<MemoryRegistry>> =
968    LazyLock::new(|| Mutex::new(HashMap::new()));
969
970fn memory_table(fleet: &Arc<Fleet>) -> Arc<MemoryCellTable> {
971    let fleet_identity = Arc::as_ptr(fleet) as usize;
972    let mut tables = lock_unpoisoned(&MEMORY_TABLES);
973    tables.retain(|_, table| table.strong_count() > 0);
974    if let Some(table) = tables.get(&fleet_identity).and_then(Weak::upgrade) {
975        return table;
976    }
977    let table = Arc::new(MemoryCellTable::new(Arc::clone(fleet)));
978    tables.insert(fleet_identity, Arc::downgrade(&table));
979    table
980}
981
982#[cfg(unix)]
983struct ShmCellTable {
984    region: ShmRegion,
985    structural_lock: Mutex<()>,
986}
987
988#[cfg(unix)]
989impl ShmCellTable {
990    fn open_or_create(name: &str) -> Result<Self> {
991        use std::{io, ptr};
992
993        let (region, _initialization_lock) =
994            ShmRegion::open_or_create_locked(name, shm_segment_size())?;
995        if region.created() {
996            unsafe {
997                ptr::write(
998                    region.as_ptr().cast::<CellStateHeader>(),
999                    CellStateHeader::new(),
1000                );
1001                let slots = region.as_ptr().add(size_of::<CellStateHeader>());
1002                ptr::write_bytes(
1003                    slots,
1004                    0,
1005                    CELL_CAPACITY * size_of::<CellSlot>(),
1006                );
1007            }
1008        } else {
1009            let header = unsafe { &*region.as_ptr().cast::<CellStateHeader>() };
1010            if header.magic != STATE_MAGIC {
1011                return Err(Error::Io(io::Error::new(
1012                    io::ErrorKind::InvalidData,
1013                    format!(
1014                        "SHM segment {name} has wrong magic 0x{:08X} (expected 0x{STATE_MAGIC:08X})",
1015                        header.magic
1016                    ),
1017                )));
1018            }
1019            if header.version != STATE_VERSION
1020                || usize::from(header.header_size)
1021                    != size_of::<CellStateHeader>()
1022                || header.capacity as usize != CELL_CAPACITY
1023                || header.slot_size as usize != size_of::<CellSlot>()
1024            {
1025                return Err(Error::Io(io::Error::new(
1026                    io::ErrorKind::InvalidData,
1027                    format!(
1028                        "SHM segment {name} has an incompatible cell-state layout"
1029                    ),
1030                )));
1031            }
1032        }
1033        Ok(Self { region, structural_lock: Mutex::new(()) })
1034    }
1035
1036    fn header(&self) -> &CellStateHeader {
1037        unsafe { &*self.region.as_ptr().cast::<CellStateHeader>() }
1038    }
1039
1040    fn slots(&self) -> &[CellSlot] {
1041        unsafe {
1042            std::slice::from_raw_parts(
1043                self.region
1044                    .as_ptr()
1045                    .add(size_of::<CellStateHeader>())
1046                    .cast::<CellSlot>(),
1047                CELL_CAPACITY,
1048            )
1049        }
1050    }
1051
1052    fn with_structure<T>(
1053        &self,
1054        operation: impl FnOnce(&[CellSlot], &AtomicU32) -> Result<T>,
1055    ) -> Result<T> {
1056        let _local = lock_unpoisoned(&self.structural_lock);
1057        let _process = self.region.lock_exclusive()?;
1058        operation(self.slots(), &self.header().hint)
1059    }
1060}
1061
1062#[cfg(unix)]
1063fn shm_segment_size() -> usize {
1064    size_of::<CellStateHeader>() + CELL_CAPACITY * size_of::<CellSlot>()
1065}
1066
1067const _: () = assert!(CELL_CAPACITY.is_power_of_two());
1068const _: () = assert!(CELL_CAPACITY <= u32::MAX as usize);
1069const _: () = assert!(size_of::<CellStateHeader>() == 64);
1070const _: () = assert!(size_of::<CellSlot>() == 64);
1071
1072#[cfg(test)]
1073mod registry_tests {
1074    use std::sync::Arc;
1075
1076    use orbit_core::Fleet;
1077
1078    use super::Cells;
1079
1080    /// The in-memory registry keys tables by the fleet's address. The
1081    /// table therefore keeps the fleet alive: otherwise a dropped fleet's
1082    /// address could be reused by a new one, which would then find and
1083    /// share the old table.
1084    #[test]
1085    fn a_memory_table_keeps_its_fleet_alive() {
1086        let fleet = Arc::new(Fleet::join("cell-registry-test", 1).unwrap());
1087        let cells = Cells::new(Arc::clone(&fleet)).unwrap();
1088        assert!(Arc::strong_count(&fleet) > 1);
1089        drop(cells);
1090        assert_eq!(Arc::strong_count(&fleet), 1);
1091    }
1092}
1093
1094#[cfg(test)]
1095mod wait_tests {
1096    use std::sync::Arc;
1097    use std::time::Duration;
1098
1099    use orbit_core::Fleet;
1100
1101    use super::{Cells, Error};
1102
1103    fn cells() -> Cells {
1104        Cells::new(Arc::new(Fleet::join("cell-wait-test", 1).unwrap())).unwrap()
1105    }
1106
1107    /// A waiter parks on the change count and a write from another thread
1108    /// wakes it; it then reads the value the write left. Ten writes while
1109    /// parked are one wake.
1110    #[test]
1111    fn a_write_wakes_a_waiter_once() {
1112        let cells = cells();
1113        let cell = cells.allocate(0_i64).unwrap();
1114        let since = cell.version().unwrap();
1115        let watcher = cell.clone();
1116        let waiter = std::thread::spawn(move || watcher.wait_changed(since));
1117
1118        std::thread::sleep(Duration::from_millis(30));
1119        for _ in 0..10 {
1120            cell.fetch_add(1).unwrap();
1121        }
1122        let now = waiter.join().unwrap().unwrap();
1123        assert!(now > since);
1124        assert_eq!(cell.load().unwrap(), 10);
1125        // Nothing new since: a wait would park; the count says so.
1126        assert_eq!(cell.version().unwrap(), cell.version().unwrap());
1127    }
1128
1129    /// Releasing the cell is a change too: a parked waiter comes back Stale
1130    /// instead of sleeping on a slot somebody else may take.
1131    #[test]
1132    fn release_wakes_a_waiter_stale() {
1133        let cells = cells();
1134        let cell = cells.allocate(7_i64).unwrap();
1135        let since = cell.version().unwrap();
1136        let watcher = cell.clone();
1137        let waiter = std::thread::spawn(move || watcher.wait_changed(since));
1138
1139        std::thread::sleep(Duration::from_millis(30));
1140        cell.release().unwrap();
1141        assert!(matches!(waiter.join().unwrap(), Err(Error::Stale(_))));
1142    }
1143}