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> {
144 let total = peer_directory_size();
145 let file = OpenOptions::new()
146 .read(true).write(true).create(true).truncate(true)
147 .open(path.as_ref())?;
148 file.set_len(total as u64)?;
149 let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
150 unsafe { init_dir_layout(mmap.as_mut_ptr()) };
151 let raw_ptr = mmap.as_mut_ptr();
152 Ok(Self { _backing: DirBacking::File(file, mmap), raw_ptr })
153 }
154
155 pub fn open(path: impl AsRef<Path>) -> Result<Self, RingError> {
158 let total = peer_directory_size();
159 let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
160 if (file.metadata()?.len() as usize) < total {
161 return Err(RingError::LayoutMismatch);
162 }
163 let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
164 let header = unsafe { &*(mmap.as_ptr() as *const DirHeader) };
165 if header.magic.load(AtomOrd::Acquire) != DIR_MAGIC {
166 return Err(RingError::LayoutMismatch);
167 }
168 let raw_ptr = mmap.as_ptr() as *mut u8;
169 Ok(Self { _backing: DirBacking::File(file, mmap), raw_ptr })
170 }
171
172 pub fn create_or_open_shm(name: &str) -> Result<Self, RingError> {
176 let mut shm = crate::shm_file::ShmFile::create_or_open_named(
177 name, peer_directory_size(),
178 ).map_err(|e| RingError::IoError(e.kind()))?;
179 if shm.len() < peer_directory_size() {
180 return Err(RingError::LayoutMismatch);
181 }
182 let raw_ptr = shm.as_mut_slice().as_mut_ptr();
183 let dir = Self { _backing: DirBacking::Shm(shm), raw_ptr };
184 let header = dir.header();
185 if header.magic
186 .compare_exchange(0, 1, AtomOrd::AcqRel, AtomOrd::Acquire)
187 .is_ok()
188 {
189 unsafe { init_dir_layout(raw_ptr) };
192 } else {
193 let mut spins = 0u32;
194 while dir.header().magic.load(AtomOrd::Acquire) != DIR_MAGIC {
195 std::hint::spin_loop();
196 spins += 1;
197 if spins > 100_000_000 {
198 return Err(RingError::LayoutMismatch);
199 }
200 }
201 }
202 Ok(dir)
203 }
204
205 #[inline]
206 fn header(&self) -> &DirHeader {
207 unsafe { &*(self.raw_ptr as *const DirHeader) }
208 }
209
210 #[inline]
211 fn bitmap_word(&self, base: usize, i: usize) -> &AtomicU64 {
212 unsafe { &*(self.raw_ptr.add(base + i * 8) as *const AtomicU64) }
213 }
214
215 #[inline]
216 fn owner_word(&self, ring: usize) -> &AtomicU64 {
217 debug_assert!(ring < PRODUCER_SLOT_CEILING);
218 unsafe { &*(self.raw_ptr.add(OFF_OWNERS + ring * 8) as *const AtomicU64) }
219 }
220
221 #[inline]
222 fn pid_word(&self, slot: usize) -> &AtomicU64 {
223 debug_assert!(slot < CONSUMER_SLOT_CEILING);
224 unsafe { &*(self.raw_ptr.add(OFF_C_PIDS + slot * 8) as *const AtomicU64) }
225 }
226
227 #[inline]
228 fn producer_pid_word(&self, slot: usize) -> &AtomicU64 {
229 debug_assert!(slot < PRODUCER_SLOT_CEILING);
230 unsafe { &*(self.raw_ptr.add(OFF_P_PIDS + slot * 8) as *const AtomicU64) }
231 }
232
233 #[inline]
236 pub fn epoch(&self) -> u64 {
237 self.header().epoch.load(AtomOrd::Acquire)
238 }
239
240 pub fn bump_epoch(&self) -> u64 {
242 self.header().epoch.fetch_add(1, AtomOrd::AcqRel) + 1
243 }
244
245 #[inline]
247 pub fn published(&self) -> usize {
248 self.header().published.load(AtomOrd::Acquire) as usize
249 }
250
251 pub fn publish_rings(&self, to: usize) {
255 self.header().published.fetch_max(to as u32, AtomOrd::AcqRel);
256 self.bump_epoch();
257 }
258
259 #[inline]
261 pub fn active_producers(&self) -> usize {
262 self.header().active_producers.load(AtomOrd::Acquire) as usize
263 }
264
265 #[inline]
267 pub fn active_consumers(&self) -> usize {
268 self.header().active_consumers.load(AtomOrd::Acquire) as usize
269 }
270
271 pub fn claim_producer_slot(&self) -> Option<usize> {
275 let slot = self.claim_bit(OFF_P_BITMAP, P_WORDS)?;
276 self.producer_pid_word(slot).store(std::process::id() as u64, AtomOrd::Release);
277 self.header().active_producers.fetch_add(1, AtomOrd::AcqRel);
278 self.bump_epoch();
279 Some(slot)
280 }
281
282 pub fn release_producer_slot(&self, slot: usize) {
285 if slot < PRODUCER_SLOT_CEILING {
286 self.producer_pid_word(slot).store(0, AtomOrd::Release);
287 }
288 if self.release_bit(OFF_P_BITMAP, P_WORDS, slot) {
289 self.header().active_producers.fetch_sub(1, AtomOrd::AcqRel);
290 self.bump_epoch();
291 }
292 }
293
294 pub fn reap_dead_peers(&self) {
300 for slot in self.claimed_slots(OFF_P_BITMAP, P_WORDS) {
301 let pid = self.producer_pid_word(slot).load(AtomOrd::Acquire);
302 if pid != 0 && pid != std::process::id() as u64
303 && !process_alive(pid as u32)
304 {
305 self.release_producer_slot(slot);
306 }
307 }
308 for slot in self.claimed_slots(OFF_C_BITMAP, C_WORDS) {
309 let pid = self.pid_word(slot).load(AtomOrd::Acquire);
310 if pid != 0 && pid != std::process::id() as u64
311 && !process_alive(pid as u32)
312 {
313 self.release_consumer_slot(slot);
314 }
315 }
316 }
317
318 fn claimed_slots(&self, base: usize, words: usize) -> Vec<usize> {
319 let mut slots = Vec::new();
320 for w in 0..words {
321 let mut word = self.bitmap_word(base, w).load(AtomOrd::Acquire);
322 while word != 0 {
323 let bit = word.trailing_zeros() as usize;
324 slots.push(w * 64 + bit);
325 word &= word - 1;
326 }
327 }
328 slots
329 }
330
331 pub fn claim_consumer_slot(&self) -> Option<usize> {
334 let slot = self.claim_bit(OFF_C_BITMAP, C_WORDS)?;
335 self.pid_word(slot).store(std::process::id() as u64, AtomOrd::Release);
336 self.header().active_consumers.fetch_add(1, AtomOrd::AcqRel);
337 self.bump_epoch();
338 Some(slot)
339 }
340
341 pub fn release_consumer_slot(&self, slot: usize) {
345 if slot >= CONSUMER_SLOT_CEILING {
346 return;
347 }
348 self.pid_word(slot).store(0, AtomOrd::Release);
349 if self.release_bit(OFF_C_BITMAP, C_WORDS, slot) {
350 self.header().active_consumers.fetch_sub(1, AtomOrd::AcqRel);
351 self.bump_epoch();
352 }
353 }
354
355 pub fn consumer_slot_claimed(&self, slot: usize) -> bool {
357 if slot >= CONSUMER_SLOT_CEILING {
358 return false;
359 }
360 let word = self.bitmap_word(OFF_C_BITMAP, slot / 64).load(AtomOrd::Acquire);
361 word & (1u64 << (slot % 64)) != 0
362 }
363
364 pub fn claimed_consumer_slots(&self) -> Vec<u16> {
368 let mut slots = Vec::new();
369 for w in 0..C_WORDS {
370 let mut word = self.bitmap_word(OFF_C_BITMAP, w).load(AtomOrd::Acquire);
371 while word != 0 {
372 let bit = word.trailing_zeros() as usize;
373 slots.push((w * 64 + bit) as u16);
374 word &= word - 1;
375 }
376 }
377 slots
378 }
379
380 #[inline]
382 pub fn ring_owner(&self, ring: usize) -> (u16, u16) {
383 unpack_owner(self.owner_word(ring).load(AtomOrd::Acquire))
384 }
385
386 pub fn try_claim_ring(&self, ring: usize, me: u16) -> bool {
391 let cur = pack_owner(OWNER_NONE, OWNER_NONE);
392 self.owner_word(ring)
393 .compare_exchange(cur, pack_owner(me, OWNER_NONE),
394 AtomOrd::AcqRel, AtomOrd::Acquire)
395 .is_ok()
396 }
397
398 pub fn request_handoff(&self, ring: usize, target: u16) {
402 let word = self.owner_word(ring);
403 let mut cur = word.load(AtomOrd::Acquire);
404 loop {
405 let (owner, _) = unpack_owner(cur);
406 if owner == target || owner == OWNER_NONE {
407 return; }
409 match word.compare_exchange(cur, pack_owner(owner, target),
410 AtomOrd::AcqRel, AtomOrd::Acquire) {
411 Ok(_) => return,
412 Err(actual) => cur = actual,
413 }
414 }
415 }
416
417 pub fn apply_handoff(&self, ring: usize, me: u16) -> Option<u16> {
421 let word = self.owner_word(ring);
422 let cur = word.load(AtomOrd::Acquire);
423 let (owner, pending) = unpack_owner(cur);
424 if owner != me || pending == OWNER_NONE {
425 return None;
426 }
427 match word.compare_exchange(cur, pack_owner(pending, OWNER_NONE),
428 AtomOrd::AcqRel, AtomOrd::Acquire) {
429 Ok(_) => {
430 self.bump_epoch();
431 Some(pending)
432 }
433 Err(_) => None,
434 }
435 }
436
437 pub fn transfer_ring(&self, ring: usize, me: u16, to: u16) {
440 let word = self.owner_word(ring);
441 let mut cur = word.load(AtomOrd::Acquire);
442 loop {
443 let (owner, _) = unpack_owner(cur);
444 if owner != me {
445 return;
446 }
447 match word.compare_exchange(cur, pack_owner(to, OWNER_NONE),
448 AtomOrd::AcqRel, AtomOrd::Acquire) {
449 Ok(_) => return,
450 Err(actual) => cur = actual,
451 }
452 }
453 }
454
455 pub fn try_takeover(&self, ring: usize, dead_owner: u16, me: u16) -> bool {
461 let slot = dead_owner as usize;
462 if self.consumer_slot_claimed(slot) {
463 let pid = self.pid_word(slot).load(AtomOrd::Acquire);
464 if pid == std::process::id() as u64 || pid == 0 || process_alive(pid as u32) {
465 return false;
466 }
467 }
468 let word = self.owner_word(ring);
469 let cur = word.load(AtomOrd::Acquire);
470 let (owner, _) = unpack_owner(cur);
471 if owner != dead_owner {
472 return false;
473 }
474 let swapped = word.compare_exchange(
475 cur, pack_owner(me, OWNER_NONE), AtomOrd::AcqRel, AtomOrd::Acquire,
476 ).is_ok();
477 if swapped {
478 self.bump_epoch();
479 }
480 swapped
481 }
482
483 fn claim_bit(&self, base: usize, words: usize) -> Option<usize> {
484 for w in 0..words {
485 let word = self.bitmap_word(base, w);
486 let mut cur = word.load(AtomOrd::Acquire);
487 loop {
488 if cur == u64::MAX {
489 break; }
491 let bit = (!cur).trailing_zeros() as usize;
492 match word.compare_exchange(cur, cur | (1u64 << bit),
493 AtomOrd::AcqRel, AtomOrd::Acquire) {
494 Ok(_) => return Some(w * 64 + bit),
495 Err(actual) => cur = actual,
496 }
497 }
498 }
499 None
500 }
501
502 fn release_bit(&self, base: usize, words: usize, slot: usize) -> bool {
503 if slot >= words * 64 {
504 return false;
505 }
506 let word = self.bitmap_word(base, slot / 64);
507 let mask = 1u64 << (slot % 64);
508 word.fetch_and(!mask, AtomOrd::AcqRel) & mask != 0
509 }
510}
511
512#[cfg(unix)]
516fn process_alive(pid: u32) -> bool {
517 let rc = unsafe { libc::kill(pid as libc::pid_t, 0) };
518 rc == 0
519 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
520}
521
522#[cfg(windows)]
523fn process_alive(pid: u32) -> bool {
524 use windows_sys::Win32::Foundation::{CloseHandle, STILL_ACTIVE};
525 use windows_sys::Win32::System::Threading::{
526 GetExitCodeProcess, OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION,
527 };
528 unsafe {
529 let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, pid);
530 if handle.is_null() {
531 return std::io::Error::last_os_error().raw_os_error()
534 == Some(5 );
535 }
536 let mut code: u32 = 0;
537 let ok = GetExitCodeProcess(handle, &mut code);
538 CloseHandle(handle);
539 ok != 0 && code == STILL_ACTIVE as u32
540 }
541}
542
543#[cfg(test)]
544mod tests {
545 use super::*;
546
547 #[test]
548 fn producer_slots_claim_release_recycle() {
549 let dir = PeerDirectory::create_anon().unwrap();
550 assert_eq!(dir.claim_producer_slot(), Some(0));
551 assert_eq!(dir.claim_producer_slot(), Some(1));
552 assert_eq!(dir.claim_producer_slot(), Some(2));
553 assert_eq!(dir.active_producers(), 3);
554 dir.release_producer_slot(1);
555 assert_eq!(dir.active_producers(), 2);
556 assert_eq!(dir.claim_producer_slot(), Some(1));
558 assert_eq!(dir.active_producers(), 3);
559 }
560
561 #[test]
562 fn epoch_bumps_on_every_topology_change() {
563 let dir = PeerDirectory::create_anon().unwrap();
564 let e0 = dir.epoch();
565 let p = dir.claim_producer_slot().unwrap();
566 assert!(dir.epoch() > e0);
567 let e1 = dir.epoch();
568 dir.release_producer_slot(p);
569 assert!(dir.epoch() > e1);
570 let e2 = dir.epoch();
571 dir.publish_rings(4);
572 assert!(dir.epoch() > e2);
573 assert_eq!(dir.published(), 4);
574 dir.publish_rings(2);
576 assert_eq!(dir.published(), 4);
577 }
578
579 #[test]
580 fn ring_ownership_claim_handoff_transfer() {
581 let dir = PeerDirectory::create_anon().unwrap();
582 assert_eq!(dir.ring_owner(0), (OWNER_NONE, OWNER_NONE));
583 assert!(dir.try_claim_ring(0, 3));
584 assert!(!dir.try_claim_ring(0, 5), "claimed ring must reject CAS");
585 assert_eq!(dir.ring_owner(0), (3, OWNER_NONE));
586
587 dir.request_handoff(0, 7);
589 assert_eq!(dir.ring_owner(0), (3, 7));
590 assert_eq!(dir.apply_handoff(0, 5), None, "non-owner cannot apply");
591 assert_eq!(dir.apply_handoff(0, 3), Some(7));
592 assert_eq!(dir.ring_owner(0), (7, OWNER_NONE));
593
594 dir.transfer_ring(0, 7, 2);
596 assert_eq!(dir.ring_owner(0), (2, OWNER_NONE));
597 }
598
599 #[test]
600 fn takeover_requires_dead_or_released_owner() {
601 let dir = PeerDirectory::create_anon().unwrap();
602 let slot = dir.claim_consumer_slot().unwrap() as u16;
603 assert!(dir.try_claim_ring(0, slot));
604 assert!(!dir.try_takeover(0, slot, 9));
606 dir.release_consumer_slot(slot as usize);
608 assert!(dir.try_takeover(0, slot, 9));
609 assert_eq!(dir.ring_owner(0), (9, OWNER_NONE));
610 }
611
612 #[test]
613 fn consumer_slots_record_pid_and_enumerate_densely() {
614 let dir = PeerDirectory::create_anon().unwrap();
615 let a = dir.claim_consumer_slot().unwrap();
616 let b = dir.claim_consumer_slot().unwrap();
617 let c = dir.claim_consumer_slot().unwrap();
618 assert_eq!((a, b, c), (0, 1, 2));
619 dir.release_consumer_slot(b);
620 assert_eq!(dir.claimed_consumer_slots(), vec![0u16, 2]);
621 assert_eq!(dir.active_consumers(), 2);
622 }
623
624 #[test]
625 fn file_backed_directory_shares_claims_across_handles() {
626 let base = std::env::temp_dir()
627 .join(format!("subetha_peerdir_{}", std::process::id()));
628 let path = base.with_extension("peers.bin");
629 let dir_a = PeerDirectory::create(&path).unwrap();
630 let p = dir_a.claim_producer_slot().unwrap();
631 dir_a.publish_rings(p + 1);
632
633 let dir_b = PeerDirectory::open(&path).unwrap();
634 assert_eq!(dir_b.active_producers(), 1);
635 assert_eq!(dir_b.published(), p + 1);
636 assert_eq!(dir_b.claim_producer_slot(), Some(p + 1));
637 assert_eq!(dir_a.active_producers(), 2, "claim visible both ways");
638
639 drop(dir_a);
640 drop(dir_b);
641 std::fs::remove_file(&path).ok();
642 }
643}