a3s_code_core/state_graph/
store.rs

1use super::GraphEventRecord;
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#[async_trait]
12pub trait GraphEventStore: Send + Sync {
13    /// Atomically replace one branch generation with a complete event log.
14    async fn save(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()>;
15    async fn load(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>>;
16    async fn delete(&self, branch_id: &str) -> Result<()>;
17}
18
19#[derive(Debug, Default)]
20pub struct MemoryGraphEventStore {
21    branches: RwLock<HashMap<String, Vec<GraphEventRecord>>>,
22}
23
24impl MemoryGraphEventStore {
25    pub fn new() -> Self {
26        Self::default()
27    }
28}
29
30#[async_trait]
31impl GraphEventStore for MemoryGraphEventStore {
32    async fn save(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()> {
33        self.branches
34            .write()
35            .await
36            .insert(branch_id.to_string(), events.to_vec());
37        Ok(())
38    }
39
40    async fn load(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>> {
41        Ok(self.branches.read().await.get(branch_id).cloned())
42    }
43
44    async fn delete(&self, branch_id: &str) -> Result<()> {
45        self.branches.write().await.remove(branch_id);
46        Ok(())
47    }
48}
49
50#[derive(Debug)]
51pub struct FileGraphEventStore {
52    root: PathBuf,
53    write_lock: Mutex<()>,
54}
55
56impl FileGraphEventStore {
57    pub fn new(root: impl Into<PathBuf>) -> Self {
58        Self {
59            root: root.into(),
60            write_lock: Mutex::new(()),
61        }
62    }
63
64    fn path(&self, branch_id: &str) -> PathBuf {
65        self.root
66            .join(format!("{}.json", sha256::digest(branch_id.as_bytes())))
67    }
68
69    fn temp_path(path: &Path) -> PathBuf {
70        path.with_extension(format!("tmp-{}", uuid::Uuid::new_v4()))
71    }
72}
73
74#[async_trait]
75impl GraphEventStore for FileGraphEventStore {
76    async fn save(&self, branch_id: &str, events: &[GraphEventRecord]) -> Result<()> {
77        let _guard = self.write_lock.lock().await;
78        tokio::fs::create_dir_all(&self.root)
79            .await
80            .with_context(|| format!("create graph event store `{}`", self.root.display()))?;
81        let path = self.path(branch_id);
82        let temp = Self::temp_path(&path);
83        let bytes = serde_json::to_vec(events).context("serialize graph event log")?;
84        if bytes.len() > MAX_GRAPH_EVENT_LOG_BYTES {
85            anyhow::bail!(
86                "graph event log is {} bytes; limit is {} bytes",
87                bytes.len(),
88                MAX_GRAPH_EVENT_LOG_BYTES
89            );
90        }
91        let result = async {
92            let mut file = tokio::fs::File::create(&temp)
93                .await
94                .with_context(|| format!("create graph event log `{}`", temp.display()))?;
95            file.write_all(&bytes)
96                .await
97                .context("write graph event log")?;
98            file.sync_all().await.context("sync graph event log")?;
99            drop(file);
100            let temp = temp.clone();
101            let path = path.clone();
102            tokio::task::spawn_blocking(move || {
103                tempfile::TempPath::try_from_path(temp)?
104                    .persist(path)
105                    .map_err(|error| error.error)
106            })
107            .await
108            .context("atomic graph event replace task failed")??;
109            Ok::<_, anyhow::Error>(())
110        }
111        .await;
112        if result.is_err() {
113            let _ = tokio::fs::remove_file(&temp).await;
114        }
115        result
116    }
117
118    async fn load(&self, branch_id: &str) -> Result<Option<Vec<GraphEventRecord>>> {
119        let path = self.path(branch_id);
120        let bytes = match tokio::fs::read(&path).await {
121            Ok(bytes) => bytes,
122            Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
123            Err(error) => {
124                return Err(error)
125                    .with_context(|| format!("read graph event log `{}`", path.display()))
126            }
127        };
128        if bytes.len() > MAX_GRAPH_EVENT_LOG_BYTES {
129            anyhow::bail!(
130                "graph event log `{}` is {} bytes; limit is {} bytes",
131                path.display(),
132                bytes.len(),
133                MAX_GRAPH_EVENT_LOG_BYTES
134            );
135        }
136        serde_json::from_slice(&bytes)
137            .with_context(|| format!("decode graph event log `{}`", path.display()))
138            .map(Some)
139    }
140
141    async fn delete(&self, branch_id: &str) -> Result<()> {
142        let _guard = self.write_lock.lock().await;
143        match tokio::fs::remove_file(self.path(branch_id)).await {
144            Ok(()) => Ok(()),
145            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
146            Err(error) => Err(error).context("delete graph event log"),
147        }
148    }
149}
a3s_code_core/state_graph/store.rs

a3s_code_core/state_graph/
store.rs
