1use std::fs::{File, OpenOptions};
41use std::path::Path;
42use std::sync::atomic::{AtomicU32, AtomicU64, Ordering as AtomOrd};
43
44use memmap2::{MmapMut, MmapOptions};
45
46use crate::shared_ring::RingError;
47
48pub const PRODUCER_SLOT_CEILING: usize = 4096;
53
54pub const CONSUMER_SLOT_CEILING: usize = 256;
57
58pub const OWNER_NONE: u16 = u16::MAX;
60
61const DIR_MAGIC: u32 = 0x5045_4552; const P_WORDS: usize = PRODUCER_SLOT_CEILING / 64;
64const C_WORDS: usize = CONSUMER_SLOT_CEILING / 64;
65
66#[repr(C)]
67struct DirHeader {
68 magic: AtomicU32,
69 _rsvd: u32,
70 epoch: AtomicU64,
72 published: AtomicU32,
74 active_producers: AtomicU32,
75 active_consumers: AtomicU32,
76 _pad: u32,
77}
78
79const OFF_P_BITMAP: usize = std::mem::size_of::<DirHeader>();
80const OFF_C_BITMAP: usize = OFF_P_BITMAP + P_WORDS * 8;
81const OFF_C_PIDS: usize = OFF_C_BITMAP + C_WORDS * 8;
82const OFF_P_PIDS: usize = OFF_C_PIDS + CONSUMER_SLOT_CEILING * 8;
83const OFF_OWNERS: usize = OFF_P_PIDS + PRODUCER_SLOT_CEILING * 8;
84
85pub const fn peer_directory_size() -> usize {
87 OFF_OWNERS + PRODUCER_SLOT_CEILING * 8
88}
89
90#[allow(dead_code)]
92enum DirBacking {
93 Anon(MmapMut),
94 File(File, MmapMut),
95 Shm(crate::shm_file::ShmFile),
96}
97
98pub struct PeerDirectory {
100 _backing: DirBacking,
101 raw_ptr: *mut u8,
102}
103
104unsafe impl Send for PeerDirectory {}
105unsafe impl Sync for PeerDirectory {}
106
107unsafe fn init_dir_layout(ptr: *mut u8) {
110 unsafe {
111 std::ptr::write_bytes(ptr, 0, peer_directory_size());
112 let owners = ptr.add(OFF_OWNERS) as *mut u64;
114 let none = pack_owner(OWNER_NONE, OWNER_NONE);
115 for i in 0..PRODUCER_SLOT_CEILING {
116 std::ptr::write(owners.add(i), none);
117 }
118 let header = &*(ptr as *const DirHeader);
119 header.magic.store(DIR_MAGIC, AtomOrd::Release);
120 }
121}
122
123#[inline]
124const fn pack_owner(owner: u16, pending: u16) -> u64 {
125 (owner as u64) | ((pending as u64) << 16)
126}
127
128#[inline]
129const fn unpack_owner(word: u64) -> (u16, u16) {
130 (word as u16, (word >> 16) as u16)
131}
132
133impl PeerDirectory {
134 pub fn create_anon() -> Result<Self, RingError> {
136 let mut mmap = MmapOptions::new().len(peer_directory_size()).map_anon()?;
137 unsafe { init_dir_layout(mmap.as_mut_ptr()) };
138 let raw_ptr = mmap.as_mut_ptr();
139 Ok(Self { _backing: DirBacking::Anon(mmap), raw_ptr })
140 }
141
142 pub fn create(path: impl AsRef<Path>) -> Result<Self, RingError> {
147 let total = peer_directory_size();
148 let (file, mut mmap) = crate::mmf_attach::create_or_attach(
149 path.as_ref(),
150 total,
151 |ptr| unsafe { init_dir_layout(ptr) },
152 |ptr| {
153 let header = unsafe { &*(ptr as *const DirHeader) };
154 header.magic.load(AtomOrd::Acquire) == DIR_MAGIC
155 },
156 )
157 .map_err(|e| RingError::IoError(e.kind()))?;
158 let raw_ptr = mmap.as_mut_ptr();
159 Ok(Self { _backing: DirBacking::File(file, mmap), raw_ptr })
160 }
161
162 pub fn reset(path: impl AsRef<Path>) -> Result<Self, RingError> {
165 let total = peer_directory_size();
166 let (file, mut mmap) =
167 crate::mmf_attach::reset(path.as_ref(), total, |ptr| unsafe { init_dir_layout(ptr) })
168 .map_err(|e| RingError::IoError(e.kind()))?;
169 let raw_ptr = mmap.as_mut_ptr();
170 Ok(Self { _backing: DirBacking::File(file, mmap), raw_ptr })
171 }
172
173 pub fn open(path: impl AsRef<Path>) -> Result<Self, RingError> {
176 let total = peer_directory_size();
177 let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
178 if (file.metadata()?.len() as usize) < total {
179 return Err(RingError::LayoutMismatch);
180 }
181 let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
182 let header = unsafe { &*(mmap.as_ptr() as *const DirHeader) };
183 if header.magic.load(AtomOrd::Acquire) != DIR_MAGIC {
184 return Err(RingError::LayoutMismatch);
185 }
186 let raw_ptr = mmap.as_ptr() as *mut u8;
187 Ok(Self { _backing: DirBacking::File(file, mmap), raw_ptr })
188 }
189
190 pub fn create_or_open_shm(name: &str) -> Result<Self, RingError> {
194 let mut shm = crate::shm_file::ShmFile::create_or_open_named(
195 name, peer_directory_size(),
196 ).map_err(|e| RingError::IoError(e.kind()))?;
197 if shm.len() < peer_directory_size() {
198 return Err(RingError::LayoutMismatch);
199 }
200 let raw_ptr = shm.as_mut_slice().as_mut_ptr();
201 let dir = Self { _backing: DirBacking::Shm(shm), raw_ptr };
202 let header = dir.header();
203 if header.magic
204 .compare_exchange(0, 1, AtomOrd::AcqRel, AtomOrd::Acquire)
205 .is_ok()
206 {
207 unsafe { init_dir_layout(raw_ptr) };
210 } else {
211 let mut spins = 0u32;
212 while dir.header().magic.load(AtomOrd::Acquire) != DIR_MAGIC {
213 std::hint::spin_loop();
214 spins += 1;
215 if spins > 100_000_000 {
216 return Err(RingError::LayoutMismatch);
217 }
218 }
219 }
220 Ok(dir)
221 }
222
223 #[inline]
224 fn header(&self) -> &DirHeader {
225 unsafe { &*(self.raw_ptr as *const DirHeader) }
226 }
227
228 #[inline]
229 fn bitmap_word(&self, base: usize, i: usize) -> &AtomicU64 {
230 unsafe { &*(self.raw_ptr.add(base + i * 8) as *const AtomicU64) }
231 }
232
233 #[inline]
234 fn owner_word(&self, ring: usize) -> &AtomicU64 {
235 debug_assert!(ring < PRODUCER_SLOT_CEILING);
236 unsafe { &*(self.raw_ptr.add(OFF_OWNERS + ring * 8) as *const AtomicU64) }
237 }
238
239 #[inline]
240 fn pid_word(&self, slot: usize) -> &AtomicU64 {
241 debug_assert!(slot < CONSUMER_SLOT_CEILING);
242 unsafe { &*(self.raw_ptr.add(OFF_C_PIDS + slot * 8) as *const AtomicU64) }
243 }
244
245 #[inline]
246 fn producer_pid_word(&self, slot: usize) -> &AtomicU64 {
247 debug_assert!(slot < PRODUCER_SLOT_CEILING);
248 unsafe { &*(self.raw_ptr.add(OFF_P_PIDS + slot * 8) as *const AtomicU64) }
249 }
250
251 #[inline]
254 pub fn epoch(&self) -> u64 {
255 self.header().epoch.load(AtomOrd::Acquire)
256 }
257
258 pub fn bump_epoch(&self) -> u64 {
260 self.header().epoch.fetch_add(1, AtomOrd::AcqRel) + 1
261 }
262
263 #[inline]
265 pub fn published(&self) -> usize {
266 self.header().published.load(AtomOrd::Acquire) as usize
267 }
268
269 pub fn publish_rings(&self, to: usize) {
273 self.header().published.fetch_max(to as u32, AtomOrd::AcqRel);
274 self.bump_epoch();
275 }
276
277 #[inline]
279 pub fn active_producers(&self) -> usize {
280 self.header().active_producers.load(AtomOrd::Acquire) as usize
281 }
282
283 #[inline]
285 pub fn active_consumers(&self) -> usize {
286 self.header().active_consumers.load(AtomOrd::Acquire) as usize
287 }
288
289 pub fn claim_producer_slot(&self) -> Option<usize> {
293 let slot = self.claim_bit(OFF_P_BITMAP, P_WORDS)?;
294 self.producer_pid_word(slot).store(std::process::id() as u64, AtomOrd::Release);
295 self.header().active_producers.fetch_add(1, AtomOrd::AcqRel);
296 self.bump_epoch();
297 Some(slot)
298 }
299
300 pub fn release_producer_slot(&self, slot: usize) {
303 if slot < PRODUCER_SLOT_CEILING {
304 self.producer_pid_word(slot).store(0, AtomOrd::Release);
305 }
306 if self.release_bit(OFF_P_BITMAP, P_WORDS, slot) {
307 self.header().active_producers.fetch_sub(1, AtomOrd::AcqRel);
308 self.bump_epoch();
309 }
310 }
311
312 pub fn reap_dead_peers(&self) {
318 for slot in self.claimed_slots(OFF_P_BITMAP, P_WORDS) {
319 let pid = self.producer_pid_word(slot).load(AtomOrd::Acquire);
320 if pid != 0 && pid != std::process::id() as u64
321 && !process_alive(pid as u32)
322 {
323 self.release_producer_slot(slot);
324 }
325 }
326 for slot in self.claimed_slots(OFF_C_BITMAP, C_WORDS) {
327 let pid = self.pid_word(slot).load(AtomOrd::Acquire);
328 if pid != 0 && pid != std::process::id() as u64
329 && !process_alive(pid as u32)
330 {
331 self.release_consumer_slot(slot);
332 }
333 }
334 }
335
336 fn claimed_slots(&self, base: usize, words: usize) -> Vec<usize> {
337 let mut slots = Vec::new();
338 for w in 0..words {
339 let mut word = self.bitmap_word(base, w).load(AtomOrd::Acquire);
340 while word != 0 {
341 let bit = word.trailing_zeros() as usize;
342 slots.push(w * 64 + bit);
343 word &= word - 1;
344 }
345 }
346 slots
347 }
348
349 pub fn claim_consumer_slot(&self) -> Option<usize> {
352 let slot = self.claim_bit(OFF_C_BITMAP, C_WORDS)?;
353 self.pid_word(slot).store(std::process::id() as u64, AtomOrd::Release);
354 self.header().active_consumers.fetch_add(1, AtomOrd::AcqRel);
355 self.bump_epoch();
356 Some(slot)
357 }
358
359 pub fn release_consumer_slot(&self, slot: usize) {
363 if slot >= CONSUMER_SLOT_CEILING {
364 return;
365 }
366 self.pid_word(slot).store(0, AtomOrd::Release);
367 if self.release_bit(OFF_C_BITMAP, C_WORDS, slot) {
368 self.header().active_consumers.fetch_sub(1, AtomOrd::AcqRel);
369 self.bump_epoch();
370 }
371 }
372
373 pub fn consumer_slot_claimed(&self, slot: usize) -> bool {
375 if slot >= CONSUMER_SLOT_CEILING {
376 return false;
377 }
378 let word = self.bitmap_word(OFF_C_BITMAP, slot / 64).load(AtomOrd::Acquire);
379 word & (1u64 << (slot % 64)) != 0
380 }
381
382 pub fn claimed_consumer_slots(&self) -> Vec<u16> {
386 let mut slots = Vec::new();
387 for w in 0..C_WORDS {
388 let mut word = self.bitmap_word(OFF_C_BITMAP, w).load(AtomOrd::Acquire);
389 while word != 0 {
390 let bit = word.trailing_zeros() as usize;
391 slots.push((w * 64 + bit) as u16);
392 word &= word - 1;
393 }
394 }
395 slots
396 }
397
398 #[inline]
400 pub fn ring_owner(&self, ring: usize) -> (u16, u16) {
401 unpack_owner(self.owner_word(ring).load(AtomOrd::Acquire))
402 }
403
404 pub fn try_claim_ring(&self, ring: usize, me: u16) -> bool {
409 let cur = pack_owner(OWNER_NONE, OWNER_NONE);
410 self.owner_word(ring)
411 .compare_exchange(cur, pack_owner(me, OWNER_NONE),
412 AtomOrd::AcqRel, AtomOrd::Acquire)
413 .is_ok()
414 }
415
416 pub fn request_handoff(&self, ring: usize, target: u16) {
420 let word = self.owner_word(ring);
421 let mut cur = word.load(AtomOrd::Acquire);
422 loop {
423 let (owner, _) = unpack_owner(cur);
424 if owner == target || owner == OWNER_NONE {
425 return; }
427 match word.compare_exchange(cur, pack_owner(owner, target),
428 AtomOrd::AcqRel, AtomOrd::Acquire) {
429 Ok(_) => return,
430 Err(actual) => cur = actual,
431 }
432 }
433 }
434
435 pub fn apply_handoff(&self, ring: usize, me: u16) -> Option<u16> {
439 let word = self.owner_word(ring);
440 let cur = word.load(AtomOrd::Acquire);
441 let (owner, pending) = unpack_owner(cur);
442 if owner != me || pending == OWNER_NONE {
443 return None;
444 }
445 match word.compare_exchange(cur, pack_owner(pending, OWNER_NONE),
446 AtomOrd::AcqRel, AtomOrd::Acquire) {
447 Ok(_) => {
448 self.bump_epoch();
449 Some(pending)
450 }
451 Err(_) => None,
452 }
453 }
454
455 pub fn transfer_ring(&self, ring: usize, me: u16, to: u16) {
458 let word = self.owner_word(ring);
459 let mut cur = word.load(AtomOrd::Acquire);
460 loop {
461 let (owner, _) = unpack_owner(cur);
462 if owner != me {
463 return;
464 }
465 match word.compare_exchange(cur, pack_owner(to, OWNER_NONE),
466 AtomOrd::AcqRel, AtomOrd::Acquire) {
467 Ok(_) => return,
468 Err(actual) => cur = actual,
469 }
470 }
471 }
472
473 pub fn try_takeover(&self, ring: usize, dead_owner: u16, me: u16) -> bool {
479 let slot = dead_owner as usize;
480 if self.consumer_slot_claimed(slot) {
481 let pid = self.pid_word(slot).load(AtomOrd::Acquire);
482 if pid == std::process::id() as u64 || pid == 0 || process_alive(pid as u32) {
483 return false;
484 }
485 }
486 let word = self.owner_word(ring);
487 let cur = word.load(AtomOrd::Acquire);
488 let (owner, _) = unpack_owner(cur);
489 if owner != dead_owner {
490 return false;
491 }
492 let swapped = word.compare_exchange(
493 cur, pack_owner(me, OWNER_NONE), AtomOrd::AcqRel, AtomOrd::Acquire,
494 ).is_ok();
495 if swapped {
496 self.bump_epoch();
497 }
498 swapped
499 }
500
501 fn claim_bit(&self, base: usize, words: usize) -> Option<usize> {
502 for w in 0..words {
503 let word = self.bitmap_word(base, w);
504 let mut cur = word.load(AtomOrd::Acquire);
505 loop {
506 if cur == u64::MAX {
507 break; }
509 let bit = (!cur).trailing_zeros() as usize;
510 match word.compare_exchange(cur, cur | (1u64 << bit),
511 AtomOrd::AcqRel, AtomOrd::Acquire) {
512 Ok(_) => return Some(w * 64 + bit),
513 Err(actual) => cur = actual,
514 }
515 }
516 }
517 None
518 }
519
520 fn release_bit(&self, base: usize, words: usize, slot: usize) -> bool {
521 if slot >= words * 64 {
522 return false;
523 }
524 let word = self.bitmap_word(base, slot / 64);
525 let mask = 1u64 << (slot % 64);
526 word.fetch_and(!mask, AtomOrd::AcqRel) & mask != 0
527 }
528}
529
530#[cfg(unix)]
534fn process_alive(pid: u32) -> bool {
535 let rc = unsafe { libc::kill(pid as libc::pid_t, 0) };
536 rc == 0
537 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
538}
539
540#[cfg(windows)]
541fn process_alive(pid: u32) -> bool {
542 use windows_sys::Win32::Foundation::{CloseHandle, STILL_ACTIVE};
543 use windows_sys::Win32::System::Threading::{
544 GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION,
545 };
546 unsafe {
547 let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
548 if handle.is_null() {
549 return std::io::Error::last_os_error().raw_os_error()
552 == Some(5 );
553 }
554 let mut code: u32 = 0;
555 let ok = GetExitCodeProcess(handle, &mut code);
556 CloseHandle(handle);
557 ok != 0 && code == STILL_ACTIVE as u32
558 }
559}
560
561#[cfg(test)]
562mod tests {
563 use super::*;
564
565 #[test]
566 fn producer_slots_claim_release_recycle() {
567 let dir = PeerDirectory::create_anon().unwrap();
568 assert_eq!(dir.claim_producer_slot(), Some(0));
569 assert_eq!(dir.claim_producer_slot(), Some(1));
570 assert_eq!(dir.claim_producer_slot(), Some(2));
571 assert_eq!(dir.active_producers(), 3);
572 dir.release_producer_slot(1);
573 assert_eq!(dir.active_producers(), 2);
574 assert_eq!(dir.claim_producer_slot(), Some(1));
576 assert_eq!(dir.active_producers(), 3);
577 }
578
579 #[test]
580 fn epoch_bumps_on_every_topology_change() {
581 let dir = PeerDirectory::create_anon().unwrap();
582 let e0 = dir.epoch();
583 let p = dir.claim_producer_slot().unwrap();
584 assert!(dir.epoch() > e0);
585 let e1 = dir.epoch();
586 dir.release_producer_slot(p);
587 assert!(dir.epoch() > e1);
588 let e2 = dir.epoch();
589 dir.publish_rings(4);
590 assert!(dir.epoch() > e2);
591 assert_eq!(dir.published(), 4);
592 dir.publish_rings(2);
594 assert_eq!(dir.published(), 4);
595 }
596
597 #[test]
598 fn ring_ownership_claim_handoff_transfer() {
599 let dir = PeerDirectory::create_anon().unwrap();
600 assert_eq!(dir.ring_owner(0), (OWNER_NONE, OWNER_NONE));
601 assert!(dir.try_claim_ring(0, 3));
602 assert!(!dir.try_claim_ring(0, 5), "claimed ring must reject CAS");
603 assert_eq!(dir.ring_owner(0), (3, OWNER_NONE));
604
605 dir.request_handoff(0, 7);
607 assert_eq!(dir.ring_owner(0), (3, 7));
608 assert_eq!(dir.apply_handoff(0, 5), None, "non-owner cannot apply");
609 assert_eq!(dir.apply_handoff(0, 3), Some(7));
610 assert_eq!(dir.ring_owner(0), (7, OWNER_NONE));
611
612 dir.transfer_ring(0, 7, 2);
614 assert_eq!(dir.ring_owner(0), (2, OWNER_NONE));
615 }
616
617 #[test]
618 fn takeover_requires_dead_or_released_owner() {
619 let dir = PeerDirectory::create_anon().unwrap();
620 let slot = dir.claim_consumer_slot().unwrap() as u16;
621 assert!(dir.try_claim_ring(0, slot));
622 assert!(!dir.try_takeover(0, slot, 9));
624 dir.release_consumer_slot(slot as usize);
626 assert!(dir.try_takeover(0, slot, 9));
627 assert_eq!(dir.ring_owner(0), (9, OWNER_NONE));
628 }
629
630 #[test]
631 fn consumer_slots_record_pid_and_enumerate_densely() {
632 let dir = PeerDirectory::create_anon().unwrap();
633 let a = dir.claim_consumer_slot().unwrap();
634 let b = dir.claim_consumer_slot().unwrap();
635 let c = dir.claim_consumer_slot().unwrap();
636 assert_eq!((a, b, c), (0, 1, 2));
637 dir.release_consumer_slot(b);
638 assert_eq!(dir.claimed_consumer_slots(), vec![0u16, 2]);
639 assert_eq!(dir.active_consumers(), 2);
640 }
641
642 #[test]
643 fn file_backed_directory_shares_claims_across_handles() {
644 let base = std::env::temp_dir()
645 .join(format!("subetha_peerdir_{}", std::process::id()));
646 let path = base.with_extension("peers.bin");
647 let dir_a = PeerDirectory::create(&path).unwrap();
648 let p = dir_a.claim_producer_slot().unwrap();
649 dir_a.publish_rings(p + 1);
650
651 let dir_b = PeerDirectory::open(&path).unwrap();
652 assert_eq!(dir_b.active_producers(), 1);
653 assert_eq!(dir_b.published(), p + 1);
654 assert_eq!(dir_b.claim_producer_slot(), Some(p + 1));
655 assert_eq!(dir_a.active_producers(), 2, "claim visible both ways");
656
657 drop(dir_a);
658 drop(dir_b);
659 std::fs::remove_file(&path).ok();
660 }
661}