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    pub fn create(path: impl AsRef<Path>, capacity: usize) -> Result<Self, HeartbeatError> {
109        assert!(capacity >= 1);
110        let total = heartbeat_file_size(capacity);
111        let file = OpenOptions::new()
112            .read(true).write(true).create(true).truncate(true)
113            .open(path.as_ref())?;
114        file.set_len(total as u64)?;
115        let mut mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
116        let hdr_ptr = mmap.as_mut_ptr() as *mut HeartbeatHeader;
117        unsafe {
118            std::ptr::write(hdr_ptr, HeartbeatHeader {
119                magic: HEARTBEAT_MAGIC,
120                capacity: capacity as u64,
121                epoch: AtomicU64::new(0),
122                _reserved: [0; 40],
123            });
124        }
125        let slots_base = unsafe {
126            mmap.as_mut_ptr().add(std::mem::size_of::<HeartbeatHeader>())
127        };
128        for i in 0..capacity {
129            let slot_ptr = unsafe {
130                slots_base.add(i * std::mem::size_of::<HeartbeatSlot>()) as *mut HeartbeatSlot
131            };
132            unsafe {
133                std::ptr::write(slot_ptr, HeartbeatSlot {
134                    pid: AtomicU32::new(EMPTY_PID),
135                    seq_version: AtomicU32::new(0),
136                    last_seen_epoch: AtomicU64::new(0),
137                    in_flight_bitmap: AtomicU64::new(0),
138                    role: AtomicU32::new(0),
139                    _pad: [0; 36],
140                });
141            }
142        }
143        Ok(Self {
144            _file: file, mmap, capacity,
145            header_sidecar: subetha_core::HandshakeHeader::new(),
146            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
147        })
148    }
149
150    pub fn open(path: impl AsRef<Path>, expected_capacity: usize) -> Result<Self, HeartbeatError> {
151        let file = OpenOptions::new().read(true).write(true).open(path.as_ref())?;
152        let total = heartbeat_file_size(expected_capacity);
153        let mmap = unsafe { MmapOptions::new().len(total).map_mut(&file)? };
154        let header = unsafe { &*(mmap.as_ptr() as *const HeartbeatHeader) };
155        if header.magic != HEARTBEAT_MAGIC || header.capacity != expected_capacity as u64 {
156            return Err(HeartbeatError::LayoutMismatch);
157        }
158        Ok(Self {
159            _file: file, mmap, capacity: expected_capacity,
160            header_sidecar: subetha_core::HandshakeHeader::new(),
161            ring_sidecar: Box::new(subetha_core::ObservationRing::new()),
162        })
163    }
164
165    pub fn capacity(&self) -> usize { self.capacity }
166
167    pub fn header(&self) -> &HeartbeatHeader {
168        unsafe { &*(self.mmap.as_ptr() as *const HeartbeatHeader) }
169    }
170
171    fn slot(&self, idx: usize) -> &HeartbeatSlot {
172        let base = unsafe {
173            self.mmap.as_ptr().add(std::mem::size_of::<HeartbeatHeader>())
174        };
175        unsafe {
176            &*(base.add(idx * std::mem::size_of::<HeartbeatSlot>()) as *const HeartbeatSlot)
177        }
178    }
179
180    /// Register the current process. Returns the slot index. CAS-claim
181    /// of the first empty slot.
182    pub fn register(&self, pid: u32) -> Result<usize, HeartbeatError> {
183        for i in 0..self.capacity {
184            let slot = self.slot(i);
185            if slot.pid.compare_exchange(
186                EMPTY_PID, pid, Ordering::AcqRel, Ordering::Acquire,
187            ).is_ok() {
188                slot.seq_version.fetch_add(1, Ordering::Release);
189                slot.last_seen_epoch.store(
190                    self.header().epoch.load(Ordering::Acquire),
191                    Ordering::Release,
192                );
193                slot.in_flight_bitmap.store(0, Ordering::Release);
194                slot.role.store(0, Ordering::Release);
195                slot.seq_version.fetch_add(1, Ordering::Release);
196                self.ring_sidecar
197                    .push_op(crate::sidecar_ops::liveness::OP_REGISTER, 0);
198                return Ok(i);
199            }
200        }
201        self.ring_sidecar
202            .push_op(crate::sidecar_ops::liveness::OP_REGISTER, 1);
203        Err(HeartbeatError::TableFull)
204    }
205
206    /// Release the slot at `idx`. Call before process exit.
207    pub fn unregister(&self, idx: usize) {
208        let slot = self.slot(idx);
209        slot.seq_version.fetch_add(1, Ordering::Release);
210        slot.in_flight_bitmap.store(0, Ordering::Release);
211        slot.pid.store(EMPTY_PID, Ordering::Release);
212        slot.seq_version.fetch_add(1, Ordering::Release);
213    }
214
215    /// Heartbeat: advance this slot's `last_seen_epoch` to match the
216    /// global epoch. Call once per scan tick.
217    pub fn beat(&self, idx: usize) {
218        let global = self.header().epoch.load(Ordering::Acquire);
219        let slot = self.slot(idx);
220        slot.last_seen_epoch.store(global, Ordering::Release);
221        self.ring_sidecar
222            .push_op(crate::sidecar_ops::liveness::OP_BEAT, 0);
223    }
224
225    /// Advance the global epoch. Watchdog calls this once per scan
226    /// interval. Returns the new epoch value.
227    pub fn tick_global_epoch(&self) -> u64 {
228        let v = self.header().epoch.fetch_add(1, Ordering::AcqRel) + 1;
229        self.ring_sidecar
230            .push_op(crate::sidecar_ops::liveness::OP_TICK_EPOCH, 0);
231        v
232    }
233
234    pub fn global_epoch(&self) -> u64 {
235        self.header().epoch.load(Ordering::Acquire)
236    }
237
238    /// Mark a work unit as in-flight for `slot_idx`.
239    pub fn mark_in_flight(&self, slot_idx: usize, bit: u8) {
240        debug_assert!((bit as usize) < IN_FLIGHT_SLOTS);
241        let slot = self.slot(slot_idx);
242        slot.in_flight_bitmap.fetch_or(1u64 << bit, Ordering::AcqRel);
243    }
244
245    pub fn clear_in_flight(&self, slot_idx: usize, bit: u8) {
246        debug_assert!((bit as usize) < IN_FLIGHT_SLOTS);
247        let slot = self.slot(slot_idx);
248        slot.in_flight_bitmap.fetch_and(!(1u64 << bit), Ordering::AcqRel);
249    }
250
251    /// Snapshot a slot via SeqLock retry. Returns `None` if the slot
252    /// is vacant.
253    pub fn snapshot(&self, idx: usize) -> Option<HeartbeatSnapshot> {
254        let slot = self.slot(idx);
255        loop {
256            let v1 = slot.seq_version.load(Ordering::Acquire);
257            if v1 & 1 != 0 { continue; }  // writer in progress
258            let pid = slot.pid.load(Ordering::Acquire);
259            let last = slot.last_seen_epoch.load(Ordering::Acquire);
260            let inflight = slot.in_flight_bitmap.load(Ordering::Acquire);
261            let role = slot.role.load(Ordering::Acquire);
262            let v2 = slot.seq_version.load(Ordering::Acquire);
263            if v1 == v2 {
264                if pid == EMPTY_PID { return None; }
265                return Some(HeartbeatSnapshot {
266                    pid, last_seen_epoch: last,
267                    in_flight_bitmap: inflight, role,
268                });
269            }
270        }
271    }
272}
273
274/// Snapshot of one slot's state. Cheap to copy.
275#[derive(Debug, Clone, Copy, PartialEq, Eq)]
276pub struct HeartbeatSnapshot {
277    pub pid: u32,
278    pub last_seen_epoch: u64,
279    pub in_flight_bitmap: u64,
280    pub role: u32,
281}
282
283/// Crate-internal accessor for the watchdog module. NOT pub-exported
284/// from the crate (only re-exported intra-crate).
285#[doc(hidden)]
286pub fn __slot_for_watchdog(table: &HeartbeatTable, idx: usize) -> &HeartbeatSlot {
287    table.slot(idx)
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293
294    fn tmp_path(name: &str) -> std::path::PathBuf {
295        let mut p = std::env::temp_dir();
296        let pid = std::process::id();
297        p.push(format!("subetha-hb-{name}-{pid}.bin"));
298        p
299    }
300
301    #[test]
302    fn register_returns_slot_indices() {
303        let p = tmp_path("register");
304        let t = HeartbeatTable::create(&p, 4).unwrap();
305        let s0 = t.register(1001).unwrap();
306        let s1 = t.register(1002).unwrap();
307        assert_ne!(s0, s1);
308        std::fs::remove_file(&p).ok();
309    }
310
311    #[test]
312    fn table_full_returns_error() {
313        let p = tmp_path("table-full");
314        let t = HeartbeatTable::create(&p, 2).unwrap();
315        let _val = t.register(1).unwrap();
316        let _val = t.register(2).unwrap();
317        assert_eq!(t.register(3).unwrap_err(), HeartbeatError::TableFull);
318        std::fs::remove_file(&p).ok();
319    }
320
321    #[test]
322    fn beat_advances_last_seen_epoch() {
323        let p = tmp_path("beat");
324        let t = HeartbeatTable::create(&p, 1).unwrap();
325        let s = t.register(99).unwrap();
326        for _ in 0..5 { t.tick_global_epoch(); }
327        let snap_before = t.snapshot(s).unwrap();
328        let global_after_tick = t.global_epoch();
329        t.beat(s);
330        let snap_after = t.snapshot(s).unwrap();
331        assert!(snap_after.last_seen_epoch > snap_before.last_seen_epoch);
332        assert_eq!(snap_after.last_seen_epoch, global_after_tick);
333        std::fs::remove_file(&p).ok();
334    }
335
336    #[test]
337    fn unregister_frees_slot_for_reuse() {
338        let p = tmp_path("unreg");
339        let t = HeartbeatTable::create(&p, 2).unwrap();
340        let s0 = t.register(11).unwrap();
341        let _s1 = t.register(22).unwrap();
342        t.unregister(s0);
343        // Now there should be a free slot.
344        let new = t.register(33).unwrap();
345        assert_eq!(new, s0);
346        std::fs::remove_file(&p).ok();
347    }
348
349    #[test]
350    fn in_flight_bitmap_mark_and_clear() {
351        let p = tmp_path("inflight");
352        let t = HeartbeatTable::create(&p, 1).unwrap();
353        let s = t.register(7).unwrap();
354        t.mark_in_flight(s, 3);
355        t.mark_in_flight(s, 5);
356        let snap = t.snapshot(s).unwrap();
357        assert_eq!(snap.in_flight_bitmap, (1u64 << 3) | (1u64 << 5));
358        t.clear_in_flight(s, 3);
359        let snap = t.snapshot(s).unwrap();
360        assert_eq!(snap.in_flight_bitmap, 1u64 << 5);
361        std::fs::remove_file(&p).ok();
362    }
363
364    #[test]
365    fn snapshot_via_seqlock_returns_consistent_data() {
366        let p = tmp_path("snap");
367        let t = HeartbeatTable::create(&p, 1).unwrap();
368        let s = t.register(42).unwrap();
369        t.tick_global_epoch();
370        t.beat(s);
371        let snap = t.snapshot(s).unwrap();
372        assert_eq!(snap.pid, 42);
373        assert!(snap.last_seen_epoch >= 1);
374        std::fs::remove_file(&p).ok();
375    }
376}
subetha_cxc/heartbeat.rs

subetha_cxc/
heartbeat.rs
