use std::collections::BTreeMap;
use std::fs::{self, File};
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use serde::{Deserialize, Serialize};
use crate::run::{PauseKind, WorkflowOutcome};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum WorkflowRunStatus {
Running,
Completed,
Paused,
BudgetExceeded,
Cancelled,
Failed,
}
impl WorkflowRunStatus {
#[must_use]
pub const fn from_outcome(outcome: &WorkflowOutcome) -> Self {
match outcome {
WorkflowOutcome::Completed { .. } => Self::Completed,
WorkflowOutcome::Paused { .. } => Self::Paused,
WorkflowOutcome::BudgetExceeded { .. } => Self::BudgetExceeded,
WorkflowOutcome::Cancelled => Self::Cancelled,
WorkflowOutcome::Failed { .. } => Self::Failed,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct WorkflowRunRecord {
pub run_id: String,
pub name: String,
#[serde(default)]
pub description: String,
pub status: WorkflowRunStatus,
pub journal_path: PathBuf,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub script_ref: Option<String>,
pub created_at_ms: u64,
pub updated_at_ms: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub pause_kind: Option<PauseKind>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub message: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub result: Option<serde_json::Value>,
}
impl WorkflowRunRecord {
#[must_use]
pub fn new_running(
run_id: impl Into<String>,
name: impl Into<String>,
journal_path: PathBuf,
) -> Self {
let now = unix_now_ms();
Self {
run_id: run_id.into(),
name: name.into(),
description: String::new(),
status: WorkflowRunStatus::Running,
journal_path,
script_ref: None,
created_at_ms: now,
updated_at_ms: now,
pause_kind: None,
message: None,
result: None,
}
}
pub fn apply_outcome(&mut self, outcome: &WorkflowOutcome) {
self.status = WorkflowRunStatus::from_outcome(outcome);
self.updated_at_ms = unix_now_ms();
match outcome {
WorkflowOutcome::Paused { kind, message } => {
self.pause_kind = Some(*kind);
self.message = Some(message.clone());
self.result = None;
}
WorkflowOutcome::BudgetExceeded { message }
| WorkflowOutcome::Failed { error: message } => {
self.pause_kind = None;
self.message = Some(message.clone());
self.result = None;
}
WorkflowOutcome::Completed { result } => {
self.pause_kind = None;
self.message = None;
self.result = Some(result.clone());
}
WorkflowOutcome::Cancelled => {
self.pause_kind = None;
self.message = None;
self.result = None;
}
}
}
}
pub trait WorkflowRunStore: Send + Sync {
fn put(&self, record: WorkflowRunRecord) -> Result<(), StoreError>;
fn get(&self, run_id: &str) -> Result<Option<WorkflowRunRecord>, StoreError>;
fn list(&self) -> Result<Vec<WorkflowRunRecord>, StoreError>;
fn delete(&self, run_id: &str) -> Result<bool, StoreError>;
}
#[derive(Debug, thiserror::Error)]
pub enum StoreError {
#[error("workflow store io: {0}")]
Io(#[from] std::io::Error),
#[error("workflow store parse: {0}")]
Parse(String),
}
#[derive(Debug, Default)]
pub struct MemoryWorkflowRunStore {
map: Mutex<BTreeMap<String, WorkflowRunRecord>>,
}
impl MemoryWorkflowRunStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
}
impl WorkflowRunStore for MemoryWorkflowRunStore {
fn put(&self, record: WorkflowRunRecord) -> Result<(), StoreError> {
self.map
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(record.run_id.clone(), record);
Ok(())
}
fn get(&self, run_id: &str) -> Result<Option<WorkflowRunRecord>, StoreError> {
Ok(self
.map
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.get(run_id)
.cloned())
}
fn list(&self) -> Result<Vec<WorkflowRunRecord>, StoreError> {
let mut rows: Vec<_> = self
.map
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.values()
.cloned()
.collect();
rows.sort_by_key(|r| std::cmp::Reverse(r.updated_at_ms));
Ok(rows)
}
fn delete(&self, run_id: &str) -> Result<bool, StoreError> {
Ok(self
.map
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(run_id)
.is_some())
}
}
#[derive(Debug, Clone)]
pub struct FileWorkflowRunStore {
root: PathBuf,
}
impl FileWorkflowRunStore {
#[must_use]
pub fn new(root: impl Into<PathBuf>) -> Self {
Self { root: root.into() }
}
#[must_use]
pub fn root(&self) -> &Path {
&self.root
}
#[must_use]
pub fn journal_path_for(&self, run_id: &str) -> PathBuf {
self.root.join("journals").join(format!("{run_id}.jsonl"))
}
fn runs_dir(&self) -> PathBuf {
self.root.join("runs")
}
fn record_path(&self, run_id: &str) -> PathBuf {
self.runs_dir().join(format!("{run_id}.json"))
}
}
impl WorkflowRunStore for FileWorkflowRunStore {
fn put(&self, record: WorkflowRunRecord) -> Result<(), StoreError> {
let dir = self.runs_dir();
fs::create_dir_all(&dir)?;
let path = self.record_path(&record.run_id);
let tmp = path.with_extension("json.tmp");
let body =
serde_json::to_vec_pretty(&record).map_err(|e| StoreError::Parse(e.to_string()))?;
{
let mut f = File::create(&tmp)?;
f.write_all(&body)?;
f.sync_data()?;
}
fs::rename(&tmp, &path)?;
Ok(())
}
fn get(&self, run_id: &str) -> Result<Option<WorkflowRunRecord>, StoreError> {
let path = self.record_path(run_id);
if !path.is_file() {
return Ok(None);
}
let bytes = fs::read(&path)?;
let rec = serde_json::from_slice(&bytes).map_err(|e| StoreError::Parse(e.to_string()))?;
Ok(Some(rec))
}
fn list(&self) -> Result<Vec<WorkflowRunRecord>, StoreError> {
let dir = self.runs_dir();
if !dir.is_dir() {
return Ok(Vec::new());
}
let mut rows = Vec::new();
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("json") {
continue;
}
let bytes = fs::read(&path)?;
let rec: WorkflowRunRecord =
serde_json::from_slice(&bytes).map_err(|e| StoreError::Parse(e.to_string()))?;
rows.push(rec);
}
rows.sort_by_key(|r| std::cmp::Reverse(r.updated_at_ms));
Ok(rows)
}
fn delete(&self, run_id: &str) -> Result<bool, StoreError> {
let path = self.record_path(run_id);
if !path.is_file() {
return Ok(false);
}
fs::remove_file(path)?;
Ok(true)
}
}
pub fn peek_jsonl_line(path: &Path) -> Result<Option<String>, StoreError> {
if !path.is_file() {
return Ok(None);
}
let f = File::open(path)?;
let mut lines = BufReader::new(f).lines();
match lines.next() {
Some(Ok(line)) => Ok(Some(line)),
Some(Err(e)) => Err(e.into()),
None => Ok(None),
}
}
fn unix_now_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| u64::try_from(d.as_millis()).unwrap_or(u64::MAX))
}
#[cfg(test)]
mod tests {
use tempfile::tempdir;
use super::*;
#[test]
fn memory_put_list_delete() {
let store = MemoryWorkflowRunStore::new();
let mut rec = WorkflowRunRecord::new_running("run_1", "demo", PathBuf::from("j.jsonl"));
store.put(rec.clone()).expect("put");
rec.apply_outcome(&WorkflowOutcome::Completed {
result: serde_json::json!({"ok": true}),
});
store.put(rec).expect("put2");
let listed = store.list().expect("list");
assert_eq!(listed.len(), 1);
assert_eq!(
listed.first().map(|r| r.status),
Some(WorkflowRunStatus::Completed)
);
assert!(store.delete("run_1").expect("del"));
assert!(store.get("run_1").expect("get").is_none());
}
#[test]
fn file_store_round_trip() {
let dir = tempdir().expect("tmp");
let store = FileWorkflowRunStore::new(dir.path());
let journal = store.journal_path_for("wf_abc");
let mut rec = WorkflowRunRecord::new_running("wf_abc", "fanout", journal);
rec.description = "test".into();
store.put(rec.clone()).expect("put");
let got = store.get("wf_abc").expect("get").expect("some");
assert_eq!(got.name, "fanout");
assert_eq!(store.list().expect("list").len(), 1);
rec.apply_outcome(&WorkflowOutcome::Failed {
error: "boom".into(),
});
store.put(rec).expect("put fail");
let got = store.get("wf_abc").expect("g2").expect("s");
assert_eq!(got.status, WorkflowRunStatus::Failed);
assert_eq!(got.message.as_deref(), Some("boom"));
}
}