crates/jera/src/agent/roster.rs
//! Durable node **roster** + heartbeat/liveness for the multi-machine scheduler.
//!
//! The [`Scheduler`](super::Scheduler) holds live [`Agent`](super::Agent)s in
//! memory; the roster is the *durable* list of nodes the fleet knows about —
//! **Loki / Odin / NUC** + any Njord worker — each with its
//! [`AgentSpec`](super::AgentSpec) (id/group/caps/cores/mem), a gRPC `endpoint`, and
//! the last time it was heard from. A node silent past the **heartbeat window** is
//! DEAD and must never be scheduled to; [`Roster::heartbeat`] refreshes the stamp,
//! [`Roster::evict_dead`] drops the stale ones.
//!
//! Pure std + serde (a JSON file is the durable store, [`Roster::load`]/
//! [`Roster::save`]) — fully testable with no network and **no clock coupling**:
//! every liveness method takes `now` (micros since epoch), so a test can prove "a
//! node stale past the window is dead" deterministically. The wire feature adds no
//! new RPC — the existing `State` probe (`agent-grpc`) is the liveness ping; the
//! consumer calls [`Roster::heartbeat`] after a successful probe.
use std::collections::BTreeMap;
use std::path::Path;
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use super::{AgentId, AgentSpec, AgentState};
/// Default liveness window in seconds: a node silent longer than this is dead.
pub const DEFAULT_HEARTBEAT_WINDOW_SECS: i64 = 30;
/// One node in the durable roster: its spec, how to reach it, and when it last
/// checked in (micros since epoch).
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct NodeEntry {
pub spec: AgentSpec,
/// Where the worker's `AgentWorker` gRPC server listens (e.g. `http://loki:7900`).
pub endpoint: String,
/// Last time this node was heard from (a heartbeat / successful `State` probe).
pub last_heartbeat_micros: i64,
}
impl NodeEntry {
/// Live iff heard from within `window_micros` of `now`.
pub fn is_live(&self, now: i64, window_micros: i64) -> bool {
now.saturating_sub(self.last_heartbeat_micros) <= window_micros
}
}
/// The durable roster — a keyed set of [`NodeEntry`] + a liveness window. Persisted
/// as JSON ([`load`](Self::load)/[`save`](Self::save)); seeded from a Loki/Odin/NUC
/// roster file. The control-plane loads it, probes each node, [`heartbeat`](Self::heartbeat)s
/// the ones that answer, and [`evict_dead`](Self::evict_dead)s the rest.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Roster {
nodes: BTreeMap<AgentId, NodeEntry>,
/// Liveness window in micros — a node silent longer than this is dead.
window_micros: i64,
}
impl Default for Roster {
fn default() -> Self {
Self::new(DEFAULT_HEARTBEAT_WINDOW_SECS)
}
}
impl Roster {
/// A fresh roster with a `window_secs` liveness window.
pub fn new(window_secs: i64) -> Self {
Self {
nodes: BTreeMap::new(),
window_micros: window_secs * 1_000_000,
}
}
/// The liveness window, in micros.
pub fn window_micros(&self) -> i64 {
self.window_micros
}
pub fn len(&self) -> usize {
self.nodes.len()
}
pub fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
/// Look up a node by id.
pub fn get(&self, id: &str) -> Option<&NodeEntry> {
self.nodes.get(id)
}
/// Iterate every known node (live or dead).
pub fn nodes(&self) -> impl Iterator<Item = &NodeEntry> {
self.nodes.values()
}
/// Register (or replace) a node, stamping `now` as its first heartbeat.
pub fn register(&mut self, spec: AgentSpec, endpoint: impl Into<String>, now: i64) {
let id = spec.id.clone();
self.nodes.insert(
id,
NodeEntry {
spec,
endpoint: endpoint.into(),
last_heartbeat_micros: now,
},
);
}
/// Refresh node `id`'s heartbeat to `now`. Returns `true` if the node exists
/// (a heartbeat for an unknown node is a no-op returning `false`, so a caller
/// can tell "resurrect vs never-registered" apart).
pub fn heartbeat(&mut self, id: &str, now: i64) -> bool {
match self.nodes.get_mut(id) {
Some(n) => {
n.last_heartbeat_micros = now;
true
}
None => false,
}
}
/// Whether node `id` is currently live (exists ∧ within the window at `now`).
pub fn is_live(&self, id: &str, now: i64) -> bool {
self.nodes
.get(id)
.is_some_and(|n| n.is_live(now, self.window_micros))
}
/// The live nodes at `now` (safe to schedule to).
pub fn live(&self, now: i64) -> Vec<&NodeEntry> {
self.nodes
.values()
.filter(|n| n.is_live(now, self.window_micros))
.collect()
}
/// The ids of dead (stale) nodes at `now`.
pub fn dead(&self, now: i64) -> Vec<AgentId> {
self.nodes
.values()
.filter(|n| !n.is_live(now, self.window_micros))
.map(|n| n.spec.id.clone())
.collect()
}
/// Evict every dead node at `now`; returns the removed ids. Emits a
/// `jera-roster` functional-status row per eviction (no-op unless a consumer
/// installed a `nornir-testmatrix` hook), so dead-node eviction is observable.
pub fn evict_dead(&mut self, now: i64) -> Vec<AgentId> {
let dead = self.dead(now);
for id in &dead {
self.nodes.remove(id);
crate::jobs::selftest::emit(
"jera-roster",
"evict_dead_node",
true,
&format!("node '{id}' silent past the heartbeat window → evicted"),
);
}
dead
}
/// Load a roster from a JSON file (the durable store).
pub fn load(path: &Path) -> Result<Self> {
let bytes =
std::fs::read(path).with_context(|| format!("read roster {}", path.display()))?;
serde_json::from_slice(&bytes).with_context(|| format!("decode roster {}", path.display()))
}
/// Save this roster to a JSON file (the durable store). Creates parent dirs.
pub fn save(&self, path: &Path) -> Result<()> {
if let Some(dir) = path.parent() {
std::fs::create_dir_all(dir)
.with_context(|| format!("create roster dir {}", dir.display()))?;
}
let bytes = serde_json::to_vec_pretty(self).context("encode roster")?;
std::fs::write(path, bytes).with_context(|| format!("write roster {}", path.display()))
}
/// Observable fleet state as data — feeds a UI "Nodes" pane / a `nornir-testmatrix`
/// cell at the consumer (id/group/endpoint/cores/mem + live-at-`now`).
pub fn state_json(&self, now: i64) -> serde_json::Value {
serde_json::json!({
"window_secs": self.window_micros / 1_000_000,
"nodes": self.nodes.values().map(|n| {
serde_json::json!({
"id": n.spec.id,
"group": n.spec.group,
"endpoint": n.endpoint,
"cores": n.spec.cores,
"mem_mb": n.spec.mem_mb,
"last_heartbeat_micros": n.last_heartbeat_micros,
"live": n.is_live(now, self.window_micros),
"state": if n.is_live(now, self.window_micros) { AgentState::Idle } else { AgentState::Offline },
})
}).collect::<Vec<_>>(),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
const SEC: i64 = 1_000_000;
fn spec(id: &str, group: &str, cores: u32) -> AgentSpec {
AgentSpec::new(id, group, ["linux"], cores, 65536)
}
#[test]
fn stale_node_is_dead_after_the_window_and_heartbeat_revives_it() {
// RED-when-broken: a node not heard from within the window MUST read dead,
// and a fresh heartbeat MUST bring it back live.
let mut r = Roster::new(30); // 30s window
let t0 = 1_000 * SEC;
r.register(spec("loki", "loki", 12), "http://loki:7900", t0);
assert!(r.is_live("loki", t0), "just-registered node is live");
// 20s later: still within the window.
assert!(r.is_live("loki", t0 + 20 * SEC), "within window ⇒ live");
// 31s later: past the window ⇒ dead.
let stale = t0 + 31 * SEC;
assert!(!r.is_live("loki", stale), "silent past the window ⇒ dead");
assert_eq!(r.dead(stale), vec!["loki".to_string()]);
assert!(r.live(stale).is_empty());
// A heartbeat at the stale instant revives it.
assert!(
r.heartbeat("loki", stale),
"heartbeat of a known node returns true"
);
assert!(r.is_live("loki", stale), "heartbeat brings it back live");
assert!(
!r.heartbeat("ghost", stale),
"heartbeat of an unknown node is false"
);
}
#[test]
fn evict_dead_removes_only_the_stale_nodes() {
let mut r = Roster::new(30);
let t0 = 100 * SEC;
r.register(spec("loki", "loki", 12), "http://loki:7900", t0);
r.register(spec("odin", "odin", 32), "http://odin:7900", t0);
// odin heartbeats later; loki stays silent.
let later = t0 + 40 * SEC;
r.heartbeat("odin", later);
// At `later`, loki (last seen t0, 40s ago) is dead; odin is live.
let evicted = r.evict_dead(later);
assert_eq!(
evicted,
vec!["loki".to_string()],
"only the stale node evicted"
);
assert!(r.get("loki").is_none(), "loki removed from the roster");
assert!(r.get("odin").is_some(), "odin retained");
assert_eq!(r.len(), 1);
}
#[test]
fn roster_saves_and_loads_durably() {
// RED-when-broken: the durable store must round-trip the whole roster.
let mut r = Roster::new(45);
let t0 = 7 * SEC;
r.register(spec("nuc", "nuc", 6), "http://nuc:7900", t0);
r.register(spec("loki", "loki", 12), "http://loki:7900", t0);
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("roster").join("fleet.json");
r.save(&path).expect("save");
let back = Roster::load(&path).expect("load");
assert_eq!(back, r, "roster round-trips through the durable JSON store");
assert_eq!(back.window_micros(), 45 * SEC);
assert_eq!(back.get("loki").unwrap().endpoint, "http://loki:7900");
}
#[test]
fn state_json_exposes_live_and_dead_as_data() {
let mut r = Roster::new(30);
let t0 = 0;
r.register(spec("loki", "loki", 12), "http://loki:7900", t0);
let now = t0 + 31 * SEC; // loki is now stale
let v = r.state_json(now);
assert_eq!(v["window_secs"], 30);
let n = &v["nodes"][0];
assert_eq!(n["id"], "loki");
assert_eq!(n["endpoint"], "http://loki:7900");
assert_eq!(n["cores"], 12);
assert_eq!(n["live"], false);
assert_eq!(n["state"], "offline");
}
}