use std::collections::BTreeSet;
use std::fs::{self, File, OpenOptions, TryLockError};
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use crate::error::TopologyError;
use crate::record::{ExecutionRecord, RecordSet};
pub const JOURNAL_SUBDIR: &str = "topology";
pub const JOURNAL_FILE: &str = "records.jsonl";
pub const SCHEMA_VERSION: u32 = 1;
#[derive(Debug, Error)]
pub enum JournalError {
#[error("topology journal I/O failed: {0}")]
Io(#[from] std::io::Error),
#[error(transparent)]
Record(#[from] TopologyError),
#[error(
"topology journal holds records from {} embedders ({}) — name one to load",
.embedders.len(),
.embedders.join(", ")
)]
MixedEmbedders { embedders: Vec<String> },
#[error("topology journal holds no records for embedder {embedder}")]
NoRecordsForEmbedder { embedder: String },
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct JournalEntry {
pub schema: u32,
pub embedder: String,
pub record: ExecutionRecord,
}
impl JournalEntry {
pub fn new(embedder: impl Into<String>, record: ExecutionRecord) -> Self {
Self {
schema: SCHEMA_VERSION,
embedder: embedder.into(),
record,
}
}
}
pub fn journal_path(state_root: &Path) -> PathBuf {
state_root.join(JOURNAL_SUBDIR).join(JOURNAL_FILE)
}
#[derive(Debug)]
pub struct RecordJournal {
path: PathBuf,
writer: BufWriter<File>,
_lock: File,
}
impl Drop for RecordJournal {
fn drop(&mut self) {
let _ = self.writer.flush();
let _ = self._lock.unlock();
}
}
impl RecordJournal {
pub fn open(path: &Path) -> Result<Self, JournalError> {
if let Some(parent) = path.parent() {
if !parent.as_os_str().is_empty() {
fs::create_dir_all(parent)?;
}
}
let lock_path = {
let mut s = path.as_os_str().to_owned();
s.push(".lock");
PathBuf::from(s)
};
let lock = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&lock_path)?;
match lock.try_lock() {
Ok(()) => {}
Err(TryLockError::WouldBlock) => {
return Err(JournalError::Io(std::io::Error::new(
std::io::ErrorKind::WouldBlock,
format!(
"topology record journal already open by another writer \
(advisory lock held on {})",
lock_path.display()
),
)));
}
Err(TryLockError::Error(e)) => return Err(JournalError::Io(e)),
}
let needs_newline = match File::open(path) {
Ok(mut existing) => {
use std::io::{Read, Seek, SeekFrom};
if existing.metadata()?.len() == 0 {
false
} else {
existing.seek(SeekFrom::End(-1))?;
let mut last = [0u8; 1];
existing.read_exact(&mut last)?;
last[0] != b'\n'
}
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => false,
Err(e) => return Err(JournalError::Io(e)),
};
let file = OpenOptions::new().create(true).append(true).open(path)?;
let mut writer = BufWriter::new(file);
if needs_newline {
writer.write_all(b"\n")?;
writer.flush()?;
}
Ok(Self {
path: path.to_path_buf(),
writer,
_lock: lock,
})
}
pub fn path(&self) -> &Path {
&self.path
}
pub fn append(&mut self, embedder: &str, record: &ExecutionRecord) -> Result<(), JournalError> {
let entry = JournalEntry::new(embedder, record.clone());
self.append_entry(&entry)
}
pub fn append_entry(&mut self, entry: &JournalEntry) -> Result<(), JournalError> {
let line = serde_json::to_string(entry).map_err(std::io::Error::other)?;
match self
.writer
.write_all(line.as_bytes())
.and_then(|()| self.writer.write_all(b"\n"))
.and_then(|()| self.writer.flush())
{
Ok(()) => Ok(()),
Err(e) => {
match self.reopen_writer() {
Ok(()) => Err(JournalError::Io(e)),
Err(reopen) => Err(JournalError::Io(reopen)),
}
}
}
}
pub fn sync(&mut self) -> Result<(), JournalError> {
self.writer.flush()?;
self.writer.get_ref().sync_all()?;
Ok(())
}
fn reopen_writer(&mut self) -> std::io::Result<()> {
let file = OpenOptions::new()
.create(true)
.append(true)
.open(&self.path)?;
self.writer = BufWriter::new(file);
Ok(())
}
pub fn load(path: &Path) -> Result<Vec<JournalEntry>, JournalError> {
let file = match File::open(path) {
Ok(f) => f,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(e) => return Err(JournalError::Io(e)),
};
let mut entries = Vec::new();
for line in BufReader::new(file).lines() {
let line = line?;
let line = line.trim();
if line.is_empty() {
continue;
}
if let Ok(entry) = serde_json::from_str::<JournalEntry>(line) {
if entry.schema <= SCHEMA_VERSION {
entries.push(entry);
}
}
}
Ok(entries)
}
pub fn embedders(path: &Path) -> Result<Vec<String>, JournalError> {
let set: BTreeSet<String> = Self::load(path)?
.into_iter()
.map(|entry| entry.embedder)
.collect();
Ok(set.into_iter().collect())
}
pub fn load_records(path: &Path, embedder: Option<&str>) -> Result<RecordSet, JournalError> {
let entries = Self::load(path)?;
if entries.is_empty() {
return Err(JournalError::Record(TopologyError::NoRecords {
kind: "journal",
}));
}
let wanted = match embedder {
Some(name) => name.to_string(),
None => {
let found: BTreeSet<&str> = entries.iter().map(|e| e.embedder.as_str()).collect();
if found.len() > 1 {
return Err(JournalError::MixedEmbedders {
embedders: found.into_iter().map(str::to_owned).collect(),
});
}
found
.into_iter()
.next()
.expect("non-empty entries hold at least one embedder")
.to_string()
}
};
let records: Vec<ExecutionRecord> = entries
.into_iter()
.filter(|e| e.embedder == wanted)
.map(|e| e.record)
.collect();
if records.is_empty() {
return Err(JournalError::NoRecordsForEmbedder { embedder: wanted });
}
Ok(RecordSet::with_embedder(records, wanted)?)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::topology::{CoordinationShape, Topology};
use crate::{SelectorConfig, TopologySelector};
fn record(task: &str, q: f32, shape: CoordinationShape, tokens: u64) -> ExecutionRecord {
ExecutionRecord::new(
task,
vec![q, 1.0 - q],
shape.topology(4).unwrap(),
1.0,
tokens,
)
}
#[test]
fn append_then_load_round_trips() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append(
"mini-lm",
&record("t1", 0.2, CoordinationShape::Debate, 900),
)
.unwrap();
journal
.append(
"mini-lm",
&record("t1", 0.2, CoordinationShape::Pipeline, 300),
)
.unwrap();
journal.sync().unwrap();
}
let entries = RecordJournal::load(&path).unwrap();
assert_eq!(entries.len(), 2);
assert!(entries.iter().all(|e| e.embedder == "mini-lm"));
assert!(entries.iter().all(|e| e.schema == SCHEMA_VERSION));
assert_eq!(entries[0].record.tokens, 900);
}
#[test]
fn the_journal_lands_under_the_state_root() {
let path = journal_path(Path::new("/state"));
assert!(path.ends_with("topology/records.jsonl"), "{path:?}");
}
#[test]
fn a_missing_journal_is_empty_not_an_error() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = dir.join("does-not-exist.jsonl");
assert!(RecordJournal::load(&path).unwrap().is_empty());
assert!(RecordJournal::embedders(&path).unwrap().is_empty());
}
#[test]
fn reopening_preserves_earlier_records() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
for tokens in [100u64, 200, 300] {
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Debate, tokens))
.unwrap();
}
let entries = RecordJournal::load(&path).unwrap();
assert_eq!(entries.len(), 3);
assert_eq!(entries[2].record.tokens, 300);
}
#[test]
fn a_torn_tail_is_skipped_and_the_next_append_starts_a_fresh_line() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Debate, 100))
.unwrap();
}
{
let mut raw = fs::read_to_string(&path).unwrap();
raw.push_str(r#"{"schema":1,"embedder":"e","record":{"task_id":"#);
fs::write(&path, raw).unwrap();
}
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Pipeline, 200))
.unwrap();
}
let entries = RecordJournal::load(&path).unwrap();
assert_eq!(entries.len(), 2, "torn line skipped, both intact ones kept");
assert_eq!(entries[1].record.tokens, 200);
}
#[test]
fn blank_lines_are_tolerated() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Debate, 100))
.unwrap();
}
let raw = fs::read_to_string(&path).unwrap();
fs::write(&path, format!("\n\n{raw}\n\n")).unwrap();
assert_eq!(RecordJournal::load(&path).unwrap().len(), 1);
}
#[test]
fn an_internally_inconsistent_topology_line_is_skipped_not_loaded() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Debate, 100))
.unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Pipeline, 200))
.unwrap();
}
{
let mut raw = fs::read_to_string(&path).unwrap();
raw.push_str(
r#"{"schema":1,"embedder":"e","record":{"task_id":"bad","query":[0.5,0.5],"#,
);
raw.push_str(r#""topology":{"n":4,"edges":[true]},"utility":1.0,"tokens":5}}"#);
raw.push('\n');
fs::write(&path, raw).unwrap();
}
let entries = RecordJournal::load(&path).unwrap();
assert_eq!(entries.len(), 2, "the inconsistent line must be skipped");
assert!(entries.iter().all(|e| e.record.task_id == "t"));
for entry in &entries {
let t = &entry.record.topology;
for i in 0..t.n() {
for j in 0..t.n() {
let _ = t.edge(i, j);
}
}
}
}
#[test]
fn a_future_schema_line_is_skipped_not_guessed_at() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("e", &record("t", 0.5, CoordinationShape::Debate, 100))
.unwrap();
}
let mut future = JournalEntry::new("e", record("t", 0.5, CoordinationShape::Pipeline, 7));
future.schema = SCHEMA_VERSION + 1;
{
let mut raw = fs::read_to_string(&path).unwrap();
raw.push_str(&serde_json::to_string(&future).unwrap());
raw.push('\n');
fs::write(&path, raw).unwrap();
}
let entries = RecordJournal::load(&path).unwrap();
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].record.tokens, 100);
}
#[test]
fn a_second_writer_is_refused_while_the_first_is_open() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
let _first = RecordJournal::open(&path).unwrap();
match RecordJournal::open(&path) {
Err(JournalError::Io(e)) => {
assert_eq!(e.kind(), std::io::ErrorKind::WouldBlock, "{e}")
}
other => panic!("expected a WouldBlock refusal, got {other:?}"),
}
}
#[test]
fn the_lock_is_released_when_the_journal_drops() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let _first = RecordJournal::open(&path).unwrap();
}
assert!(RecordJournal::open(&path).is_ok());
}
#[test]
fn load_records_labels_the_set_with_its_embedder() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
for i in 0..4 {
let task = format!("t{i}");
journal
.append(
"mini-lm",
&record(&task, 0.3, CoordinationShape::Debate, 900),
)
.unwrap();
journal
.append(
"mini-lm",
&record(&task, 0.3, CoordinationShape::Pipeline, 300),
)
.unwrap();
}
}
let set = RecordJournal::load_records(&path, None).unwrap();
assert_eq!(set.len(), 8);
assert_eq!(set.embedder(), Some("mini-lm"));
}
#[test]
fn a_mixed_journal_will_not_fold_into_one_set_by_accident() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("mini-lm", &record("t", 0.3, CoordinationShape::Debate, 900))
.unwrap();
journal
.append(
"bge-small",
&record("t", 0.3, CoordinationShape::Pipeline, 300),
)
.unwrap();
}
match RecordJournal::load_records(&path, None) {
Err(JournalError::MixedEmbedders { embedders }) => {
assert_eq!(embedders, vec!["bge-small", "mini-lm"]);
}
other => panic!("expected MixedEmbedders, got {other:?}"),
}
assert_eq!(
RecordJournal::embedders(&path).unwrap(),
vec!["bge-small", "mini-lm"]
);
let set = RecordJournal::load_records(&path, Some("mini-lm")).unwrap();
assert_eq!(set.len(), 1);
assert_eq!(set.embedder(), Some("mini-lm"));
}
#[test]
fn asking_for_an_absent_embedder_is_an_error() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal
.append("mini-lm", &record("t", 0.3, CoordinationShape::Debate, 900))
.unwrap();
}
assert!(matches!(
RecordJournal::load_records(&path, Some("bge-small")),
Err(JournalError::NoRecordsForEmbedder { .. })
));
}
#[test]
fn an_empty_journal_reports_no_records() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
RecordJournal::open(&path).unwrap();
assert!(matches!(
RecordJournal::load_records(&path, None),
Err(JournalError::Record(TopologyError::NoRecords { .. }))
));
}
#[test]
fn a_selector_fitted_from_the_journal_refuses_a_foreign_query() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
{
let mut journal = RecordJournal::open(&path).unwrap();
for i in 0..6 {
let task = format!("t{i}");
let q = i as f32 / 6.0;
for (shape, tokens) in [
(CoordinationShape::Debate, 600),
(CoordinationShape::Pipeline, 1800),
] {
journal
.append("mini-lm", &record(&task, q, shape, tokens))
.unwrap();
}
}
}
let set = RecordJournal::load_records(&path, None).unwrap();
let selector = TopologySelector::fit(&set, &SelectorConfig::default()).unwrap();
assert_eq!(selector.embedder(), Some("mini-lm"));
assert!(selector.select_with(&[0.5, 0.5], "mini-lm").is_ok());
match selector.select_with(&[0.5, 0.5], "bge-small") {
Err(TopologyError::EmbedderMismatch { fitted, query }) => {
assert_eq!(fitted, "mini-lm");
assert_eq!(query, "bge-small");
}
other => panic!("expected EmbedderMismatch, got {other:?}"),
}
assert!(selector.select(&[0.5, 0.5]).is_ok());
}
#[test]
fn an_unlabeled_selector_cannot_be_checked_so_the_query_passes() {
let set = RecordSet::new(vec![
record("t", 0.3, CoordinationShape::Debate, 900),
record("t", 0.3, CoordinationShape::Pipeline, 300),
])
.unwrap();
let selector = TopologySelector::fit(&set, &SelectorConfig::default()).unwrap();
assert_eq!(selector.embedder(), None);
assert!(selector.select_with(&[0.5, 0.5], "anything").is_ok());
}
#[test]
fn topologies_survive_the_json_round_trip_intact() {
let dir = tempfile::tempdir().unwrap();
let dir = dir.path();
let path = journal_path(dir);
let original = record("t", 0.25, CoordinationShape::Supervisor, 1234);
{
let mut journal = RecordJournal::open(&path).unwrap();
journal.append("e", &original).unwrap();
}
let loaded = RecordJournal::load(&path).unwrap();
assert_eq!(loaded[0].record, original);
assert_eq!(
loaded[0].record.topology,
Topology::star(4, 3).unwrap(),
"supervisor is a star hubbed on the last agent"
);
}
}