orbit_cell/

text.rs

//! Fixed-width text cells: the string arena beside the scalar table.
//!
//! A text cell holds up to [`CELL_TEXT_MAX`] bytes of UTF-8 in shared memory.
//! Sixty-four bits can be swapped in one instruction; a string cannot, so a
//! cell publishes an inactive bank with one full-width revision. A bounded
//! reader validates that revision after copying; contention is explicit.
//! The last committed bank survives a writer exit before publication.

use std::collections::HashMap;
use std::fmt;
use std::str::FromStr;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::{Arc, LazyLock, Mutex, Weak};

use orbit_core::Fleet;
#[cfg(unix)]
use orbit_core::shm::{ShmRegion, ring_segment_name};

use crate::lifetime::{self, Lease, Owners};
use crate::{
    Error, Result, SLOT_EMPTY, SLOT_OCCUPIED, SLOT_RETIRED, lock_unpoisoned,
};

mod slot;
use slot::TextSlot;
#[cfg(unix)]
mod owners;
#[cfg(test)]
pub(crate) mod tests;

/// Reserved Orbit SHM kind for the text cell table.
pub const CELL_TEXT_KIND: u8 = 186;
/// Number of text cells one fleet generation can hold at once.
///
/// Compile-time geometry: `ORBIT_CELL_TEXT_CAPACITY`. Must be a power of two.
pub const CELL_TEXT_CAPACITY: usize = orbit_core::compile::usize_from_env(
    option_env!("ORBIT_CELL_TEXT_CAPACITY"),
    1_024,
);
/// Bytes one text cell can hold.
///
/// Compile-time geometry: `ORBIT_CELL_TEXT_MAX`. It sizes every slot, so it is
/// part of the same wire contract as the capacity.
pub const CELL_TEXT_MAX: usize = orbit_core::compile::usize_from_env(
    option_env!("ORBIT_CELL_TEXT_MAX"),
    240,
);

#[cfg(unix)]
const TEXT_MAGIC: u32 = 0x43_54_58_54; // "CTXT"
#[cfg(unix)]
const TEXT_VERSION: u16 = 5;

/// The address of one text cell; prints as `text:<slot>:<generation>`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct TextId {
    index: u32,
    generation: u32,
}

impl TextId {
    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 TextId {
    fn fmt(
        &self,
        f: &mut fmt::Formatter<'_>,
    ) -> fmt::Result {
        write!(f, "text:{}:{}", self.index, self.generation)
    }
}

impl FromStr for TextId {
    type Err = Error;

    fn from_str(text: &str) -> Result<Self> {
        let malformed = || Error::Malformed(text.to_owned());
        let rest = text.strip_prefix("text:").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())?,
        })
    }
}

/// One text cell. Cheap to clone; every clone is the same place.
#[derive(Clone)]
pub struct Text {
    table: TextBackend,
    id: TextId,
    _lease: Arc<Lease>,
}

impl Text {
    pub fn id(&self) -> TextId {
        self.id
    }

    /// A bounded copy of the current text. Returns [`Error::Contended`] if
    /// another publication reuses its bank while the copy is in progress.
    pub fn load(&self) -> Result<String> {
        let mut bytes = vec![0_u8; CELL_TEXT_MAX];
        let len = self.slot()?.copy(self.id, &mut bytes)?;
        bytes.truncate(len);
        String::from_utf8(bytes)
            .map_err(|_| Error::Corrupt(self.id.to_string()))
    }

    pub fn len(&self) -> Result<usize> {
        self.slot()?.copy(self.id, &mut [])
    }

    pub fn is_empty(&self) -> Result<bool> {
        Ok(self.len()? == 0)
    }

    /// Replace the text. A live writer causes [`Error::Contended`], not a spin.
    pub fn store(
        &self,
        text: &str,
    ) -> Result<()> {
        if text.len() > CELL_TEXT_MAX {
            return Err(Error::TooLong { len: text.len(), max: CELL_TEXT_MAX });
        }
        let slot = self.slot()?;
        slot.writer(self.table.writer_token())?
            .write(text.as_bytes(), false)?;
        Ok(())
    }

    /// Append by publishing a complete new bank. A refused write changes
    /// neither the committed value nor its revision.
    pub fn append(
        &self,
        text: &str,
    ) -> Result<usize> {
        self.slot()?
            .writer(self.table.writer_token())?
            .write(text.as_bytes(), true)
    }

    /// Retire this generation immediately. The physical slot remains pinned
    /// until every handle to the old generation has been dropped.
    pub fn release(self) -> Result<()> {
        self.table.release(self.id)
    }

