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