agentmux_srv/backend/storage/filestore/
zone_txn.rs1use rusqlite::{params, OptionalExtension, Transaction};
17
18use super::core::FileStore;
19use super::counter::{new_gen, read_bytes, read_row, AppendMode};
20use crate::backend::storage::error::StoreError;
21
22pub struct ZoneTxn<'t> {
24 tx: &'t Transaction<'t>,
25 zone_id: String,
26 now: i64,
27 touched: Vec<String>,
28}
29
30impl ZoneTxn<'_> {
31 pub fn read(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
33 read_in(self.tx, &self.zone_id, name)
34 }
35
36 pub fn put(&mut self, name: &str, data: &[u8]) -> Result<(), StoreError> {
38 self.ensure(name)?;
39 FileStore::replace_content_in(self.tx, &self.zone_id, name, data, self.now)?;
40 self.touch(name);
41 Ok(())
42 }
43
44 pub fn put_if_absent(&mut self, name: &str, data: &[u8]) -> Result<bool, StoreError> {
47 if read_row(self.tx, &self.zone_id, name)?.is_some() {
48 return Ok(false);
49 }
50 self.put(name, data)?;
51 Ok(true)
52 }
53
54 pub fn append_lines(&mut self, name: &str, data: &[u8]) -> Result<(), StoreError> {
56 self.ensure(name)?;
57 FileStore::append_in_tx(self.tx, &self.zone_id, name, data, AppendMode::Lines, self.now)?;
58 self.touch(name);
59 Ok(())
60 }
61
62 fn ensure(&self, name: &str) -> Result<(), StoreError> {
63 let exists = self
64 .tx
65 .query_row(
66 "SELECT 1 FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
67 params![self.zone_id, name],
68 |_| Ok(()),
69 )
70 .optional()?
71 .is_some();
72 if !exists {
73 self.tx.execute(
74 "INSERT INTO db_wave_file (zoneid, name, size, createdts, modts, opts, meta,
75 gen, lines, lines_size, lines_tail, lines_modts, lines_rev)
76 VALUES (?1, ?2, 0, ?3, ?3, '{}', '{}', ?4, 0, 0, 0, ?3, 0)",
77 params![self.zone_id, name, self.now, new_gen()],
78 )?;
79 }
80 Ok(())
81 }
82
83 fn touch(&mut self, name: &str) {
84 if !self.touched.iter().any(|n| n == name) {
85 self.touched.push(name.to_string());
86 }
87 }
88}
89
90fn read_in(tx: &Transaction<'_>, zone_id: &str, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
92 let Some(row) = read_row(tx, zone_id, name)? else { return Ok(None) };
93 if row.size == 0 {
94 return Ok(Some(Vec::new()));
95 }
96 Ok(Some(read_bytes(tx, zone_id, name, 0, row.size)?))
97}
98
99impl FileStore {
100 #[track_caller]
109 pub fn zone_txn<T>(
110 &self,
111 zone_id: &str,
112 f: impl FnOnce(&mut ZoneTxn<'_>) -> Result<T, StoreError>,
113 ) -> Result<T, StoreError> {
114 let now = agentmux_common::time::now_ms();
115 let mut touched = Vec::new();
116 let out = self.write_txn(|tx| {
117 let mut z = ZoneTxn { tx, zone_id: zone_id.to_string(), now, touched: Vec::new() };
118 let out = f(&mut z)?;
119 touched = z.touched;
120 Ok(out)
121 })?;
122 let names: Vec<&str> = touched.iter().map(String::as_str).collect();
123 self.forget_cached(zone_id, &names);
124 Ok(out)
125 }
126
127 #[track_caller]
131 pub fn read_files_consistent(&self, zone_id: &str, names: &[&str]) -> Result<Vec<Option<Vec<u8>>>, StoreError> {
132 self.read_txn(|tx| names.iter().map(|name| read_in(tx, zone_id, name)).collect())
133 }
134}
135
136#[cfg(test)]
137mod tests {
138 use super::*;
139
140 const ZONE: &str = "agent-uid:u1:memory";
141
142 #[test]
143 fn writes_commit_together_and_reads_see_them() {
144 let fs = FileStore::open_in_memory().unwrap();
145 fs.zone_txn(ZONE, |z| {
146 z.put("heads.json", b"{\"a\":1}")?;
147 z.append_lines("log.jsonl", b"{\"e\":1}\n")?;
148 z.append_lines("log.jsonl", b"{\"e\":2}\n")
149 })
150 .unwrap();
151 let got = fs.read_files_consistent(ZONE, &["heads.json", "log.jsonl", "missing"]).unwrap();
152 assert_eq!(got[0].as_deref(), Some(&b"{\"a\":1}"[..]));
153 assert_eq!(got[1].as_deref(), Some(&b"{\"e\":1}\n{\"e\":2}\n"[..]));
154 assert_eq!(got[2], None);
155 }
156
157 #[test]
158 fn an_error_rolls_back_every_write() {
159 let fs = FileStore::open_in_memory().unwrap();
160 fs.zone_txn(ZONE, |z| z.put("heads.json", b"old")).unwrap();
161 let r: Result<(), _> = fs.zone_txn(ZONE, |z| {
162 z.put("heads.json", b"new")?;
163 z.append_lines("log.jsonl", b"x\n")?;
164 Err(StoreError::Other("stale parent".into()))
165 });
166 assert!(r.is_err());
167 let got = fs.read_files_consistent(ZONE, &["heads.json", "log.jsonl"]).unwrap();
168 assert_eq!(got[0].as_deref(), Some(&b"old"[..]));
169 assert_eq!(got[1], None, "the log append rolled back with it");
170 }
171
172 #[test]
173 fn put_if_absent_never_changes_existing_content() {
174 let fs = FileStore::open_in_memory().unwrap();
175 let created = fs.zone_txn(ZONE, |z| z.put_if_absent("blob/abc", b"one")).unwrap();
176 let again = fs.zone_txn(ZONE, |z| z.put_if_absent("blob/abc", b"two")).unwrap();
177 assert!(created && !again);
178 assert_eq!(fs.read_files_consistent(ZONE, &["blob/abc"]).unwrap()[0].as_deref(), Some(&b"one"[..]));
179 }
180
181 #[test]
185 fn a_second_process_sees_a_rewrite_at_its_real_size() {
186 let dir = tempfile::tempdir().unwrap();
187 let path = dir.path().join("filestore.db");
188 let a = FileStore::open(&path).unwrap();
189 let b = FileStore::open(&path).unwrap();
190 b.make_file(ZONE, "heads.json", Default::default(), Default::default()).unwrap();
193 a.zone_txn(ZONE, |z| z.put("heads.json", b"a much longer body")).unwrap();
194 assert_eq!(b.read_file(ZONE, "heads.json").unwrap().as_deref(), Some(&b""[..]), "the cached read is stale");
195 let got = b.read_files_consistent(ZONE, &["heads.json"]).unwrap();
196 assert_eq!(got[0].as_deref(), Some(&b"a much longer body"[..]));
197 }
198
199 #[test]
202 fn concurrent_read_modify_writes_both_land() {
203 let dir = tempfile::tempdir().unwrap();
204 let path = dir.path().join("filestore.db");
205 let a = std::sync::Arc::new(FileStore::open(&path).unwrap());
206 let b = std::sync::Arc::new(FileStore::open(&path).unwrap());
207 a.zone_txn(ZONE, |z| z.put("n", b"0")).unwrap();
208 let bump = |fs: std::sync::Arc<FileStore>| {
209 std::thread::spawn(move || {
210 for _ in 0..25 {
211 fs.zone_txn(ZONE, |z| {
212 let n: u32 = String::from_utf8(z.read("n")?.unwrap()).unwrap().parse().unwrap();
213 z.put("n", (n + 1).to_string().as_bytes())
214 })
215 .unwrap();
216 }
217 })
218 };
219 let (ta, tb) = (bump(a.clone()), bump(b.clone()));
220 ta.join().unwrap();
221 tb.join().unwrap();
222 assert_eq!(a.read_files_consistent(ZONE, &["n"]).unwrap()[0].as_deref(), Some(&b"50"[..]));
223 }
224}