    /// The full-width committed revision token. Compare for equality; do not
    /// truncate it or interpret its bank bit as an in-progress marker.
    pub fn version(&self) -> Result<u64> {
        Ok(self.slot()?.committed.load(Ordering::Acquire))
    }

    /// Park until the text has been written since `since`, then return the
    /// sequence now. A released cell wakes every waiter with
    /// [`Error::StaleText`]. Blocking; an async runtime wraps it.
    pub fn wait_changed(
        &self,
        since: u64,
    ) -> Result<u64> {
        loop {
            let slot = self.slot()?;
            let now = slot.committed.load(Ordering::SeqCst);
            if now != since {
                return Ok(now);
            }
            slot.wake.wait_if(|| {
                slot.state.load(Ordering::SeqCst) == SLOT_OCCUPIED
                    && slot.committed.load(Ordering::SeqCst) == since
            })?;
        }
    }

    /// [`Self::wait_changed`], but a second word can end it. `Ok(None)`
    /// means the interrupt arrived, not that time ran out — nothing here
    /// measures time. The canceller sets `interrupt`, then calls
    /// [`Self::wake_waiters`]. The wake advances a separate parking generation;
    /// it does not change the text or its committed revision.
    pub fn wait_changed_until(
        &self,
        since: u64,
        interrupt: &core::sync::atomic::AtomicBool,
    ) -> Result<Option<u64>> {
        loop {
            if interrupt.load(Ordering::SeqCst) {
                return Ok(None);
            }

            let slot = self.slot()?;
            let now = slot.committed.load(Ordering::SeqCst);
            if now != since {
                return Ok(Some(now));
            }

            slot.wake.wait_if(|| {
                slot.state.load(Ordering::SeqCst) == SLOT_OCCUPIED
                    && slot.committed.load(Ordering::SeqCst) == since
                    && !interrupt.load(Ordering::SeqCst)
            })?;
        }
    }

    /// Wake everyone parked on this text cell without appending to it.
    /// Only useful after the word the waiter also watches has been set.
    pub fn wake_waiters(&self) -> Result<()> {
        self.slot()?.wake.notify();
        Ok(())
    }

    fn slot(&self) -> Result<&TextSlot> {
        self.table.slot(self.id)
    }
}

#[derive(Clone)]
pub(crate) enum TextBackend {
    InMemory(Arc<MemoryTextTable>),
    #[cfg(unix)]
    Shm(Arc<ShmTextTable>),
}

impl TextBackend {
    pub(crate) fn new(fleet: &Arc<Fleet>) -> Result<Self> {
        Ok(if fleet.is_shm() {
            #[cfg(unix)]
            {
                Self::Shm(shm_table(
                    &ring_segment_name(fleet.name(), CELL_TEXT_KIND),
                    lifetime::owner_capacity(fleet),
                )?)
            }
            #[cfg(not(unix))]
            unreachable!("non-Unix fleets cannot use POSIX SHM")
        } else {
            Self::InMemory(memory_table(fleet))
        })
    }

    pub(crate) fn allocate(
        &self,
        text: &str,
    ) -> Result<Text> {
        if text.len() > CELL_TEXT_MAX {
            return Err(Error::TooLong { len: text.len(), max: CELL_TEXT_MAX });
        }
        let (id, lease) = self.with_structure(|slots, hint, owners| {
            let capacity = slots.len();
            let start = hint.load(Ordering::Relaxed) as usize;
            let allocate = || -> Result<Option<(TextId, Arc<Lease>)>> {
                for offset in 0..capacity {
                    let index = (start + offset) & (capacity - 1);
                    let slot = &slots[index];
                    if matches!(
                        slot.state.load(Ordering::Acquire),
                        SLOT_EMPTY | SLOT_RETIRED
                    ) && slot.generation.load(Ordering::Relaxed) != u32::MAX
                        && slot.writer.load(Ordering::Acquire) == 0
                        && !owners.pinned(index)
                    {
                        let lease = owners.lease(index)?;
                        let generation = slot.install(text.as_bytes());
                        hint.store(
                            ((index + 1) & (capacity - 1)) as u32,
                            Ordering::Relaxed,
                        );
                        return Ok(Some((
                            TextId { index: index as u32, generation },
                            lease,
                        )));
                    }
                }
                Ok(None)
            };
            if let Some(allocated) = allocate()? {
                return Ok(allocated);
            }
            if self.reclaim_dead(owners)?
                && let Some(allocated) = allocate()?
            {
                return Ok(allocated);
            }
            Err(Error::Full { capacity })
        })?;
        Ok(Text { table: self.clone(), id, _lease: lease })
    }

