harness_loop/orchestrator/
store.rs

1//! `RunStore` — persists Run + Job state so a Run can resume after a crash.
2//!
3//! Persistence is what separates "fancy async callbacks" from a recoverable
4//! task system: the store is the source of truth. The orchestrator saves the
5//! whole [`Run`] after every state transition; on resume, Jobs caught mid-
6//! flight (`Running`/`Queued`/`Retrying`) are reset to `Pending` and the Run
7//! continues from its succeeded results.
8
9use crate::orchestrator::run::Run;
10use async_trait::async_trait;
11use std::collections::HashMap;
12use std::path::PathBuf;
13use std::sync::Mutex;
14
15#[derive(Debug, thiserror::Error)]
16#[non_exhaustive]
17pub enum StoreError {
18    #[error("run store io: {0}")]
19    Io(String),
20    #[error("run store serde: {0}")]
21    Serde(String),
22}
23
24#[async_trait]
25pub trait RunStore: Send + Sync {
26    async fn save(&self, run: &Run) -> Result<(), StoreError>;
27    async fn load(&self, run_id: &str) -> Result<Option<Run>, StoreError>;
28}
29
30/// In-memory store — for tests and ephemeral Runs.
31#[derive(Default)]
32pub struct InMemoryRunStore {
33    runs: Mutex<HashMap<String, Run>>,
34}
35
36impl InMemoryRunStore {
37    pub fn new() -> Self {
38        Self::default()
39    }
40}
41
42#[async_trait]
43impl RunStore for InMemoryRunStore {
44    async fn save(&self, run: &Run) -> Result<(), StoreError> {
45        self.runs
46            .lock()
47            .unwrap()
48            .insert(run.id.clone(), run.clone());
49        Ok(())
50    }
51    async fn load(&self, run_id: &str) -> Result<Option<Run>, StoreError> {
52        Ok(self.runs.lock().unwrap().get(run_id).cloned())
53    }
54}
55
56/// One JSON file per Run under `dir` (`<dir>/<run_id>.json`).
57pub struct FileRunStore {
58    dir: PathBuf,
59}
60
61impl FileRunStore {
62    pub fn open(dir: impl Into<PathBuf>) -> Result<Self, StoreError> {
63        let dir = dir.into();
64        std::fs::create_dir_all(&dir).map_err(|e| StoreError::Io(e.to_string()))?;
65        Ok(Self { dir })
66    }
67    fn path(&self, run_id: &str) -> PathBuf {
68        // run_id is caller-controlled; keep it to a safe filename.
69        let safe: String = run_id
70            .chars()
71            .map(|c| {
72                if c.is_alphanumeric() || c == '-' || c == '_' {
73                    c
74                } else {
75                    '_'
76                }
77            })
78            .collect();
79        self.dir.join(format!("{safe}.json"))
80    }
81}
82
83#[async_trait]
84impl RunStore for FileRunStore {
85    async fn save(&self, run: &Run) -> Result<(), StoreError> {
86        let json = serde_json::to_vec_pretty(run).map_err(|e| StoreError::Serde(e.to_string()))?;
87        std::fs::write(self.path(&run.id), json).map_err(|e| StoreError::Io(e.to_string()))
88    }
89    async fn load(&self, run_id: &str) -> Result<Option<Run>, StoreError> {
90        let p = self.path(run_id);
91        if !p.exists() {
92            return Ok(None);
93        }
94        let bytes = std::fs::read(&p).map_err(|e| StoreError::Io(e.to_string()))?;
95        let run = serde_json::from_slice(&bytes).map_err(|e| StoreError::Serde(e.to_string()))?;
96        Ok(Some(run))
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::orchestrator::dag::Dag;
104    use crate::orchestrator::job::Job;
105
106    #[tokio::test]
107    async fn file_store_roundtrips() {
108        let dir = std::env::temp_dir().join(format!("harness-orch-{}", std::process::id()));
109        let store = FileRunStore::open(&dir).unwrap();
110        let run = Run::new("run-1", "goal", Dag::from_jobs([Job::new("a", "x")]));
111        store.save(&run).await.unwrap();
112        let back = store.load("run-1").await.unwrap().unwrap();
113        assert_eq!(back.id, "run-1");
114        assert_eq!(back.dag.len(), 1);
115        assert!(store.load("missing").await.unwrap().is_none());
116        let _ = std::fs::remove_dir_all(&dir);
117    }
118}
harness_loop/orchestrator/store.rs

harness_loop/orchestrator/
store.rs
