Skip to main content

agentmux_srv/backend/storage/filestore/
zone_txn.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Read-modify-write of small files in one zone, as one transaction across
5//! processes (SPEC_MEMORY_FOLLOWS_THE_AGENT_2026_09_24.md ยง2.1.1).
6//!
7//! The global transcript store is one database file that every srv instance
8//! on the machine opens. A caller that reads a file, decides, and writes it
9//! back (the memory record's `heads.json` next to its `log.jsonl`) must do
10//! all of that inside one `BEGIN IMMEDIATE` transaction, or two processes
11//! both read the same state and one's write is lost. Reads here are sized
12//! from the database row, never this process's cache, which another
13//! process's write leaves stale; the touched files are dropped from the
14//! cache after the commit.
15
16use 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
22/// One zone, inside one write transaction. See [`FileStore::zone_txn`].
23pub struct ZoneTxn<'t> {
24    tx: &'t Transaction<'t>,
25    zone_id: String,
26    now: i64,
27    touched: Vec<String>,
28}
29
30impl ZoneTxn<'_> {
31    /// A file's whole content as committed, or `None` if it doesn't exist.
32    pub fn read(&self, name: &str) -> Result<Option<Vec<u8>>, StoreError> {
33        read_in(self.tx, &self.zone_id, name)
34    }
35
36    /// Replace a file's content, creating the file if it doesn't exist.
37    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    /// Create a file with `data` unless it already exists; `true` if it was
45    /// created. For content-addressed files, whose content never changes.
46    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    /// Append complete `\n`-terminated lines, creating the file if needed.
55    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
90/// A file's whole content, sized from its row in `tx`.
91fn 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    /// Run `f` over `zone_id` as one `BEGIN IMMEDIATE` transaction: its reads
101    /// see the committed state and its writes commit together, or not at all
102    /// when `f` returns `Err`. Serialized against every process sharing the
103    /// database.
104    ///
105    /// `f` runs with this store's connection locked: it must not call any
106    /// other `FileStore` method (the lock isn't reentrant โ€” that deadlocks).
107    /// Use the [`ZoneTxn`] it is given.
108    #[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    /// Read several files of one zone as one consistent snapshot, sized from
128    /// the database โ€” never from this process's cache, which another
129    /// process's write leaves stale.
130    #[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    /// Two handles on one database file stand in for two srv processes: a
182    /// read through one sees the other's rewrite at its real size, even with
183    /// the old size cached.
184    #[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 creates the file, so b's cache holds its row (size 0) โ€” the
191        // stale state a size-from-cache read would trust.
192        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    /// A read-decide-write in one process can't interleave with another's:
200    /// both increments land.
201    #[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}