    pub(crate) fn open(
        &self,
        id: TextId,
    ) -> Result<Text> {
        let lease = self.with_structure(|slots, _, owners| {
            slots
                .get(id.index as usize)
                .filter(|slot| slot.is(id))
                .ok_or(Error::StaleText(id))?;
            owners.lease(id.index as usize)
        })?;
        Ok(Text { table: self.clone(), id, _lease: lease })
    }

    pub(crate) fn release(
        &self,
        id: TextId,
    ) -> Result<()> {
        self.with_structure(|slots, _, _| {
            let slot = slots
                .get(id.index as usize)
                .filter(|slot| slot.is(id))
                .ok_or(Error::StaleText(id))?;
            slot.state.store(SLOT_RETIRED, Ordering::SeqCst);
            slot.wake.changed();
            Ok(())
        })
    }

    pub(crate) fn is_live(
        &self,
        id: TextId,
    ) -> bool {
        self.slot(id).is_ok()
    }

    pub(crate) fn reset_all(&self) -> Result<()> {
        self.with_structure(|slots, hint, _| {
            for slot in slots {
                slot.state.store(SLOT_RETIRED, Ordering::SeqCst);
                slot.wake.changed();
            }
            hint.store(0, Ordering::Relaxed);
            Ok(())
        })
    }

    #[cfg(unix)]
    pub(crate) fn unlink(&self) -> Result<()> {
        match self {
            Self::InMemory(_) => self.reset_all(),
            Self::Shm(table) => table.shared.region.unlink().map_err(Error::Io),
        }
    }

    fn writer_token(&self) -> u64 {
        match self {
            Self::InMemory(_) => 1,
            #[cfg(unix)]
            Self::Shm(table) => table.token,
        }
    }

    fn reclaim_dead(
        &self,
        owners: &Owners,
    ) -> Result<bool> {
        owners.reclaim_dead_with(|owner| {
            #[cfg(unix)]
            if let Self::Shm(table) = self {
                owners::recover(&table.shared.region, owner)?;
            }
            #[cfg(not(unix))]
            let _ = owner;
            Ok(())
        })
    }

    pub(crate) fn recover_writers(&self) -> Result<usize> {
        let mut recovered = 0;
        self.with_structure(|_, _, pins| {
            pins.reclaim_dead_with(|owner| {
                #[cfg(unix)]
                if let Self::Shm(table) = self {
                    recovered += owners::recover(&table.shared.region, owner)?;
                }
                #[cfg(not(unix))]
                let _ = owner;
                Ok(())
            })?;
            Ok(())
        })?;
        Ok(recovered)
    }

    fn slots(&self) -> &[TextSlot] {
        match self {
            Self::InMemory(table) => &table.slots,
            #[cfg(unix)]
            Self::Shm(table) => table.slots(),
        }
    }

    fn slot(
        &self,
        id: TextId,
    ) -> Result<&TextSlot> {
        self.slots()
            .get(id.index as usize)
            .filter(|slot| slot.is(id))
            .ok_or(Error::StaleText(id))
    }

    fn with_structure<T>(
        &self,
        operation: impl FnOnce(&[TextSlot], &AtomicU32, &Owners) -> Result<T>,
    ) -> Result<T> {
        match self {
            Self::InMemory(table) => {
                let _local = lock_unpoisoned(&table.structural_lock);
                operation(&table.slots, &table.hint, &table.owners)
            }
            #[cfg(unix)]
            Self::Shm(table) => table.with_structure(operation),
        }
    }
}

#[repr(C, align(64))]
struct TextStateHeader {
    #[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,
    owners: u32,
    _reserved: [u8; 40],
}

impl TextStateHeader {
    fn new(owners: usize) -> Self {
        Self {
            #[cfg(unix)]
            magic: TEXT_MAGIC,
            #[cfg(not(unix))]
            _magic: 0,
            #[cfg(unix)]
            version: TEXT_VERSION,
            #[cfg(not(unix))]
            _version: 0,
            header_size: size_of::<Self>() as u16,
            capacity: CELL_TEXT_CAPACITY as u32,
            slot_size: size_of::<TextSlot>() as u32,
            hint: AtomicU32::new(0),
            owners: owners as u32,
            _reserved: [0; 40],
        }
    }
}

pub(crate) struct MemoryTextTable {
    /// Held so the fleet's address, which keys the registry, cannot be
    /// reused by another fleet while this table lives.
    _fleet: Arc<Fleet>,
    slots: Vec<TextSlot>,
    hint: AtomicU32,
    structural_lock: Mutex<()>,
    owners: Owners,
}

