1use std::sync::Arc;
5use std::sync::atomic::{AtomicU64, Ordering};
6
7use bytes::Bytes;
8use dashmap::DashMap;
9
10use crate::OrbitTyped;
11use crate::error::{Error, Result};
12use crate::id::NetId64;
13#[cfg(any(target_os = "linux", target_os = "freebsd"))]
14use crate::ring::RingEventFd;
15#[cfg(unix)]
16use crate::ring::shm::{ShmRing, ShmRingRegistry};
17use crate::ring::{Frame, Ring, RingRegistry, RingTopology};
18
19mod cursor;
20pub use cursor::{FleetLaneCursor, FleetLanePoll};
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
29#[repr(transparent)]
30pub struct NodeId(pub u16);
31
32impl NodeId {
33 pub const ZERO: Self = Self(0);
34
35 pub const fn new(value: u16) -> Self {
36 Self(value)
37 }
38
39 pub const fn get(self) -> u16 {
40 self.0
41 }
42}
43
44impl std::fmt::Display for NodeId {
45 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46 write!(f, "node:{}", self.0)
47 }
48}
49
50#[derive(Clone)]
53pub struct Fleet {
54 inner: Arc<FleetInner>,
55}
56
57struct FleetInner {
58 name: &'static str,
59 fleet_capacity: u16,
60 node_id: NodeId,
61 id_counters: DashMap<u8, Arc<AtomicU64>>,
66 backing: RingBacking,
70}
71
72enum RingBacking {
76 InMemory(RingRegistry),
80 #[cfg(unix)]
84 Shm(ShmRingRegistry),
85}
86
87impl Fleet {
88 pub fn join(name: &'static str, fleet_capacity: u16) -> Result<Self> {
91 Self::join_as(name, fleet_capacity, NodeId::ZERO)
92 }
93
94 pub fn join_as(name: &'static str, fleet_capacity: u16, node_id: NodeId) -> Result<Self> {
96 if fleet_capacity == 0 {
97 return Err(Error::EmptyFleet);
98 }
99 if node_id.get() >= fleet_capacity {
100 return Err(Error::NodeOutsideFleet {
101 node_id: node_id.get(),
102 fleet_capacity,
103 });
104 }
105 Ok(Self {
106 inner: Arc::new(FleetInner {
107 name,
108 fleet_capacity,
109 node_id,
110 id_counters: DashMap::new(),
111 backing: RingBacking::InMemory(RingRegistry::new(fleet_capacity)),
112 }),
113 })
114 }
115
116 #[cfg(unix)]
128 pub fn join_shm(name: &'static str, fleet_capacity: u16) -> Result<Self> {
129 Self::join_shm_as(name, fleet_capacity, NodeId::ZERO)
130 }
131
132 #[cfg(unix)]
138 pub fn join_shm_as(name: &'static str, fleet_capacity: u16, node_id: NodeId) -> Result<Self> {
139 if fleet_capacity == 0 {
140 return Err(Error::EmptyFleet);
141 }
142 if node_id.get() >= fleet_capacity {
143 return Err(Error::NodeOutsideFleet {
144 node_id: node_id.get(),
145 fleet_capacity,
146 });
147 }
148 Ok(Self {
149 inner: Arc::new(FleetInner {
150 name,
151 fleet_capacity,
152 node_id,
153 id_counters: DashMap::new(),
154 backing: RingBacking::Shm(ShmRingRegistry::new(name, fleet_capacity)),
155 }),
156 })
157 }
158
159 pub fn name(&self) -> &'static str {
160 self.inner.name
161 }
162
163 pub fn fleet_capacity(&self) -> u16 {
165 self.inner.fleet_capacity
166 }
167
168 pub fn node_id(&self) -> NodeId {
169 self.inner.node_id
170 }
171
172 pub fn next_id<T: OrbitTyped>(&self) -> NetId64 {
180 let counter_arc = self
181 .inner
182 .id_counters
183 .entry(T::KIND)
184 .or_insert_with(|| Arc::new(AtomicU64::new(0)))
185 .clone();
186 let counter = counter_arc.fetch_add(1, Ordering::Relaxed);
187 NetId64::make(T::KIND, self.node_id().get(), counter)
188 }
189
190 pub fn is_shm(&self) -> bool {
193 #[cfg(unix)]
194 {
195 matches!(self.inner.backing, RingBacking::Shm(_))
196 }
197 #[cfg(not(unix))]
198 {
199 false
200 }
201 }
202
203 pub fn ring<T: OrbitTyped>(&self) -> Arc<Ring> {
211 match &self.inner.backing {
212 RingBacking::InMemory(r) => r.get_or_create::<T>(),
213 #[cfg(unix)]
214 RingBacking::Shm(_) => {
215 panic!("Fleet::ring called on SHM-backed fleet — use Fleet::shm_ring instead");
216 }
217 }
218 }
219
220 #[cfg(unix)]
232 pub fn shm_ring<T: OrbitTyped>(&self) -> std::io::Result<Arc<ShmRing>> {
233 match &self.inner.backing {
234 RingBacking::Shm(r) => r.get_or_create_for::<T>(),
235 RingBacking::InMemory(_) => {
236 panic!("Fleet::shm_ring called on in-memory fleet — use Fleet::ring instead");
237 }
238 }
239 }
240
241 pub fn publish<T: OrbitTyped>(&self, frame_kind: u8, ver: u64, payload: Bytes) -> NetId64 {
251 match &self.inner.backing {
252 RingBacking::InMemory(r) => {
253 let ring = r.get_or_create::<T>();
254 ring.write(self.node_id(), frame_kind, ver, payload)
255 }
256 #[cfg(unix)]
257 RingBacking::Shm(r) => {
258 let ring = r
259 .get_or_create_for::<T>()
260 .expect("SHM ring open failed — fleet unusable");
261 ring.write(self.node_id(), frame_kind, ver, payload)
262 .expect("SHM ring write failed")
263 }
264 }
265 }
266
267 pub fn publish_batch<T: OrbitTyped>(
280 &self,
281 frame_kind: u8,
282 ver: u64,
283 payloads: Vec<Bytes>,
284 ) -> Vec<NetId64> {
285 match &self.inner.backing {
286 RingBacking::InMemory(r) => {
287 let ring = r.get_or_create::<T>();
288 ring.write_batch(self.node_id(), frame_kind, ver, payloads)
289 }
290 #[cfg(unix)]
291 RingBacking::Shm(r) => {
292 let ring = r
293 .get_or_create_for::<T>()
294 .expect("SHM ring open failed — fleet unusable");
295 ring.write_batch(self.node_id(), frame_kind, ver, payloads)
296 .expect("SHM ring batch write failed")
297 }
298 }
299 }
300
301 pub fn read(&self, id: NetId64) -> Option<Frame> {
305 match &self.inner.backing {
306 RingBacking::InMemory(r) => r.lookup(id.kind())?.read(id),
307 #[cfg(unix)]
308 RingBacking::Shm(r) => r.lookup(id.kind())?.read(id),
309 }
310 }
311
312 pub fn read_head<T: OrbitTyped>(&self) -> Option<Frame> {
315 if T::RING_SPEC.topology == RingTopology::PerNode {
316 let head = self.lane_head::<T>(self.node_id());
317 return (head > 0)
318 .then(|| self.read_lane_at::<T>(self.node_id(), head - 1))
319 .flatten();
320 }
321 match &self.inner.backing {
322 RingBacking::InMemory(r) => {
323 let ring = r.get_or_create::<T>();
324 ring.read_head()
325 }
326 #[cfg(unix)]
327 RingBacking::Shm(r) => {
328 let ring = r.get_or_create_for::<T>().ok()?;
329 ring.read_head()
330 }
331 }
332 }
333
334 pub fn head<T: OrbitTyped>(&self) -> u64 {
343 if T::RING_SPEC.topology == RingTopology::PerNode {
344 return self.lane_head::<T>(self.node_id());
345 }
346 match &self.inner.backing {
347 RingBacking::InMemory(r) => r.get_or_create::<T>().head(),
348 #[cfg(unix)]
349 RingBacking::Shm(r) => r
350 .get_or_create_for::<T>()
351 .map(|ring| ring.head())
352 .unwrap_or(0),
353 }
354 }
355
356 pub fn read_at<T: OrbitTyped>(&self, counter: u64) -> Option<Frame> {
364 if T::RING_SPEC.topology == RingTopology::PerNode {
365 return self.read_lane_at::<T>(self.node_id(), counter);
366 }
367 match &self.inner.backing {
368 RingBacking::InMemory(r) => r.get_or_create::<T>().read_at(counter),
369 #[cfg(unix)]
370 RingBacking::Shm(r) => r.get_or_create_for::<T>().ok()?.read_at(counter),
371 }
372 }
373
374 pub(crate) fn read_state_at<T: OrbitTyped>(
375 &self,
376 counter: u64,
377 ) -> crate::ring::cursor::RingRead {
378 if T::RING_SPEC.topology == RingTopology::PerNode {
379 return self.read_lane_state_at::<T>(self.node_id(), counter);
380 }
381 match &self.inner.backing {
382 RingBacking::InMemory(r) => r.get_or_create::<T>().read_state_at(counter),
383 #[cfg(unix)]
384 RingBacking::Shm(r) => r
385 .get_or_create_for::<T>()
386 .map(|ring| ring.read_state_at(counter))
387 .unwrap_or(crate::ring::cursor::RingRead::Unavailable),
388 }
389 }
390
391 pub fn lane_head<T: OrbitTyped>(&self, node_id: NodeId) -> u64 {
395 match &self.inner.backing {
396 RingBacking::InMemory(r) => r.get_or_create::<T>().lane_head(node_id),
397 #[cfg(unix)]
398 RingBacking::Shm(r) => r
399 .get_or_create_for::<T>()
400 .map(|ring| ring.lane_head(node_id))
401 .unwrap_or(0),
402 }
403 }
404
405 pub fn read_lane_at<T: OrbitTyped>(&self, node_id: NodeId, counter: u64) -> Option<Frame> {
407 match &self.inner.backing {
408 RingBacking::InMemory(r) => r.get_or_create::<T>().read_lane_at(node_id, counter),
409 #[cfg(unix)]
410 RingBacking::Shm(r) => r
411 .get_or_create_for::<T>()
412 .ok()?
413 .read_lane_at(node_id, counter),
414 }
415 }
416
417 pub(crate) fn read_lane_state_at<T: OrbitTyped>(
418 &self,
419 node_id: NodeId,
420 counter: u64,
421 ) -> crate::ring::cursor::RingRead {
422 match &self.inner.backing {
423 RingBacking::InMemory(r) => r.get_or_create::<T>().read_lane_state_at(node_id, counter),
424 #[cfg(unix)]
425 RingBacking::Shm(r) => r
426 .get_or_create_for::<T>()
427 .map(|ring| ring.read_lane_state_at(node_id, counter))
428 .unwrap_or(crate::ring::cursor::RingRead::Unavailable),
429 }
430 }
431
432 pub fn ring_capacity<T: OrbitTyped>(&self) -> usize {
436 match &self.inner.backing {
437 RingBacking::InMemory(r) => r.get_or_create::<T>().capacity(),
438 #[cfg(unix)]
439 RingBacking::Shm(r) => r
440 .get_or_create_for::<T>()
441 .map(|ring| ring.capacity())
442 .unwrap_or(T::RING_SPEC.capacity),
443 }
444 }
445
446 pub fn next_ring_version<T: OrbitTyped>(&self) -> u64 {
452 match &self.inner.backing {
453 RingBacking::InMemory(r) => r.get_or_create::<T>().next_version(),
454 #[cfg(unix)]
455 RingBacking::Shm(r) => r
456 .get_or_create_for::<T>()
457 .expect("SHM ring open failed — fleet unusable")
458 .next_version(),
459 }
460 }
461
462 pub fn current_ring_version<T: OrbitTyped>(&self) -> u64 {
464 match &self.inner.backing {
465 RingBacking::InMemory(r) => r.get_or_create::<T>().current_version(),
466 #[cfg(unix)]
467 RingBacking::Shm(r) => r
468 .get_or_create_for::<T>()
469 .expect("SHM ring open failed — fleet unusable")
470 .current_version(),
471 }
472 }
473
474 pub fn reset_ring<T: OrbitTyped>(&self) -> std::io::Result<()> {
482 match &self.inner.backing {
483 RingBacking::InMemory(r) => {
484 r.get_or_create::<T>().reset();
485 Ok(())
486 }
487 #[cfg(unix)]
488 RingBacking::Shm(r) => {
489 r.get_or_create_for::<T>()?.reset();
490 Ok(())
491 }
492 }
493 }
494
495 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
496 pub fn ring_event_fd<T: OrbitTyped>(&self) -> std::io::Result<RingEventFd> {
502 match &self.inner.backing {
503 RingBacking::Shm(rings) => RingEventFd::new(rings.get_or_create_for::<T>()?),
504 RingBacking::InMemory(_) => Err(std::io::Error::new(
505 std::io::ErrorKind::Unsupported,
506 "Orbit eventfd requires a shared-memory fleet",
507 )),
508 }
509 }
510
511 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
512 pub fn publish_notified<T: OrbitTyped>(
514 &self,
515 frame_kind: u8,
516 ver: u64,
517 payload: Bytes,
518 ) -> std::io::Result<NetId64> {
519 match &self.inner.backing {
520 RingBacking::Shm(rings) => {
521 let ring = rings.get_or_create_for::<T>()?;
522 let id = ring.write(self.node_id(), frame_kind, ver, payload)?;
523 RingEventFd::notify(&ring)?;
524 Ok(id)
525 }
526 RingBacking::InMemory(rings) => {
528 Ok(rings
529 .get_or_create::<T>()
530 .write(self.node_id(), frame_kind, ver, payload))
531 }
532 }
533 }
534
535 #[cfg(any(target_os = "linux", target_os = "freebsd"))]
536 pub fn publish_batch_notified<T: OrbitTyped>(
539 &self,
540 frame_kind: u8,
541 ver: u64,
542 payloads: Vec<Bytes>,
543 ) -> std::io::Result<Vec<NetId64>> {
544 match &self.inner.backing {
545 RingBacking::Shm(rings) => {
546 let ring = rings.get_or_create_for::<T>()?;
547 let ids = ring.write_batch(self.node_id(), frame_kind, ver, payloads)?;
548 if !ids.is_empty() {
549 RingEventFd::notify(&ring)?;
550 }
551 Ok(ids)
552 }
553 RingBacking::InMemory(rings) => Ok(rings.get_or_create::<T>().write_batch(
554 self.node_id(),
555 frame_kind,
556 ver,
557 payloads,
558 )),
559 }
560 }
561}
562
563impl std::fmt::Debug for Fleet {
564 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
565 f.debug_struct("Fleet")
566 .field("name", &self.inner.name)
567 .field("fleet_capacity", &self.inner.fleet_capacity)
568 .field("node_id", &self.inner.node_id)
569 .field("id_counters", &self.inner.id_counters.len())
570 .finish()
571 }
572}