use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use bytes::Bytes;
use dashmap::DashMap;
use crate::OrbitTyped;
use crate::error::{Error, Result};
use crate::id::NetId64;
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
use crate::ring::RingEventFd;
#[cfg(unix)]
use crate::ring::shm::{ShmRing, ShmRingRegistry};
use crate::ring::{Frame, Ring, RingRegistry, RingTopology};
mod cursor;
pub use cursor::{FleetLaneCursor, FleetLanePoll};
#[cfg(unix)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FleetObserver {
name: String,
uid: u32,
}
#[cfg(unix)]
impl FleetObserver {
pub fn attach_existing(name: impl Into<String>) -> std::io::Result<Self> {
let uid = unsafe { libc::geteuid() };
Self::attach_existing_for_uid(name, uid)
}
pub fn attach_existing_for_uid(name: impl Into<String>, uid: u32) -> std::io::Result<Self> {
let name = name.into();
if name.is_empty() {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"fleet name must not be empty",
));
}
if name.contains('/') || name.contains('\0') {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"fleet name must not contain '/' or a NUL byte",
));
}
Ok(Self { name, uid })
}
pub fn name(&self) -> &str {
&self.name
}
pub fn uid(&self) -> u32 {
self.uid
}
pub fn ring(&self, kind: u8) -> std::io::Result<crate::ring::shm::ShmRingView> {
crate::ring::shm::ShmRingView::attach_existing_for_uid(&self.name, kind, self.uid)
}
pub fn typed_ring<T: OrbitTyped>(&self) -> std::io::Result<crate::ring::shm::ShmRingView> {
let view = self.ring(T::KIND)?;
if view.metadata().spec != T::RING_SPEC {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
format!(
"OrbitTyped KIND {} declares {:?}; existing spec is {:?}",
T::KIND,
T::RING_SPEC,
view.metadata().spec
),
));
}
Ok(view)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(transparent)]
pub struct NodeId(pub u16);
impl NodeId {
pub const ZERO: Self = Self(0);
pub const fn new(value: u16) -> Self {
Self(value)
}
pub const fn get(self) -> u16 {
self.0
}
}
impl std::fmt::Display for NodeId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "node:{}", self.0)
}
}
#[derive(Clone)]
pub struct Fleet {
inner: Arc<FleetInner>,
}
struct FleetInner {
name: &'static str,
fleet_capacity: u16,
node_id: NodeId,
id_counters: DashMap<u8, Arc<AtomicU64>>,
backing: RingBacking,
#[cfg(unix)]
#[allow(dead_code)]
membership: Option<crate::shm::FleetMembership>,
}
enum RingBacking {
InMemory(RingRegistry),
#[cfg(unix)]
Shm(ShmRingRegistry),
}
impl Fleet {
pub fn join(name: &'static str, fleet_capacity: u16) -> Result<Self> {
Self::join_as(name, fleet_capacity, NodeId::ZERO)
}
pub fn join_as(name: &'static str, fleet_capacity: u16, node_id: NodeId) -> Result<Self> {
if fleet_capacity == 0 {
return Err(Error::EmptyFleet);
}
if node_id.get() >= fleet_capacity {
return Err(Error::NodeOutsideFleet {
node_id: node_id.get(),
fleet_capacity,
});
}
Ok(Self {
inner: Arc::new(FleetInner {
name,
fleet_capacity,
node_id,
id_counters: DashMap::new(),
backing: RingBacking::InMemory(RingRegistry::new(fleet_capacity)),
#[cfg(unix)]
membership: None,
}),
})
}
#[cfg(unix)]
pub fn join_shm(name: &'static str, fleet_capacity: u16) -> Result<Self> {
Self::join_shm_as(name, fleet_capacity, NodeId::ZERO)
}
#[cfg(unix)]
pub fn join_shm_as(name: &'static str, fleet_capacity: u16, node_id: NodeId) -> Result<Self> {
if fleet_capacity == 0 {
return Err(Error::EmptyFleet);
}
if node_id.get() >= fleet_capacity {
return Err(Error::NodeOutsideFleet {
node_id: node_id.get(),
fleet_capacity,
});
}
let membership = crate::shm::join_fleet_membership(name).map_err(Error::Io)?;
Ok(Self {
inner: Arc::new(FleetInner {
name,
fleet_capacity,
node_id,
id_counters: DashMap::new(),
backing: RingBacking::Shm(ShmRingRegistry::new(name, fleet_capacity)),
membership: Some(membership),
}),
})
}
pub fn name(&self) -> &'static str {
self.inner.name
}
pub fn fleet_capacity(&self) -> u16 {
self.inner.fleet_capacity
}
pub fn node_id(&self) -> NodeId {
self.inner.node_id
}
pub fn next_id<T: OrbitTyped>(&self) -> NetId64 {
let counter_arc = self
.inner
.id_counters
.entry(T::KIND)
.or_insert_with(|| Arc::new(AtomicU64::new(0)))
.clone();
let counter = counter_arc.fetch_add(1, Ordering::Relaxed);
NetId64::make(T::KIND, self.node_id().get(), counter)
}
pub fn is_shm(&self) -> bool {
#[cfg(unix)]
{
matches!(self.inner.backing, RingBacking::Shm(_))
}
#[cfg(not(unix))]
{
false
}
}
pub fn ring<T: OrbitTyped>(&self) -> Arc<Ring> {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>(),
#[cfg(unix)]
RingBacking::Shm(_) => {
panic!("Fleet::ring called on SHM-backed fleet — use Fleet::shm_ring instead");
}
}
}
#[cfg(unix)]
pub fn shm_ring<T: OrbitTyped>(&self) -> std::io::Result<Arc<ShmRing>> {
match &self.inner.backing {
RingBacking::Shm(r) => r.get_or_create_for::<T>(),
RingBacking::InMemory(_) => {
panic!("Fleet::shm_ring called on in-memory fleet — use Fleet::ring instead");
}
}
}
pub fn publish<T: OrbitTyped>(&self, frame_kind: u8, ver: u64, payload: Bytes) -> NetId64 {
match &self.inner.backing {
RingBacking::InMemory(r) => {
let ring = r.get_or_create::<T>();
ring.write(self.node_id(), frame_kind, ver, payload)
}
#[cfg(unix)]
RingBacking::Shm(r) => {
let ring = r
.get_or_create_for::<T>()
.expect("SHM ring open failed — fleet unusable");
ring.write(self.node_id(), frame_kind, ver, payload)
.expect("SHM ring write failed")
}
}
}
pub fn publish_batch<T: OrbitTyped>(
&self,
frame_kind: u8,
ver: u64,
payloads: Vec<Bytes>,
) -> Vec<NetId64> {
match &self.inner.backing {
RingBacking::InMemory(r) => {
let ring = r.get_or_create::<T>();
ring.write_batch(self.node_id(), frame_kind, ver, payloads)
}
#[cfg(unix)]
RingBacking::Shm(r) => {
let ring = r
.get_or_create_for::<T>()
.expect("SHM ring open failed — fleet unusable");
ring.write_batch(self.node_id(), frame_kind, ver, payloads)
.expect("SHM ring batch write failed")
}
}
}
pub fn read(&self, id: NetId64) -> Option<Frame> {
match &self.inner.backing {
RingBacking::InMemory(r) => r.lookup(id.kind())?.read(id),
#[cfg(unix)]
RingBacking::Shm(r) => r.lookup(id.kind())?.read(id),
}
}
pub fn read_head<T: OrbitTyped>(&self) -> Option<Frame> {
if T::RING_SPEC.topology == RingTopology::PerNode {
let head = self.lane_head::<T>(self.node_id());
return (head > 0)
.then(|| self.read_lane_at::<T>(self.node_id(), head - 1))
.flatten();
}
match &self.inner.backing {
RingBacking::InMemory(r) => {
let ring = r.get_or_create::<T>();
ring.read_head()
}
#[cfg(unix)]
RingBacking::Shm(r) => {
let ring = r.get_or_create_for::<T>().ok()?;
ring.read_head()
}
}
}
pub fn head<T: OrbitTyped>(&self) -> u64 {
if T::RING_SPEC.topology == RingTopology::PerNode {
return self.lane_head::<T>(self.node_id());
}
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().head(),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.map(|ring| ring.head())
.unwrap_or(0),
}
}
pub fn read_at<T: OrbitTyped>(&self, counter: u64) -> Option<Frame> {
if T::RING_SPEC.topology == RingTopology::PerNode {
return self.read_lane_at::<T>(self.node_id(), counter);
}
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().read_at(counter),
#[cfg(unix)]
RingBacking::Shm(r) => r.get_or_create_for::<T>().ok()?.read_at(counter),
}
}
pub(crate) fn read_state_at<T: OrbitTyped>(
&self,
counter: u64,
) -> crate::ring::cursor::RingRead {
if T::RING_SPEC.topology == RingTopology::PerNode {
return self.read_lane_state_at::<T>(self.node_id(), counter);
}
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().read_state_at(counter),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.map(|ring| ring.read_state_at(counter))
.unwrap_or(crate::ring::cursor::RingRead::Unavailable),
}
}
pub fn lane_head<T: OrbitTyped>(&self, node_id: NodeId) -> u64 {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().lane_head(node_id),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.map(|ring| ring.lane_head(node_id))
.unwrap_or(0),
}
}
pub fn read_lane_at<T: OrbitTyped>(&self, node_id: NodeId, counter: u64) -> Option<Frame> {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().read_lane_at(node_id, counter),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.ok()?
.read_lane_at(node_id, counter),
}
}
pub(crate) fn read_lane_state_at<T: OrbitTyped>(
&self,
node_id: NodeId,
counter: u64,
) -> crate::ring::cursor::RingRead {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().read_lane_state_at(node_id, counter),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.map(|ring| ring.read_lane_state_at(node_id, counter))
.unwrap_or(crate::ring::cursor::RingRead::Unavailable),
}
}
pub fn ring_capacity<T: OrbitTyped>(&self) -> usize {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().capacity(),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.map(|ring| ring.capacity())
.unwrap_or(T::RING_SPEC.capacity),
}
}
pub fn next_ring_version<T: OrbitTyped>(&self) -> u64 {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().next_version(),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.expect("SHM ring open failed — fleet unusable")
.next_version(),
}
}
pub fn current_ring_version<T: OrbitTyped>(&self) -> u64 {
match &self.inner.backing {
RingBacking::InMemory(r) => r.get_or_create::<T>().current_version(),
#[cfg(unix)]
RingBacking::Shm(r) => r
.get_or_create_for::<T>()
.expect("SHM ring open failed — fleet unusable")
.current_version(),
}
}
pub fn reset_ring<T: OrbitTyped>(&self) -> std::io::Result<()> {
match &self.inner.backing {
RingBacking::InMemory(r) => {
r.get_or_create::<T>().reset();
Ok(())
}
#[cfg(unix)]
RingBacking::Shm(r) => {
r.get_or_create_for::<T>()?.reset();
Ok(())
}
}
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub fn ring_event_fd<T: OrbitTyped>(&self) -> std::io::Result<RingEventFd> {
match &self.inner.backing {
RingBacking::Shm(rings) => RingEventFd::new(rings.get_or_create_for::<T>()?),
RingBacking::InMemory(_) => Err(std::io::Error::new(
std::io::ErrorKind::Unsupported,
"Orbit eventfd requires a shared-memory fleet",
)),
}
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub fn publish_notified<T: OrbitTyped>(
&self,
frame_kind: u8,
ver: u64,
payload: Bytes,
) -> std::io::Result<NetId64> {
match &self.inner.backing {
RingBacking::Shm(rings) => {
let ring = rings.get_or_create_for::<T>()?;
let id = ring.write(self.node_id(), frame_kind, ver, payload)?;
RingEventFd::notify(&ring)?;
Ok(id)
}
RingBacking::InMemory(rings) => {
Ok(rings
.get_or_create::<T>()
.write(self.node_id(), frame_kind, ver, payload))
}
}
}
#[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
pub fn publish_batch_notified<T: OrbitTyped>(
&self,
frame_kind: u8,
ver: u64,
payloads: Vec<Bytes>,
) -> std::io::Result<Vec<NetId64>> {
match &self.inner.backing {
RingBacking::Shm(rings) => {
let ring = rings.get_or_create_for::<T>()?;
let ids = ring.write_batch(self.node_id(), frame_kind, ver, payloads)?;
if !ids.is_empty() {
RingEventFd::notify(&ring)?;
}
Ok(ids)
}
RingBacking::InMemory(rings) => Ok(rings.get_or_create::<T>().write_batch(
self.node_id(),
frame_kind,
ver,
payloads,
)),
}
}
}
impl std::fmt::Debug for Fleet {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Fleet")
.field("name", &self.inner.name)
.field("fleet_capacity", &self.inner.fleet_capacity)
.field("node_id", &self.inner.node_id)
.field("id_counters", &self.inner.id_counters.len())
.finish()
}
}