use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use bytes::Bytes;
use crate::id::NetId64;
use crate::{NodeId, OrbitTyped};
pub mod cursor;
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
mod readiness;
#[cfg(unix)]
pub mod shm;
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub use readiness::{ParkedRingEventFd, RingEventFd};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum RingTopology {
Shared = 0,
PerNode = 1,
SharedOrdered = 2
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RingSpec {
pub capacity: usize,
pub payload_capacity: usize,
pub topology: RingTopology
}
impl RingSpec {
pub const fn new(
capacity: usize,
payload_capacity: usize
) -> Self {
Self { capacity, payload_capacity, topology: RingTopology::Shared }
}
pub const fn per_node(
capacity: usize,
payload_capacity: usize
) -> Self {
Self { capacity, payload_capacity, topology: RingTopology::PerNode }
}
pub const fn shared_ordered(
capacity: usize,
payload_capacity: usize
) -> Self {
Self { capacity, payload_capacity, topology: RingTopology::SharedOrdered }
}
pub(crate) fn assert_valid(self) {
assert!(self.capacity > 0, "ring capacity must be > 0");
assert!(self.capacity.is_power_of_two(), "ring capacity must be a power of two");
assert!(
self.payload_capacity <= u32::MAX as usize,
"ring payload capacity must fit in u32"
);
}
}
struct RingLane {
write_pos: AtomicU64,
write_lock: Mutex<()>,
slots: Vec<RwLock<Option<Frame>>>
}
impl RingLane {
fn new(capacity: usize) -> Self {
let mut slots = Vec::with_capacity(capacity);
for _ in 0..capacity {
slots.push(RwLock::new(None));
}
Self { write_pos: AtomicU64::new(0), write_lock: Mutex::new(()), slots }
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Frame {
pub id: NetId64,
pub kind: u8,
pub ver: u64,
pub payload: Bytes
}
pub struct Ring {
kind: u8,
capacity: usize,
payload_capacity: usize,
topology: RingTopology,
version_counter: AtomicU64,
lanes: Vec<RingLane>
}
impl Ring {
pub fn new<T: OrbitTyped>() -> Self {
Self::new_for_fleet::<T>(1)
}
pub fn new_for_fleet<T: OrbitTyped>(fleet_capacity: u16) -> Self {
assert!(fleet_capacity > 0, "ring fleet capacity must be > 0");
let spec = T::RING_SPEC;
spec.assert_valid();
let capacity = spec.capacity;
let lane_count = match spec.topology {
RingTopology::Shared | RingTopology::SharedOrdered => 1,
RingTopology::PerNode => usize::from(fleet_capacity)
};
let mut lanes = Vec::with_capacity(lane_count);
for _ in 0..lane_count {
lanes.push(RingLane::new(capacity));
}
Self {
kind: T::KIND,
capacity,
payload_capacity: spec.payload_capacity,
topology: spec.topology,
version_counter: AtomicU64::new(0),
lanes
}
}
pub fn kind(&self) -> u8 {
self.kind
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn payload_capacity(&self) -> usize {
self.payload_capacity
}
pub fn spec(&self) -> RingSpec {
RingSpec {
capacity: self.capacity,
payload_capacity: self.payload_capacity,
topology: self.topology
}
}
pub fn head(&self) -> u64 {
self.lanes[0].write_pos.load(Ordering::Acquire)
}
pub fn lane_count(&self) -> usize {
self.lanes.len()
}
pub fn lane_head(
&self,
node_id: NodeId
) -> u64 {
self.lane(node_id).write_pos.load(Ordering::Acquire)
}
pub fn next_version(&self) -> u64 {
self.version_counter
.fetch_add(1, Ordering::AcqRel)
.checked_add(1)
.expect("ring semantic version exhausted")
}
pub fn current_version(&self) -> u64 {
self.version_counter.load(Ordering::Acquire)
}
pub fn write(
&self,
node_id: NodeId,
frame_kind: u8,
ver: u64,
payload: Bytes
) -> NetId64 {
assert!(
payload.len() <= self.payload_capacity,
"payload {} > ring payload capacity {}",
payload.len(),
self.payload_capacity
);
let lane = self.lane(node_id);
match self.topology {
RingTopology::Shared => {
let counter = lane.write_pos.fetch_add(1, Ordering::AcqRel);
self.write_frame(lane, node_id, counter, frame_kind, ver, payload)
}
RingTopology::PerNode | RingTopology::SharedOrdered => {
let _write = lane.write_lock.lock().unwrap_or_else(|error| error.into_inner());
let counter = lane.write_pos.load(Ordering::Relaxed);
let id = self.write_frame(lane, node_id, counter, frame_kind, ver, payload);
lane.write_pos.store(counter.wrapping_add(1), Ordering::Release);
id
}
}
}
pub fn write_batch(
&self,
node_id: NodeId,
frame_kind: u8,
ver: u64,
payloads: Vec<Bytes>
) -> Vec<NetId64> {
assert!(
payloads.len() <= self.capacity,
"batch {} > ring capacity {}",
payloads.len(),
self.capacity
);
for payload in &payloads {
assert!(
payload.len() <= self.payload_capacity,
"payload {} > ring payload capacity {}",
payload.len(),
self.payload_capacity
);
}
if payloads.is_empty() {
return Vec::new();
}
let lane = self.lane(node_id);
match self.topology {
RingTopology::Shared => {
let start = lane.write_pos.fetch_add(payloads.len() as u64, Ordering::AcqRel);
payloads
.into_iter()
.enumerate()
.map(|(offset, payload)| {
self.write_frame(
lane,
node_id,
start.wrapping_add(offset as u64),
frame_kind,
ver,
payload
)
})
.collect()
}
RingTopology::PerNode | RingTopology::SharedOrdered => {
let _write = lane.write_lock.lock().unwrap_or_else(|error| error.into_inner());
let start = lane.write_pos.load(Ordering::Relaxed);
let ids = payloads
.into_iter()
.enumerate()
.map(|(offset, payload)| {
self.write_frame(
lane,
node_id,
start.wrapping_add(offset as u64),
frame_kind,
ver,
payload
)
})
.collect::<Vec<_>>();
lane.write_pos.store(start.wrapping_add(ids.len() as u64), Ordering::Release);
ids
}
}
}
pub fn read(
&self,
id: NetId64
) -> Option<Frame> {
if id.kind() != self.kind {
return None;
}
let lane = self.lane_for_frame(id)?;
let slot_idx = (id.counter() as usize) % self.capacity;
let guard = lane.slots[slot_idx].read().expect("ring slot poisoned");
match &*guard {
Some(f) if f.id == id => Some(f.clone()),
_ => None
}
}
pub fn read_head(&self) -> Option<Frame> {
let head = self.head();
if head == 0 {
return None;
}
let slot_idx = ((head - 1) as usize) % self.capacity;
self.lanes[0].slots[slot_idx].read().expect("ring slot poisoned").clone()
}
pub fn read_at(
&self,
counter: u64
) -> Option<Frame> {
let slot_idx = (counter as usize) % self.capacity;
self.lanes[0].slots[slot_idx].read().expect("ring slot poisoned").clone()
}
pub(crate) fn read_state_at(
&self,
counter: u64
) -> cursor::RingRead {
match self.read_at(counter) {
Some(frame) if frame.id.counter() == counter => cursor::RingRead::Ready(frame),
Some(frame) if frame.id.counter() > counter => cursor::RingRead::Unavailable,
Some(_) | None => cursor::RingRead::Pending
}
}
pub(crate) fn read_lane_at(
&self,
node_id: NodeId,
counter: u64
) -> Option<Frame> {
let lane = self.lane(node_id);
let slot_idx = (counter as usize) % self.capacity;
lane.slots[slot_idx].read().expect("ring slot poisoned").clone()
}
pub(crate) fn read_lane_state_at(
&self,
node_id: NodeId,
counter: u64
) -> cursor::RingRead {
match self.read_lane_at(node_id, counter) {
Some(frame) if frame.id.counter() == counter => cursor::RingRead::Ready(frame),
Some(frame) if frame.id.counter() > counter => cursor::RingRead::Unavailable,
Some(_) | None if self.topology != RingTopology::Shared => {
cursor::RingRead::Unavailable
}
Some(_) | None => cursor::RingRead::Pending
}
}
pub fn reset(&self) {
for lane in &self.lanes {
for slot in &lane.slots {
*slot.write().expect("ring slot poisoned") = None;
}
lane.write_pos.store(0, Ordering::Release);
}
self.version_counter.store(0, Ordering::Release);
}
fn lane(
&self,
node_id: NodeId
) -> &RingLane {
let index = match self.topology {
RingTopology::Shared | RingTopology::SharedOrdered => 0,
RingTopology::PerNode => usize::from(node_id.get())
};
self.lanes.get(index).unwrap_or_else(|| {
panic!("node {} is outside ring lane count {}", node_id.get(), self.lanes.len())
})
}
fn lane_for_frame(
&self,
id: NetId64
) -> Option<&RingLane> {
let index = match self.topology {
RingTopology::Shared | RingTopology::SharedOrdered => 0,
RingTopology::PerNode => usize::from(id.node())
};
self.lanes.get(index)
}
fn write_frame(
&self,
lane: &RingLane,
node_id: NodeId,
counter: u64,
frame_kind: u8,
ver: u64,
payload: Bytes
) -> NetId64 {
let id = NetId64::make(self.kind, node_id.get(), counter);
let slot_idx = (counter as usize) % self.capacity;
let frame = Frame { id, kind: frame_kind, ver, payload };
let mut guard = lane.slots[slot_idx].write().expect("ring slot poisoned");
*guard = Some(frame);
id
}
}
impl cursor::RingFrameSource for Ring {
fn kind(&self) -> u8 {
Ring::kind(self)
}
fn head(&self) -> u64 {
Ring::head(self)
}
fn capacity(&self) -> usize {
Ring::capacity(self)
}
fn read_at(
&self,
counter: u64
) -> Option<Frame> {
Ring::read_at(self, counter)
}
fn read_state_at(
&self,
counter: u64
) -> cursor::RingRead {
Ring::read_state_at(self, counter)
}
}
#[cfg(unix)]
impl cursor::RingFrameSource for shm::ShmRing {
fn kind(&self) -> u8 {
shm::ShmRing::kind(self)
}
fn head(&self) -> u64 {
shm::ShmRing::head(self)
}
fn capacity(&self) -> usize {
shm::ShmRing::capacity(self)
}
fn read_at(
&self,
counter: u64
) -> Option<Frame> {
shm::ShmRing::read_at(self, counter)
}
fn read_state_at(
&self,
counter: u64
) -> cursor::RingRead {
shm::ShmRing::read_state_at(self, counter)
}
}
impl std::fmt::Debug for Ring {
fn fmt(
&self,
f: &mut std::fmt::Formatter<'_>
) -> std::fmt::Result {
f.debug_struct("Ring")
.field("kind", &self.kind)
.field("capacity", &self.capacity)
.field("payload_capacity", &self.payload_capacity)
.field("topology", &self.topology)
.field("lane_count", &self.lanes.len())
.field("head", &self.head())
.finish()
}
}
pub(crate) struct RingRegistry {
fleet_capacity: u16,
rings: dashmap::DashMap<u8, Arc<Ring>>
}
impl RingRegistry {
pub fn new(fleet_capacity: u16) -> Self {
Self { fleet_capacity, rings: dashmap::DashMap::new() }
}
pub fn get_or_create<T: OrbitTyped>(&self) -> Arc<Ring> {
let ring = self
.rings
.entry(T::KIND)
.or_insert_with(|| Arc::new(Ring::new_for_fleet::<T>(self.fleet_capacity)))
.clone();
assert_eq!(
ring.spec(),
T::RING_SPEC,
"OrbitTyped KIND {} was reused with a different ring spec",
T::KIND
);
ring
}
pub fn lookup(
&self,
kind: u8
) -> Option<Arc<Ring>> {
self.rings.get(&kind).map(|e| e.clone())
}
}