a3s_code_core/state_graph/
store.rs

1use super::{GraphEventRecord, GraphRuntime};
2use anyhow::{Context, Result};
3use async_trait::async_trait;
4use std::collections::HashMap;
5use std::path::{Path, PathBuf};
6use tokio::io::AsyncWriteExt;
7use tokio::sync::{Mutex, RwLock};
8
9const MAX_GRAPH_EVENT_LOG_BYTES: usize = 256 * 1024 * 1024;
10
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum GraphSaveOutcome {
13    Saved,
14    Conflict { actual_head: Option<String> },
15}
16
17pub fn graph_event_head(events: &[GraphEventRecord]) -> Option<&str> {
18    events.last().map(|record| record.record_hash.as_str())
19}
20
21#[async_trait]
22pub trait GraphEventStore: Send + Sync {
23    /// Atomically replace one branch generation. Intended for a single writer.
24    async fn save(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()>;
25    /// Compare the persisted head and publish an extending generation atomically.
26    async fn save_if_head(
27        &self,
28        _branch_id: &str,
29        _expected_head: Option<&str>,
30        _events: &[GraphEventRecord],
31    ) -> Result<GraphSaveOutcome> {
32        anyhow::bail!("graph event store does not support compare-and-swap")
33    }
34    async fn load(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>>;
35    async fn delete(&self, branch_id: &str) -> Result<()>;
36}
37
38#[derive(Debug, Default)]
39pub struct MemoryGraphEventStore {
40    branches: RwLock<HashMap<String, Vec<GraphEventRecord>>>,
41}
42
43impl MemoryGraphEventStore {
44    pub fn new() -> Self {
45        Self::default()
46    }
47}
48
49#[async_trait]
50impl GraphEventStore for MemoryGraphEventStore {
51    async fn save(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()> {
52        self.branches
53            .write()
54            .await
55            .insert(branch_id.to_string(), events.to_vec());
56        Ok(())
57    }
58
59    async fn save_if_head(
60        &self,
61        branch_id: &str,
62        expected_head: Option<&str>,
63        events: &[GraphEventRecord],
64    ) -> Result<GraphSaveOutcome> {
65        validate_graph_events(events)?;
66        let mut branches = self.branches.write().await;
67        let current = branches.get(branch_id);
68        let actual_head = current.and_then(|records| graph_event_head(records));
69        if actual_head != expected_head {
70            return Ok(GraphSaveOutcome::Conflict {
71                actual_head: actual_head.map(str::to_string),
72            });
73        }
74        ensure_extends(current.map(Vec::as_slice), events)?;
75        branches.insert(branch_id.to_string(), events.to_vec());
76        Ok(GraphSaveOutcome::Saved)
77    }
78
79    async fn load(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>> {
80        Ok(self.branches.read().await.get(branch_id).cloned())
81    }
82
83    async fn delete(&self, branch_id: &str) -> Result<()> {
84        self.branches.write().await.remove(branch_id);
85        Ok(())
86    }
87}
88
89#[derive(Debug)]
90pub struct FileGraphEventStore {
91    root: PathBuf,
92    write_lock: Mutex<()>,
93}
94
95impl FileGraphEventStore {
96    pub fn new(root: impl Into<PathBuf>) -> Self {
97        Self {
98            root: root.into(),
99            write_lock: Mutex::new(()),
100        }
101    }
102
103    fn path(&self, branch_id: &str) -> PathBuf {
104        self.root
105            .join(format!("{}.json", sha256::digest(branch_id.as_bytes())))
106    }
107
108    fn temp_path(path: &Path) -> PathBuf {
109        path.with_extension(format!("tmp-{}", uuid::Uuid::new_v4()))
110    }
111
112    async fn acquire_process_lock(&self) -> Result<std::fs::File> {
113        tokio::fs::create_dir_all(&self.root)
114            .await
115            .with_context(|| format!("create graph event store `{}`", self.root.display()))?;
116        let path = self.root.join(".graph-store.lock");
117        tokio::task::spawn_blocking(move || {
118            use fs2::FileExt;
119            let file = std::fs::OpenOptions::new()
120                .create(true)
121                .truncate(false)
122                .read(true)
123                .write(true)
124                .open(&path)
125                .with_context(|| format!("open graph store lock `{}`", path.display()))?;
126            file.lock_exclusive()
127                .with_context(|| format!("lock graph store `{}`", path.display()))?;
128            Ok(file)
129        })
130        .await
131        .context("graph store lock task failed")?
132    }
133
134    async fn load_unlocked(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>> {
135        read_graph_events(&self.path(branch_id)).await
136    }
137
138    async fn publish_unlocked(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()> {
139        let path = self.path(branch_id);
140        let temp = Self::temp_path(&path);
141        let bytes = encode_graph_events(events)?;
142        let result = async {
143            let mut file = tokio::fs::File::create(&temp)
144                .await
145                .with_context(|| format!("create graph event log `{}`", temp.display()))?;
146            file.write_all(&bytes)
147                .await
148                .context("write graph event log")?;
149            file.sync_all().await.context("sync graph event log")?;
150            drop(file);
151            let temp = temp.clone();
152            let path = path.clone();
153            tokio::task::spawn_blocking(move || {
154                tempfile::TempPath::try_from_path(temp)?
155                    .persist(path)
156                    .map_err(|error| error.error)
157            })
158            .await
159            .context("atomic graph event replace task failed")??;
160            Ok::<_, anyhow::Error>(())
161        }
162        .await;
163        if result.is_err() {
164            let _ = tokio::fs::remove_file(&temp).await;
165        }
166        result
167    }
168}
169
170#[async_trait]
171impl GraphEventStore for FileGraphEventStore {
172    async fn save(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()> {
173        let _guard = self.write_lock.lock().await;
174        let _process_guard = self.acquire_process_lock().await?;
175        self.publish_unlocked(branch_id, events).await
176    }
177
178    async fn save_if_head(
179        &self,
180        branch_id: &str,
181        expected_head: Option<&str>,
182        events: &[GraphEventRecord],
183    ) -> Result<GraphSaveOutcome> {
184        validate_graph_events(events)?;
185        let _guard = self.write_lock.lock().await;
186        let _process_guard = self.acquire_process_lock().await?;
187        let current = self.load_unlocked(branch_id).await?;
188        let actual_head = current.as_deref().and_then(graph_event_head);
189        if actual_head != expected_head {
190            return Ok(GraphSaveOutcome::Conflict {
191                actual_head: actual_head.map(str::to_string),
192            });
193        }
194        ensure_extends(current.as_deref(), events)?;
195        self.publish_unlocked(branch_id, events).await?;
196        Ok(GraphSaveOutcome::Saved)
197    }
198
199    async fn load(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>> {
200        self.load_unlocked(branch_id).await
201    }
202
203    async fn delete(&self, branch_id: &str) -> Result<()> {
204        let _guard = self.write_lock.lock().await;
205        let _process_guard = self.acquire_process_lock().await?;
206        match tokio::fs::remove_file(self.path(branch_id)).await {
207            Ok(()) => Ok(()),
208            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
209            Err(error) => Err(error).context("delete graph event log"),
210        }
211    }
212}
213
214fn validate_graph_events(events: &[GraphEventRecord]) -> Result<()> {
215    GraphRuntime::strict_replay(events).context("validate graph event generation")?;
216    Ok(())
217}
218
219fn ensure_extends(
220    current: Option<&[GraphEventRecord]>,
221    candidate: &[GraphEventRecord],
222) -> Result<()> {
223    let Some(current) = current else {
224        return Ok(());
225    };
226    if candidate.len() < current.len() || candidate[..current.len()] != *current {
227        anyhow::bail!("candidate graph generation does not extend the persisted history");
228    }
229    Ok(())
230}
231
232fn encode_graph_events(events: &[GraphEventRecord]) -> Result<Vec<u8>> {
233    let bytes = serde_json::to_vec(events).context("serialize graph event log")?;
234    if bytes.len() > MAX_GRAPH_EVENT_LOG_BYTES {
235        anyhow::bail!(
236            "graph event log is {} bytes; limit is {} bytes",
237            bytes.len(),
238            MAX_GRAPH_EVENT_LOG_BYTES
239        );
240    }
241    Ok(bytes)
242}
243
244async fn read_graph_events(path: &Path) -> Result<Option<Vec<GraphEventRecord>>> {
245    let bytes = match tokio::fs::read(path).await {
246        Ok(bytes) => bytes,
247        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
248        Err(error) => {
249            return Err(error).with_context(|| format!("read graph event log `{}`", path.display()))
250        }
251    };
252    if bytes.len() > MAX_GRAPH_EVENT_LOG_BYTES {
253        anyhow::bail!(
254            "graph event log `{}` is {} bytes; limit is {} bytes",
255            path.display(),
256            bytes.len(),
257            MAX_GRAPH_EVENT_LOG_BYTES
258        );
259    }
260    serde_json::from_slice(&bytes)
261        .with_context(|| format!("decode graph event log `{}`", path.display()))
262        .map(Some)
263}
a3s_code_core/state_graph/store.rs

a3s_code_core/state_graph/
store.rs
