use std::collections::HashMap;
use std::mem::size_of;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, LazyLock, Mutex, Weak};
use orbit_core::readiness::{Readiness, Signal};
#[cfg(unix)]
use orbit_core::shm::{ShmRegion, ring_segment_name};
use orbit_core::{Fleet, OrbitEpoch};
use crate::layout::{
Direction, Doorbell, FLAG_READER_GONE, FLAG_RESET, Geometry, Header, SIDE_CLAIMED, SIDE_DEAD,
SLOT_EMPTY, SLOT_LIVE, Slot
};
use crate::wake::{Doorstep, Driver, Interest, Registry};
use crate::{Error, Incarnation, Result, StreamSpec, lock_unpoisoned};
enum Backing {
Memory(AlignedBytes),
#[cfg(unix)]
Shm(ShmRegion)
}
struct AlignedBytes {
ptr: *mut u8,
layout: std::alloc::Layout
}
impl AlignedBytes {
fn zeroed(size: usize) -> Self {
let layout = std::alloc::Layout::from_size_align(size, 64).expect("segment layout");
let ptr = unsafe { std::alloc::alloc_zeroed(layout) };
assert!(!ptr.is_null(), "stream table allocation failed");
Self { ptr, layout }
}
}
impl Drop for AlignedBytes {
fn drop(&mut self) {
unsafe { std::alloc::dealloc(self.ptr, self.layout) }
}
}
unsafe impl Send for AlignedBytes {}
unsafe impl Sync for AlignedBytes {}
pub(crate) struct Table {
_fleet: Arc<Fleet>,
backing: Backing,
geometry: Geometry,
kind: u8,
node: u16,
incarnation: u64,
allocate: Mutex<usize>,
pub(crate) registry: Registry,
driver: Mutex<Option<Driving>>,
#[cfg(unix)]
bell: Option<crate::bell::Bell>,
readiness: std::sync::OnceLock<Signal>
}
impl Table {
fn base(&self) -> *mut u8 {
match &self.backing {
Backing::Memory(bytes) => bytes.ptr,
#[cfg(unix)]
Backing::Shm(region) => region.as_ptr()
}
}
pub(crate) fn is_shared(&self) -> bool {
match &self.backing {
Backing::Memory(_) => false,
#[cfg(unix)]
Backing::Shm(_) => true
}
}
pub(crate) fn geometry(&self) -> &Geometry {
&self.geometry
}
pub(crate) fn kind(&self) -> u8 {
self.kind
}
pub(crate) fn node(&self) -> u16 {
self.node
}
pub(crate) fn incarnation(&self) -> u64 {
self.incarnation
}
fn header(&self) -> &Header {
unsafe { &*self.base().cast::<Header>() }
}
pub(crate) fn epoch(&self) -> u64 {
self.header().epoch.load(Ordering::Acquire)
}
#[cfg(all(unix, feature = "tokio"))]
pub(crate) fn own_doorbell(&self) -> &Doorbell {
self.doorbell(usize::from(self.node))
}
fn doorbell(
&self,
node: usize
) -> &Doorbell {
debug_assert!(node < self.geometry.fleet_capacity);
unsafe {
&*self
.base()
.add(self.geometry.doorbells_offset + node * size_of::<Doorbell>())
.cast::<Doorbell>()
}
}
fn bitmap_word(
&self,
offset: usize,
node: usize,
word: usize
) -> &AtomicU64 {
debug_assert!(node < self.geometry.fleet_capacity && word < self.geometry.bitmap_words);
unsafe {
&*self
.base()
.add(offset + (node * self.geometry.bitmap_words + word) * size_of::<AtomicU64>())
.cast::<AtomicU64>()
}
}
fn pending_word(
&self,
node: usize,
word: usize
) -> &AtomicU64 {
self.bitmap_word(self.geometry.pending_offset, node, word)
}
fn offer_word(
&self,
node: usize,
word: usize
) -> &AtomicU64 {
self.bitmap_word(self.geometry.offers_offset, node, word)
}
pub(crate) fn slots(&self) -> &[Slot] {
unsafe {
std::slice::from_raw_parts(
self.base().add(self.geometry.slots_offset).cast::<Slot>(),
self.geometry.total_slots
)
}
}
pub(crate) fn buffer(
&self,
slot: usize,
direction: usize
) -> *mut u8 {
unsafe { self.base().add(self.geometry.buffer_offset(slot, direction)) }
}
pub(crate) fn allocate(&self) -> Result<(usize, u32)> {
let mut hint = lock_unpoisoned(&self.allocate);
let lane_capacity = self.geometry.lane_capacity;
let lane_start = usize::from(self.node) * lane_capacity;
let slots = self.slots();
for offset in 0..lane_capacity {
let index = lane_start + ((*hint + offset) & (lane_capacity - 1));
let slot = &slots[index];
if slot.state.load(Ordering::Acquire) == SLOT_EMPTY
&& let Some(generation) = slot.install(self.node, self.incarnation)
{
*hint = (index - lane_start + 1) & (lane_capacity - 1);
return Ok((index, generation));
}
}
Err(Error::Full { capacity: lane_capacity })
}
pub(crate) fn notify(
&self,
index: usize,
slot: &Slot,
direction: &Direction,
ring: Ring
) {
direction.changes.fetch_add(1, Ordering::SeqCst);
if direction.waiters.load(Ordering::SeqCst) > 0 {
crate::wake_on(&direction.changes);
}
let sides: &[usize] = match ring {
Ring::Nobody => return,
Ring::Side(side) => &[side],
Ring::Both => &[0, 1]
};
if !self.is_shared() {
self.registry.wake(index);
self.signal_readiness();
return;
}
let mut rung = usize::MAX;
for &side in sides {
let node = usize::from(slot.node[side].load(Ordering::Acquire));
if node < self.geometry.fleet_capacity && node != rung {
self.ring(node, Some(index));
rung = node;
}
}
}
fn ring(
&self,
node: usize,
pending: Option<usize>
) {
if let Some(index) = pending {
self.pending_word(node, index / 64).fetch_or(1 << (index % 64), Ordering::SeqCst);
}
let doorbell = self.doorbell(node);
doorbell.generation.fetch_add(1, Ordering::SeqCst);
let listening = doorbell.listening.load(Ordering::SeqCst);
#[cfg(unix)]
let listening = {
use crate::bell::ARMED;
if listening & ARMED != 0
&& doorbell.listening.fetch_and(!ARMED, Ordering::SeqCst) & ARMED != 0
&& let Some(bell) = &self.bell
{
bell.ring(node);
}
listening & !ARMED
};
if listening > 0 {
crate::wake_on(&doorbell.generation);
}
}
pub(crate) fn offer(
&self,
index: usize,
node: u16
) -> Result<()> {
let node = usize::from(node);
if node >= self.geometry.fleet_capacity {
return Err(Error::Malformed(format!("node {node} is outside the fleet")));
}
let node = if self.is_shared() { node } else { usize::from(self.node) };
self.offer_word(node, index / 64).fetch_or(1 << (index % 64), Ordering::SeqCst);
self.ring(node, None);
if !self.is_shared() {
self.registry.wake_offer();
self.signal_readiness();
}
Ok(())
}
pub(crate) fn take_offer(&self) -> Option<usize> {
let node = usize::from(self.node);
for word in 0..self.geometry.bitmap_words {
let place = self.offer_word(node, word);
let mut bits = place.load(Ordering::SeqCst);
while bits != 0 {
let bit = bits.trailing_zeros() as u64;
let mask = 1_u64 << bit;
if place.fetch_and(!mask, Ordering::SeqCst) & mask != 0 {
let index = word * 64 + bit as usize;
if index < self.geometry.total_slots {
return Some(index);
}
}
bits &= !mask;
}
}
None
}
fn has_offer(&self) -> bool {
let node = usize::from(self.node);
(0..self.geometry.bitmap_words)
.any(|word| self.offer_word(node, word).load(Ordering::SeqCst) != 0)
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub(crate) fn wait_offer_timeout(
&self,
timeout: std::time::Duration
) -> Result<bool> {
let doorbell = self.doorbell(usize::from(self.node));
if self.has_offer() {
return Ok(true);
}
doorbell.listening.fetch_add(1, Ordering::SeqCst);
let seen = doorbell.generation.load(Ordering::SeqCst);
let outcome = if self.has_offer() {
Ok(true)
} else {
crate::wait_on_timeout(&doorbell.generation, seen, timeout)
};
doorbell.listening.fetch_sub(1, Ordering::SeqCst);
outcome
}
pub(crate) fn wait_offer(&self) -> Result<()> {
let doorbell = self.doorbell(usize::from(self.node));
loop {
if self.has_offer() {
return Ok(());
}
doorbell.listening.fetch_add(1, Ordering::SeqCst);
let seen = doorbell.generation.load(Ordering::SeqCst);
let outcome =
if self.has_offer() { Ok(()) } else { crate::wait_on(&doorbell.generation, seen) };
doorbell.listening.fetch_sub(1, Ordering::SeqCst);
outcome?;
}
}
pub(crate) fn register(
self: &Arc<Self>,
slot: usize,
interest: Interest,
waker: &std::task::Waker
) -> Result<()> {
self.registry.register(slot, interest, waker);
self.ensure_driver()
}
pub(crate) fn register_offer(
self: &Arc<Self>,
waker: &std::task::Waker
) -> Result<()> {
self.registry.register_offer(waker);
self.ensure_driver()
}
pub(crate) fn take_readiness(self: &Arc<Self>) -> Result<Readiness> {
let (readiness, signal) = orbit_core::readiness::pair()?;
self.readiness.set(signal).map_err(|_| {
Error::Malformed(
"this process already took the stream table's readiness descriptor".to_owned()
)
})?;
self.ensure_driver()?;
Ok(readiness)
}
fn signal_readiness(&self) {
if let Some(signal) = self.readiness.get() {
let _ = signal.signal();
}
}
fn ensure_driver(self: &Arc<Self>) -> Result<()> {
if self.is_shared() {
let mut driver = lock_unpoisoned(&self.driver);
if driver.is_none() {
#[cfg(all(unix, feature = "tokio"))]
if let (Some(bell), Ok(runtime)) =
(&self.bell, tokio::runtime::Handle::try_current())
{
*driver = Some(Driving::Reactor(
crate::bell::ReactorDriver::start(self, bell, &runtime)
.map_err(Error::Io)?
));
return Ok(());
}
*driver = Some(Driving::Thread(Driver::start(
Arc::as_ptr(self),
format!("orbit-stream-{}-driver", self.node)
)?));
}
}
Ok(())
}
pub(crate) fn node_dead(
&self,
node: u16,
incarnation: u64
) {
self.end_sides(node, |held| held == incarnation);
let doorbell = self.doorbell(usize::from(node));
if doorbell.incarnation.load(Ordering::SeqCst) == incarnation {
doorbell.listening.store(0, Ordering::SeqCst);
doorbell.incarnation.store(0, Ordering::SeqCst);
}
}
pub(crate) fn supersede(&self) {
let mine = self.incarnation;
self.end_sides(self.node, |held| held != mine);
let doorbell = self.doorbell(usize::from(self.node));
if doorbell.incarnation.load(Ordering::SeqCst) != mine {
doorbell.listening.store(0, Ordering::SeqCst);
doorbell.incarnation.store(0, Ordering::SeqCst);
}
}
fn end_sides(
&self,
node: u16,
dead: impl Fn(u64) -> bool
) {
for (index, slot) in self.slots().iter().enumerate() {
if slot.state.load(Ordering::Acquire) != SLOT_LIVE {
continue;
}
let mut touched = false;
for side in 0..2 {
if slot.claimed[side].load(Ordering::SeqCst) != SIDE_CLAIMED
|| slot.node[side].load(Ordering::Acquire) != node
|| !dead(slot.incarnation[side].load(Ordering::Acquire))
{
continue;
}
if slot.claimed[side]
.compare_exchange(SIDE_CLAIMED, SIDE_DEAD, Ordering::SeqCst, Ordering::SeqCst)
.is_err()
{
continue;
}
touched = true;
let writes = &slot.directions[side];
let reads = &slot.directions[1 - side];
writes.flags.fetch_or(FLAG_RESET, Ordering::SeqCst);
reads.flags.fetch_or(FLAG_READER_GONE, Ordering::SeqCst);
self.notify(index, slot, writes, Ring::Both);
self.notify(index, slot, reads, Ring::Both);
}
if touched
&& slot.claimed.iter().all(|side| side.load(Ordering::SeqCst) != SIDE_CLAIMED)
{
let _ = slot.state.compare_exchange(
SLOT_LIVE,
SLOT_EMPTY,
Ordering::SeqCst,
Ordering::SeqCst
);
}
}
}
pub(crate) fn reset_all(&self) {
let mut hint = lock_unpoisoned(&self.allocate);
*hint = 0;
let header = self.header();
let epoch = header.epoch.load(Ordering::Acquire);
header.epoch.store(next_epoch(epoch), Ordering::Release);
for slot in self.slots() {
slot.state.store(SLOT_EMPTY, Ordering::Release);
slot.generation.store(0, Ordering::Relaxed);
for direction in &slot.directions {
direction.changes.fetch_add(1, Ordering::SeqCst);
crate::wake_on(&direction.changes);
}
}
for node in 0..self.geometry.fleet_capacity {
for word in 0..self.geometry.bitmap_words {
self.pending_word(node, word).store(0, Ordering::Relaxed);
self.offer_word(node, word).store(0, Ordering::Relaxed);
}
}
self.registry.clear();
}
#[cfg(unix)]
pub(crate) fn unlink(&self) -> Result<()> {
match &self.backing {
Backing::Memory(_) => {
self.reset_all();
Ok(())
}
Backing::Shm(region) => region.unlink().map_err(Error::Io)
}
}
}
fn next_epoch(previous: u64) -> u64 {
OrbitEpoch::now().as_unix_ms().max(previous + 1)
}
impl Doorstep for Table {
fn generation(&self) -> &AtomicU32 {
&self.doorbell(usize::from(self.node)).generation
}
fn listening(
&self,
delta: i32
) {
let doorbell = self.doorbell(usize::from(self.node));
if delta > 0 {
doorbell.incarnation.store(self.incarnation, Ordering::SeqCst);
doorbell.listening.fetch_add(1, Ordering::SeqCst);
} else {
doorbell.listening.fetch_sub(1, Ordering::SeqCst);
}
}
fn drain(&self) {
let node = usize::from(self.node);
for word in 0..self.geometry.bitmap_words {
let mut bits = self.pending_word(node, word).swap(0, Ordering::SeqCst);
while bits != 0 {
let bit = bits.trailing_zeros() as usize;
bits &= bits - 1;
let index = word * 64 + bit;
if index < self.geometry.total_slots {
self.registry.wake(index);
}
}
}
if self.has_offer() {
self.registry.wake_offer();
}
self.signal_readiness();
}
}
impl Drop for Table {
fn drop(&mut self) {
match lock_unpoisoned(&self.driver).take() {
Some(Driving::Thread(mut driver)) => driver.stop(self.generation()),
#[cfg(all(unix, feature = "tokio"))]
Some(Driving::Reactor(driver)) => {
if let Some(bell) = &self.bell {
driver.stop(self.own_doorbell(), bell);
}
}
None => {}
}
}
}
enum Driving {
Thread(Driver),
#[cfg(all(unix, feature = "tokio"))]
Reactor(crate::bell::ReactorDriver)
}
#[derive(Hash, PartialEq, Eq)]
enum Key {
Memory(usize, u8),
#[cfg(unix)]
Shm(String, u16)
}
static TABLES: LazyLock<Mutex<HashMap<Key, Weak<Table>>>> =
LazyLock::new(|| Mutex::new(HashMap::new()));
pub(crate) fn open(
fleet: &Arc<Fleet>,
incarnation: Incarnation,
spec: StreamSpec
) -> Result<Arc<Table>> {
if !crate::waits_supported() {
return Err(Error::Io(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"orbit-stream needs a platform that can wait on a shared word: Linux, FreeBSD, or macOS 14.4 or later"
)));
}
let key = if fleet.is_shm() {
#[cfg(unix)]
{
Key::Shm(ring_segment_name(fleet.name(), spec.kind), fleet.node_id().get())
}
#[cfg(not(unix))]
unreachable!("non-Unix fleets cannot use POSIX SHM")
} else {
Key::Memory(Arc::as_ptr(fleet) as usize, spec.kind)
};
let mut tables = lock_unpoisoned(&TABLES);
tables.retain(|_, table| table.strong_count() > 0);
if let Some(table) = tables.get(&key).and_then(Weak::upgrade) {
if table.incarnation != incarnation.get() {
return Err(Error::Malformed(format!(
"this process already opened the stream table as incarnation {}",
table.incarnation
)));
}
if table.geometry.lane_capacity != spec.lane_capacity
|| table.geometry.buffer_bytes != spec.buffer_bytes
{
return Err(Error::Malformed(format!(
"this process already opened kind {} with lane_capacity={} buffer_bytes={}",
spec.kind, table.geometry.lane_capacity, table.geometry.buffer_bytes
)));
}
return Ok(table);
}
let geometry = Geometry::new(fleet.fleet_capacity(), spec);
let backing = match &key {
Key::Memory(..) => {
let bytes = AlignedBytes::zeroed(geometry.segment_size);
unsafe {
std::ptr::write(
bytes.ptr.cast::<Header>(),
Header::new(fleet.fleet_capacity(), &geometry, next_epoch(0))
)
};
Backing::Memory(bytes)
}
#[cfg(unix)]
Key::Shm(name, _) => Backing::Shm(open_shm(name, fleet.fleet_capacity(), &geometry)?)
};
let table = Arc::new(Table {
_fleet: Arc::clone(fleet),
backing,
geometry,
kind: spec.kind,
node: fleet.node_id().get(),
incarnation: incarnation.get(),
allocate: Mutex::new(0),
registry: Registry::new(geometry.total_slots),
driver: Mutex::new(None),
#[cfg(unix)]
bell: match &key {
Key::Shm(name, _) => Some(crate::bell::Bell::new(name).map_err(Error::Io)?),
Key::Memory(..) => None
},
readiness: std::sync::OnceLock::new()
});
tables.insert(key, Arc::downgrade(&table));
Ok(table)
}
#[cfg(unix)]
fn open_shm(
name: &str,
fleet_capacity: u16,
geometry: &Geometry
) -> Result<ShmRegion> {
use std::io;
let (region, _initialization_lock) =
ShmRegion::open_or_create_locked(name, geometry.segment_size)?;
if region.created() {
unsafe {
std::ptr::write(
region.as_ptr().cast::<Header>(),
Header::new(fleet_capacity, geometry, next_epoch(0))
);
}
} else {
if region.len() < geometry.segment_size {
return Err(Error::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!("SHM segment {name} is smaller than the stream table")
)));
}
let header = unsafe { &*region.as_ptr().cast::<Header>() };
if header.magic != crate::layout::MAGIC {
return Err(Error::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"SHM segment {name} has wrong magic 0x{:08X} (expected 0x{:08X})",
header.magic,
crate::layout::MAGIC
)
)));
}
if !header.compatible(fleet_capacity, geometry) {
return Err(Error::Io(io::Error::new(
io::ErrorKind::InvalidData,
format!("SHM segment {name} has an incompatible stream-table layout")
)));
}
}
Ok(region)
}
pub fn segment_size(fleet_capacity: u16) -> usize {
segment_size_for(fleet_capacity, StreamSpec::DEFAULT)
}
pub fn segment_size_for(
fleet_capacity: u16,
spec: StreamSpec
) -> usize {
Geometry::new(fleet_capacity, spec).segment_size
}
#[derive(Clone, Copy)]
pub(crate) enum Ring {
Nobody,
Side(usize),
Both
}