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 notification_waiters: AtomicU32,
98 _reserved: [u8; 20]
101}
102
103#[repr(C, align(64))]
104struct ShmLaneHeader {
105 write_pos: AtomicU64,
106 _reserved: [u8; 56]
107}
108
109#[repr(C)]
111struct ShmSlotHeader {
112 seq: AtomicU64,
116 id: AtomicU64,
118 ver: AtomicU64,
120 payload_len: AtomicU32,
122 kind: AtomicU8,
124 _reserved: [u8; 3]
125}
126
127const SLOT_HEADER_SIZE: usize = std::mem::size_of::<ShmSlotHeader>();
137const HEADER_SIZE: usize = std::mem::size_of::<ShmRingHeader>();
138const LANE_HEADER_SIZE: usize = std::mem::size_of::<ShmLaneHeader>();
139
140const _: () = assert!(HEADER_SIZE == 64);
141const _: () = assert!(LANE_HEADER_SIZE == 64);
142const _: () = assert!(SLOT_HEADER_SIZE == 32);
143
144fn invalid_input(message: impl Into<String>) -> std::io::Error {
145 std::io::Error::new(std::io::ErrorKind::InvalidInput, message.into())
146}
147
148fn lane_count_for(
149 spec: RingSpec,
150 fleet_capacity: u16
151) -> std::io::Result<usize> {
152 if fleet_capacity == 0 {
153 return Err(invalid_input("ShmRing fleet capacity must be > 0"));
154 }
155 Ok(match spec.topology {
156 RingTopology::Shared | RingTopology::SharedOrdered => 1,
157 RingTopology::PerNode => usize::from(fleet_capacity)
158 })
159}
160
161fn checked_layout(
162 spec: RingSpec,
163 fleet_capacity: u16
164) -> std::io::Result<(usize, usize, usize)> {
165 if spec.capacity == 0 {
166 return Err(invalid_input("ShmRing capacity must be > 0"));
167 }
168 if !spec.capacity.is_power_of_two() {
169 return Err(invalid_input("ShmRing capacity must be a power of two"));
170 }
171 if spec.payload_capacity > u32::MAX as usize {
172 return Err(invalid_input("ShmRing payload capacity must fit in u32"));
173 }
174
175 let unaligned = SLOT_HEADER_SIZE
176 .checked_add(spec.payload_capacity)
177 .ok_or_else(|| invalid_input("ShmRing slot size overflow"))?;
178 let slot_stride = unaligned
179 .checked_add(SLOT_ALIGNMENT - 1)
180 .map(|value| value & !(SLOT_ALIGNMENT - 1))
181 .ok_or_else(|| invalid_input("ShmRing slot stride overflow"))?;
182 if slot_stride > u32::MAX as usize {
183 return Err(invalid_input("ShmRing slot stride must fit in u32"));
184 }
185 let lane_count = lane_count_for(spec, fleet_capacity)?;
186 let lane_headers_size = lane_count
187 .checked_mul(LANE_HEADER_SIZE)
188 .ok_or_else(|| invalid_input("ShmRing lane header size overflow"))?;
189 let slots_offset = HEADER_SIZE
190 .checked_add(lane_headers_size)
191 .ok_or_else(|| invalid_input("ShmRing slots offset overflow"))?;
192 let slots_per_lane = spec
193 .capacity
194 .checked_mul(slot_stride)
195 .ok_or_else(|| invalid_input("ShmRing lane size overflow"))?;
196 let slots_size = lane_count
197 .checked_mul(slots_per_lane)
198 .ok_or_else(|| invalid_input("ShmRing slots size overflow"))?;
199 let segment_size = slots_offset
200 .checked_add(slots_size)
201 .ok_or_else(|| invalid_input("ShmRing segment size overflow"))?;
202 Ok((slot_stride, slots_offset, segment_size))
203}
204
205pub fn segment_size_for_spec(spec: RingSpec) -> std::io::Result<usize> {
207 segment_size_for_spec_and_fleet(spec, 1)
208}
209
210pub fn segment_size_for_spec_and_fleet(
212 spec: RingSpec,
213 fleet_capacity: u16
214) -> std::io::Result<usize> {
215 checked_layout(spec, fleet_capacity).map(|(_, _, segment_size)| segment_size)
216}
217
218pub struct ShmRing {
229 region: ShmRegion,
230 kind: u8,
231 capacity: usize,
232 payload_capacity: usize,
233 topology: RingTopology,
234 lane_count: usize,
235 slot_stride: usize,
236 slots_offset: usize,
237 lane_stride: usize,
238 write_locks: Vec<Mutex<()>>
239}
240
241impl ShmRing {
242 pub fn open_or_create_with_policy(
245 fleet_name: &str,
246 kind: u8,
247 spec: RingSpec,
248 policy: shm::ShmAccessPolicy
249 ) -> std::io::Result<Self> {
250 Self::open_or_create_for_fleet_with_policy(fleet_name, kind, spec, 1, policy)
251 }
252
253 pub fn open_or_create(
257 fleet_name: &str,
258 kind: u8,
259 spec: RingSpec
260 ) -> std::io::Result<Self> {
261 Self::open_or_create_for_fleet(fleet_name, kind, spec, 1)
262 }
263
264 pub fn open_or_create_for_fleet(
267 fleet_name: &str,
268 kind: u8,
269 spec: RingSpec,
270 fleet_capacity: u16
271 ) -> std::io::Result<Self> {
272 Self::open_or_create_for_fleet_with_policy(
273 fleet_name,
274 kind,
275 spec,
276 fleet_capacity,
277 shm::ShmAccessPolicy::default()
278 )
279 }
280
281 pub fn open_or_create_for_fleet_with_policy(
285 fleet_name: &str,
286 kind: u8,
287 spec: RingSpec,
288 fleet_capacity: u16,
289 policy: shm::ShmAccessPolicy
290 ) -> std::io::Result<Self> {
291 let lane_count = lane_count_for(spec, fleet_capacity)?;
292 let (slot_stride, slots_offset, size) = checked_layout(spec, fleet_capacity)?;
293 let lane_stride = spec
294 .capacity
295 .checked_mul(slot_stride)
296 .ok_or_else(|| invalid_input("ShmRing lane stride overflow"))?;
297 let name = shm::ring_segment_name(fleet_name, kind);
298
299 let (region, _initialization_lock) =
309 ShmRegion::open_or_create_locked_with_policy(&name, size, policy)?;
310
311 if region.created() {
314 unsafe {
317 let header_ptr = region.as_ptr() as *mut ShmRingHeader;
318 ptr::write(
319 header_ptr,
320 ShmRingHeader {
321 version_counter: AtomicU64::new(0),
322 capacity: spec.capacity as u64,
323 magic: MAGIC,
324 version: VERSION,
325 payload_capacity: spec.payload_capacity as u32,
326 slot_stride: slot_stride as u32,
327 kind,
328 topology: spec.topology as u8,
329 lane_count: lane_count as u16,
330 notification_generation: AtomicU32::new(0),
331 notification_waiters: AtomicU32::new(0),
332 _reserved: [0; 20]
333 }
334 );
335
336 for lane in 0..lane_count {
337 let lane_ptr = region.as_ptr().add(HEADER_SIZE + lane * LANE_HEADER_SIZE)
338 as *mut ShmLaneHeader;
339 ptr::write(
340 lane_ptr,
341 ShmLaneHeader { write_pos: AtomicU64::new(0), _reserved: [0; 56] }
342 );
343 }
344
345 let slots_ptr = region.as_ptr().add(slots_offset);
350 ptr::write_bytes(slots_ptr, 0, lane_count * lane_stride);
351 }
352 } else {
353 let header = unsafe { &*(region.as_ptr() as *const ShmRingHeader) };
356 if header.magic != MAGIC {
357 return Err(std::io::Error::new(
358 std::io::ErrorKind::InvalidData,
359 format!(
360 "SHM segment {} has wrong magic 0x{:08X} (expected 0x{:08X})",
361 name, header.magic, MAGIC
362 )
363 ));
364 }
365 if header.version != VERSION {
366 return Err(std::io::Error::new(
367 std::io::ErrorKind::InvalidData,
368 format!("SHM segment {} version {} != local {}", name, header.version, VERSION)
369 ));
370 }
371 if header.kind != kind {
372 return Err(std::io::Error::new(
373 std::io::ErrorKind::InvalidData,
374 format!("SHM segment {} kind {} != requested {}", name, header.kind, kind)
375 ));
376 }
377 if header.topology != spec.topology as u8 {
378 return Err(std::io::Error::new(
379 std::io::ErrorKind::InvalidData,
380 format!(
381 "SHM segment {} topology {} != requested {}",
382 name, header.topology, spec.topology as u8
383 )
384 ));
385 }
386 if header.lane_count as usize != lane_count {
387 return Err(std::io::Error::new(
388 std::io::ErrorKind::InvalidData,
389 format!(
390 "SHM segment {} lane count {} != requested {}",
391 name, header.lane_count, lane_count
392 )
393 ));
394 }
395 if header.capacity as usize != spec.capacity {
396 return Err(std::io::Error::new(
397 std::io::ErrorKind::InvalidData,
398 format!(
399 "SHM segment {} capacity {} != requested {}",
400 name, header.capacity, spec.capacity
401 )
402 ));
403 }
404 if header.payload_capacity as usize != spec.payload_capacity {
405 return Err(std::io::Error::new(
406 std::io::ErrorKind::InvalidData,
407 format!(
408 "SHM segment {} payload capacity {} != requested {}",
409 name, header.payload_capacity, spec.payload_capacity
410 )
411 ));
412 }
413 if header.slot_stride as usize != slot_stride {
414 return Err(std::io::Error::new(
415 std::io::ErrorKind::InvalidData,
416 format!(
417 "SHM segment {} slot stride {} != requested {}",
418 name, header.slot_stride, slot_stride
419 )
420 ));
421 }
422 }
423
424 Ok(Self {
425 region,
426 kind,
427 capacity: spec.capacity,
428 payload_capacity: spec.payload_capacity,
429 topology: spec.topology,
430 lane_count,
431 slot_stride,
432 slots_offset,
433 lane_stride,
434 write_locks: (0..lane_count).map(|_| Mutex::new(())).collect()
435 })
436 }
437
438 pub fn created(&self) -> bool {
441 self.region.created()
442 }
443
444 pub fn unlink(&self) -> std::io::Result<()> {
448 self.region.unlink()
449 }
450
451 pub fn kind(&self) -> u8 {
453 self.kind
454 }
455
456 pub fn capacity(&self) -> usize {
458 self.capacity
459 }
460
461 pub fn lane_count(&self) -> usize {
463 self.lane_count
464 }
465
466 pub fn payload_capacity(&self) -> usize {
468 self.payload_capacity
469 }
470
471 pub fn spec(&self) -> RingSpec {
472 RingSpec {
473 capacity: self.capacity,
474 payload_capacity: self.payload_capacity,
475 topology: self.topology
476 }
477 }
478
479 fn lane_header(
480 &self,
481 lane: usize
482 ) -> &ShmLaneHeader {
483 debug_assert!(lane < self.lane_count);
484 unsafe {
485 &*(self.region.as_ptr().add(HEADER_SIZE + lane * LANE_HEADER_SIZE)
486 as *const ShmLaneHeader)
487 }
488 }
489
490 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
491 pub(crate) fn notification_generation(&self) -> &AtomicU32 {
492 unsafe { &(*(self.region.as_ptr() as *const ShmRingHeader)).notification_generation }
495 }
496
497 #[cfg(any(target_os = "linux", target_os = "freebsd", target_os = "macos"))]
498 pub(crate) fn notification_waiters(&self) -> &AtomicU32 {
499 unsafe { &(*(self.region.as_ptr() as *const ShmRingHeader)).notification_waiters }
501 }
502
503 fn slot_ptr(
504 &self,
505 lane: usize,
506 idx: usize
507 ) -> *mut ShmSlotHeader {
508 debug_assert!(lane < self.lane_count);
509 debug_assert!(idx < self.capacity);
510 unsafe {
513 let base = self.region.as_ptr().add(self.slots_offset + lane * self.lane_stride);
514 base.add(idx * self.slot_stride).cast::<ShmSlotHeader>()
515 }
516 }
517
518 unsafe fn payload_ptr(slot_ptr: *mut ShmSlotHeader) -> *mut AtomicU8 {
529 unsafe { slot_ptr.cast::<u8>().add(SLOT_HEADER_SIZE).cast::<AtomicU8>() }
530 }
531
532 pub fn head(&self) -> u64 {
534 self.lane_header(0).write_pos.load(Ordering::Acquire)
535 }
536
537 pub fn lane_head(
539 &self,
540 node_id: NodeId
541 ) -> u64 {
542 let lane = self.lane_index(node_id);
543 self.lane_header(lane).write_pos.load(Ordering::Acquire)
544 }
545
546 pub fn next_version(&self) -> u64 {
549 let header = unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) };
550 header
551 .version_counter
552 .fetch_add(1, Ordering::AcqRel)
553 .checked_add(1)
554 .expect("SHM ring semantic version exhausted")
555 }
556
557 pub fn current_version(&self) -> u64 {
559 let header = unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) };
560 header.version_counter.load(Ordering::Acquire)
561 }
562
563 pub fn write(
566 &self,
567 node_id: NodeId,
568 frame_kind: u8,
569 ver: u64,
570 payload: Bytes
571 ) -> std::io::Result<NetId64> {
572 if payload.len() > self.payload_capacity {
573 return Err(std::io::Error::new(
574 std::io::ErrorKind::InvalidInput,
575 format!(
576 "payload {} > ring payload capacity {}",
577 payload.len(),
578 self.payload_capacity
579 )
580 ));
581 }
582
583 let lane = self.lane_index(node_id);
584 match self.topology {
585 RingTopology::Shared => {
586 let counter = self.lane_header(lane).write_pos.fetch_add(1, Ordering::AcqRel);
587 Ok(self.write_slot(lane, node_id, counter, frame_kind, ver, &payload))
588 }
589 RingTopology::PerNode => {
590 let _write =
591 self.write_locks[lane].lock().unwrap_or_else(|error| error.into_inner());
592 let lane_header = self.lane_header(lane);
593 let counter = lane_header.write_pos.load(Ordering::Relaxed);
594 let id = self.write_slot(lane, node_id, counter, frame_kind, ver, &payload);
595 lane_header.write_pos.store(counter.wrapping_add(1), Ordering::Release);
596 Ok(id)
597 }
598 RingTopology::SharedOrdered => {
599 let _write =
600 self.write_locks[lane].lock().unwrap_or_else(|error| error.into_inner());
601 let _cross_process = self.region.lock_exclusive()?;
602 let lane_header = self.lane_header(lane);
603 let counter = lane_header.write_pos.load(Ordering::Relaxed);
604 let id = self.write_slot(lane, node_id, counter, frame_kind, ver, &payload);
605 lane_header.write_pos.store(counter.wrapping_add(1), Ordering::Release);
606 Ok(id)
607 }
608 }
609 }
610
611 pub fn write_batch(
616 &self,
617 node_id: NodeId,
618 frame_kind: u8,
619 ver: u64,
620 payloads: Vec<Bytes>
621 ) -> std::io::Result<Vec<NetId64>> {
622 if payloads.len() > self.capacity {
623 return Err(std::io::Error::new(
624 std::io::ErrorKind::InvalidInput,
625 format!("batch {} > ring capacity {}", payloads.len(), self.capacity)
626 ));
627 }
628 if let Some(payload) = payloads.iter().find(|payload| payload.len() > self.payload_capacity)
629 {
630 return Err(std::io::Error::new(
631 std::io::ErrorKind::InvalidInput,
632 format!(
633 "payload {} > ring payload capacity {}",
634 payload.len(),
635 self.payload_capacity
636 )
637 ));
638 }
639 if payloads.is_empty() {
640 return Ok(Vec::new());
641 }
642
643 let lane = self.lane_index(node_id);
644 let write_slots = |start: u64, payloads: Vec<Bytes>| {
645 payloads
646 .into_iter()
647 .enumerate()
648 .map(|(offset, payload)| {
649 self.write_slot(
650 lane,
651 node_id,
652 start.wrapping_add(offset as u64),
653 frame_kind,
654 ver,
655 &payload
656 )
657 })
658 .collect::<Vec<_>>()
659 };
660
661 match self.topology {
662 RingTopology::Shared => {
663 let start = self
664 .lane_header(lane)
665 .write_pos
666 .fetch_add(payloads.len() as u64, Ordering::AcqRel);
667 Ok(write_slots(start, payloads))
668 }
669 RingTopology::PerNode => {
670 let _write =
671 self.write_locks[lane].lock().unwrap_or_else(|error| error.into_inner());
672 let lane_header = self.lane_header(lane);
673 let start = lane_header.write_pos.load(Ordering::Relaxed);
674 let ids = write_slots(start, payloads);
675 lane_header
676 .write_pos
677 .store(start.wrapping_add(ids.len() as u64), Ordering::Release);
678 Ok(ids)
679 }
680 RingTopology::SharedOrdered => {
681 let _write =
682 self.write_locks[lane].lock().unwrap_or_else(|error| error.into_inner());
683 let _cross_process = self.region.lock_exclusive()?;
684 let lane_header = self.lane_header(lane);
685 let start = lane_header.write_pos.load(Ordering::Relaxed);
686 let ids = write_slots(start, payloads);
687 lane_header
688 .write_pos
689 .store(start.wrapping_add(ids.len() as u64), Ordering::Release);
690 Ok(ids)
691 }
692 }
693 }
694
695 pub fn read(
699 &self,
700 id: NetId64
701 ) -> Option<Frame> {
702 if id.kind() != self.kind {
703 return None;
704 }
705 let lane = self.lane_index_for_frame(id)?;
706 let counter = id.counter();
707 let slot_idx = (counter as usize) & (self.capacity - 1);
708 let slot_ptr = self.slot_ptr(lane, slot_idx);
709
710 for _ in 0..3 {
712 let Some(frame) = (unsafe { read_committed_frame(slot_ptr, self.payload_capacity) })
713 else {
714 continue;
715 };
716 if frame.id.counter() == counter {
717 return Some(frame);
718 } else {
719 return None;
721 }
722 }
723 None
724 }
725
726 pub fn read_head(&self) -> Option<Frame> {
729 let head = self.head();
730 if head == 0 {
731 return None;
732 }
733 let counter = head - 1;
734 let slot_idx = (counter as usize) & (self.capacity - 1);
735 let slot_ptr = self.slot_ptr(0, slot_idx);
736
737 for _ in 0..3 {
738 if let Some(frame) = unsafe { read_committed_frame(slot_ptr, self.payload_capacity) } {
739 return Some(frame);
740 }
741 }
742 None
743 }
744
745 pub fn read_at(
750 &self,
751 counter: u64
752 ) -> Option<Frame> {
753 self.read_lane_index_at(0, counter)
754 }
755
756 pub fn read_lane_at(
758 &self,
759 node_id: NodeId,
760 counter: u64
761 ) -> Option<Frame> {
762 let lane = self.lane_index(node_id);
763 self.read_lane_index_at(lane, counter)
764 }
765
766 fn read_lane_index_at(
767 &self,
768 lane: usize,
769 counter: u64
770 ) -> Option<Frame> {
771 let slot_idx = (counter as usize) & (self.capacity - 1);
772 let slot_ptr = self.slot_ptr(lane, slot_idx);
773 for _ in 0..3 {
774 if let Some(frame) = unsafe { read_committed_frame(slot_ptr, self.payload_capacity) } {
775 return Some(frame);
776 }
777 }
778 None
779 }
780
781 pub(crate) fn read_state_at(
782 &self,
783 counter: u64
784 ) -> crate::ring::cursor::RingRead {
785 self.read_lane_index_state_at(0, counter)
786 }
787
788 pub(crate) fn read_lane_state_at(
789 &self,
790 node_id: NodeId,
791 counter: u64
792 ) -> crate::ring::cursor::RingRead {
793 let lane = self.lane_index(node_id);
794 self.read_lane_index_state_at(lane, counter)
795 }
796
797 fn read_lane_index_state_at(
798 &self,
799 lane: usize,
800 counter: u64
801 ) -> crate::ring::cursor::RingRead {
802 use crate::ring::cursor::RingRead;
803
804 let slot_idx = (counter as usize) & (self.capacity - 1);
805 let slot_ptr = self.slot_ptr(lane, slot_idx);
806 let expected_committed =
807 counter.checked_mul(2).and_then(|value| value.checked_add(2)).expect("seq overflow");
808
809 for _ in 0..3 {
810 let sequence = unsafe { &*slot_ptr }.seq.load(Ordering::Acquire);
811 if sequence < expected_committed {
812 return if self.topology != RingTopology::Shared {
813 RingRead::Unavailable
814 } else {
815 RingRead::Pending
816 };
817 }
818 if sequence > expected_committed {
819 return RingRead::Unavailable;
820 }
821 if let Some(frame) = unsafe { read_committed_frame(slot_ptr, self.payload_capacity) } {
822 return if frame.id.counter() == counter {
823 RingRead::Ready(frame)
824 } else {
825 RingRead::Unavailable
826 };
827 }
828 }
829
830 let sequence = unsafe { &*slot_ptr }.seq.load(Ordering::Acquire);
831 if sequence < expected_committed {
832 if self.topology != RingTopology::Shared {
833 RingRead::Unavailable
834 } else {
835 RingRead::Pending
836 }
837 } else {
838 RingRead::Unavailable
839 }
840 }
841
842 pub fn reset(&self) {
848 unsafe {
851 let slots_ptr = self.region.as_ptr().add(self.slots_offset);
852 ptr::write_bytes(slots_ptr, 0, self.lane_count * self.lane_stride);
853 }
854 for lane in 0..self.lane_count {
855 self.lane_header(lane).write_pos.store(0, Ordering::Release);
856 }
857 let header = unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) };
858 header.version_counter.store(0, Ordering::Release);
859 }
860
861 fn lane_index(
862 &self,
863 node_id: NodeId
864 ) -> usize {
865 let lane = match self.topology {
866 RingTopology::Shared | RingTopology::SharedOrdered => 0,
867 RingTopology::PerNode => usize::from(node_id.get())
868 };
869 assert!(
870 lane < self.lane_count,
871 "node {} is outside SHM ring lane count {}",
872 node_id.get(),
873 self.lane_count
874 );
875 lane
876 }
877
878 fn lane_index_for_frame(
879 &self,
880 id: NetId64
881 ) -> Option<usize> {
882 let lane = match self.topology {
883 RingTopology::Shared | RingTopology::SharedOrdered => 0,
884 RingTopology::PerNode => usize::from(id.node())
885 };
886 (lane < self.lane_count).then_some(lane)
887 }
888
889 fn write_slot(
890 &self,
891 lane: usize,
892 node_id: NodeId,
893 counter: u64,
894 frame_kind: u8,
895 ver: u64,
896 payload: &[u8]
897 ) -> NetId64 {
898 let id = NetId64::make(self.kind, node_id.get(), counter);
899 let slot_idx = (counter as usize) & (self.capacity - 1);
900 let slot_ptr = self.slot_ptr(lane, slot_idx);
901
902 unsafe {
919 let slot = &*slot_ptr;
920 let mid_seq = counter
921 .checked_mul(2)
922 .and_then(|value| value.checked_add(1))
923 .expect("seq overflow");
924 let final_seq = mid_seq.wrapping_add(1);
925
926 slot.seq.store(mid_seq, Ordering::Relaxed);
927 fence(Ordering::Release);
928 slot.id.store(id.raw(), Ordering::Relaxed);
929 slot.ver.store(ver, Ordering::Relaxed);
930 slot.payload_len.store(payload.len() as u32, Ordering::Relaxed);
931 slot.kind.store(frame_kind, Ordering::Relaxed);
932 let bytes = Self::payload_ptr(slot_ptr);
933 for (index, byte) in payload.iter().enumerate() {
934 (*bytes.add(index)).store(*byte, Ordering::Relaxed);
935 }
936 slot.seq.store(final_seq, Ordering::Release);
937 }
938
939 id
940 }
941}
942
943#[derive(Clone, Debug, PartialEq, Eq)]
953pub struct ShmRingMetadata {
954 pub name: String,
956 pub segment_size: usize,
958 pub format_version: u32,
960 pub kind: u8,
962 pub spec: RingSpec,
964 pub lane_count: usize
966}
967
968pub struct ShmRingView {
976 region: ShmRegion,
977 metadata: ShmRingMetadata,
978 slot_stride: usize,
979 slots_offset: usize,
980 lane_stride: usize
981}
982
983impl std::fmt::Debug for ShmRingView {
984 fn fmt(
985 &self,
986 formatter: &mut std::fmt::Formatter<'_>
987 ) -> std::fmt::Result {
988 formatter.debug_struct("ShmRingView").field("metadata", &self.metadata).finish()
989 }
990}
991
992impl ShmRingView {
993 pub fn attach_existing_with_policy(
995 fleet_name: &str,
996 kind: u8,
997 policy: shm::ShmAccessPolicy
998 ) -> std::io::Result<Self> {
999 Self::attach_existing_for_uid_with_policy(
1000 fleet_name,
1001 kind,
1002 unsafe { libc::geteuid() },
1003 policy
1004 )
1005 }
1006
1007 pub fn attach_existing(
1009 fleet_name: &str,
1010 kind: u8
1011 ) -> std::io::Result<Self> {
1012 let uid = unsafe { libc::geteuid() };
1014 Self::attach_existing_for_uid(fleet_name, kind, uid)
1015 }
1016
1017 pub fn attach_existing_for_uid(
1022 fleet_name: &str,
1023 kind: u8,
1024 uid: u32
1025 ) -> std::io::Result<Self> {
1026 Self::attach_existing_for_uid_with_policy(
1027 fleet_name,
1028 kind,
1029 uid,
1030 shm::ShmAccessPolicy::default()
1031 )
1032 }
1033
1034 pub fn attach_existing_for_uid_with_policy(
1037 fleet_name: &str,
1038 kind: u8,
1039 uid: u32,
1040 policy: shm::ShmAccessPolicy
1041 ) -> std::io::Result<Self> {
1042 let name = shm::ring_segment_name_for_uid(fleet_name, kind, uid);
1043 let region = ShmRegion::open_existing_read_only(&name, HEADER_SIZE, uid, policy)?;
1044
1045 let header = unsafe { &*(region.as_ptr() as *const ShmRingHeader) };
1048 if header.magic != MAGIC {
1049 return Err(std::io::Error::new(
1050 std::io::ErrorKind::InvalidData,
1051 format!(
1052 "SHM segment {} has wrong magic 0x{:08X} (expected 0x{:08X})",
1053 name, header.magic, MAGIC
1054 )
1055 ));
1056 }
1057 if header.version != VERSION {
1058 return Err(std::io::Error::new(
1059 std::io::ErrorKind::InvalidData,
1060 format!("SHM segment {} version {} != local {}", name, header.version, VERSION)
1061 ));
1062 }
1063 if header.kind != kind {
1064 return Err(std::io::Error::new(
1065 std::io::ErrorKind::InvalidData,
1066 format!("SHM segment {} kind {} != requested {}", name, header.kind, kind)
1067 ));
1068 }
1069
1070 let topology = match header.topology {
1071 value if value == RingTopology::Shared as u8 => RingTopology::Shared,
1072 value if value == RingTopology::PerNode as u8 => RingTopology::PerNode,
1073 value if value == RingTopology::SharedOrdered as u8 => RingTopology::SharedOrdered,
1074 value => {
1075 return Err(std::io::Error::new(
1076 std::io::ErrorKind::InvalidData,
1077 format!("SHM segment {name} has unknown topology {value}")
1078 ));
1079 }
1080 };
1081 let capacity = usize::try_from(header.capacity).map_err(|_| {
1082 std::io::Error::new(
1083 std::io::ErrorKind::InvalidData,
1084 format!(
1085 "SHM segment {name} capacity {} does not fit this platform",
1086 header.capacity
1087 )
1088 )
1089 })?;
1090 let lane_count = usize::from(header.lane_count);
1091 if lane_count == 0 {
1092 return Err(std::io::Error::new(
1093 std::io::ErrorKind::InvalidData,
1094 format!("SHM segment {name} has zero lanes")
1095 ));
1096 }
1097 if topology != RingTopology::PerNode && lane_count != 1 {
1098 return Err(std::io::Error::new(
1099 std::io::ErrorKind::InvalidData,
1100 format!(
1101 "SHM segment {name} topology {topology:?} requires one lane, found {lane_count}"
1102 )
1103 ));
1104 }
1105 let spec =
1106 RingSpec { capacity, payload_capacity: header.payload_capacity as usize, topology };
1107 let fleet_capacity = if topology == RingTopology::PerNode { header.lane_count } else { 1 };
1108 let (slot_stride, slots_offset, expected_size) = checked_layout(spec, fleet_capacity)
1109 .map_err(|error| {
1110 std::io::Error::new(
1111 std::io::ErrorKind::InvalidData,
1112 format!("SHM segment {name} has invalid geometry: {error}")
1113 )
1114 })?;
1115 if header.slot_stride as usize != slot_stride {
1116 return Err(std::io::Error::new(
1117 std::io::ErrorKind::InvalidData,
1118 format!(
1119 "SHM segment {} slot stride {} != geometry {}",
1120 name, header.slot_stride, slot_stride
1121 )
1122 ));
1123 }
1124 if region.len() < expected_size {
1125 return Err(std::io::Error::new(
1126 std::io::ErrorKind::InvalidData,
1127 format!(
1128 "SHM segment {} size {} is smaller than persisted geometry {}",
1129 name,
1130 region.len(),
1131 expected_size
1132 )
1133 ));
1134 }
1135 let lane_stride = capacity
1136 .checked_mul(slot_stride)
1137 .ok_or_else(|| invalid_input(format!("SHM segment {name} lane stride overflow")))?;
1138
1139 Ok(Self {
1140 metadata: ShmRingMetadata {
1141 name,
1142 segment_size: region.len(),
1143 format_version: header.version,
1144 kind,
1145 spec,
1146 lane_count
1147 },
1148 region,
1149 slot_stride,
1150 slots_offset,
1151 lane_stride
1152 })
1153 }
1154
1155 pub fn attach_typed<T: crate::OrbitTyped>(fleet_name: &str) -> std::io::Result<Self> {
1157 let view = Self::attach_existing(fleet_name, T::KIND)?;
1158 if view.metadata.spec != T::RING_SPEC {
1159 return Err(std::io::Error::new(
1160 std::io::ErrorKind::InvalidData,
1161 format!(
1162 "OrbitTyped KIND {} declares {:?}; existing spec is {:?}",
1163 T::KIND,
1164 T::RING_SPEC,
1165 view.metadata.spec
1166 )
1167 ));
1168 }
1169 Ok(view)
1170 }
1171
1172 pub fn metadata(&self) -> &ShmRingMetadata {
1173 &self.metadata
1174 }
1175
1176 pub fn current_version(&self) -> u64 {
1178 self.header().version_counter.load(Ordering::Acquire)
1179 }
1180
1181 pub fn notification_generation(&self) -> u32 {
1183 self.header().notification_generation.load(Ordering::Acquire)
1184 }
1185
1186 pub fn lane(
1188 &self,
1189 lane: usize
1190 ) -> std::io::Result<ShmRingLaneView<'_>> {
1191 if lane >= self.metadata.lane_count {
1192 return Err(std::io::Error::new(
1193 std::io::ErrorKind::InvalidInput,
1194 format!("lane {lane} is outside SHM ring lane count {}", self.metadata.lane_count)
1195 ));
1196 }
1197 Ok(ShmRingLaneView { ring: self, lane })
1198 }
1199
1200 fn header(&self) -> &ShmRingHeader {
1201 unsafe { &*(self.region.as_ptr() as *const ShmRingHeader) }
1203 }
1204
1205 fn lane_header(
1206 &self,
1207 lane: usize
1208 ) -> &ShmLaneHeader {
1209 debug_assert!(lane < self.metadata.lane_count);
1210 unsafe {
1212 &*(self.region.as_ptr().add(HEADER_SIZE + lane * LANE_HEADER_SIZE)
1213 as *const ShmLaneHeader)
1214 }
1215 }
1216
1217 fn slot_ptr(
1218 &self,
1219 lane: usize,
1220 counter: u64
1221 ) -> *mut ShmSlotHeader {
1222 debug_assert!(lane < self.metadata.lane_count);
1223 let slot = (counter as usize) & (self.metadata.spec.capacity - 1);
1224 unsafe {
1227 self.region
1228 .as_ptr()
1229 .add(self.slots_offset + lane * self.lane_stride + slot * self.slot_stride)
1230 .cast::<ShmSlotHeader>()
1231 }
1232 }
1233
1234 fn read_lane_state_at(
1235 &self,
1236 lane: usize,
1237 counter: u64
1238 ) -> crate::ring::cursor::RingRead {
1239 use crate::ring::cursor::RingRead;
1240
1241 let slot_ptr = self.slot_ptr(lane, counter);
1242 let Some(expected_committed) =
1243 counter.checked_mul(2).and_then(|value| value.checked_add(2))
1244 else {
1245 return RingRead::Unavailable;
1246 };
1247
1248 for _ in 0..3 {
1249 let sequence = unsafe { &*slot_ptr }.seq.load(Ordering::Acquire);
1251 if sequence < expected_committed {
1252 return if self.metadata.spec.topology == RingTopology::Shared {
1253 RingRead::Pending
1254 } else {
1255 RingRead::Unavailable
1256 };
1257 }
1258 if sequence > expected_committed {
1259 return RingRead::Unavailable;
1260 }
1261 if let Some(frame) =
1263 unsafe { read_committed_frame(slot_ptr, self.metadata.spec.payload_capacity) }
1264 {
1265 let lane_matches = self.metadata.spec.topology != RingTopology::PerNode
1266 || usize::from(frame.id.node()) == lane;
1267 return if frame.id.kind() == self.metadata.kind
1268 && frame.id.counter() == counter
1269 && lane_matches
1270 {
1271 RingRead::Ready(frame)
1272 } else {
1273 RingRead::Unavailable
1274 };
1275 }
1276 }
1277 RingRead::Unavailable
1278 }
1279}
1280
1281#[derive(Clone, Copy, Debug)]
1283pub struct ShmRingLaneView<'a> {
1284 ring: &'a ShmRingView,
1285 lane: usize
1286}
1287
1288impl ShmRingLaneView<'_> {
1289 pub fn index(self) -> usize {
1290 self.lane
1291 }
1292
1293 pub fn head(self) -> u64 {
1295 self.ring.lane_header(self.lane).write_pos.load(Ordering::Acquire)
1296 }
1297
1298 pub fn retained_range(self) -> std::ops::Range<u64> {
1300 let head = self.head();
1301 head.saturating_sub(self.ring.metadata.spec.capacity as u64)..head
1302 }
1303
1304 pub fn read_at(
1306 self,
1307 counter: u64
1308 ) -> Option<Frame> {
1309 match self.ring.read_lane_state_at(self.lane, counter) {
1310 crate::ring::cursor::RingRead::Ready(frame) => Some(frame),
1311 crate::ring::cursor::RingRead::Pending | crate::ring::cursor::RingRead::Unavailable => {
1312 None
1313 }
1314 }
1315 }
1316
1317 pub fn read_head(self) -> Option<Frame> {
1319 let head = self.head();
1320 head.checked_sub(1).and_then(|counter| self.read_at(counter))
1321 }
1322}
1323
1324impl crate::ring::cursor::RingFrameSource for ShmRingLaneView<'_> {
1325 fn kind(&self) -> u8 {
1326 self.ring.metadata.kind
1327 }
1328
1329 fn head(&self) -> u64 {
1330 (*self).head()
1331 }
1332
1333 fn capacity(&self) -> usize {
1334 self.ring.metadata.spec.capacity
1335 }
1336
1337 fn read_at(
1338 &self,
1339 counter: u64
1340 ) -> Option<Frame> {
1341 (*self).read_at(counter)
1342 }
1343
1344 fn read_state_at(
1345 &self,
1346 counter: u64
1347 ) -> crate::ring::cursor::RingRead {
1348 self.ring.read_lane_state_at(self.lane, counter)
1349 }
1350}
1351
1352pub struct ShmRingRegistry {
1360 fleet_name: String,
1361 fleet_capacity: u16,
1362 policies: std::collections::BTreeMap<u8, shm::ShmAccessPolicy>,
1363 rings: dashmap::DashMap<u8, std::sync::Arc<ShmRing>>
1364}
1365
1366impl ShmRingRegistry {
1367 pub fn new(
1368 fleet_name: impl Into<String>,
1369 fleet_capacity: u16
1370 ) -> Self {
1371 Self::with_policies(fleet_name, fleet_capacity, [])
1372 }
1373
1374 pub fn with_policies(
1377 fleet_name: impl Into<String>,
1378 fleet_capacity: u16,
1379 policies: impl IntoIterator<Item = (u8, shm::ShmAccessPolicy)>
1380 ) -> Self {
1381 Self {
1382 fleet_name: fleet_name.into(),
1383 fleet_capacity,
1384 policies: policies.into_iter().collect(),
1385 rings: dashmap::DashMap::new()
1386 }
1387 }
1388
1389 pub fn get_or_create_for<T: crate::OrbitTyped>(
1393 &self
1394 ) -> std::io::Result<std::sync::Arc<ShmRing>> {
1395 if let Some(entry) = self.rings.get(&T::KIND) {
1396 if entry.spec() != T::RING_SPEC {
1397 return Err(std::io::Error::new(
1398 std::io::ErrorKind::InvalidData,
1399 format!(
1400 "OrbitTyped KIND {} was reused with ring spec {:?}; existing spec is {:?}",
1401 T::KIND,
1402 T::RING_SPEC,
1403 entry.spec()
1404 )
1405 ));
1406 }
1407 return Ok(entry.clone());
1408 }
1409 let ring = std::sync::Arc::new(ShmRing::open_or_create_for_fleet_with_policy(
1410 &self.fleet_name,
1411 T::KIND,
1412 T::RING_SPEC,
1413 self.fleet_capacity,
1414 self.policies.get(&T::KIND).copied().unwrap_or_default()
1415 )?);
1416 let entry = self.rings.entry(T::KIND).or_insert_with(|| ring.clone());
1417 if entry.spec() != T::RING_SPEC {
1418 return Err(std::io::Error::new(
1419 std::io::ErrorKind::InvalidData,
1420 format!(
1421 "OrbitTyped KIND {} raced with ring spec {:?}; installed spec is {:?}",
1422 T::KIND,
1423 T::RING_SPEC,
1424 entry.spec()
1425 )
1426 ));
1427 }
1428 Ok(entry.clone())
1429 }
1430
1431 pub fn lookup(
1434 &self,
1435 kind: u8
1436 ) -> Option<std::sync::Arc<ShmRing>> {
1437 self.rings.get(&kind).map(|e| e.clone())
1438 }
1439}
1440
1441unsafe fn read_committed_frame(
1449 slot_ptr: *mut ShmSlotHeader,
1450 payload_capacity: usize
1451) -> Option<Frame> {
1452 let slot = unsafe { &*slot_ptr };
1453 let seq_pre = slot.seq.load(Ordering::Acquire);
1454 if seq_pre == 0 {
1455 return None;
1457 }
1458 if seq_pre & 1 == 1 {
1459 return None;
1461 }
1462
1463 let id = NetId64::from_raw(slot.id.load(Ordering::Relaxed));
1467 let kind = slot.kind.load(Ordering::Relaxed);
1468 let ver = slot.ver.load(Ordering::Relaxed);
1469 let payload_len = slot.payload_len.load(Ordering::Relaxed) as usize;
1470 if payload_len > payload_capacity {
1471 return None;
1473 }
1474 let payload_src = unsafe { ShmRing::payload_ptr(slot_ptr) };
1479 let mut payload_buf = Vec::with_capacity(payload_len);
1480 for index in 0..payload_len {
1481 payload_buf.push(unsafe { (*payload_src.add(index)).load(Ordering::Relaxed) });
1482 }
1483
1484 fence(Ordering::Acquire);
1492 let seq_post = slot.seq.load(Ordering::Relaxed);
1493 if seq_pre != seq_post {
1494 return None;
1496 }
1497
1498 Some(Frame { id, kind, ver, payload: Bytes::from(payload_buf) })
1499}
1500
1501#[cfg(test)]
1502mod tests {
1503 use std::sync::atomic::{AtomicU64, Ordering};
1504
1505 use nix::sys::wait::{WaitStatus, waitpid};
1506 use nix::unistd::{ForkResult, fork};
1507
1508 use super::*;
1509 use crate::ring::cursor::{RingCursor, poll_ring};
1510
1511 #[test]
1512 fn cursor_retries_a_claimed_slot_after_it_commits() {
1513 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1514
1515 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1516 let fleet_name = format!("p{:x}{test_id:x}", std::process::id());
1517 let ring = ShmRing::open_or_create(&fleet_name, 199, RingSpec::new(4, 16))
1518 .expect("create test ring");
1519 ring.reset();
1520 let slot_ptr = ring.slot_ptr(0, 0);
1521
1522 ring.lane_header(0).write_pos.store(1, Ordering::Release);
1523 unsafe { &*slot_ptr }.seq.store(1, Ordering::Release);
1524
1525 let mut cursor = RingCursor::from_start();
1526 let pending = poll_ring(&ring, &mut cursor);
1527 assert!(pending.is_empty());
1528 assert_eq!(cursor.next_counter(), 0);
1529
1530 let payload = b"ready";
1531 unsafe {
1532 let slot = &*slot_ptr;
1533 slot.id.store(NetId64::make(199, NodeId::ZERO.get(), 0).raw(), Ordering::Relaxed);
1534 slot.ver.store(7, Ordering::Relaxed);
1535 slot.payload_len.store(payload.len() as u32, Ordering::Relaxed);
1536 slot.kind.store(1, Ordering::Relaxed);
1537 let bytes = ShmRing::payload_ptr(slot_ptr);
1538 for (index, byte) in payload.iter().enumerate() {
1539 (*bytes.add(index)).store(*byte, Ordering::Relaxed);
1540 }
1541 }
1542 unsafe { &*slot_ptr }.seq.store(2, Ordering::Release);
1543
1544 let committed = poll_ring(&ring, &mut cursor);
1545 assert_eq!(committed.frames.len(), 1);
1546 assert_eq!(&committed.frames[0].payload[..], payload);
1547 assert_eq!(cursor.next_counter(), 1);
1548
1549 ring.unlink().expect("unlink test ring");
1550 }
1551
1552 #[test]
1553 fn per_node_head_ignores_a_writer_that_dies_before_commit() {
1554 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1555
1556 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1557 let fleet_name = format!("d{:x}{test_id:x}", std::process::id());
1558 let spec = RingSpec::per_node(4, 16);
1559 let abandoned = ShmRing::open_or_create_for_fleet(&fleet_name, 198, spec, 2)
1560 .expect("create per-node test ring");
1561 abandoned.reset();
1562
1563 let slot_ptr = abandoned.slot_ptr(1, 0);
1564 unsafe { &*slot_ptr }.seq.store(1, Ordering::Release);
1565 assert_eq!(abandoned.lane_head(NodeId::new(1)), 0);
1566
1567 let replacement = ShmRing::open_or_create_for_fleet(&fleet_name, 198, spec, 2)
1568 .expect("replacement attaches");
1569 let id = replacement
1570 .write(NodeId::new(1), 1, 9, Bytes::from_static(b"recovered"))
1571 .expect("replacement commits");
1572
1573 assert_eq!(id.counter(), 0);
1574 assert_eq!(replacement.lane_head(NodeId::new(1)), 1);
1575 assert_eq!(&replacement.read(id).expect("frame visible").payload[..], b"recovered");
1576
1577 replacement.unlink().expect("unlink test ring");
1578 }
1579
1580 #[test]
1581 fn per_node_lane_serializes_concurrent_local_publishers() {
1582 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1583
1584 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1585 let fleet_name = format!("c{:x}{test_id:x}", std::process::id());
1586 let ring = std::sync::Arc::new(
1587 ShmRing::open_or_create_for_fleet(&fleet_name, 197, RingSpec::per_node(512, 0), 2)
1588 .expect("create concurrent writer ring")
1589 );
1590 ring.reset();
1591
1592 let mut writers = Vec::new();
1593 for _ in 0..4 {
1594 let ring = ring.clone();
1595 writers.push(std::thread::spawn(move || {
1596 (0..64)
1597 .map(|_| {
1598 ring.write(NodeId::new(1), 1, 0, Bytes::new()).expect("publish").counter()
1599 })
1600 .collect::<Vec<_>>()
1601 }));
1602 }
1603
1604 let mut counters = writers
1605 .into_iter()
1606 .flat_map(|writer| writer.join().expect("writer joins"))
1607 .collect::<Vec<_>>();
1608 counters.sort_unstable();
1609 assert_eq!(counters, (0..256).collect::<Vec<_>>());
1610 assert_eq!(ring.lane_head(NodeId::new(1)), 256);
1611
1612 ring.unlink().expect("unlink test ring");
1613 }
1614
1615 #[test]
1616 fn shared_ordered_recovers_after_a_locked_writer_process_dies() {
1617 static TEST_ID: AtomicU64 = AtomicU64::new(0);
1618
1619 let test_id = TEST_ID.fetch_add(1, Ordering::Relaxed);
1620 let fleet_name = format!("o{:x}{test_id:x}", std::process::id());
1621 let spec = RingSpec::shared_ordered(4, 16);
1622 let ring = ShmRing::open_or_create_for_fleet(&fleet_name, 196, spec, 2)
1623 .expect("create shared-ordered test ring");
1624 ring.reset();
1625
1626 match unsafe { fork() }.expect("fork test writer") {
1627 ForkResult::Child => {
1628 let _lock = ring.region.lock_exclusive().expect("child locks ring");
1632 let slot_ptr = ring.slot_ptr(0, 0);
1633 unsafe { &*slot_ptr }.seq.store(1, Ordering::Release);
1634
1635 unsafe { libc::_exit(0) };
1638 }
1639 ForkResult::Parent { child } => {
1640 let status = waitpid(child, None).expect("wait for abandoned writer");
1641 assert!(matches!(status, WaitStatus::Exited(_, 0)));
1642
1643 let replacement = ShmRing::open_or_create_for_fleet(&fleet_name, 196, spec, 2)
1648 .expect("replacement attaches");
1649 let recovered_lock = replacement
1650 .region
1651 .try_lock_exclusive()
1652 .expect("kernel released dead writer lock");
1653 drop(recovered_lock);
1654
1655 let id = replacement
1656 .write(NodeId::new(1), 1, 9, Bytes::from_static(b"recovered"))
1657 .expect("replacement commits counter zero");
1658 assert_eq!(id.counter(), 0);
1659 assert_eq!(replacement.head(), 1);
1660 assert_eq!(
1661 &replacement.read(id).expect("recovered frame").payload[..],
1662 b"recovered"
1663 );
1664
1665 replacement.unlink().expect("unlink shared-ordered test ring");
1666 }
1667 }
1668 }
1669}