use std::collections::HashMap;
use std::fmt;
use std::marker::PhantomData;
use std::str::FromStr;
use std::sync::atomic::{AtomicU8, AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, LazyLock, Mutex, MutexGuard, Weak};
use orbit_core::Fleet;
#[cfg(unix)]
use orbit_core::shm::{ShmRegion, ring_segment_name};
mod text;
pub use text::{CELL_TEXT_CAPACITY, CELL_TEXT_KIND, CELL_TEXT_MAX, Text, TextId};
pub const CELL_STATE_KIND: u8 = 233;
pub const CELL_CAPACITY: usize =
orbit_core::compile::usize_from_env(option_env!("ORBIT_CELL_CAPACITY"), 4_096);
pub(crate) const SLOT_EMPTY: u8 = 0;
pub(crate) const SLOT_OCCUPIED: u8 = 1;
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub(crate) fn waits_supported() -> bool {
orbit_core::sync::supported()
}
#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
pub(crate) fn waits_supported() -> bool {
false
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub(crate) fn wait_on(
word: &AtomicU32,
expected: u32
) -> Result<()> {
orbit_core::sync::wait_word(word, expected).map_err(Error::Io)
}
#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
pub(crate) fn wait_on(
_word: &AtomicU32,
_expected: u32
) -> Result<()> {
Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"orbit-cell needs a platform that can wait on a shared word"
)))
}
pub(crate) fn wake_on(word: &AtomicU32) {
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
let _ = orbit_core::sync::wake_word(word);
#[cfg(not(any(target_os = "linux", target_os = "freebsd", target_os = "macos")))]
let _ = word;
}
#[cfg(unix)]
const STATE_MAGIC: u32 = 0x43_45_4C_4C; #[cfg(unix)]
const STATE_VERSION: u16 = 2;
pub type Result<T, E = Error> = std::result::Result<T, E>;
#[non_exhaustive]
#[derive(Debug)]
pub enum Error {
Stale(CellId),
StaleText(TextId),
TooLong {
len: usize,
max: usize
},
Corrupt(String),
TypeMismatch {
id: CellId,
expected: &'static str,
found: &'static str
},
Overflow,
Full {
capacity: usize
},
Malformed(String),
Io(std::io::Error)
}
impl fmt::Display for Error {
fn fmt(
&self,
f: &mut fmt::Formatter<'_>
) -> fmt::Result {
match self {
Self::Stale(id) => write!(f, "cell {id} has been released"),
Self::StaleText(id) => write!(f, "text cell {id} has been released"),
Self::TooLong { len, max } => {
write!(f, "text does not fit the cell: len={len} max={max}")
}
Self::Corrupt(id) => write!(f, "text cell {id} holds bytes that are not UTF-8"),
Self::TypeMismatch { id, expected, found } => {
write!(f, "cell {id} holds {found}, not {expected}")
}
Self::Overflow => f.write_str("cell value is outside the integer range"),
Self::Full { capacity } => write!(f, "cell table is full: capacity={capacity}"),
Self::Malformed(text) => write!(f, "not a cell id: {text:?}"),
Self::Io(error) => write!(f, "Orbit cell io error: {error}")
}
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(error) => Some(error),
_ => None
}
}
}
impl From<std::io::Error> for Error {
fn from(value: std::io::Error) -> Self {
Self::Io(value)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct CellId {
index: u32,
generation: u32
}
impl CellId {
pub const fn index(self) -> u32 {
self.index
}
pub const fn generation(self) -> u32 {
self.generation
}
pub const fn to_bits(self) -> u64 {
((self.index as u64) << 32) | self.generation as u64
}
pub const fn from_bits(bits: u64) -> Self {
Self { index: (bits >> 32) as u32, generation: bits as u32 }
}
}
impl fmt::Display for CellId {
fn fmt(
&self,
f: &mut fmt::Formatter<'_>
) -> fmt::Result {
write!(f, "cell:{}:{}", self.index, self.generation)
}
}
impl FromStr for CellId {
type Err = Error;
fn from_str(text: &str) -> Result<Self> {
let malformed = || Error::Malformed(text.to_owned());
let rest = text.strip_prefix("cell:").ok_or_else(malformed)?;
let (index, generation) = rest.split_once(':').ok_or_else(malformed)?;
Ok(Self {
index: index.parse().map_err(|_| malformed())?,
generation: generation.parse().map_err(|_| malformed())?
})
}
}
pub trait CellType: Copy + Send + Sync + 'static {
const TAG: u8;
const NAME: &'static str;
fn to_bits(self) -> u64;
fn from_bits(bits: u64) -> Self;
}
impl CellType for i64 {
const TAG: u8 = 1;
const NAME: &'static str = "int";
fn to_bits(self) -> u64 {
self as u64
}
fn from_bits(bits: u64) -> Self {
bits as i64
}
}
impl CellType for u64 {
const TAG: u8 = 2;
const NAME: &'static str = "unsigned";
fn to_bits(self) -> u64 {
self
}
fn from_bits(bits: u64) -> Self {
bits
}
}
impl CellType for f64 {
const TAG: u8 = 3;
const NAME: &'static str = "float";
fn to_bits(self) -> u64 {
self.to_bits()
}
fn from_bits(bits: u64) -> Self {
f64::from_bits(bits)
}
}
impl CellType for bool {
const TAG: u8 = 4;
const NAME: &'static str = "bool";
fn to_bits(self) -> u64 {
u64::from(self)
}
fn from_bits(bits: u64) -> Self {
bits != 0
}
}
fn type_name(tag: u8) -> &'static str {
match tag {
1 => i64::NAME,
2 => u64::NAME,
3 => f64::NAME,
4 => bool::NAME,
_ => "unknown"
}
}
#[derive(Clone)]
pub struct Cells {
backend: CellBackend,
text: text::TextBackend
}
#[derive(Clone)]
enum CellBackend {
InMemory(Arc<MemoryCellTable>),
#[cfg(unix)]
Shm(Arc<ShmCellTable>)
}
impl Cells {
pub fn new(fleet: Arc<Fleet>) -> Result<Self> {
if !crate::waits_supported() {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"orbit-cell needs a platform that can wait on a shared word: Linux, FreeBSD, or macOS 14.4 or later"
)));
}
let backend = if fleet.is_shm() {
#[cfg(unix)]
{
CellBackend::Shm(Arc::new(ShmCellTable::open_or_create(&ring_segment_name(
fleet.name(),
CELL_STATE_KIND
))?))
}
#[cfg(not(unix))]
unreachable!("non-Unix fleets cannot use POSIX SHM")
} else {
CellBackend::InMemory(memory_table(&fleet))
};
let text = text::TextBackend::new(&fleet)?;
Ok(Self { backend, text })
}
pub fn allocate_text(
&self,
initial: &str
) -> Result<Text> {
self.text.allocate(initial)
}
pub fn open_text(
&self,
id: TextId
) -> Result<Text> {
self.text.open(id)
}
pub fn release_text(
&self,
id: TextId
) -> Result<()> {
self.text.release(id)
}
pub fn is_text_live(
&self,
id: TextId
) -> bool {
self.text.is_live(id)
}
pub fn allocate<T: CellType>(
&self,
initial: T
) -> Result<Orbital<T>> {
let id = self.with_structure(|slots, hint| {
let capacity = slots.len();
let start = hint.load(Ordering::Relaxed) as usize;
for offset in 0..capacity {
let index = (start + offset) & (capacity - 1);
let slot = &slots[index];
if slot.state.load(Ordering::Acquire) == SLOT_EMPTY {
let generation = slot.install(T::TAG, initial.to_bits());
hint.store(((index + 1) & (capacity - 1)) as u32, Ordering::Relaxed);
return Ok(CellId { index: index as u32, generation });
}
}
Err(Error::Full { capacity })
})?;
Ok(Orbital { cells: self.clone(), id, _type: PhantomData })
}
pub fn open<T: CellType>(
&self,
id: CellId
) -> Result<Orbital<T>> {
let slot = self.slot(id)?;
let found = slot.tag.load(Ordering::Relaxed);
if found != T::TAG {
return Err(Error::TypeMismatch { id, expected: T::NAME, found: type_name(found) });
}
Ok(Orbital { cells: self.clone(), id, _type: PhantomData })
}
pub fn release(
&self,
id: CellId
) -> Result<()> {
self.with_structure(|slots, _| {
let slot =
slots.get(id.index as usize).filter(|slot| slot.is(id)).ok_or(Error::Stale(id))?;
slot.state.store(SLOT_EMPTY, Ordering::Release);
slot.changed();
Ok(())
})
}
pub fn is_live(
&self,
id: CellId
) -> bool {
self.slot(id).is_ok()
}
pub fn reset_all(&self) -> Result<()> {
self.with_structure(|slots, hint| {
for slot in slots {
slot.state.store(SLOT_EMPTY, Ordering::Release);
slot.changed();
}
hint.store(0, Ordering::Relaxed);
Ok(())
})?;
self.text.reset_all()
}
#[cfg(unix)]
pub fn unlink(&self) -> Result<()> {
match &self.backend {
CellBackend::InMemory(_) => self.reset_all()?,
CellBackend::Shm(table) => table.region.unlink().map_err(Error::Io)?
}
self.text.unlink()
}
fn slots(&self) -> &[CellSlot] {
match &self.backend {
CellBackend::InMemory(table) => &table.slots,
#[cfg(unix)]
CellBackend::Shm(table) => table.slots()
}
}
fn slot(
&self,
id: CellId
) -> Result<&CellSlot> {
self.slots().get(id.index as usize).filter(|slot| slot.is(id)).ok_or(Error::Stale(id))
}
fn with_structure<T>(
&self,
operation: impl FnOnce(&[CellSlot], &AtomicU32) -> Result<T>
) -> Result<T> {
match &self.backend {
CellBackend::InMemory(table) => {
let _local = lock_unpoisoned(&table.structural_lock);
operation(&table.slots, &table.hint)
}
#[cfg(unix)]
CellBackend::Shm(table) => table.with_structure(operation)
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Claim {
Took { taken: i64, after: i64 },
Exhausted
}
#[derive(Clone)]
pub struct Orbital<T: CellType> {
cells: Cells,
id: CellId,
_type: PhantomData<T>
}
impl<T: CellType> Orbital<T> {
pub fn id(&self) -> CellId {
self.id
}
pub fn load(&self) -> Result<T> {
Ok(T::from_bits(self.slot()?.value.load(Ordering::Acquire)))
}
pub fn store(
&self,
value: T
) -> Result<()> {
let slot = self.slot()?;
slot.value.store(value.to_bits(), Ordering::Release);
slot.changed();
Ok(())
}
pub fn swap(
&self,
value: T
) -> Result<T> {
let slot = self.slot()?;
let previous = slot.value.swap(value.to_bits(), Ordering::AcqRel);
slot.changed();
Ok(T::from_bits(previous))
}
pub fn version(&self) -> Result<u32> {
Ok(self.slot()?.changes.load(Ordering::Acquire))
}
pub fn wait_changed(
&self,
since: u32
) -> Result<u32> {
loop {
let slot = self.slot()?;
let now = slot.changes.load(Ordering::SeqCst);
if now != since {
return Ok(now);
}
slot.waiters.fetch_add(1, Ordering::SeqCst);
let outcome = if slot.changes.load(Ordering::SeqCst) == since {
wait_on(&slot.changes, since)
} else {
Ok(())
};
slot.waiters.fetch_sub(1, Ordering::SeqCst);
outcome?;
}
}
pub fn compare_exchange(
&self,
current: T,
new: T
) -> Result<std::result::Result<T, T>> {
let slot = self.slot()?;
let exchanged = slot.value.compare_exchange(
current.to_bits(),
new.to_bits(),
Ordering::AcqRel,
Ordering::Acquire
);
if exchanged.is_ok() {
slot.changed();
}
Ok(exchanged.map(T::from_bits).map_err(T::from_bits))
}
pub fn wait_changed_until(
&self,
since: u32,
interrupt: &core::sync::atomic::AtomicBool
) -> Result<Option<u32>> {
loop {
if interrupt.load(Ordering::SeqCst) {
return Ok(None);
}
let slot = self.slot()?;
let now = slot.changes.load(Ordering::SeqCst);
if now != since {
return Ok(Some(now));
}
slot.waiters.fetch_add(1, Ordering::SeqCst);
let outcome = if slot.changes.load(Ordering::SeqCst) == since
&& !interrupt.load(Ordering::SeqCst)
{
wait_on(&slot.changes, since)
} else {
Ok(())
};
slot.waiters.fetch_sub(1, Ordering::SeqCst);
outcome?;
}
}
pub fn wake_waiters(&self) -> Result<()> {
wake_on(&self.slot()?.changes);
Ok(())
}
pub fn release(self) -> Result<()> {
self.cells.release(self.id)
}
fn slot(&self) -> Result<&CellSlot> {
self.cells.slot(self.id)
}
}
impl Orbital<i64> {
pub fn fetch_add(
&self,
by: i64
) -> Result<i64> {
let slot = self.slot()?;
slot.value
.try_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
(bits as i64).checked_add(by).map(|value| value as u64)
})
.map(|previous| {
slot.changed();
(previous as i64) + by
})
.map_err(|_| Error::Overflow)
}
pub fn fetch_sub(
&self,
by: i64
) -> Result<i64> {
self.fetch_add(by.checked_neg().ok_or(Error::Overflow)?)
}
pub fn add_until(
&self,
by: i64,
limit: i64
) -> Result<Claim> {
if by <= 0 {
return Err(Error::Overflow);
}
let slot = self.slot()?;
loop {
let current = slot.value.load(Ordering::Acquire) as i64;
if current >= limit {
return Ok(Claim::Exhausted);
}
let taken = by.min(limit - current);
let after = current + taken;
if slot
.value
.compare_exchange(current as u64, after as u64, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
slot.changed();
return Ok(Claim::Took { taken, after });
}
}
}
}
impl Orbital<u64> {
pub fn fetch_add(
&self,
by: u64
) -> Result<u64> {
let slot = self.slot()?;
slot.value
.try_update(Ordering::AcqRel, Ordering::Acquire, |bits| bits.checked_add(by))
.map(|previous| {
slot.changed();
previous + by
})
.map_err(|_| Error::Overflow)
}
pub fn fetch_sub(
&self,
by: u64
) -> Result<u64> {
let slot = self.slot()?;
slot.value
.try_update(Ordering::AcqRel, Ordering::Acquire, |bits| bits.checked_sub(by))
.map(|previous| {
slot.changed();
previous - by
})
.map_err(|_| Error::Overflow)
}
}
impl Orbital<f64> {
pub fn fetch_add(
&self,
by: f64
) -> Result<f64> {
let slot = self.slot()?;
let previous = slot
.value
.try_update(Ordering::AcqRel, Ordering::Acquire, |bits| {
Some((f64::from_bits(bits) + by).to_bits())
})
.unwrap_or_else(|bits| bits);
slot.changed();
Ok(f64::from_bits(previous) + by)
}
}
impl Orbital<bool> {
pub fn set(&self) -> Result<bool> {
self.swap(true)
}
pub fn clear(&self) -> Result<bool> {
self.swap(false)
}
}
#[repr(C, align(64))]
struct CellStateHeader {
#[cfg(unix)]
magic: u32,
#[cfg(not(unix))]
_magic: u32,
#[cfg(unix)]
version: u16,
#[cfg(not(unix))]
_version: u16,
header_size: u16,
capacity: u32,
slot_size: u32,
hint: AtomicU32,
_reserved: [u8; 44]
}
impl CellStateHeader {
fn new() -> Self {
Self {
#[cfg(unix)]
magic: STATE_MAGIC,
#[cfg(not(unix))]
_magic: 0,
#[cfg(unix)]
version: STATE_VERSION,
#[cfg(not(unix))]
_version: 0,
header_size: size_of::<Self>() as u16,
capacity: CELL_CAPACITY as u32,
slot_size: size_of::<CellSlot>() as u32,
hint: AtomicU32::new(0),
_reserved: [0; 44]
}
}
}
#[repr(C, align(64))]
struct CellSlot {
state: AtomicU8,
tag: AtomicU8,
_reserved: [u8; 2],
generation: AtomicU32,
value: AtomicU64,
changes: AtomicU32,
waiters: AtomicU32,
_padding: [u8; 40]
}
impl CellSlot {
fn empty() -> Self {
Self {
state: AtomicU8::new(SLOT_EMPTY),
tag: AtomicU8::new(0),
_reserved: [0; 2],
generation: AtomicU32::new(0),
value: AtomicU64::new(0),
changes: AtomicU32::new(0),
waiters: AtomicU32::new(0),
_padding: [0; 40]
}
}
fn changed(&self) {
self.changes.fetch_add(1, Ordering::SeqCst);
if self.waiters.load(Ordering::SeqCst) > 0 {
wake_on(&self.changes);
}
}
fn is(
&self,
id: CellId
) -> bool {
self.state.load(Ordering::Acquire) == SLOT_OCCUPIED
&& self.generation.load(Ordering::Relaxed) == id.generation
}
fn install(
&self,
tag: u8,
value: u64
) -> u32 {
let generation = self.generation.load(Ordering::Relaxed).wrapping_add(1).max(1);
self.tag.store(tag, Ordering::Relaxed);
self.value.store(value, Ordering::Relaxed);
self.changes.store(0, Ordering::Relaxed);
self.waiters.store(0, Ordering::Relaxed);
self.generation.store(generation, Ordering::Relaxed);
self.state.store(SLOT_OCCUPIED, Ordering::Release);
generation
}
}
pub(crate) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(|error| error.into_inner())
}
struct MemoryCellTable {
_fleet: Arc<Fleet>,
slots: Vec<CellSlot>,
hint: AtomicU32,
structural_lock: Mutex<()>
}
impl MemoryCellTable {
fn new(fleet: Arc<Fleet>) -> Self {
Self {
_fleet: fleet,
slots: (0..CELL_CAPACITY).map(|_| CellSlot::empty()).collect(),
hint: AtomicU32::new(0),
structural_lock: Mutex::new(())
}
}
}
type MemoryRegistry = HashMap<usize, Weak<MemoryCellTable>>;
static MEMORY_TABLES: LazyLock<Mutex<MemoryRegistry>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
fn memory_table(fleet: &Arc<Fleet>) -> Arc<MemoryCellTable> {
let fleet_identity = Arc::as_ptr(fleet) as usize;
let mut tables = lock_unpoisoned(&MEMORY_TABLES);
tables.retain(|_, table| table.strong_count() > 0);
if let Some(table) = tables.get(&fleet_identity).and_then(Weak::upgrade) {
return table;
}
let table = Arc::new(MemoryCellTable::new(Arc::clone(fleet)));
tables.insert(fleet_identity, Arc::downgrade(&table));
table
}
#[cfg(unix)]
struct ShmCellTable {
region: ShmRegion,
structural_lock: Mutex<()>
}
#[cfg(unix)]
impl ShmCellTable {
fn open_or_create(name: &str) -> Result<Self> {
use std::{io, ptr};
let (region, _initialization_lock) =
ShmRegion::open_or_create_locked(name, shm_segment_size())?;
if region.created() {
unsafe {
ptr::write(region.as_ptr().cast::<CellStateHeader>(), CellStateHeader::new());
let slots = region.as_ptr().add(size_of::<CellStateHeader>());
ptr::write_bytes(slots, 0, CELL_CAPACITY * size_of::<CellSlot>());
}
} else {
let header = unsafe { &*region.as_ptr().cast::<CellStateHeader>() };
if header.magic != STATE_MAGIC {
return Err(Error::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"SHM segment {name} has wrong magic 0x{:08X} (expected 0x{STATE_MAGIC:08X})",
header.magic
)
)));
}
if header.version != STATE_VERSION
|| usize::from(header.header_size) != size_of::<CellStateHeader>()
|| header.capacity as usize != CELL_CAPACITY
|| header.slot_size as usize != size_of::<CellSlot>()
{
return Err(Error::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!("SHM segment {name} has an incompatible cell-state layout")
)));
}
}
Ok(Self { region, structural_lock: Mutex::new(()) })
}
fn header(&self) -> &CellStateHeader {
unsafe { &*self.region.as_ptr().cast::<CellStateHeader>() }
}
fn slots(&self) -> &[CellSlot] {
unsafe {
std::slice::from_raw_parts(
self.region.as_ptr().add(size_of::<CellStateHeader>()).cast::<CellSlot>(),
CELL_CAPACITY
)
}
}
fn with_structure<T>(
&self,
operation: impl FnOnce(&[CellSlot], &AtomicU32) -> Result<T>
) -> Result<T> {
let _local = lock_unpoisoned(&self.structural_lock);
let _process = self.region.lock_exclusive()?;
operation(self.slots(), &self.header().hint)
}
}
#[cfg(unix)]
fn shm_segment_size() -> usize {
size_of::<CellStateHeader>() + CELL_CAPACITY * size_of::<CellSlot>()
}
const _: () = assert!(CELL_CAPACITY.is_power_of_two());
const _: () = assert!(CELL_CAPACITY <= u32::MAX as usize);
const _: () = assert!(size_of::<CellStateHeader>() == 64);
const _: () = assert!(size_of::<CellSlot>() == 64);
#[cfg(test)]
mod registry_tests {
use std::sync::Arc;
use orbit_core::Fleet;
use super::Cells;
#[test]
fn a_memory_table_keeps_its_fleet_alive() {
let fleet = Arc::new(Fleet::join("cell-registry-test", 1).unwrap());
let cells = Cells::new(Arc::clone(&fleet)).unwrap();
assert!(Arc::strong_count(&fleet) > 1);
drop(cells);
assert_eq!(Arc::strong_count(&fleet), 1);
}
}
#[cfg(test)]
mod wait_tests {
use std::sync::Arc;
use std::time::Duration;
use orbit_core::Fleet;
use super::{Cells, Error};
fn cells() -> Cells {
Cells::new(Arc::new(Fleet::join("cell-wait-test", 1).unwrap())).unwrap()
}
#[test]
fn a_write_wakes_a_waiter_once() {
let cells = cells();
let cell = cells.allocate(0_i64).unwrap();
let since = cell.version().unwrap();
let watcher = cell.clone();
let waiter = std::thread::spawn(move || watcher.wait_changed(since));
std::thread::sleep(Duration::from_millis(30));
for _ in 0..10 {
cell.fetch_add(1).unwrap();
}
let now = waiter.join().unwrap().unwrap();
assert!(now > since);
assert_eq!(cell.load().unwrap(), 10);
assert_eq!(cell.version().unwrap(), cell.version().unwrap());
}
#[test]
fn release_wakes_a_waiter_stale() {
let cells = cells();
let cell = cells.allocate(7_i64).unwrap();
let since = cell.version().unwrap();
let watcher = cell.clone();
let waiter = std::thread::spawn(move || watcher.wait_changed(since));
std::thread::sleep(Duration::from_millis(30));
cell.release().unwrap();
assert!(matches!(waiter.join().unwrap(), Err(Error::Stale(_))));
}
}