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;
pub const CELL_TEXT_KIND: u8 = 186;
pub const CELL_TEXT_CAPACITY: usize = orbit_core::compile::usize_from_env(
option_env!("ORBIT_CELL_TEXT_CAPACITY"),
1_024,
);
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; #[cfg(unix)]
const TEXT_VERSION: u16 = 5;
#[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())?,
})
}
}
#[derive(Clone)]
pub struct Text {
table: TextBackend,
id: TextId,
_lease: Arc<Lease>,
}
impl Text {
pub fn id(&self) -> TextId {
self.id
}
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 load_into<'a>(
&self,
buffer: &'a mut [u8],
) -> Result<&'a str> {
let len = self.slot()?.copy(self.id, buffer)?;
if len > buffer.len() {
return Err(Error::BufferTooSmall {
required: len,
available: buffer.len(),
});
}
std::str::from_utf8(&buffer[..len])
.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)
}
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(())
}
pub fn append(
&self,
text: &str,
) -> Result<usize> {
self.slot()?
.writer(self.table.writer_token())?
.write(text.as_bytes(), true)
}
pub fn release(self) -> Result<()> {
self.table.release(self.id)
}
pub fn version(&self) -> Result<u64> {
Ok(self.slot()?.committed.load(Ordering::Acquire))
}
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
})?;
}
}
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)
})?;
}
}
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 {
_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);