impl MemoryTextTable {
    fn new(fleet: Arc<Fleet>) -> Self {
        Self {
            _fleet: fleet,
            slots: (0..CELL_TEXT_CAPACITY).map(|_| TextSlot::empty()).collect(),
            hint: AtomicU32::new(0),
            structural_lock: Mutex::new(()),
            owners: Owners::memory(CELL_TEXT_CAPACITY),
        }
    }
}

type MemoryRegistry = HashMap<usize, Weak<MemoryTextTable>>;

static MEMORY_TABLES: LazyLock<Mutex<MemoryRegistry>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

fn memory_table(fleet: &Arc<Fleet>) -> Arc<MemoryTextTable> {
    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(MemoryTextTable::new(Arc::clone(fleet)));
    tables.insert(fleet_identity, Arc::downgrade(&table));
    table
}

#[cfg(unix)]
pub(crate) struct ShmTextTable {
    shared: Arc<lifetime::Shared>,
    token: u64,
    structural_lock: Mutex<()>,
}

#[cfg(unix)]
impl ShmTextTable {
    fn open_or_create(
        name: &str,
        owners: usize,
    ) -> Result<Self> {
        use std::{io, ptr};

        let (region, _initialization_lock) =
            ShmRegion::open_or_create_locked(name, shm_segment_size(owners)?)?;
        if region.created() {
            unsafe {
                ptr::write(
                    region.as_ptr().cast::<TextStateHeader>(),
                    TextStateHeader::new(owners),
                );
                let slots = region.as_ptr().add(size_of::<TextStateHeader>());
                ptr::write_bytes(
                    slots,
                    0,
                    shm_segment_size(owners)? - size_of::<TextStateHeader>(),
                );
            }
        } else {
            let header = unsafe { &*region.as_ptr().cast::<TextStateHeader>() };
            if header.magic != TEXT_MAGIC {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!(
                        "SHM segment {name} has wrong magic 0x{:08X} (expected 0x{TEXT_MAGIC:08X})",
                        header.magic
                    ),
                )));
            }
            if header.version != TEXT_VERSION
                || usize::from(header.header_size)
                    != size_of::<TextStateHeader>()
                || header.capacity as usize != CELL_TEXT_CAPACITY
                || header.slot_size as usize != size_of::<TextSlot>()
                || header.owners as usize != owners
            {
                return Err(Error::Io(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!(
                        "SHM segment {name} has an incompatible text-cell layout"
                    ),
                )));
            }
        }
        let shared = lifetime::Shared::claim_locked_with(
            region,
            name,
            owners,
            CELL_TEXT_CAPACITY,
            owners::pins_offset(owners)?,
            |owner, region| owners::claim(region, owner),
        )?;
        let token = owners::token(&shared.region, shared.owner());
        Ok(Self { shared, token, structural_lock: Mutex::new(()) })
    }

    fn header(&self) -> &TextStateHeader {
        unsafe { &*self.shared.region.as_ptr().cast::<TextStateHeader>() }
    }

    fn slots(&self) -> &[TextSlot] {
        owners::slots(&self.shared.region)
    }

    fn with_structure<T>(
        &self,
        operation: impl FnOnce(&[TextSlot], &AtomicU32, &Owners) -> Result<T>,
    ) -> Result<T> {
        let owners = Owners::Shared(self.shared.clone());
        owners.check_process()?;
        let _local = lock_unpoisoned(&self.structural_lock);
        let _process = self.shared.region.lock_exclusive()?;
        operation(self.slots(), &self.header().hint, &owners)
    }
}

#[cfg(unix)]
fn shm_segment_size(owners: usize) -> Result<usize> {
    lifetime::segment_size(
        owners::pins_offset(owners)?,
        CELL_TEXT_CAPACITY,
        owners,
    )
}

#[cfg(unix)]
fn shm_table(
    name: &str,
    owners: usize,
) -> Result<Arc<ShmTextTable>> {
    static TABLES: LazyLock<Mutex<HashMap<String, Weak<ShmTextTable>>>> =
        LazyLock::new(|| Mutex::new(HashMap::new()));
    let mut tables = lock_unpoisoned(&TABLES);
    tables.retain(|_, table| table.strong_count() != 0);
    if let Some(table) = tables.get(name).and_then(Weak::upgrade)
        && table.shared.is_current(owners)
    {
        return Ok(table);
    }
    let table = Arc::new(ShmTextTable::open_or_create(name, owners)?);
    tables.insert(name.to_owned(), Arc::downgrade(&table));
    Ok(table)
}

const _: () = assert!(CELL_TEXT_CAPACITY.is_power_of_two());
const _: () = assert!(CELL_TEXT_MAX > 0 && CELL_TEXT_MAX <= u16::MAX as usize);
const _: () = assert!(size_of::<TextStateHeader>() == 64);