From a1ff95fa61536f6880941ee506ce2796c9a0d662 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Wed, 7 Oct 2026 06:13:25 -0700 Subject: [PATCH] agent: memory store lock on file_lock; mark a lock held before taking it The memory store's `.memory.lock` was a raw flock outside file_lock: not name-matched, not in the held set, so a model write could rename over it and let two memory mutations run at once. StoreLock is now file_lock::try_lock(Target::File(/.memory)): an OFD lock on `.memory.beyond-lock` plus the legacy flock on `.memory.lock` older binaries take, both in the held set and refused by write_atomic. An old binary's `.memory.lock` with no record beside it is matched by name. try_lock now marks each lock file held before it tries to lock it, and a failed attempt drops its marks (a Marks guard, kept by FileLock until after it unlocks), closing the window where a held Itself key read as not held. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JimHGjsfk2Ktm5GxyZJKKk --- crates/agent/ARCHITECTURE.md | 18 +++-- crates/agent/src/file_lock.rs | 131 +++++++++++++++++++++--------- crates/agent/src/memory/file.rs | 137 ++++++++++++++++++++------------ 3 files changed, 194 insertions(+), 92 deletions(-) diff --git a/crates/agent/ARCHITECTURE.md b/crates/agent/ARCHITECTURE.md index 3abc5b56..d8fff7c3 100644 --- a/crates/agent/ARCHITECTURE.md +++ b/crates/agent/ARCHITECTURE.md @@ -3086,8 +3086,9 @@ coherent across a tenant's sessions. ### Locks `acquire_session_lock(path)` is `file_lock::try_lock` on the session — `Target::Dir` for a session -directory, `Target::File` for a `.jsonl` (the MCP manifest's write lock is the same primitive on the -manifest file). `Ok(None)` means someone holds it right now. It is **liveness only** — correctness is +directory, `Target::File` for a `.jsonl`. The MCP manifest's write lock and the memory store's lock +are the same primitive, `Target::File` on the manifest file and on `/.memory`; the memory +store's legacy `flock` is the `.memory.lock` older binaries took. `Ok(None)` means someone holds it right now. It is **liveness only** — correctness is the epoch fence, which needs no lock — so a lock lost to a crash or a partition costs a retry, never history. @@ -3127,19 +3128,22 @@ attempts, measured). So a target is a lock file only when one of three things sa - **Its name**, compared without regard to ASCII case: a record lock file, or a legacy `lock` / `.lock` beside one. -- **It is a lock this process holds.** `try_lock` puts the (device, inode) of each lock file it holds - in a small counted set and `Drop` takes them out once released; `write_atomic` looks the target's +- **It is a lock this process holds.** `try_lock` puts the (device, inode) of each lock file in a + small counted set **before** it tries to lock it, takes it out again if the attempt fails, and + `Drop` takes it out only after unlocking — so no moment exists when a lock is held but reads as not + held (a brief false "held" mid-acquire is the safe direction). `write_atomic` looks the target's `stat` up there. That covers a journal key held through `Target::Itself` under any name, a hard link to one of our own lock files, and a spelling a case-insensitive filesystem folds onto one. - **An old binary's legacy lock with no record lock file beside it**, by where it sits: a `lock` in a session directory (one holding `000001.jsonl`, which is never deleted), or a `.lock` beside a - session file (`*.jsonl`) or the MCP manifest. Whether anyone holds it is not asked. + session file (`*.jsonl`), the MCP manifest or the memory store's `.memory`. Whether anyone holds it + is not asked. Anything else is an ordinary file and an edit of it goes through: a SQLite database another program holds mid-transaction, a daemon's pid file, an unrelated `Cargo.lock`, a `lock` in a directory that holds no session. (Verified on a casefold ext4 directory as well as here.) A worktree seed leaves -record lock files behind (a copy of a lock is no lock). A failed session start's directory is taken back through -`file_lock::dir_holds_only_lock_files` and `FileLock::release_and_remove_files`, which **closes** the +record lock files behind (a copy of a lock is no lock). A failed session start's directory is taken +back through `file_lock::dir_holds_only_lock_files` and `FileLock::release_and_remove_files`, which **closes** the lock files' descriptors before the directory is removed: an NFS client turns the unlink of a file it still has open into a rename to `.nfs*`, which stays until the last close and makes `remove_dir` fail with ENOTEMPTY (`serve_ws::tests::an_empty_session_directory_is_taken_back_with_its_lock` fails on a diff --git a/crates/agent/src/file_lock.rs b/crates/agent/src/file_lock.rs index e1064126..f4f3c460 100644 --- a/crates/agent/src/file_lock.rs +++ b/crates/agent/src/file_lock.rs @@ -2,8 +2,8 @@ //! //! The kernel releases the lock when its holder exits, however it exits, so there is no staleness to //! judge and no lockfile to break: a crashed holder's leftover file is simply unlocked. Used for a -//! session's liveness lock (`session_store::acquire_session_lock`) and the MCP manifest cache's write -//! lock (`tools::mcp_manifest`). Non-blocking: [`try_lock`] answers at once (bar a short retry, below), +//! session's liveness lock (`session_store::acquire_session_lock`), the MCP manifest cache's write +//! lock (`tools::mcp_manifest`) and the memory store's lock (`memory::file`). Non-blocking: [`try_lock`] answers at once (bar a short retry, below), //! and a caller that wants to wait polls it off the async runtime. //! //! **The lock is an open file description (OFD) lock** on Linux (`fcntl(F_OFD_SETLK)`, whole file, @@ -119,7 +119,8 @@ pub fn is_record_lock_file(path: &Path) -> bool { /// own lock files. /// - **An old binary's legacy lock with no record lock file beside it**, by where it sits: a `lock` /// in a session directory (one holding `000001.jsonl`, which is never deleted), or a `.lock` -/// beside a session file (`*.jsonl`) or the MCP manifest. Whether anyone holds it is not asked. +/// beside a session file (`*.jsonl`), the MCP manifest or the memory store's `.memory`. Whether +/// anyone holds it is not asked. /// /// Anything else — a SQLite database another program holds open and locked, a daemon's pid file — is /// an ordinary file, and an edit of it goes through. @@ -144,6 +145,7 @@ fn is_lock_file_by_name(path: &Path) -> bool { }; stem.ends_with(".jsonl") || stem == "mcp-manifest.json" + || stem == ".memory" || path .with_file_name(format!("{}{RECORD_SUFFIX}", &name[..stem.len()])) .exists() @@ -151,8 +153,8 @@ fn is_lock_file_by_name(path: &Path) -> bool { type FileId = (u64, u64); -/// The (device, inode) of every lock file this process holds, counted: [`try_lock`] adds a lock's -/// files once it holds them, and dropping the lock takes them out again once it is released. +/// The (device, inode) of every lock file this process holds or is acquiring, counted (see +/// [`Marks`]). #[cfg(unix)] fn held() -> std::sync::MutexGuard<'static, std::collections::HashMap> { static HELD: std::sync::OnceLock>> = @@ -168,28 +170,43 @@ fn file_id(meta: &std::fs::Metadata) -> FileId { (meta.dev(), meta.ino()) } -#[cfg(unix)] -fn mark_held(files: &[&File], add: bool) { - let ids: Vec = files - .iter() - .filter_map(|f| f.metadata().ok()) - .map(|m| file_id(&m)) - .collect(); - let mut held = held(); - for id in ids { - if add { - *held.entry(id).or_default() += 1; - } else if let std::collections::hash_map::Entry::Occupied(mut e) = held.entry(id) { - *e.get_mut() -= 1; - if *e.get() == 0 { - e.remove(); - } +/// The held-set entries one acquire made, taken out again when it is dropped. An acquire marks each +/// lock file **before** it tries to lock it, and an acquire that fails drops its marks — so there is +/// no moment when a lock is held but its file reads as not held. (A brief false "held" while an +/// acquire is in flight is the safe direction.) A [`FileLock`] keeps its marks until it is dropped, +/// which is after it has unlocked. +#[derive(Default)] +struct Marks(Vec); + +impl Marks { + fn mark(&mut self, file: &File) { + #[cfg(unix)] + if let Ok(meta) = file.metadata() { + let id = file_id(&meta); + *held().entry(id).or_default() += 1; + self.0.push(id); } + #[cfg(not(unix))] + let _ = file; } } -#[cfg(not(unix))] -fn mark_held(_: &[&File], _: bool) {} +impl Drop for Marks { + fn drop(&mut self) { + #[cfg(unix)] + if !self.0.is_empty() { + let mut held = held(); + for id in self.0.drain(..) { + if let std::collections::hash_map::Entry::Occupied(mut e) = held.entry(id) { + *e.get_mut() -= 1; + if *e.get() == 0 { + e.remove(); + } + } + } + } + } +} /// Whether the file at `path` is one this process holds a lock on — by `stat`, nothing opened. fn is_held_here(path: &Path) -> bool { @@ -224,6 +241,8 @@ pub struct FileLock { legacy: Option, record_path: PathBuf, legacy_path: Option, + /// Dropped after [`FileLock`]'s own `Drop` has unlocked, so the set never misses a held lock. + _marks: Marks, } impl FileLock { @@ -234,10 +253,6 @@ impl FileLock { &self.record } - fn files(&self) -> Vec<&File> { - std::iter::once(&self.record).chain(&self.legacy).collect() - } - /// Release the lock and unlink its files — for a lock directory being taken back. What matters is /// that the descriptors are closed before the directory is removed: an NFS client turns the unlink /// of a file it still has open into a rename to `.nfs*`, which stays until the last close and keeps @@ -258,8 +273,6 @@ impl Drop for FileLock { if let Some(legacy) = &self.legacy { let _ = legacy.unlock(); } - // Out of the held set only once unlocked, so the set never misses a lock still held. - mark_held(&self.files(), false); } } @@ -274,7 +287,9 @@ pub fn try_lock(target: Target<'_>) -> std::io::Result> { let record_path = target.record_path(); let legacy_path = target.legacy_path(); for _ in 0..RETRIES { + let mut marks = Marks::default(); let record = open_lock_file(&record_path)?; + marks.mark(&record); #[cfg(test)] tests::between_open_and_lock(&record_path); if !try_lock_description(&record)? { @@ -285,7 +300,11 @@ pub fn try_lock(target: Target<'_>) -> std::io::Result> { unlock(&record); continue; } - let legacy = match legacy_path.as_deref().map(take_legacy_flock).transpose()? { + let legacy = match legacy_path + .as_deref() + .map(|path| take_legacy_flock(path, &mut marks)) + .transpose()? + { None => None, Some(Legacy::Taken(file)) => Some(file), // An old binary holds it: held. @@ -298,14 +317,13 @@ pub fn try_lock(target: Target<'_>) -> std::io::Result> { continue; } }; - let lock = FileLock { + return Ok(Some(FileLock { record, legacy, record_path, legacy_path, - }; - mark_held(&lock.files(), true); - return Ok(Some(lock)); + _marks: marks, + })); } Ok(None) } @@ -377,8 +395,9 @@ enum Legacy { } /// The legacy `flock` older binaries take (see the module doc), retried for [`LEGACY_RETRY`]. -fn take_legacy_flock(path: &Path) -> std::io::Result { +fn take_legacy_flock(path: &Path, marks: &mut Marks) -> std::io::Result { let file = open_lock_file(path)?; + marks.mark(&file); let deadline = Instant::now() + LEGACY_RETRY; loop { match file.try_lock() { @@ -624,6 +643,42 @@ pub(crate) mod tests { assert!(try_lock(Target::Itself(&key)).unwrap().is_some()); } + /// No moment exists when a lock is held but its file reads as not held: an acquire marks the lock + /// file before it tries to lock it. A failed attempt takes its mark out again, so a key a + /// competing acquire could not get is ordinary once its holder lets go. + #[cfg(unix)] + #[test] + fn a_lock_file_reads_as_held_from_before_it_is_locked_until_after_it_is_released() { + let dir = tempfile::tempdir().unwrap(); + let key = dir.path().join("journal.key"); + let seen = std::rc::Rc::new(std::cell::Cell::new(None)); + let saw = seen.clone(); + BETWEEN.with(|hook| { + *hook.borrow_mut() = Some(Box::new(move |path: &Path| { + saw.set(Some(is_lock_file(path))); + })); + }); + let held = try_lock(Target::Itself(&key)).unwrap().unwrap(); + assert_eq!(seen.get(), Some(true), "held before the lock is taken"); + seen.set(None); + assert!( + try_lock(Target::Itself(&key)).unwrap().is_none(), + "a competing acquire fails" + ); + BETWEEN.with(|hook| *hook.borrow_mut() = None); + assert_eq!( + seen.get(), + Some(true), + "and was held throughout its attempt" + ); + assert!(is_lock_file(&key)); + drop(held); + assert!( + !is_lock_file(&key), + "the failed attempt took its mark out again" + ); + } + /// A lock file is recognized by its name in any case — a record lock file, a legacy one beside a /// record lock file, or an old binary's legacy one where a session or the MCP manifest keeps it — /// and nothing else is: an unrelated `Cargo.lock`, or a `lock` in a directory that holds no @@ -649,7 +704,11 @@ pub(crate) mod tests { "a lock in a session directory" ); assert!(is_lock_file(&old.join("Lock")), "in any case"); - for beside in ["1700000000_abc.jsonl.lock", "mcp-manifest.json.lock"] { + for beside in [ + "1700000000_abc.jsonl.lock", + "mcp-manifest.json.lock", + ".memory.lock", + ] { assert!(is_lock_file(&dir.path().join(beside)), "{beside}"); } std::fs::write(dir.path().join("Cargo.lock"), "").unwrap(); diff --git a/crates/agent/src/memory/file.rs b/crates/agent/src/memory/file.rs index 81567432..c40b5783 100644 --- a/crates/agent/src/memory/file.rs +++ b/crates/agent/src/memory/file.rs @@ -173,16 +173,12 @@ impl Store { } } - /// The store-wide lock path guarding every mutation (one lock for the whole store keeps cross-file - /// operations like `rename` and index updates consistent). - fn lock_path(&self) -> PathBuf { - self.dir.join(".memory.lock") - } - - /// Acquire the store lock, having ensured the store directory exists. + /// Acquire the store-wide lock guarding every mutation (one lock for the whole store keeps + /// cross-file operations like `rename` and index updates consistent), having ensured the store + /// directory exists. fn lock(&self) -> Result { fs::create_dir_all(&self.dir).map_err(|e| MemoryError::Backend(e.to_string()))?; - StoreLock::acquire(&self.lock_path()).map_err(|e| MemoryError::Backend(e.to_string())) + StoreLock::acquire(&self.dir, LOCK_TIMEOUT).map_err(|e| MemoryError::Backend(e.to_string())) } /// Decode what a document's bytes say — sealed or not. @@ -544,57 +540,41 @@ impl MemoryBackend for FileBackend { } } -// ---- StoreLock: a cross-process advisory lock, duplicated per this crate's store convention -------- +// ---- StoreLock: the store-wide lock, a `file_lock` lock ----------------------------------------------- const LOCK_TIMEOUT: Duration = Duration::from_secs(5); const LOCK_RETRY_INTERVAL: Duration = Duration::from_millis(20); -/// Held for the length of one mutation. Released by closing the descriptor on `Drop`, so a panicked or -/// early-returning holder frees it — and so does a holder that is killed outright. -/// -/// The lock is the file descriptor, not the file's existence. The previous `create_new`-as-a-mutex -/// scheme had to guess when an abandoned lock file was stale (a fixed 10-second age), which is both a -/// stall for anyone waiting and a silent lock break for anyone slower than the guess. `flock` has no -/// such ambiguity: the kernel releases it when the holder's descriptor closes, crash included. The lock -/// file itself is left behind deliberately — unlinking it while holding the lock is the classic race, -/// where a waiter ends up locking an inode that is no longer the one at the path. +/// Held for the length of one mutation: [`crate::file_lock`]'s lock on `/.memory` — an OFD +/// lock on `.memory.beyond-lock`, plus, for binaries before it, the `flock` on `.memory.lock` they +/// took. The kernel releases it when the holder exits, crash included, so there is no staleness to +/// guess. Being a `file_lock` lock is what keeps it a lock: no other descriptor (a model `read` of +/// the file) releases it, and `write_atomic` refuses to rename over either file — a model `write` or +/// `edit` that replaced `.memory.lock` would leave the next mutation locking a different inode, and +/// two mutations running at once. struct StoreLock { - _file: fs::File, + _lock: crate::file_lock::FileLock, } impl StoreLock { - fn acquire(lock_path: &Path) -> io::Result { - if let Some(parent) = lock_path.parent() { - fs::create_dir_all(parent)?; - } - // Read+write, because a network filesystem may emulate `flock` with POSIX record locks, which - // require a writable descriptor. - let mut opts = fs::OpenOptions::new(); - opts.read(true).write(true).create(true).truncate(false); - #[cfg(unix)] - { - use std::os::unix::fs::OpenOptionsExt; - opts.mode(0o600); - } - let file = opts.open(lock_path)?; - let deadline = Instant::now() + LOCK_TIMEOUT; + /// Poll for the lock (the store directory must exist) until `timeout`. + fn acquire(dir: &Path, timeout: Duration) -> io::Result { + let key = dir.join(".memory"); + let deadline = Instant::now() + timeout; loop { - match file.try_lock() { - Ok(()) => return Ok(Self { _file: file }), - Err(fs::TryLockError::WouldBlock) => { - if Instant::now() >= deadline { - return Err(io::Error::new( - ErrorKind::TimedOut, - format!( - "timed out waiting for memory store lock at {}", - lock_path.display() - ), - )); - } - std::thread::sleep(LOCK_RETRY_INTERVAL); - } - Err(fs::TryLockError::Error(e)) => return Err(e), + if let Some(lock) = crate::file_lock::try_lock(crate::file_lock::Target::File(&key))? { + return Ok(Self { _lock: lock }); } + if Instant::now() >= deadline { + return Err(io::Error::new( + ErrorKind::TimedOut, + format!( + "timed out waiting for memory store lock in {}", + dir.display() + ), + )); + } + std::thread::sleep(LOCK_RETRY_INTERVAL); } } } @@ -682,6 +662,65 @@ mod tests { assert!(matches!(zero, MemoryError::NotUnique { count: 0, .. })); } + /// A model `write` to the store's lock file while a mutation holds the lock is refused — were it + /// renamed over, the next mutation would lock a different inode and run alongside this one. So a + /// second mutation still waits: here it times out rather than getting the lock. + #[cfg(unix)] + #[test] + fn the_store_lock_file_is_not_replaced_and_mutations_stay_serialized() { + let dir = tempfile::tempdir().unwrap(); + let store = dir.path().join("memory"); + fs::create_dir_all(&store).unwrap(); + let held = StoreLock::acquire(&store, LOCK_TIMEOUT).unwrap(); + for name in [".memory.lock", ".memory.beyond-lock", ".MEMORY.LOCK"] { + let path = store.join(name); + if name == ".MEMORY.LOCK" && !path.exists() { + continue; // case-sensitive filesystem: a different, ordinary file + } + assert!( + crate::tools::write_atomic(path.to_str().unwrap(), b"x").is_err(), + "{name} must not be replaced" + ); + } + let second = StoreLock::acquire(&store, Duration::from_millis(100)); + assert!( + second.is_err_and(|e| e.kind() == ErrorKind::TimedOut), + "a second mutation waits for the first" + ); + drop(held); + StoreLock::acquire(&store, Duration::from_millis(100)).expect("free once released"); + } + + /// Concurrent mutations through the backend are serialized: none of them loses another's update. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn concurrent_mutations_lose_no_update() { + let (_d, b) = backend(); + b.create(&p("/memories/log.md"), "start\n").await.unwrap(); + let b = Arc::new(b); + let tasks: Vec<_> = (0..16) + .map(|i| { + let b = b.clone(); + tokio::spawn(async move { + b.insert(&p("/memories/log.md"), 1, &format!("line {i}")) + .await + .unwrap(); + }) + }) + .collect(); + for t in tasks { + t.await.unwrap(); + } + let View::Document(text) = b.view(&p("/memories/log.md"), None).await.unwrap() else { + panic!("not a document") + }; + for i in 0..16 { + assert!( + text.contains(&format!("line {i}\n")), + "line {i} lost:\n{text}" + ); + } + } + #[tokio::test] async fn insert_after_line() { let (_d, b) = backend();