1#![cfg(unix)]
56
57use std::ptr;
58use std::sync::Mutex;
59use std::sync::atomic::{AtomicU8, AtomicU32, AtomicU64, Ordering, fence};
60
61use bytes::Bytes;
62
63use crate::NodeId;
64use crate::id::NetId64;
65use crate::ring::{Frame, RingSpec, RingTopology};
66use crate::shm::{self, ShmRegion};
67
68const MAGIC: u32 = 0x4F524254; const VERSION: u32 = 3;
74const SLOT_ALIGNMENT: usize = 64;
75
76#[repr(C, align(64))]
78struct ShmRingHeader {
79 version_counter: AtomicU64,
80 capacity: u64,
81 magic: u32,
82 version: u32,
83 payload_capacity: u32,
84 slot_stride: u32,
85 kind: u8,
86 topology: u8,
87 lane_count: u16,
88 notification_generation: AtomicU32,
92 _reserved: [u8; 24],
95}
96
97#[repr(C, align(64))]
98struct ShmLaneHeader {
99 write_pos: AtomicU64,
100 _reserved: [u8; 56],
101}
102
103#[repr(C)]
105struct ShmSlotHeader {
106 seq: AtomicU64,
110 id: AtomicU64,
112 ver: AtomicU64,
114 payload_len: AtomicU32,
116 kind: AtomicU8,
118 _reserved: [u8; 3],
119}
120
121const SLOT_HEADER_SIZE: usize = std::mem::size_of::<ShmSlotHeader>();
131const HEADER_SIZE: usize = std::mem::size_of::<ShmRingHeader>();
132const LANE_HEADER_SIZE: usize = std::mem::size_of::<ShmLaneHeader>();
133
134const _: () = assert!(HEADER_SIZE == 64);
135const _: () = assert!(LANE_HEADER_SIZE == 64);
136const _: () = assert!(SLOT_HEADER_SIZE == 32);
137
138fn invalid_input(message: impl Into<String>) -> std::io::Error {
139 std::io::Error::new(std::io::ErrorKind::InvalidInput, message.into())
140}
141
142fn lane_count_for(spec: RingSpec, fleet_capacity: u16) -> std::io::Result<usize> {
143 if fleet_capacity == 0 {
144 return Err(invalid_input("ShmRing fleet capacity must be > 0"));
145 }
146 Ok(match spec.topology {
147 RingTopology::Shared | RingTopology::SharedOrdered => 1,
148 RingTopology::PerNode => usize::from(fleet_capacity),
149 })
150}
151
152fn checked_layout(spec: RingSpec, fleet_capacity: u16) -> std::io::Result<(usize, usize, usize)> {
153 if spec.capacity == 0 {
154 return Err(invalid_input("ShmRing capacity must be > 0"));
155 }
156 if !spec.capacity.is_power_of_two() {
157 return Err(invalid_input("ShmRing capacity must be a power of two"));
158 }
159 if spec.payload_capacity > u32::MAX as usize {
160 return Err(invalid_input("ShmRing payload capacity must fit in u32"));
161 }
162
163 let unaligned = SLOT_HEADER_SIZE
164 .checked_add(spec.payload_capacity)
165 .ok_or_else(|| invalid_input("ShmRing slot size overflow"))?;
166 let slot_stride = unaligned
167 .checked_add(SLOT_ALIGNMENT - 1)
168 .map(|value| value & !(SLOT_ALIGNMENT - 1))
169 .ok_or_else(|| invalid_input("ShmRing slot stride overflow"))?;
170 if slot_stride > u32::MAX as usize {
171 return Err(invalid_input("ShmRing slot stride must fit in u32"));
172 }
173 let lane_count = lane_count_for(spec, fleet_capacity)?;
174 let lane_headers_size = lane_count
175 .checked_mul(LANE_HEADER_SIZE)
176 .ok_or_else(|| invalid_input("ShmRing lane header size overflow"))?;
177 let slots_offset = HEADER_SIZE
178 .checked_add(lane_headers_size)
179 .ok_or_else(|| invalid_input("ShmRing slots offset overflow"))?;
180 let slots_per_lane = spec
181 .capacity
182 .checked_mul(slot_stride)
183 .ok_or_else(|| invalid_input("ShmRing lane size overflow"))?;
184 let slots_size = lane_count
185 .checked_mul(slots_per_lane)
186 .ok_or_else(|| invalid_input("ShmRing slots size overflow"))?;
187 let segment_size = slots_offset
188 .checked_add(slots_size)
189 .ok_or_else(|| invalid_input("ShmRing segment size overflow"))?;
190 Ok((slot_stride, slots_offset, segment_size))
191}
192
193pub fn segment_size_for_spec(spec: RingSpec) -> std::io::Result<usize> {
195 segment_size_for_spec_and_fleet(spec, 1)
196}
197
198pub fn segment_size_for_spec_and_fleet(
200 spec: RingSpec,
201 fleet_capacity: u16,
202) -> std::io::Result<usize> {
203 checked_layout(spec, fleet_capacity).map(|(_, _, segment_size)| segment_size)
204}
205
206pub struct ShmRing {
217 region: ShmRegion,
218 kind: u8,
219 capacity: usize,
220 payload_capacity: usize,
221 topology: RingTopology,
222 lane_count: usize,
223 slot_stride: usize,
224 slots_offset: usize,
225 lane_stride: usize,
226 write_locks: Vec<Mutex<()>>,
227}
228
229impl ShmRing {
230 pub fn open_or_create(fleet_name: &str, kind: u8, spec: RingSpec) -> std::io::Result<Self> {
234 Self::open_or_create_for_fleet(fleet_name, kind, spec, 1)
235 }
236
237 pub fn open_or_create_for_fleet(
240 fleet_name: &str,
241 kind: u8,
242 spec: RingSpec,
243 fleet_capacity: u16,
244 ) -> std::io::Result<Self> {
245 let lane_count = lane_count_for(spec, fleet_capacity)?;
246 let (slot_stride, slots_offset, size) = checked_layout(spec, fleet_capacity)?;
247 let lane_stride = spec
248 .capacity
249 .checked_mul(slot_stride)
250 .ok_or_else(|| invalid_input("ShmRing lane stride overflow"))?;
251 let name = shm::ring_segment_name(fleet_name, kind);
252
253 let (region, _initialization_lock) = ShmRegion::open_or_create_locked(&name, size)?;
263
264 if region.created() {
267 unsafe {
270 let header_ptr = region.as_ptr() as *mut ShmRingHeader;
271 ptr::write(
272 header_ptr,
273 ShmRingHeader {
274 version_counter: AtomicU64::new(0),
275 capacity: spec.capacity as u64,
276 magic: MAGIC,
277 version: VERSION,
278 payload_capacity: spec.payload_capacity as u32,
279 slot_stride: slot_stride as u32,
280 kind,
281 topology: spec.topology as u8,
282 lane_count: lane_count as u16,
283 notification_generation: AtomicU32::new(0),
284 _reserved: [0; 24],
285 },
286 );
287
288 for lane in 0..lane_count {
289 let lane_ptr = region.as_ptr().add(HEADER_SIZE + lane * LANE_HEADER_SIZE)
290 as *mut ShmLaneHeader;
291 ptr::write(
292 lane_ptr,
293 ShmLaneHeader {
294 write_pos: AtomicU64::new(0),
295 _reserved: [0; 56],
296 },
297 );
298 }
299
300 let slots_ptr = region.as_ptr().add(slots_offset);
305 ptr::write_bytes(slots_ptr, 0, lane_count * lane_stride);
306 }
307 } else {
308 let header = unsafe { &*(region.as_ptr() as *const ShmRingHeader) };
311 if header.magic != MAGIC {
312 return Err(std::io::Error::new(
313 std::io::ErrorKind::InvalidData,
314 format!(
315 "SHM segment {} has wrong magic 0x{:08X} (expected 0x{:08X})",
316 name, header.magic, MAGIC
317 ),
318 ));
319 }
320 if header.version != VERSION {
321 return Err(std::io::Error::new(
322 std::io::ErrorKind::InvalidData,
323 format!(
324 "SHM segment {} version {} != local {}",
325 name, header.version, VERSION
326 ),
327 ));
328 }
329 if header.kind != kind {
330 return Err(std::io::Error::new(
331 std::io::ErrorKind::InvalidData,
332 format!(
333 "SHM segment {} kind {} != requested {}",
334 name, header.kind, kind
335 ),
336 ));
337 }
338 if header.topology != spec.topology as u8 {
339 return Err(std::io::Error::new(
340 std::io::ErrorKind::InvalidData,
341 format!(
342 "SHM segment {} topology {} != requested {}",
343 name, header.topology, spec.topology as u8
344 ),
345 ));
346 }
347 if header.lane_count as usize != lane_count {
348 return Err(std::io::Error::new(
349 std::io::ErrorKind::InvalidData,
350 format!(
351 "SHM segment {} lane count {} != requested {}",
352 name, header.lane_count, lane_count
353 ),
354 ));
355 }
356 if header.capacity as usize != spec.capacity {
357 return Err(std::io::Error::new(
358 std::io::ErrorKind::InvalidData,
359 format!(
360 "SHM segment {} capacity {} != requested {}",
361 name, header.capacity, spec.capacity
362 ),
363 ));
364 }
365 if header.payload_capacity as usize != spec.payload_capacity {
366 return Err(std::io::Error::new(
367 std::io::ErrorKind::InvalidData,
368 format!(
369 "SHM segment {} payload capacity {} != requested {}",
370 name, header.payload_capacity, spec.payload_capacity
371 ),
372 ));
373 }
374 if header.slot_stride as usize != slot_stride {
375 return Err(std::io::Error::new(
376 std::io::ErrorKind::InvalidData,
377 format!(
378 "SHM segment {} slot stride {} != requested {}",
379 name, header.slot_stride, slot_stride
380 ),
381 ));
382 }
383 }
384
385 Ok(Self {
386 region,
387 kind,
388 capacity: spec.capacity,
389 payload_capacity: spec.payload_capacity,
390 topology: spec.topology,
391 lane_count,
392 slot_stride,
393 slots_offset,
394 lane_stride,
395 write_locks: (0..lane_count).map(|_| Mutex::new(())).collect(),
396 })
397 }
398
399 pub fn created(&self) -> bool {
402 self.region.created()
403 }
404
405 pub fn unlink(&self) -> std::io::Result<()> {
409 self.region.unlink()
410 }
411
412 pub fn kind(&self) -> u8 {
414 self.kind
415 }
416
417 pub fn capacity(&self) -> usize {
419 self.capacity
420 }
421
422 pub fn lane_count(&self) -> usize {
424 self.lane_count
425 }
426
427 pub fn payload_capacity(&self) -> usize {
429 self.payload_capacity
430 }
431
432 pub fn spec(&self) -> RingSpec {
433 RingSpec {
434 capacity: self.capacity,
435 payload_capacity: self.payload_capacity,
436 topology: self.topology,
437 }
438 }
439
440 fn lane_header(&self, lane: usize) -> &ShmLaneHeader {
441 debug_assert!(lane < self.lane_count);
442 unsafe {
443 &*(self
444 .region
445 .as_ptr()
446 .add(HEADER_SIZE + lane * LANE_HEADER_SIZE) as *const ShmLaneHeader)
447 }
448 }
449
450 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
451 pub(crate) fn notification_generation(&self) -> &AtomicU32 {
452 unsafe { &(*(self.region.as_ptr() as *const ShmRingHeader)).notification_generation }
455 }
456
457 fn slot_ptr(&self, lane: usize, idx: usize) -> *mut ShmSlotHeader {
458 debug_assert!(lane < self.lane_count);
459 debug_assert!(idx < self.capacity);
460 unsafe {
463 let base = self
464 .region
465 .as_ptr()
466 .add(self.slots_offset + lane * self.lane_stride);
467 base.add(idx * self.slot_stride).cast::<ShmSlotHeader>()
468 }
469 }
470
471 unsafe fn payload_ptr(slot_ptr: *mut ShmSlotHeader) -> *mut AtomicU8 {
482 unsafe {
483 slot_ptr
484 .cast::<u8>()
485 .add(SLOT_HEADER_SIZE)
486 .cast::<AtomicU8>()
487 }
488 }
489
490 pub fn head(&self) -> u64 {
492 self.lane_header(0).write_pos.load(Ordering::Acquire)
493 }
494
495 pub fn lane_head(&self, node_id: NodeId) -> u64 {
497 let lane = self.lane_index(node_id);
498 self.lane_header(lane).write_pos.load(Ordering::Acquire)
499 }
500
501 pub fn next_version(&self) -> u64 {
504 let header = unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) };
505 header
506 .version_counter
507 .fetch_add(1, Ordering::AcqRel)
508 .checked_add(1)
509 .expect("SHM ring semantic version exhausted")
510 }
511
512 pub fn current_version(&self) -> u64 {
514 let header = unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) };
515 header.version_counter.load(Ordering::Acquire)
516 }
517
518 pub fn write(
521 &self,
522 node_id: NodeId,
523 frame_kind: u8,
524 ver: u64,
525 payload: Bytes,
526 ) -> std::io::Result<NetId64> {
527 if payload.len() > self.payload_capacity {
528 return Err(std::io::Error::new(
529 std::io::ErrorKind::InvalidInput,
530 format!(
531 "payload {} > ring payload capacity {}",
532 payload.len(),
533 self.payload_capacity
534 ),
535 ));
536 }
537
538 let lane = self.lane_index(node_id);
539 match self.topology {
540 RingTopology::Shared => {
541 let counter = self
542 .lane_header(lane)
543 .write_pos
544 .fetch_add(1, Ordering::AcqRel);
545 Ok(self.write_slot(lane, node_id, counter, frame_kind, ver, &payload))
546 }
547 RingTopology::PerNode => {
548 let _write = self.write_locks[lane]
549 .lock()
550 .unwrap_or_else(|error| error.into_inner());
551 let lane_header = self.lane_header(lane);
552 let counter = lane_header.write_pos.load(Ordering::Relaxed);
553 let id = self.write_slot(lane, node_id, counter, frame_kind, ver, &payload);
554 lane_header
555 .write_pos
556 .store(counter.wrapping_add(1), Ordering::Release);
557 Ok(id)
558 }
559 RingTopology::SharedOrdered => {
560 let _write = self.write_locks[lane]
561 .lock()
562 .unwrap_or_else(|error| error.into_inner());
563 let _cross_process = self.region.lock_exclusive()?;
564 let lane_header = self.lane_header(lane);
565 let counter = lane_header.write_pos.load(Ordering::Relaxed);
566 let id = self.write_slot(lane, node_id, counter, frame_kind, ver, &payload);
567 lane_header
568 .write_pos
569 .store(counter.wrapping_add(1), Ordering::Release);
570 Ok(id)
571 }
572 }
573 }
574
575 pub fn write_batch(
580 &self,
581 node_id: NodeId,
582 frame_kind: u8,
583 ver: u64,
584 payloads: Vec<Bytes>,
585 ) -> std::io::Result<Vec<NetId64>> {
586 if payloads.len() > self.capacity {
587 return Err(std::io::Error::new(
588 std::io::ErrorKind::InvalidInput,
589 format!("batch {} > ring capacity {}", payloads.len(), self.capacity),
590 ));
591 }
592 if let Some(payload) = payloads
593 .iter()
594 .find(|payload| payload.len() > self.payload_capacity)
595 {
596 return Err(std::io::Error::new(
597 std::io::ErrorKind::InvalidInput,
598 format!(
599 "payload {} > ring payload capacity {}",
600 payload.len(),
601 self.payload_capacity
602 ),
603 ));
604 }
605 if payloads.is_empty() {
606 return Ok(Vec::new());
607 }
608
609 let lane = self.lane_index(node_id);
610 let write_slots = |start: u64, payloads: Vec<Bytes>| {
611 payloads
612 .into_iter()
613 .enumerate()
614 .map(|(offset, payload)| {
615 self.write_slot(
616 lane,
617 node_id,
618 start.wrapping_add(offset as u64),
619 frame_kind,
620 ver,
621 &payload,
622 )
623 })
624 .collect::<Vec<_>>()
625 };
626
627 match self.topology {
628 RingTopology::Shared => {
629 let start = self
630 .lane_header(lane)
631 .write_pos
632 .fetch_add(payloads.len() as u64, Ordering::AcqRel);
633 Ok(write_slots(start, payloads))
634 }
635 RingTopology::PerNode => {
636 let _write = self.write_locks[lane]
637 .lock()
638 .unwrap_or_else(|error| error.into_inner());
639 let lane_header = self.lane_header(lane);
640 let start = lane_header.write_pos.load(Ordering::Relaxed);
641 let ids = write_slots(start, payloads);
642 lane_header
643 .write_pos
644 .store(start.wrapping_add(ids.len() as u64), Ordering::Release);
645 Ok(ids)
646 }
647 RingTopology::SharedOrdered => {
648 let _write = self.write_locks[lane]
649 .lock()
650 .unwrap_or_else(|error| error.into_inner());
651 let _cross_process = self.region.lock_exclusive()?;
652 let lane_header = self.lane_header(lane);
653 let start = lane_header.write_pos.load(Ordering::Relaxed);
654 let ids = write_slots(start, payloads);
655 lane_header
656 .write_pos
657 .store(start.wrapping_add(ids.len() as u64), Ordering::Release);
658 Ok(ids)
659 }
660 }
661 }
662
663 pub fn read(&self, id: NetId64) -> Option<Frame> {
667 if id.kind() != self.kind {
668 return None;
669 }
670 let lane = self.lane_index_for_frame(id)?;
671 let counter = id.counter();
672 let slot_idx = (counter as usize) & (self.capacity - 1);
673 let slot_ptr = self.slot_ptr(lane, slot_idx);
674
675 for _ in 0..3 {
677 let Some(frame) = (unsafe { read_committed_frame(slot_ptr, self.payload_capacity) })
678 else {
679 continue;
680 };
681 if frame.id.counter() == counter {
682 return Some(frame);
683 } else {
684 return None;
686 }
687 }
688 None
689 }
690
691 pub fn read_head(&self) -> Option<Frame> {
694 let head = self.head();
695 if head == 0 {
696 return None;
697 }
698 let counter = head - 1;
699 let slot_idx = (counter as usize) & (self.capacity - 1);
700 let slot_ptr = self.slot_ptr(0, slot_idx);
701
702 for _ in 0..3 {
703 if let Some(frame) = unsafe { read_committed_frame(slot_ptr, self.payload_capacity) } {
704 return Some(frame);
705 }
706 }
707 None
708 }
709
710 pub fn read_at(&self, counter: u64) -> Option<Frame> {
715 self.read_lane_index_at(0, counter)
716 }
717
718 pub fn read_lane_at(&self, node_id: NodeId, counter: u64) -> Option<Frame> {
720 let lane = self.lane_index(node_id);
721 self.read_lane_index_at(lane, counter)
722 }
723
724 fn read_lane_index_at(&self, lane: usize, counter: u64) -> Option<Frame> {
725 let slot_idx = (counter as usize) & (self.capacity - 1);
726 let slot_ptr = self.slot_ptr(lane, slot_idx);
727 for _ in 0..3 {
728 if let Some(frame) = unsafe { read_committed_frame(slot_ptr, self.payload_capacity) } {
729 return Some(frame);
730 }
731 }
732 None
733 }
734
735 pub(crate) fn read_state_at(&self, counter: u64) -> crate::ring::cursor::RingRead {
736 self.read_lane_index_state_at(0, counter)
737 }
738
739 pub(crate) fn read_lane_state_at(
740 &self,
741 node_id: NodeId,
742 counter: u64,
743 ) -> crate::ring::cursor::RingRead {
744 let lane = self.lane_index(node_id);
745 self.read_lane_index_state_at(lane, counter)
746 }
747
748 fn read_lane_index_state_at(&self, lane: usize, counter: u64) -> crate::ring::cursor::RingRead {
749 use crate::ring::cursor::RingRead;
750
751 let slot_idx = (counter as usize) & (self.capacity - 1);
752 let slot_ptr = self.slot_ptr(lane, slot_idx);
753 let expected_committed = counter
754 .checked_mul(2)
755 .and_then(|value| value.checked_add(2))
756 .expect("seq overflow");
757
758 for _ in 0..3 {
759 let sequence = unsafe { &*slot_ptr }.seq.load(Ordering::Acquire);
760 if sequence < expected_committed {
761 return if self.topology != RingTopology::Shared {
762 RingRead::Unavailable
763 } else {
764 RingRead::Pending
765 };
766 }
767 if sequence > expected_committed {
768 return RingRead::Unavailable;
769 }
770 if let Some(frame) = unsafe { read_committed_frame(slot_ptr, self.payload_capacity) } {
771 return if frame.id.counter() == counter {
772 RingRead::Ready(frame)
773 } else {
774 RingRead::Unavailable
775 };
776 }
777 }
778
779 let sequence = unsafe { &*slot_ptr }.seq.load(Ordering::Acquire);
780 if sequence < expected_committed {
781 if self.topology != RingTopology::Shared {
782 RingRead::Unavailable
783 } else {
784 RingRead::Pending
785 }
786 } else {
787 RingRead::Unavailable
788 }
789 }
790
791 pub fn reset(&self) {
797 unsafe {
800 let slots_ptr = self.region.as_ptr().add(self.slots_offset);
801 ptr::write_bytes(slots_ptr, 0, self.lane_count * self.lane_stride);
802 }
803 for lane in 0..self.lane_count {
804 self.lane_header(lane).write_pos.store(0, Ordering::Release);
805 }
806 let header = unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) };
807 header.version_counter.store(0, Ordering::Release);
808 }
809
810 fn lane_index(&self, node_id: NodeId) -> usize {
811 let lane = match self.topology {
812 RingTopology::Shared | RingTopology::SharedOrdered => 0,
813 RingTopology::PerNode => usize::from(node_id.get()),
814 };
815 assert!(
816 lane < self.lane_count,
817 "node {} is outside SHM ring lane count {}",
818 node_id.get(),
819 self.lane_count
820 );
821 lane
822 }
823
824 fn lane_index_for_frame(&self, id: NetId64) -> Option<usize> {
825 let lane = match self.topology {
826 RingTopology::Shared | RingTopology::SharedOrdered => 0,
827 RingTopology::PerNode => usize::from(id.node()),
828 };
829 (lane < self.lane_count).then_some(lane)
830 }
831
832 fn write_slot(
833 &self,
834 lane: usize,
835 node_id: NodeId,
836 counter: u64,
837 frame_kind: u8,
838 ver: u64,
839 payload: &[u8],
840 ) -> NetId64 {
841 let id = NetId64::make(self.kind, node_id.get(), counter);
842 let slot_idx = (counter as usize) & (self.capacity - 1);
843 let slot_ptr = self.slot_ptr(lane, slot_idx);
844
845 unsafe {
862 let slot = &*slot_ptr;
863 let mid_seq = counter
864 .checked_mul(2)
865 .and_then(|value| value.checked_add(1))
866 .expect("seq overflow");
867 let final_seq = mid_seq.wrapping_add(1);
868
869 slot.seq.store(mid_seq, Ordering::Relaxed);
870 fence(Ordering::Release);
871 slot.id.store(id.raw(), Ordering::Relaxed);
872 slot.ver.store(ver, Ordering::Relaxed);
873 slot.payload_len
874 .store(payload.len() as u32, Ordering::Relaxed);
875 slot.kind.store(frame_kind, Ordering::Relaxed);
876 let bytes = Self::payload_ptr(slot_ptr);
877 for (index, byte) in payload.iter().enumerate() {
878 (*bytes.add(index)).store(*byte, Ordering::Relaxed);
879 }
880 slot.seq.store(final_seq, Ordering::Release);
881 }
882
883 id
884 }
885}
886
887#[derive(Clone, Debug, PartialEq, Eq)]
897pub struct ShmRingMetadata {
898 pub name: String,
900 pub segment_size: usize,
902 pub format_version: u32,
904 pub kind: u8,
906 pub spec: RingSpec,
908 pub lane_count: usize,
910}
911
912pub struct ShmRingView {
920 region: ShmRegion,
921 metadata: ShmRingMetadata,
922 slot_stride: usize,
923 slots_offset: usize,
924 lane_stride: usize,
925}
926
927impl std::fmt::Debug for ShmRingView {
928 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
929 formatter
930 .debug_struct("ShmRingView")
931 .field("metadata", &self.metadata)
932 .finish()
933 }
934}
935
936impl ShmRingView {
937 pub fn attach_existing(fleet_name: &str, kind: u8) -> std::io::Result<Self> {
939 let uid = unsafe { libc::geteuid() };
941 Self::attach_existing_for_uid(fleet_name, kind, uid)
942 }
943
944 pub fn attach_existing_for_uid(fleet_name: &str, kind: u8, uid: u32) -> std::io::Result<Self> {
949 let name = shm::ring_segment_name_for_uid(fleet_name, kind, uid);
950 let region = ShmRegion::open_existing_read_only(&name, HEADER_SIZE)?;
951
952 let header = unsafe { &*(region.as_ptr() as *const ShmRingHeader) };
955 if header.magic != MAGIC {
956 return Err(std::io::Error::new(
957 std::io::ErrorKind::InvalidData,
958 format!(
959 "SHM segment {} has wrong magic 0x{:08X} (expected 0x{:08X})",
960 name, header.magic, MAGIC
961 ),
962 ));
963 }
964 if header.version != VERSION {
965 return Err(std::io::Error::new(
966 std::io::ErrorKind::InvalidData,
967 format!(
968 "SHM segment {} version {} != local {}",
969 name, header.version, VERSION
970 ),
971 ));
972 }
973 if header.kind != kind {
974 return Err(std::io::Error::new(
975 std::io::ErrorKind::InvalidData,
976 format!(
977 "SHM segment {} kind {} != requested {}",
978 name, header.kind, kind
979 ),
980 ));
981 }
982
983 let topology = match header.topology {
984 value if value == RingTopology::Shared as u8 => RingTopology::Shared,
985 value if value == RingTopology::PerNode as u8 => RingTopology::PerNode,
986 value if value == RingTopology::SharedOrdered as u8 => RingTopology::SharedOrdered,
987 value => {
988 return Err(std::io::Error::new(
989 std::io::ErrorKind::InvalidData,
990 format!("SHM segment {name} has unknown topology {value}"),
991 ));
992 }
993 };
994 let capacity = usize::try_from(header.capacity).map_err(|_| {
995 std::io::Error::new(
996 std::io::ErrorKind::InvalidData,
997 format!(
998 "SHM segment {name} capacity {} does not fit this platform",
999 header.capacity
1000 ),
1001 )
1002 })?;
1003 let lane_count = usize::from(header.lane_count);
1004 if lane_count == 0 {
1005 return Err(std::io::Error::new(
1006 std::io::ErrorKind::InvalidData,
1007 format!("SHM segment {name} has zero lanes"),
1008 ));
1009 }
1010 if topology != RingTopology::PerNode && lane_count != 1 {
1011 return Err(std::io::Error::new(
1012 std::io::ErrorKind::InvalidData,
1013 format!(
1014 "SHM segment {name} topology {topology:?} requires one lane, found {lane_count}"
1015 ),
1016 ));
1017 }
1018 let spec = RingSpec {
1019 capacity,
1020 payload_capacity: header.payload_capacity as usize,
1021 topology,
1022 };
1023 let fleet_capacity = if topology == RingTopology::PerNode {
1024 header.lane_count
1025 } else {
1026 1
1027 };
1028 let (slot_stride, slots_offset, expected_size) = checked_layout(spec, fleet_capacity)
1029 .map_err(|error| {
1030 std::io::Error::new(
1031 std::io::ErrorKind::InvalidData,
1032 format!("SHM segment {name} has invalid geometry: {error}"),
1033 )
1034 })?;
1035 if header.slot_stride as usize != slot_stride {
1036 return Err(std::io::Error::new(
1037 std::io::ErrorKind::InvalidData,
1038 format!(
1039 "SHM segment {} slot stride {} != geometry {}",
1040 name, header.slot_stride, slot_stride
1041 ),
1042 ));
1043 }
1044 if region.len() < expected_size {
1045 return Err(std::io::Error::new(
1046 std::io::ErrorKind::InvalidData,
1047 format!(
1048 "SHM segment {} size {} is smaller than persisted geometry {}",
1049 name,
1050 region.len(),
1051 expected_size
1052 ),
1053 ));
1054 }
1055 let lane_stride = capacity
1056 .checked_mul(slot_stride)
1057 .ok_or_else(|| invalid_input(format!("SHM segment {name} lane stride overflow")))?;
1058
1059 Ok(Self {
1060 metadata: ShmRingMetadata {
1061 name,
1062 segment_size: region.len(),
1063 format_version: header.version,
1064 kind,
1065 spec,
1066 lane_count,
1067 },
1068 region,
1069 slot_stride,
1070 slots_offset,
1071 lane_stride,
1072 })
1073 }
1074
1075 pub fn attach_typed<T: crate::OrbitTyped>(fleet_name: &str) -> std::io::Result<Self> {
1077 let view = Self::attach_existing(fleet_name, T::KIND)?;
1078 if view.metadata.spec != T::RING_SPEC {
1079 return Err(std::io::Error::new(
1080 std::io::ErrorKind::InvalidData,
1081 format!(
1082 "OrbitTyped KIND {} declares {:?}; existing spec is {:?}",
1083 T::KIND,
1084 T::RING_SPEC,
1085 view.metadata.spec
1086 ),
1087 ));
1088 }
1089 Ok(view)
1090 }
1091
1092 pub fn metadata(&self) -> &ShmRingMetadata {
1093 &self.metadata
1094 }
1095
1096 pub fn current_version(&self) -> u64 {
1098 self.header().version_counter.load(Ordering::Acquire)
1099 }
1100
1101 pub fn notification_generation(&self) -> u32 {
1103 self.header()
1104 .notification_generation
1105 .load(Ordering::Acquire)
1106 }
1107
1108 pub fn lane(&self, lane: usize) -> std::io::Result<ShmRingLaneView<'_>> {
1110 if lane >= self.metadata.lane_count {
1111 return Err(std::io::Error::new(
1112 std::io::ErrorKind::InvalidInput,
1113 format!(
1114 "lane {lane} is outside SHM ring lane count {}",
1115 self.metadata.lane_count
1116 ),
1117 ));
1118 }
1119 Ok(ShmRingLaneView { ring: self, lane })
1120 }
1121
1122 fn header(&self) -> &ShmRingHeader {
1123 unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) }
1125 }
1126
1127 fn lane_header(&self, lane: usize) -> &ShmLaneHeader {
1128 debug_assert!(lane < self.metadata.lane_count);
1129 unsafe {
1131 &*(self
1132 .region
1133 .as_ptr()
1134 .add(HEADER_SIZE + lane * LANE_HEADER_SIZE) as *const ShmLaneHeader)
1135 }
1136 }
1137
1138 fn slot_ptr(&self, lane: usize, counter: u64) -> *mut ShmSlotHeader {
1139 debug_assert!(lane < self.metadata.lane_count);
1140 let slot = (counter as usize) & (self.metadata.spec.capacity - 1);
1141 unsafe {
1144 self.region
1145 .as_ptr()
1146 .add(self.slots_offset + lane * self.lane_stride + slot * self.slot_stride)
1147 .cast::<ShmSlotHeader>()
1148 }
1149 }
1150
1151 fn read_lane_state_at(&self, lane: usize, counter: u64) -> crate::ring::cursor::RingRead {
1152 use crate::ring::cursor::RingRead;
1153
1154 let slot_ptr = self.slot_ptr(lane, counter);
1155 let Some(expected_committed) = counter
1156 .checked_mul(2)
1157 .and_then(|value| value.checked_add(2))
1158 else {
1159 return RingRead::Unavailable;
1160 };
1161
1162 for _ in 0..3 {
1163 let sequence = unsafe { &*slot_ptr }.seq.load(Ordering::Acquire);
1165 if sequence < expected_committed {
1166 return if self.metadata.spec.topology == RingTopology::Shared {
1167 RingRead::Pending
1168 } else {
1169 RingRead::Unavailable
1170 };
1171 }
1172 if sequence > expected_committed {
1173 return RingRead::Unavailable;
1174 }
1175 if let Some(frame) =
1177 unsafe { read_committed_frame(slot_ptr, self.metadata.spec.payload_capacity) }
1178 {
1179 let lane_matches = self.metadata.spec.topology != RingTopology::PerNode
1180 || usize::from(frame.id.node()) == lane;
1181 return if frame.id.kind() == self.metadata.kind
1182 && frame.id.counter() == counter
1183 && lane_matches
1184 {
1185 RingRead::Ready(frame)
1186 } else {
1187 RingRead::Unavailable
1188 };
1189 }
1190 }
1191 RingRead::Unavailable
1192 }
1193}
1194
1195#[derive(Clone, Copy, Debug)]
1197pub struct ShmRingLaneView<'a> {
1198 ring: &'a ShmRingView,
1199 lane: usize,
1200}
1201
1202impl ShmRingLaneView<'_> {
1203 pub fn index(self) -> usize {
1204 self.lane
1205 }
1206
1207 pub fn head(self) -> u64 {
1209 self.ring
1210 .lane_header(self.lane)
1211 .write_pos
1212 .load(Ordering::Acquire)
1213 }
1214
1215 pub fn retained_range(self) -> std::ops::Range<u64> {
1217 let head = self.head();
1218 head.saturating_sub(self.ring.metadata.spec.capacity as u64)..head
1219 }
1220
1221 pub fn read_at(self, counter: u64) -> Option<Frame> {
1223 match self.ring.read_lane_state_at(self.lane, counter) {
1224 crate::ring::cursor::RingRead::Ready(frame) => Some(frame),
1225 crate::ring::cursor::RingRead::Pending | crate::ring::cursor::RingRead::Unavailable => {
1226 None
1227 }
1228 }
1229 }
1230
1231 pub fn read_head(self) -> Option<Frame> {
1233 let head = self.head();
1234 head.checked_sub(1)
1235 .and_then(|counter| self.read_at(counter))
1236 }
1237}
1238
1239impl crate::ring::cursor::RingFrameSource for ShmRingLaneView<'_> {
1240 fn kind(&self) -> u8 {
1241 self.ring.metadata.kind
1242 }
1243
1244 fn head(&self) -> u64 {
1245 (*self).head()
1246 }
1247
1248 fn capacity(&self) -> usize {
1249 self.ring.metadata.spec.capacity
1250 }
1251
1252 fn read_at(&self, counter: u64) -> Option<Frame> {
1253 (*self).read_at(counter)
1254 }
1255
1256 fn read_state_at(&self, counter: u64) -> crate::ring::cursor::RingRead {
1257 self.ring.read_lane_state_at(self.lane, counter)
1258 }
1259}
1260
1261pub struct ShmRingRegistry {
1269 fleet_name: String,
1270 fleet_capacity: u16,
1271 rings: dashmap::DashMap<u8, std::sync::Arc<ShmRing>>,
1272}
1273
1274impl ShmRingRegistry {
1275 pub fn new(fleet_name: impl Into<String>, fleet_capacity: u16) -> Self {
1276 Self {
1277 fleet_name: fleet_name.into(),
1278 fleet_capacity,
1279 rings: dashmap::DashMap::new(),
1280 }
1281 }
1282
1283 pub fn get_or_create_for<T: crate::OrbitTyped>(
1287 &self,
1288 ) -> std::io::Result<std::sync::Arc<ShmRing>> {
1289 if let Some(entry) = self.rings.get(&T::KIND) {
1290 if entry.spec() != T::RING_SPEC {
1291 return Err(std::io::Error::new(
1292 std::io::ErrorKind::InvalidData,
1293 format!(
1294 "OrbitTyped KIND {} was reused with ring spec {:?}; existing spec is {:?}",
1295 T::KIND,
1296 T::RING_SPEC,
1297 entry.spec()
1298 ),
1299 ));
1300 }
1301 return Ok(entry.clone());
1302 }
1303 let ring = std::sync::Arc::new(ShmRing::open_or_create_for_fleet(
1304 &self.fleet_name,
1305 T::KIND,
1306 T::RING_SPEC,
1307 self.fleet_capacity,
1308 )?);
1309 let entry = self.rings.entry(T::KIND).or_insert_with(|| ring.clone());
1310 if entry.spec() != T::RING_SPEC {
1311 return Err(std::io::Error::new(
1312 std::io::ErrorKind::InvalidData,
1313 format!(
1314 "OrbitTyped KIND {} raced with ring spec {:?}; installed spec is {:?}",
1315 T::KIND,
1316 T::RING_SPEC,
1317 entry.spec()
1318 ),
1319 ));
1320 }
1321 Ok(entry.clone())
1322 }
1323
1324 pub fn lookup(&self, kind: u8) -> Option<std::sync::Arc<ShmRing>> {
1327 self.rings.get(&kind).map(|e| e.clone())
1328 }
1329}
1330
1331unsafe fn read_committed_frame(
1339 slot_ptr: *mut ShmSlotHeader,
1340 payload_capacity: usize,
1341) -> Option<Frame> {
1342 let slot = unsafe { &*slot_ptr };
1343 let seq_pre = slot.seq.load(Ordering::Acquire);
1344 if seq_pre == 0 {
1345 return None;
1347 }
1348 if seq_pre & 1 == 1 {
1349 return None;
1351 }
1352
1353 let id = NetId64::from_raw(slot.id.load(Ordering::Relaxed));
1357 let kind = slot.kind.load(Ordering::Relaxed);
1358 let ver = slot.ver.load(Ordering::Relaxed);
1359 let payload_len = slot.payload_len.load(Ordering::Relaxed) as usize;
1360 if payload_len > payload_capacity {
1361 return None;
1363 }
1364 let payload_src = unsafe { ShmRing::payload_ptr(slot_ptr) };
1369 let mut payload_buf = Vec::with_capacity(payload_len);
1370 for index in 0..payload_len {
1371 payload_buf.push(unsafe { (*payload_src.add(index)).load(Ordering::Relaxed) });
1372 }
1373
1374 fence(Ordering::Acquire);
1382 let seq_post = slot.seq.load(Ordering::Relaxed);
1383 if seq_pre != seq_post {
1384 return None;
1386 }
1387
1388 Some(Frame {
1389 id,
1390 kind,
1391 ver,
1392 payload: Bytes::from(payload_buf),
1393 })
1394}
1395
1396#[cfg(test)]
1397mod tests {
1398 use std::sync::atomic::{AtomicU64, Ordering};
1399
1400 use nix::sys::wait::{WaitStatus, waitpid};
1401 use nix::unistd::{ForkResult, fork};
1402
1403 use super::*;
1404 use crate::ring::cursor::{RingCursor, poll_ring};
1405
1406 #[test]
1407 fn cursor_retries_a_claimed_slot_after_it_commits() {
1408 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1409
1410 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1411 let fleet_name = format!("p{:x}{test_id:x}", std::process::id());
1412 let ring = ShmRing::open_or_create(&fleet_name, 199, RingSpec::new(4, 16))
1413 .expect("create test ring");
1414 ring.reset();
1415 let slot_ptr = ring.slot_ptr(0, 0);
1416
1417 ring.lane_header(0).write_pos.store(1, Ordering::Release);
1418 unsafe { &*slot_ptr }.seq.store(1, Ordering::Release);
1419
1420 let mut cursor = RingCursor::from_start();
1421 let pending = poll_ring(&ring, &mut cursor);
1422 assert!(pending.is_empty());
1423 assert_eq!(cursor.next_counter(), 0);
1424
1425 let payload = b"ready";
1426 unsafe {
1427 let slot = &*slot_ptr;
1428 slot.id.store(
1429 NetId64::make(199, NodeId::ZERO.get(), 0).raw(),
1430 Ordering::Relaxed,
1431 );
1432 slot.ver.store(7, Ordering::Relaxed);
1433 slot.payload_len
1434 .store(payload.len() as u32, Ordering::Relaxed);
1435 slot.kind.store(1, Ordering::Relaxed);
1436 let bytes = ShmRing::payload_ptr(slot_ptr);
1437 for (index, byte) in payload.iter().enumerate() {
1438 (*bytes.add(index)).store(*byte, Ordering::Relaxed);
1439 }
1440 }
1441 unsafe { &*slot_ptr }.seq.store(2, Ordering::Release);
1442
1443 let committed = poll_ring(&ring, &mut cursor);
1444 assert_eq!(committed.frames.len(), 1);
1445 assert_eq!(&committed.frames[0].payload[..], payload);
1446 assert_eq!(cursor.next_counter(), 1);
1447
1448 ring.unlink().expect("unlink test ring");
1449 }
1450
1451 #[test]
1452 fn per_node_head_ignores_a_writer_that_dies_before_commit() {
1453 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1454
1455 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1456 let fleet_name = format!("d{:x}{test_id:x}", std::process::id());
1457 let spec = RingSpec::per_node(4, 16);
1458 let abandoned = ShmRing::open_or_create_for_fleet(&fleet_name, 198, spec, 2)
1459 .expect("create per-node test ring");
1460 abandoned.reset();
1461
1462 let slot_ptr = abandoned.slot_ptr(1, 0);
1463 unsafe { &*slot_ptr }.seq.store(1, Ordering::Release);
1464 assert_eq!(abandoned.lane_head(NodeId::new(1)), 0);
1465
1466 let replacement = ShmRing::open_or_create_for_fleet(&fleet_name, 198, spec, 2)
1467 .expect("replacement attaches");
1468 let id = replacement
1469 .write(NodeId::new(1), 1, 9, Bytes::from_static(b"recovered"))
1470 .expect("replacement commits");
1471
1472 assert_eq!(id.counter(), 0);
1473 assert_eq!(replacement.lane_head(NodeId::new(1)), 1);
1474 assert_eq!(
1475 &replacement.read(id).expect("frame visible").payload[..],
1476 b"recovered"
1477 );
1478
1479 replacement.unlink().expect("unlink test ring");
1480 }
1481
1482 #[test]
1483 fn per_node_lane_serializes_concurrent_local_publishers() {
1484 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1485
1486 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1487 let fleet_name = format!("c{:x}{test_id:x}", std::process::id());
1488 let ring = std::sync::Arc::new(
1489 ShmRing::open_or_create_for_fleet(&fleet_name, 197, RingSpec::per_node(512, 0), 2)
1490 .expect("create concurrent writer ring"),
1491 );
1492 ring.reset();
1493
1494 let mut writers = Vec::new();
1495 for _ in 0..4 {
1496 let ring = ring.clone();
1497 writers.push(std::thread::spawn(move || {
1498 (0..64)
1499 .map(|_| {
1500 ring.write(NodeId::new(1), 1, 0, Bytes::new())
1501 .expect("publish")
1502 .counter()
1503 })
1504 .collect::<Vec<_>>()
1505 }));
1506 }
1507
1508 let mut counters = writers
1509 .into_iter()
1510 .flat_map(|writer| writer.join().expect("writer joins"))
1511 .collect::<Vec<_>>();
1512 counters.sort_unstable();
1513 assert_eq!(counters, (0..256).collect::<Vec<_>>());
1514 assert_eq!(ring.lane_head(NodeId::new(1)), 256);
1515
1516 ring.unlink().expect("unlink test ring");
1517 }
1518
1519 #[test]
1520 fn shared_ordered_recovers_after_a_locked_writer_process_dies() {
1521 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1522
1523 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1524 let fleet_name = format!("o{:x}{test_id:x}", std::process::id());
1525 let spec = RingSpec::shared_ordered(4, 16);
1526 let ring = ShmRing::open_or_create_for_fleet(&fleet_name, 196, spec, 2)
1527 .expect("create shared-ordered test ring");
1528 ring.reset();
1529
1530 match unsafe { fork() }.expect("fork test writer") {
1531 ForkResult::Child => {
1532 let _lock = ring.region.lock_exclusive().expect("child locks ring");
1536 let slot_ptr = ring.slot_ptr(0, 0);
1537 unsafe { &*slot_ptr }.seq.store(1, Ordering::Release);
1538
1539 unsafe { libc::_exit(0) };
1542 }
1543 ForkResult::Parent { child } => {
1544 let status = waitpid(child, None).expect("wait for abandoned writer");
1545 assert!(matches!(status, WaitStatus::Exited(_, 0)));
1546
1547 let replacement = ShmRing::open_or_create_for_fleet(&fleet_name, 196, spec, 2)
1552 .expect("replacement attaches");
1553 let recovered_lock = replacement
1554 .region
1555 .try_lock_exclusive()
1556 .expect("kernel released dead writer lock");
1557 drop(recovered_lock);
1558
1559 let id = replacement
1560 .write(NodeId::new(1), 1, 9, Bytes::from_static(b"recovered"))
1561 .expect("replacement commits counter zero");
1562 assert_eq!(id.counter(), 0);
1563 assert_eq!(replacement.head(), 1);
1564 assert_eq!(
1565 &replacement.read(id).expect("recovered frame").payload[..],
1566 b"recovered"
1567 );
1568
1569 replacement
1570 .unlink()
1571 .expect("unlink shared-ordered test ring");
1572 }
1573 }
1574 }
1575}