use std::collections::HashMap;
use std::path::PathBuf;
use std::time::{Duration, Instant};
use chrono::{DateTime, Utc};
use async_trait::async_trait;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;
use crate::event::Event;
#[derive(Debug, thiserror::Error)]
pub enum JournalError {
#[error("sequence conflict for run {run_id}: expected {expected}, got {got}")]
Conflict { run_id: String, expected: u64, got: u64 },
#[error("io error: {0}")]
Io(#[from] std::io::Error),
#[error("corrupt journal entry: {0}")]
Corrupt(String),
#[error("journal backend error: {0}")]
Backend(String),
}
#[derive(Debug, Clone, PartialEq)]
pub struct LeaseHolder {
pub owner: String,
pub expires_at: DateTime<Utc>,
}
#[async_trait]
pub trait Journal: Send + Sync {
async fn append(&self, event: &Event) -> Result<(), JournalError>;
async fn load(&self, run_id: &str) -> Result<Vec<Event>, JournalError>;
async fn acquire_lease(
&self,
_run_id: &str,
_owner: &str,
_ttl: Duration,
) -> Result<Result<(), LeaseHolder>, JournalError> {
Ok(Ok(()))
}
async fn renew_lease(&self, _run_id: &str, _owner: &str, _ttl: Duration) -> Result<bool, JournalError> {
Ok(true)
}
async fn release_lease(&self, _run_id: &str, _owner: &str) -> Result<(), JournalError> {
Ok(())
}
}
#[derive(Default)]
struct Leases(std::sync::Mutex<HashMap<String, (String, Instant)>>);
impl Leases {
fn acquire(&self, run_id: &str, owner: &str, ttl: Duration) -> Result<(), LeaseHolder> {
let mut map = self.0.lock().unwrap();
let now = Instant::now();
if let Some((holder, until)) = map.get(run_id) {
if holder != owner && *until > now {
let left = chrono::Duration::from_std(*until - now).unwrap_or_default();
return Err(LeaseHolder { owner: holder.clone(), expires_at: Utc::now() + left });
}
}
map.insert(run_id.to_owned(), (owner.to_owned(), now + ttl));
Ok(())
}
fn renew(&self, run_id: &str, owner: &str, ttl: Duration) -> bool {
let mut map = self.0.lock().unwrap();
match map.get_mut(run_id) {
Some((holder, until)) if holder == owner => {
*until = Instant::now() + ttl;
true
}
_ => false,
}
}
fn release(&self, run_id: &str, owner: &str) {
let mut map = self.0.lock().unwrap();
if map.get(run_id).is_some_and(|(holder, _)| holder == owner) {
map.remove(run_id);
}
}
}
#[derive(Default)]
pub struct InMemoryJournal {
runs: Mutex<HashMap<String, Vec<Event>>>,
leases: Leases,
}
impl InMemoryJournal {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl Journal for InMemoryJournal {
async fn append(&self, event: &Event) -> Result<(), JournalError> {
let mut runs = self.runs.lock().await;
let events = runs.entry(event.run_id.clone()).or_default();
let expected = events.len() as u64;
if event.seq != expected {
return Err(JournalError::Conflict { run_id: event.run_id.clone(), expected, got: event.seq });
}
events.push(event.clone());
Ok(())
}
async fn load(&self, run_id: &str) -> Result<Vec<Event>, JournalError> {
Ok(self.runs.lock().await.get(run_id).cloned().unwrap_or_default())
}
async fn acquire_lease(
&self,
run_id: &str,
owner: &str,
ttl: Duration,
) -> Result<Result<(), LeaseHolder>, JournalError> {
Ok(self.leases.acquire(run_id, owner, ttl))
}
async fn renew_lease(&self, run_id: &str, owner: &str, ttl: Duration) -> Result<bool, JournalError> {
Ok(self.leases.renew(run_id, owner, ttl))
}
async fn release_lease(&self, run_id: &str, owner: &str) -> Result<(), JournalError> {
self.leases.release(run_id, owner);
Ok(())
}
}
pub struct JsonlJournal {
dir: PathBuf,
next_seq: Mutex<HashMap<String, u64>>,
leases: Leases,
}
impl JsonlJournal {
pub fn new(dir: impl Into<PathBuf>) -> Result<Self, JournalError> {
let dir = dir.into();
std::fs::create_dir_all(&dir)?;
Ok(Self { dir, next_seq: Mutex::new(HashMap::new()), leases: Leases::default() })
}
pub fn path_for(&self, run_id: &str) -> PathBuf {
self.dir.join(format!("{run_id}.jsonl"))
}
pub fn read_file(path: &std::path::Path) -> Result<Vec<Event>, JournalError> {
let text = std::fs::read_to_string(path)?;
parse_lines(&text)
}
}
fn parse_lines(text: &str) -> Result<Vec<Event>, JournalError> {
text.lines()
.filter(|l| !l.trim().is_empty())
.map(|l| serde_json::from_str(l).map_err(|e| JournalError::Corrupt(e.to_string())))
.collect()
}
#[async_trait]
impl Journal for JsonlJournal {
async fn append(&self, event: &Event) -> Result<(), JournalError> {
let mut seqs = self.next_seq.lock().await;
let expected = match seqs.get(&event.run_id) {
Some(n) => *n,
None => self.load(&event.run_id).await?.len() as u64,
};
if event.seq != expected {
return Err(JournalError::Conflict { run_id: event.run_id.clone(), expected, got: event.seq });
}
let mut line = serde_json::to_vec(event).map_err(|e| JournalError::Corrupt(e.to_string()))?;
line.push(b'\n');
let mut file =
tokio::fs::OpenOptions::new().create(true).append(true).open(self.path_for(&event.run_id)).await?;
file.write_all(&line).await?;
file.sync_data().await?;
seqs.insert(event.run_id.clone(), expected + 1);
Ok(())
}
async fn load(&self, run_id: &str) -> Result<Vec<Event>, JournalError> {
match tokio::fs::read_to_string(self.path_for(run_id)).await {
Ok(text) => parse_lines(&text),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Vec::new()),
Err(e) => Err(e.into()),
}
}
async fn acquire_lease(
&self,
run_id: &str,
owner: &str,
ttl: Duration,
) -> Result<Result<(), LeaseHolder>, JournalError> {
Ok(self.leases.acquire(run_id, owner, ttl))
}
async fn renew_lease(&self, run_id: &str, owner: &str, ttl: Duration) -> Result<bool, JournalError> {
Ok(self.leases.renew(run_id, owner, ttl))
}
async fn release_lease(&self, run_id: &str, owner: &str) -> Result<(), JournalError> {
self.leases.release(run_id, owner);
Ok(())
}
}