use std::collections::{BTreeMap, HashMap};
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::PathBuf;
use std::sync::Mutex;
use super::events::{CheckpointData, hash_to_hex};
use super::identity::PhaseIdentity;
use super::storage::{Checkpoint, OpCounts, PhaseEntry, PhaseStatus, now_rfc3339};
pub const CHECKPOINT_VERSION: u32 = 1;
pub struct CheckpointWriter {
path: PathBuf,
enabled: bool,
inner: Mutex<Inner>,
run_reached_end: std::sync::atomic::AtomicBool,
_lock_fd: Option<LockHandle>,
}
struct LockHandle(#[allow(dead_code)] File);
struct Inner {
doc: Checkpoint,
index: HashMap<String, usize>,
file: File,
}
impl CheckpointWriter {
pub fn new(path: PathBuf, session: String, started_at: String, invocation: u32) -> Self {
let doc = Checkpoint {
version: CHECKPOINT_VERSION,
session: session.clone(),
started_at: started_at.clone(),
checkpoint_at: started_at.clone(),
invocation,
phases: Vec::new(),
};
let lock = acquire_flock(&path);
let file = open_append(&path);
let writer = Self {
path,
enabled: true,
inner: Mutex::new(Inner {
doc,
index: HashMap::new(),
file,
}),
run_reached_end: std::sync::atomic::AtomicBool::new(false),
_lock_fd: lock,
};
writer.append_event(CheckpointData::SessionStart {
at: now_rfc3339(),
version: CHECKPOINT_VERSION,
session,
started_at,
invocation,
});
writer
}
pub fn from_existing(
path: PathBuf,
mut doc: Checkpoint,
new_checkpoint_at: String,
new_invocation: u32,
) -> Self {
doc.checkpoint_at = new_checkpoint_at;
doc.invocation = new_invocation;
let session = doc.session.clone();
let started_at = doc.started_at.clone();
let index = build_index(&doc.phases);
let lock = acquire_flock(&path);
let file = open_append(&path);
let writer = Self {
path,
enabled: true,
inner: Mutex::new(Inner { doc, index, file }),
run_reached_end: std::sync::atomic::AtomicBool::new(false),
_lock_fd: lock,
};
writer.append_event(CheckpointData::SessionStart {
at: now_rfc3339(),
version: CHECKPOINT_VERSION,
session,
started_at,
invocation: new_invocation,
});
writer
}
pub fn disabled(path: PathBuf) -> Self {
let doc = Checkpoint {
version: CHECKPOINT_VERSION,
session: String::new(),
started_at: String::new(),
checkpoint_at: String::new(),
invocation: 0,
phases: Vec::new(),
};
let file = open_append(std::path::Path::new("/dev/null"));
Self {
path,
enabled: false,
inner: Mutex::new(Inner {
doc,
index: HashMap::new(),
file,
}),
run_reached_end: std::sync::atomic::AtomicBool::new(false),
_lock_fd: None,
}
}
pub fn declare_phase(&self, identity: PhaseIdentity, skip_eligible: bool) {
let event = {
let mut g = self.inner.lock().unwrap();
let key = identity_key(&identity);
if g.index.contains_key(&key) {
return;
}
let entry = PhaseEntry {
identity: identity.clone(),
skip_eligible,
params_consumed: None,
status: PhaseStatus::Pending,
duration_secs: None,
op_counts: None,
cursor_state: None,
error: None,
};
g.doc.phases.push(entry);
let idx = g.doc.phases.len() - 1;
g.index.insert(key, idx);
CheckpointData::PhaseDeclared {
at: now_rfc3339(),
identity,
skip_eligible,
}
};
self.append_event(event);
}
pub fn phase_started(&self, identity: &PhaseIdentity) {
let updated = self.with_entry(identity, |e| {
e.status = PhaseStatus::Running;
e.error = None;
});
if updated {
self.append_event(CheckpointData::PhaseStarted {
at: now_rfc3339(),
identity: identity.clone(),
});
}
}
pub fn phase_completed(&self, identity: &PhaseIdentity, duration_secs: f64) {
let final_counts = {
let mut g = self.inner.lock().unwrap();
let key = identity_key(identity);
if let Some(&idx) = g.index.get(&key) {
let entry = &mut g.doc.phases[idx];
let counts = entry.op_counts.clone().unwrap_or_default();
entry.status = PhaseStatus::Completed;
entry.duration_secs = Some(duration_secs);
entry.op_counts = Some(counts.clone());
entry.cursor_state = None;
entry.error = None;
counts
} else {
return;
}
};
self.append_event(CheckpointData::PhaseCompleted {
at: now_rfc3339(),
identity: identity.clone(),
duration_secs,
op_counts: final_counts,
});
}
pub fn phase_failed(&self, identity: &PhaseIdentity, error: &str) {
let counts = {
let err_owned = error.to_string();
let mut g = self.inner.lock().unwrap();
let key = identity_key(identity);
if let Some(&idx) = g.index.get(&key) {
let entry = &mut g.doc.phases[idx];
entry.status = PhaseStatus::Failed;
entry.error = Some(err_owned);
entry.cursor_state = None;
entry.op_counts.clone()
} else {
return;
}
};
self.append_event(CheckpointData::PhaseFailed {
at: now_rfc3339(),
identity: identity.clone(),
error: error.to_string(),
op_counts: counts,
});
}
pub fn update_op_counts(&self, identity: &PhaseIdentity, counts: OpCounts) {
let cursor_state = {
let mut g = self.inner.lock().unwrap();
let key = identity_key(identity);
if let Some(&idx) = g.index.get(&key) {
let entry = &mut g.doc.phases[idx];
entry.op_counts = Some(counts.clone());
entry.cursor_state.clone()
} else {
return;
}
};
self.append_event(CheckpointData::PhaseProgress {
at: now_rfc3339(),
identity: identity.clone(),
op_counts: counts,
cursor_state,
});
}
pub fn update_phase_hash(
&self,
identity: &PhaseIdentity,
hash: [u8; 32],
params_consumed: Option<String>,
) {
let updated = self.with_entry(identity, |e| {
e.identity.phase_hash = Some(hash);
e.params_consumed = params_consumed.clone();
});
if updated {
self.append_event(CheckpointData::PhaseHash {
at: now_rfc3339(),
identity: identity.clone(),
hash_hex: hash_to_hex(&hash),
params_consumed,
});
}
}
pub fn update_cursor(&self, identity: &PhaseIdentity, cursor_state: serde_json::Value) {
let counts = {
let mut g = self.inner.lock().unwrap();
let key = identity_key(identity);
if let Some(&idx) = g.index.get(&key) {
let entry = &mut g.doc.phases[idx];
entry.cursor_state = Some(cursor_state.clone());
entry.op_counts.clone().unwrap_or_default()
} else {
return;
}
};
self.append_event(CheckpointData::PhaseProgress {
at: now_rfc3339(),
identity: identity.clone(),
op_counts: counts,
cursor_state: Some(cursor_state),
});
}
pub fn emit_scope_enter(
&self,
kind: &str,
coords: BTreeMap<String, serde_json::Value>,
path: Vec<BTreeMap<String, serde_json::Value>>,
) {
self.append_event(CheckpointData::ScopeEnter {
at: now_rfc3339(),
kind: kind.to_string(),
coords,
path,
});
}
pub fn emit_scope_exit(
&self,
kind: &str,
coords: BTreeMap<String, serde_json::Value>,
path: Vec<BTreeMap<String, serde_json::Value>>,
outcome: &str,
) {
self.append_event(CheckpointData::ScopeExit {
at: now_rfc3339(),
kind: kind.to_string(),
coords,
path,
outcome: outcome.to_string(),
});
}
pub fn flush(&self) -> Result<(), String> {
if !self.enabled {
return Ok(());
}
let g = self.inner.lock().unwrap();
match g.file.sync_data() {
Ok(()) => Ok(()),
Err(e) => Err(format!("fdatasync {}: {e}", self.path.display())),
}
}
pub fn snapshot(&self) -> Checkpoint {
self.inner.lock().unwrap().doc.clone()
}
pub fn mark_run_reached_end(&self) {
self.run_reached_end
.store(true, std::sync::atomic::Ordering::Relaxed);
}
pub fn resume_hint(&self) -> Option<String> {
if !self.enabled {
return None;
}
let ended = self
.run_reached_end
.load(std::sync::atomic::Ordering::Relaxed);
let cp = self.snapshot();
let recoverable = cp.phases.iter().any(|e| {
e.skip_eligible
&& match e.status {
PhaseStatus::Completed => false,
PhaseStatus::Failed | PhaseStatus::Running => true,
PhaseStatus::Pending => !ended,
}
});
if !recoverable {
return None;
}
Some(format!(
"This session has resumable phases that didn't complete.\n \
To continue from where it stopped:\n \
nmbrs run <workload> --session-dir {} (already set if you exported \
SESSION_DIRECTORY) --resume\n \
To pin the session name for repeatable resumes:\n \
nmbrs run <workload> --session {} (then add --resume next time)",
self.path
.parent()
.map(|p| p.display().to_string())
.unwrap_or_default(),
cp.session,
))
}
pub fn path(&self) -> &std::path::Path {
&self.path
}
fn append_event(&self, event: CheckpointData) {
if !self.enabled {
return;
}
let mut g = self.inner.lock().unwrap();
let mut line = match serde_json::to_string(&event) {
Ok(s) => s,
Err(e) => {
eprintln!("checkpoint: serialise event failed: {e}; dropping record",);
return;
}
};
line.push('\n');
if let Err(e) = g.file.write_all(line.as_bytes()) {
eprintln!("checkpoint: append to {}: {e}", self.path.display(),);
}
}
fn with_entry<F: FnOnce(&mut PhaseEntry)>(&self, identity: &PhaseIdentity, f: F) -> bool {
let mut g = self.inner.lock().unwrap();
let key = identity_key(identity);
if let Some(&idx) = g.index.get(&key) {
f(&mut g.doc.phases[idx]);
true
} else {
false
}
}
}
fn open_append(path: &std::path::Path) -> File {
if let Some(parent) = path.parent()
&& let Err(e) = std::fs::create_dir_all(parent)
{
panic!(
"checkpoint: create parent dir {} for {}: {e}",
parent.display(),
path.display(),
);
}
OpenOptions::new()
.create(true)
.append(true)
.open(path)
.unwrap_or_else(|e| panic!("checkpoint: open append {} failed: {e}", path.display(),))
}
pub(crate) fn identity_key(identity: &PhaseIdentity) -> String {
let path_json = serde_json::to_string(&identity.yaml_path).unwrap_or_else(|_| String::new());
format!("{path_json}\x1f{}", identity.coords)
}
fn acquire_flock(checkpoint_path: &std::path::Path) -> Option<LockHandle> {
let parent = checkpoint_path.parent()?;
if let Err(e) = std::fs::create_dir_all(parent) {
eprintln!(
"warning: could not create checkpoint dir {}: {e} (concurrent-resume protection skipped)",
parent.display(),
);
return None;
}
let lock_path = checkpoint_path.with_extension("lock");
let file = match OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(&lock_path)
{
Ok(f) => f,
Err(e) => {
eprintln!(
"warning: could not open lockfile {}: {e} (concurrent-resume protection skipped)",
lock_path.display(),
);
return None;
}
};
match file.try_lock() {
Ok(()) => Some(LockHandle(file)),
Err(std::fs::TryLockError::WouldBlock) => {
panic!(
"checkpoint: another process holds the resume lock at {} \
(concurrent `nmbrs run --resume` against the same session?). \
If you're certain no other process is running, remove the \
lockfile and retry.",
lock_path.display(),
);
}
Err(std::fs::TryLockError::Error(e)) => {
eprintln!(
"warning: lock on {} failed: {e} (concurrent-resume protection skipped)",
lock_path.display(),
);
None
}
}
}
fn build_index(phases: &[PhaseEntry]) -> HashMap<String, usize> {
let mut m = HashMap::with_capacity(phases.len());
for (i, e) in phases.iter().enumerate() {
m.insert(identity_key(&e.identity), i);
}
m
}
#[cfg(test)]
mod tests {
use super::*;
use crate::checkpoint::{PathSegment, PhaseIdentity};
fn ident(name: &str, coords: &str) -> PhaseIdentity {
PhaseIdentity {
yaml_path: vec![
PathSegment::Scenario("s".into()),
PathSegment::Phase(name.into()),
],
coords: coords.into(),
phase_hash: Some([0xcd; 32]),
}
}
fn tempdir() -> std::path::PathBuf {
let d = std::env::temp_dir().join(format!(
"nmbrs-checkpoint-writer-{}",
crate::scratch_suffix()
));
std::fs::create_dir_all(&d).unwrap();
d
}
#[test]
fn declare_then_complete_then_flush() {
let dir = tempdir();
let path = dir.join("checkpoint.jsonl");
let w = CheckpointWriter::new(
path.clone(),
"sess".into(),
"2026-01-01T00:00:00Z".into(),
1,
);
let id = ident("schema", "");
w.declare_phase(id.clone(), true);
w.phase_started(&id);
w.phase_completed(&id, 1.5);
w.flush().expect("flush");
let snap = w.snapshot();
assert_eq!(snap.phases.len(), 1);
assert_eq!(snap.phases[0].status, PhaseStatus::Completed);
assert_eq!(snap.phases[0].duration_secs, Some(1.5));
let raw = std::fs::read_to_string(&path).expect("read log");
let lines: Vec<&str> = raw.lines().collect();
assert_eq!(
lines.len(),
4,
"expected session_start, phase_declared, phase_started, phase_completed"
);
assert!(lines[0].contains("\"type\":\"session_start\""));
assert!(lines[1].contains("\"type\":\"phase_declared\""));
assert!(lines[2].contains("\"type\":\"phase_started\""));
assert!(lines[3].contains("\"type\":\"phase_completed\""));
}
#[test]
fn resume_hint_respects_run_end_boundary() {
let dir = tempdir();
let w = CheckpointWriter::new(
dir.join("checkpoint.jsonl"),
"sess".into(),
"2026-01-01T00:00:00Z".into(),
1,
);
let ran = ident("tier", "(part=0)");
let excluded = ident("tier", "(part=17)");
w.declare_phase(ran.clone(), true);
w.declare_phase(excluded.clone(), true);
w.phase_started(&ran);
w.phase_completed(&ran, 1.0);
assert!(
w.resume_hint().is_some(),
"an interrupted run must advise resuming pending phases"
);
w.mark_run_reached_end();
assert!(
w.resume_hint().is_none(),
"a run that reached its end must not advise resuming \
predicate-excluded phases"
);
let failed = ident("tier", "(part=3)");
w.declare_phase(failed.clone(), true);
w.phase_started(&failed);
w.phase_failed(&failed, "boom");
assert!(
w.resume_hint().is_some(),
"failed phases must keep the hint even on a clean end"
);
}
#[test]
fn disabled_writer_persists_nothing() {
let dir = tempdir();
let path = dir.join("checkpoint.jsonl");
let w = CheckpointWriter::disabled(path.clone());
let id = ident("teardown", "(table=changeme_default)");
w.declare_phase(id.clone(), true);
w.phase_started(&id);
w.phase_completed(&id, 1.0);
w.flush().expect("flush is a harmless no-op");
assert!(
!path.exists(),
"dry-run must not create a checkpoint file at {}",
path.display()
);
assert!(
w.resume_hint().is_none(),
"dry-run must never advertise a resumable session"
);
}
#[test]
fn redundant_declare_is_idempotent() {
let dir = tempdir();
let path = dir.join("c.jsonl");
let w = CheckpointWriter::new(path.clone(), "s".into(), "t".into(), 1);
let id = ident("p", "(k=1)");
w.declare_phase(id.clone(), true);
w.declare_phase(id.clone(), false); let snap = w.snapshot();
assert_eq!(snap.phases.len(), 1);
assert!(snap.phases[0].skip_eligible, "first declare wins");
let raw = std::fs::read_to_string(&path).expect("read");
let count = raw
.lines()
.filter(|l| l.contains("\"type\":\"phase_declared\""))
.count();
assert_eq!(count, 1, "second declare must not emit a duplicate event");
}
#[test]
fn from_existing_emits_fresh_session_start() {
let dir = tempdir();
let path = dir.join("c.jsonl");
let saved = {
let w =
CheckpointWriter::new(path.clone(), "s".into(), "2026-01-01T00:00:00Z".into(), 1);
let id = ident("schema", "");
w.declare_phase(id.clone(), true);
w.phase_completed(&id, 0.5);
w.flush().expect("flush");
w.snapshot()
};
let w2 =
CheckpointWriter::from_existing(path.clone(), saved, "2026-01-01T00:01:00Z".into(), 2);
let snap = w2.snapshot();
assert_eq!(snap.invocation, 2);
assert_eq!(snap.phases.len(), 1);
assert_eq!(snap.phases[0].status, PhaseStatus::Completed);
let raw = std::fs::read_to_string(&path).expect("read");
let count = raw
.lines()
.filter(|l| l.contains("\"type\":\"session_start\""))
.count();
assert_eq!(
count, 2,
"resume must append a fresh session_start, not rewrite"
);
}
#[test]
fn scope_enter_exit_pairs_for_two_deep_for_each() {
let dir = tempdir();
let path = dir.join("c.jsonl");
let w = CheckpointWriter::new(path.clone(), "s".into(), "2026-01-01T00:00:00Z".into(), 1);
let outer_coord = |xv: u64| -> BTreeMap<String, serde_json::Value> {
let mut m = BTreeMap::new();
m.insert("x".into(), serde_json::Value::from(xv));
m
};
let inner_coord = |yv: &str| -> BTreeMap<String, serde_json::Value> {
let mut m = BTreeMap::new();
m.insert("y".into(), serde_json::Value::from(yv));
m
};
for x in [1u64, 2u64] {
w.emit_scope_enter("for_each", outer_coord(x), Vec::new());
for y in ["a", "b"] {
w.emit_scope_enter("for_each", inner_coord(y), vec![outer_coord(x)]);
w.emit_scope_exit(
"for_each",
inner_coord(y),
vec![outer_coord(x)],
"completed",
);
}
w.emit_scope_exit("for_each", outer_coord(x), Vec::new(), "completed");
}
w.flush().expect("flush");
let raw = std::fs::read_to_string(&path).expect("read log");
let scope_events: Vec<serde_json::Value> = raw
.lines()
.map(|l| serde_json::from_str::<serde_json::Value>(l).expect("parse line"))
.filter(|v| {
let t = v.get("type").and_then(|t| t.as_str()).unwrap_or("");
t == "scope_enter" || t == "scope_exit"
})
.collect();
assert_eq!(
scope_events.len(),
12,
"expected 12 scope events, got {}",
scope_events.len()
);
let kind = |v: &serde_json::Value| {
v.get("type")
.and_then(|t| t.as_str())
.unwrap_or("")
.to_string()
};
let x_at = |v: &serde_json::Value, idx: &str| -> Option<u64> {
v.pointer(idx).and_then(|n| n.as_u64())
};
let y_at = |v: &serde_json::Value, idx: &str| -> Option<String> {
v.pointer(idx)
.and_then(|n| n.as_str())
.map(|s| s.to_string())
};
assert_eq!(kind(&scope_events[0]), "scope_enter");
assert_eq!(x_at(&scope_events[0], "/coords/x"), Some(1));
assert!(
scope_events[0]
.pointer("/path")
.and_then(|p| p.as_array())
.map(|a| a.is_empty())
.unwrap_or(false),
"outer enter must have empty path"
);
assert_eq!(kind(&scope_events[1]), "scope_enter");
assert_eq!(y_at(&scope_events[1], "/coords/y"), Some("a".to_string()));
assert_eq!(x_at(&scope_events[1], "/path/0/x"), Some(1));
assert_eq!(kind(&scope_events[2]), "scope_exit");
assert_eq!(
scope_events[2].pointer("/outcome").and_then(|s| s.as_str()),
Some("completed")
);
assert_eq!(y_at(&scope_events[3], "/coords/y"), Some("b".to_string()));
assert_eq!(kind(&scope_events[4]), "scope_exit");
assert_eq!(kind(&scope_events[5]), "scope_exit");
assert_eq!(x_at(&scope_events[5], "/coords/x"), Some(1));
assert_eq!(
scope_events[5].pointer("/outcome").and_then(|s| s.as_str()),
Some("completed")
);
assert_eq!(kind(&scope_events[6]), "scope_enter");
assert_eq!(x_at(&scope_events[6], "/coords/x"), Some(2));
assert_eq!(x_at(&scope_events[7], "/path/0/x"), Some(2));
assert_eq!(y_at(&scope_events[7], "/coords/y"), Some("a".to_string()));
assert_eq!(kind(&scope_events[11]), "scope_exit");
assert_eq!(x_at(&scope_events[11], "/coords/x"), Some(2));
let folded = super::super::storage::read(&path)
.expect("read folds")
.expect("non-empty");
assert!(
folded.phases.is_empty(),
"no phases declared, fold should be empty"
);
}
#[test]
fn scope_exit_outcome_distinguishes_interrupted_from_completed() {
let dir = tempdir();
let path = dir.join("c.jsonl");
let w = CheckpointWriter::new(path.clone(), "s".into(), "t".into(), 1);
let mut coords = BTreeMap::new();
coords.insert("k".into(), serde_json::Value::from(7u64));
w.emit_scope_enter("do_while", coords.clone(), Vec::new());
w.emit_scope_exit("do_while", coords, Vec::new(), "interrupted");
w.flush().expect("flush");
let raw = std::fs::read_to_string(&path).expect("read");
let exit_line = raw
.lines()
.find(|l| l.contains("\"type\":\"scope_exit\""))
.expect("scope_exit line");
let ev: serde_json::Value = serde_json::from_str(exit_line).unwrap();
assert_eq!(
ev.pointer("/kind").and_then(|s| s.as_str()),
Some("do_while")
);
assert_eq!(
ev.pointer("/outcome").and_then(|s| s.as_str()),
Some("interrupted")
);
}
#[test]
fn flock_blocks_concurrent_writer_on_same_path() {
let dir = tempdir();
let path = dir.join("c.jsonl");
let _w = CheckpointWriter::new(path.clone(), "s".into(), "t".into(), 1);
let result = std::panic::catch_unwind(|| {
let _w2 = CheckpointWriter::new(path.clone(), "s".into(), "t".into(), 1);
});
assert!(
result.is_err(),
"second writer should panic on flock contention"
);
}
}