subetha_cxc/
heartbeat.rs

1//! Per-process heartbeat slots stored in an MMF.
2//!
3//! Each participating process owns one [`HeartbeatSlot`] in the
4//! shared table. On each scan tick, the process advances its slot's
5//! `last_seen_epoch`. A watchdog (separate module) compares the
6//! slot's epoch against a global `epoch` counter; if the process
7//! hasn't advanced its heartbeat within the configured grace, its
8//! work is presumed dead and reclaimed.
9//!
10//! Layout:
11//! ```text
12//! +-----------------------------+
13//! | HeartbeatHeader (64B)       |
14//! |   - magic, capacity, epoch  |
15//! +-----------------------------+
16//! | HeartbeatSlot[0]  (64B)     |
17//! | HeartbeatSlot[1]  (64B)     |
18//! | ...                         |
19//! | HeartbeatSlot[N - 1]        |
20//! +-----------------------------+
21//! ```
22//!
23//! Each slot is one cache line so cross-process writes to different
24//! slots never false-share.
25
26use std::fs::{File, OpenOptions};
27use std::path::Path;
28use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
29
30use memmap2::{MmapMut, MmapOptions};
31
32pub const HEARTBEAT_MAGIC: u64 = 0x4150_4D46_4842_4154;
33
34/// Unused slot ID (no pid).
35pub const EMPTY_PID: u32 = 0;
36
37/// Maximum number of in-flight work items a slot tracks. Each bit in
38/// `in_flight_bitmap` represents one work unit; on failover those
39/// bits are reclaimable.
40pub const IN_FLIGHT_SLOTS: usize = 64;
41
42#[repr(C, align(64))]
43pub struct HeartbeatHeader {
44    pub magic: u64,
45    pub capacity: u64,
46    /// Global epoch counter; the watchdog advances this each scan.
47    pub epoch: AtomicU64,
48    _reserved: [u8; 40],
49}
50
51#[repr(C, align(64))]
52pub struct HeartbeatSlot {
53    /// Owning process id. 0 = vacant.
54    pub pid: AtomicU32,
55    /// Sequence-lock generation; bumped on each meaningful write so
56    /// readers can detect torn writes (read with seqlock retry).
57    pub seq_version: AtomicU32,
58    /// Last global epoch at which this process incremented its
59    /// heartbeat. Watchdog reclaims when `global.epoch -
60    /// last_seen_epoch > grace_epochs`.
61    pub last_seen_epoch: AtomicU64,
62    /// Bitmap of work units currently assigned to this process.
63    /// Watchdog reclaims set bits on failover.
64    pub in_flight_bitmap: AtomicU64,
65    /// Process role: 0 = worker, 1 = coordinator.
66    pub role: AtomicU32,
67    _pad: [u8; 36],
68}
69
70/// Total file size for a heartbeat table with `capacity` slots.
71pub const fn heartbeat_file_size(capacity: usize) -> usize {
72    std::mem::size_of::<HeartbeatHeader>() + capacity * std::mem::size_of::<HeartbeatSlot>()
73}
74
75/// Cross-process heartbeat registry. Each process opens this and
76/// reserves one slot via [`HeartbeatTable::register`].
77pub struct HeartbeatTable {
78    _file: File,
79    mmap: MmapMut,
80    capacity: usize,
81    header_sidecar: subetha_core::HandshakeHeader,
82    ring_sidecar: Box<subetha_core::ObservationRing>,
83}
84
85unsafe impl Send for HeartbeatTable {}
86unsafe impl Sync for HeartbeatTable {}
87
88impl subetha_sidecar::AdaptiveInstance for HeartbeatTable {
89    fn header(&self) -> &subetha_core::HandshakeHeader { &self.header_sidecar }
90    fn ring(&self) -> &subetha_core::ObservationRing { &self.ring_sidecar }
91    fn make_policy(&self) -> Box<dyn subetha_sidecar::Policy> {
92        Box::new(subetha_sidecar::NoMigrationPolicy)
93    }
94}
95
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub enum HeartbeatError {
98    LayoutMismatch,
99    TableFull,
100    IoError(std::io::ErrorKind),
101}
102
103impl From<std::io::Error> for HeartbeatError {
104    fn from(e: std::io::Error) -> Self { Self::IoError(e.kind()) }
105}
106
107impl HeartbeatTable {
108    /// Obtain the table at `path`, initializing it only when the path does
109    /// not yet exist. Attaching leaves live registrations in place; a table
110    /// built with a different capacity is a `LayoutMismatch`. Use
111    /// [`reset`](Self::reset) to deliberately clear it.
112    pub fn create(path: impl AsRef<Path>, capacity: usize) -> Result<Self, HeartbeatError> {
113        assert!(capacity >= 1);
114        let total = heartbeat_file_size(capacity);
115        let (file, mmap) = crate::mmf_attach::create_or_attach(
116            path.as_ref(),
117            total,
118            |ptr| unsafe { Self::init_region(ptr, capacity) },
119            |ptr| unsafe { (*(ptr as *const HeartbeatHeader)).magic == HEARTBEAT_MAGIC },
120        )?;
121        let hdr = unsafe { &*(mmap.as_ptr() as *const HeartbeatHeader) };
122        if hdr.capacity != capacity as u64 {
123            return Err(HeartbeatError::LayoutMismatch);
124        }
125        Ok(Self {
126            _file: file, mmap, capacity,
127            header_sidecar: subetha_core::HandshakeHeader::new(),
128            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
129        })
130    }
131
132    /// Reinitialise the table at `path`, discarding every live registration.
133    /// For a caller that knows it owns the path.
134    pub fn reset(path: impl AsRef<Path>, capacity: usize) -> Result<Self, HeartbeatError> {
135        assert!(capacity >= 1);
136        let total = heartbeat_file_size(capacity);
137        let (file, mmap) = crate::mmf_attach::reset(path.as_ref(), total, |ptr| unsafe {
138            Self::init_region(ptr, capacity)
139        })?;
140        Ok(Self {
141            _file: file, mmap, capacity,
142            header_sidecar: subetha_core::HandshakeHeader::new(),
143            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
144        })
145    }
146
147    /// Lay out a fresh table: slots first, magic last, because attachers spin
148    /// on the magic.
149    ///
150    /// # Safety
151    /// `ptr` addresses at least `heartbeat_file_size(capacity)` writable
152    /// zeroed bytes.
153    unsafe fn init_region(ptr: *mut u8, capacity: usize) {
154        let hdr_ptr = ptr as *mut HeartbeatHeader;
155        unsafe {
156            std::ptr::write(hdr_ptr, HeartbeatHeader {
157                magic: 0,
158                capacity: capacity as u64,
159                epoch: AtomicU64::new(0),
160                _reserved: [0; 40],
161            });
162            let slots_base = ptr.add(std::mem::size_of::<HeartbeatHeader>());
163            for i in 0..capacity {
164                let slot_ptr =
165                    slots_base.add(i * std::mem::size_of::<HeartbeatSlot>()) as *mut HeartbeatSlot;
166                std::ptr::write(slot_ptr, HeartbeatSlot {
167                    pid: AtomicU32::new(EMPTY_PID),
168                    seq_version: AtomicU32::new(0),
169                    last_seen_epoch: AtomicU64::new(0),
170                    in_flight_bitmap: AtomicU64::new(0),
171                    role: AtomicU32::new(0),
172                    _pad: [0; 36],
173                });
174            }
175            std::ptr::write_volatile(std::ptr::addr_of_mut!((*hdr_ptr).magic), HEARTBEAT_MAGIC);
176        }
177    }
178
179    pub fn open(path: impl AsRef<Path>, expected_capacity: usize) -> Result<Self, HeartbeatError> {
180        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
181        let total = heartbeat_file_size(expected_capacity);
182        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
183        let header = unsafe { &*(mmap.as_ptr() as *const HeartbeatHeader) };
184        if header.magic != HEARTBEAT_MAGIC || header.capacity != expected_capacity as u64 {
185            return Err(HeartbeatError::LayoutMismatch);
186        }
187        Ok(Self {
188            _file: file, mmap, capacity: expected_capacity,
189            header_sidecar: subetha_core::HandshakeHeader::new(),
190            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
191        })
192    }
193
194    pub fn capacity(&self) -> usize { self.capacity }
195
196    pub fn header(&self) -> &HeartbeatHeader {
197        unsafe { &*(self.mmap.as_ptr() as *const HeartbeatHeader) }
198    }
199
200    fn slot(&self, idx: usize) -> &HeartbeatSlot {
201        let base = unsafe {
202            self.mmap.as_ptr().add(std::mem::size_of::<HeartbeatHeader>())
203        };
204        unsafe {
205            &*(base.add(idx * std::mem::size_of::<HeartbeatSlot>()) as *const HeartbeatSlot)
206        }
207    }
208
209    /// Register the current process. Returns the slot index. CAS-claim
210    /// of the first empty slot.
211    pub fn register(&self, pid: u32) -> Result<usize, HeartbeatError> {
212        for i in 0..self.capacity {
213            let slot = self.slot(i);
214            if slot.pid.compare_exchange(
215                EMPTY_PID, pid, Ordering::AcqRel, Ordering::Acquire,
216            ).is_ok() {
217                slot.seq_version.fetch_add(1, Ordering::Release);
218                slot.last_seen_epoch.store(
219                    self.header().epoch.load(Ordering::Acquire),
220                    Ordering::Release,
221                );
222                slot.in_flight_bitmap.store(0, Ordering::Release);
223                slot.role.store(0, Ordering::Release);
224                slot.seq_version.fetch_add(1, Ordering::Release);
225                self.ring_sidecar
226                    .push_op(crate::sidecar_ops::liveness::OP_REGISTER, 0);
227                return Ok(i);
228            }
229        }
230        self.ring_sidecar
231            .push_op(crate::sidecar_ops::liveness::OP_REGISTER, 1);
232        Err(HeartbeatError::TableFull)
233    }
234
235    /// Release the slot at `idx`. Call before process exit.
236    pub fn unregister(&self, idx: usize) {
237        let slot = self.slot(idx);
238        slot.seq_version.fetch_add(1, Ordering::Release);
239        slot.in_flight_bitmap.store(0, Ordering::Release);
240        slot.pid.store(EMPTY_PID, Ordering::Release);
241        slot.seq_version.fetch_add(1, Ordering::Release);
242    }
243
244    /// Heartbeat: advance this slot's `last_seen_epoch` to match the
245    /// global epoch. Call once per scan tick.
246    pub fn beat(&self, idx: usize) {
247        let global = self.header().epoch.load(Ordering::Acquire);
248        let slot = self.slot(idx);
249        slot.last_seen_epoch.store(global, Ordering::Release);
250        self.ring_sidecar
251            .push_op(crate::sidecar_ops::liveness::OP_BEAT, 0);
252    }
253
254    /// Advance the global epoch. Watchdog calls this once per scan
255    /// interval. Returns the new epoch value.
256    pub fn tick_global_epoch(&self) -> u64 {
257        let v = self.header().epoch.fetch_add(1, Ordering::AcqRel) + 1;
258        self.ring_sidecar
259            .push_op(crate::sidecar_ops::liveness::OP_TICK_EPOCH, 0);
260        v
261    }
262
263    pub fn global_epoch(&self) -> u64 {
264        self.header().epoch.load(Ordering::Acquire)
265    }
266
267    /// Mark a work unit as in-flight for `slot_idx`.
268    pub fn mark_in_flight(&self, slot_idx: usize, bit: u8) {
269        debug_assert!((bit as usize) < IN_FLIGHT_SLOTS);
270        let slot = self.slot(slot_idx);
271        slot.in_flight_bitmap.fetch_or(1u64 << bit, Ordering::AcqRel);
272    }
273
274    pub fn clear_in_flight(&self, slot_idx: usize, bit: u8) {
275        debug_assert!((bit as usize) < IN_FLIGHT_SLOTS);
276        let slot = self.slot(slot_idx);
277        slot.in_flight_bitmap.fetch_and(!(1u64 << bit), Ordering::AcqRel);
278    }
279
280    /// Snapshot a slot via SeqLock retry. Returns `None` if the slot
281    /// is vacant.
282    pub fn snapshot(&self, idx: usize) -> Option<HeartbeatSnapshot> {
283        let slot = self.slot(idx);
284        loop {
285            let v1 = slot.seq_version.load(Ordering::Acquire);
286            if v1 & 1 != 0 { continue; }  // writer in progress
287            let pid = slot.pid.load(Ordering::Acquire);
288            let last = slot.last_seen_epoch.load(Ordering::Acquire);
289            let inflight = slot.in_flight_bitmap.load(Ordering::Acquire);
290            let role = slot.role.load(Ordering::Acquire);
291            let v2 = slot.seq_version.load(Ordering::Acquire);
292            if v1 == v2 {
293                if pid == EMPTY_PID { return None; }
294                return Some(HeartbeatSnapshot {
295                    pid, last_seen_epoch: last,
296                    in_flight_bitmap: inflight, role,
297                });
298            }
299        }
300    }
301}
302
303/// Snapshot of one slot's state. Cheap to copy.
304#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305pub struct HeartbeatSnapshot {
306    pub pid: u32,
307    pub last_seen_epoch: u64,
308    pub in_flight_bitmap: u64,
309    pub role: u32,
310}
311
312/// Crate-internal accessor for the watchdog module. NOT pub-exported
313/// from the crate (only re-exported intra-crate).
314#[doc(hidden)]
315pub fn __slot_for_watchdog(table: &HeartbeatTable, idx: usize) -> &HeartbeatSlot {
316    table.slot(idx)
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    fn tmp_path(name: &str) -> std::path::PathBuf {
324        let mut p = std::env::temp_dir();
325        let pid = std::process::id();
326        p.push(format!("subetha-hb-{name}-{pid}.bin"));
327        p
328    }
329
330    #[test]
331    fn register_returns_slot_indices() {
332        let p = tmp_path("register");
333        let t = HeartbeatTable::create(&p, 4).unwrap();
334        let s0 = t.register(1001).unwrap();
335        let s1 = t.register(1002).unwrap();
336        assert_ne!(s0, s1);
337        std::fs::remove_file(&p).ok();
338    }
339
340    #[test]
341    fn table_full_returns_error() {
342        let p = tmp_path("table-full");
343        let t = HeartbeatTable::create(&p, 2).unwrap();
344        let _val = t.register(1).unwrap();
345        let _val = t.register(2).unwrap();
346        assert_eq!(t.register(3).unwrap_err(), HeartbeatError::TableFull);
347        std::fs::remove_file(&p).ok();
348    }
349
350    #[test]
351    fn beat_advances_last_seen_epoch() {
352        let p = tmp_path("beat");
353        let t = HeartbeatTable::create(&p, 1).unwrap();
354        let s = t.register(99).unwrap();
355        for _ in 0..5 { t.tick_global_epoch(); }
356        let snap_before = t.snapshot(s).unwrap();
357        let global_after_tick = t.global_epoch();
358        t.beat(s);
359        let snap_after = t.snapshot(s).unwrap();
360        assert!(snap_after.last_seen_epoch > snap_before.last_seen_epoch);
361        assert_eq!(snap_after.last_seen_epoch, global_after_tick);
362        std::fs::remove_file(&p).ok();
363    }
364
365    #[test]
366    fn unregister_frees_slot_for_reuse() {
367        let p = tmp_path("unreg");
368        let t = HeartbeatTable::create(&p, 2).unwrap();
369        let s0 = t.register(11).unwrap();
370        let _s1 = t.register(22).unwrap();
371        t.unregister(s0);
372        // Now there should be a free slot.
373        let new = t.register(33).unwrap();
374        assert_eq!(new, s0);
375        std::fs::remove_file(&p).ok();
376    }
377
378    #[test]
379    fn in_flight_bitmap_mark_and_clear() {
380        let p = tmp_path("inflight");
381        let t = HeartbeatTable::create(&p, 1).unwrap();
382        let s = t.register(7).unwrap();
383        t.mark_in_flight(s, 3);
384        t.mark_in_flight(s, 5);
385        let snap = t.snapshot(s).unwrap();
386        assert_eq!(snap.in_flight_bitmap, (1u64 << 3) | (1u64 << 5));
387        t.clear_in_flight(s, 3);
388        let snap = t.snapshot(s).unwrap();
389        assert_eq!(snap.in_flight_bitmap, 1u64 << 5);
390        std::fs::remove_file(&p).ok();
391    }
392
393    #[test]
394    fn snapshot_via_seqlock_returns_consistent_data() {
395        let p = tmp_path("snap");
396        let t = HeartbeatTable::create(&p, 1).unwrap();
397        let s = t.register(42).unwrap();
398        t.tick_global_epoch();
399        t.beat(s);
400        let snap = t.snapshot(s).unwrap();
401        assert_eq!(snap.pid, 42);
402        assert!(snap.last_seen_epoch >= 1);
403        std::fs::remove_file(&p).ok();
404    }
405}
subetha_cxc/heartbeat.rs

subetha_cxc/
heartbeat.rs
