From 8ccbd580dd17543d92dbe57b4cec28713f5c405d Mon Sep 17 00:00:00 2001 From: Shanu Date: Wed, 19 Aug 2026 17:11:14 +0530 Subject: [PATCH 1/2] Close five defects the final production review found A seven-dimension adversarial review of main. Ten blockers were reported; I verified each against the code rather than taking the report's word, refuted one, and fixed five. The rest are recorded in the PR body with what is known and what is not. 1. A namespace beginning with `/` escaped the workspace. `sanitize_namespace` keeps `/` so namespaces can be hierarchical, but never stripped a LEADING one -- and `Path::join` with an absolute path DISCARDS the base, so `memory_dir/namespaces/` vanished and the namespace addressed anywhere on the filesystem. `clear_namespace` calls `remove_dir_all` on that path. Proved before fixing: "/Users/me/Documents" resolved to exactly that directory. (`..` was already neutral -- `.` is not in the allow-list.) Leading slashes are stripped now, with a test over six hostile inputs. 2. Concurrent first-writes of one key orphaned chunks, and `forget` left them. The row was always safe -- `ON CONFLICT(namespace, key) DO UPDATE` keeps exactly one -- but that clause does not update `document_id`, and each writer had already written `vector_chunks` under its own random id. The loser's chunks became unreachable from the row, so `forget` (which deletes by the row's id) left them and recall kept returning content the caller had deleted. The id is derived from (namespace, key) now, so both writers agree. 3. `accept_source_items` reported write failures as success. A failed `put_doc` incremented `skipped` and returned `Ok`. The contract defines `skipped` as "units the driver recognised as already present", so a locked database, a full disk or a dead embedder read as a successful no-op and the sync caller marked the items done. It propagates now, naming how many of how many were written first. 4. The remote adapters buffered whole response bodies. `json()`/`text()` read everything before any size check, and all three engines take an operator-supplied endpoint. `tinymemory-sources` already had the answer -- `read_body_capped`, with a doc comment explaining this exact OOM -- and it had simply never been applied here. Ported: a 64 MiB cap, enforced while the bytes arrive so an absent or understated Content-Length cannot get around it. 5. mem0's listing silently truncated at 1000. That single unpaginated request is the ONLY enumeration path in the adapter, so past the ceiling `get(ns, key)` returned `Ok(None)` -- "no such entry" -- for records that exist. It now refuses at its own boundary instead of answering wrongly. Paginating properly needs mem0's paging parameters verified against a live service; guessing them would have traded a loud failure for a quiet one. REFUTED, with the reasoning worth keeping: the reported SSRF bypass via bracketed IPv6 (`http://[::ffff:169.254.169.254]/`) does not exist -- I ran it. `url` normalises IPv4-mapped literals to hex (`[::ffff:a9fe:a9fe]`), so no bracketed host ever contains a dot and the single-label rule blocks every one. But it blocks them by ACCIDENT: brackets make the `Ipv6Addr` parse fail, so `is_blocked_host`'s IPv6 branch is dead code. Left as-is here rather than changed on a hunch; noted in the PR. cargo test --workspace: 1417 passed, 0 failed cargo clippy --workspace --all-targets: clean scripts/ci/engine-containment.sh: holds scripts/ci/dependency-budget.sh: 40 crates, ceiling 50 --- Cargo.lock | 1 + README.md | 34 +++++++--- adapters/remote/Cargo.toml | 7 ++- adapters/remote/src/common.rs | 55 +++++++++++++--- adapters/remote/src/mem0.rs | 41 +++++++++++- adapters/tinycortex/src/engine/mod.rs | 17 ++++- core/src/store/namespace_store/documents.rs | 70 +++++++++++++++++++-- core/src/store/namespace_store/init.rs | 46 +++++++++++++- 8 files changed, 244 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4aef392..11b430d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1957,6 +1957,7 @@ dependencies = [ "anyhow", "async-trait", "axum", + "futures", "reqwest", "serde", "serde_json", diff --git a/README.md b/README.md index d543f32..1799f3e 100644 --- a/README.md +++ b/README.md @@ -67,21 +67,41 @@ let provider = Arc::new(tinymemory::remote::supermemory_provider(backend)); The remote adapter reaches only crates.io dependencies, so cargo resolves it without any `[patch]` entries. -**The embedded engine (TinyCortex) — three patch entries:** +**The embedded engine (TinyCortex) — vendor this repository as a submodule.** + +The remote recipe above works by git because the remote adapter reaches only +published crates. The embedded engine does not: it pulls `tinycortex`, +`tinycortex-api` and `tinyagents`, none of which are published, and +`tinycortex-api` takes `tinymemory-api` *by git*, which cargo will resolve as a +second copy of a crate this workspace also provides by path. Patching that away +needs the crates on disk, so the embedded path is a submodule dependency until +these crates are published: + +```sh +git submodule add https://github.com/tinyhumansai/tinymemory vendor/tinymemory +git -C vendor/tinymemory submodule update --init --recursive +``` ```toml [dependencies] -tinymemory = { git = "https://github.com/tinyhumansai/tinymemory", features = ["tinycortex"] } +tinymemory = { path = "vendor/tinymemory", features = ["tinycortex"] } -# The engine and its api are unpublished; without these, cargo resolves a -# second copy of each from the network and type identities split at the seam. +# All four are required. The first three are unpublished crates the engine +# needs; the fourth collapses `tinycortex-api`'s git dependency on +# `tinymemory-api` onto the copy in this tree — without it two distinct +# `tinymemory_api::MemoryEntry` types exist and the seam stops type-checking. [patch.crates-io] -tinycortex = { git = "https://github.com/tinyhumansai/tinycortex" } -tinycortex-api = { git = "https://github.com/tinyhumansai/tinycortex" } +tinycortex = { path = "vendor/tinymemory/vendor/tinycortex" } +tinycortex-api = { path = "vendor/tinymemory/vendor/tinycortex/api" } +tinyagents = { path = "vendor/tinymemory/vendor/tinyagents" } [patch."https://github.com/tinyhumansai/tinymemory"] -tinymemory-api = { git = "https://github.com/tinyhumansai/tinymemory" } +tinymemory-api = { path = "vendor/tinymemory/api" } ``` +This exact patch set is what the reference consumer in `examples/` and the +repository's own root manifest use; a build missing any of the four fails at +resolution, before compiling a line. + ```rust,ignore use std::sync::Arc; use tinymemory::tinycortex::{provider, InMemoryMemoryStore}; diff --git a/adapters/remote/Cargo.toml b/adapters/remote/Cargo.toml index 2690528..3f8fcbd 100644 --- a/adapters/remote/Cargo.toml +++ b/adapters/remote/Cargo.toml @@ -16,9 +16,14 @@ async-trait = "0.1" # The storage trait deliberately uses opaque backend errors. anyhow = "1" # Native self-hosted APIs are HTTP/JSON; multipart is required by Cognee. -reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls"] } +# `stream` is for `bytes_stream()`: response bodies are read against a byte +# cap rather than buffered whole, because the endpoint is operator-supplied +# and a broken or hostile one must not be able to OOM the host. +reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "stream"] } # Remote records are translated through a private, lossless envelope. serde = { version = "1", features = ["derive"] } +# Streaming a capped body needs a Stream combinator. +futures = "0.3" serde_json = "1" # Supermemory custom ids are bounded, so namespace/key identities use SHA-256. sha2 = "0.10" diff --git a/adapters/remote/src/common.rs b/adapters/remote/src/common.rs index 96b42d7..0675f2d 100644 --- a/adapters/remote/src/common.rs +++ b/adapters/remote/src/common.rs @@ -40,6 +40,50 @@ impl std::fmt::Debug for HttpClient { } } +/// Largest response body any hosted engine may return. +/// +/// The endpoint is operator-supplied (`SupermemoryMemory::api`, +/// `Mem0Memory::new`, `CogneeMemory::self_hosted` all take an arbitrary URL), +/// so a broken or hostile server must not be able to exhaust the host's +/// memory. 64 MiB is far above any real memory payload -- the largest thing +/// these APIs return is a page of records -- and far below a size that +/// threatens a process. +const MAX_RESPONSE_BYTES: u64 = 64 * 1024 * 1024; + +/// Read a response body, failing once it exceeds [`MAX_RESPONSE_BYTES`]. +/// +/// `Response::json()`/`text()` buffer the whole body before any size check, so +/// a server that omits or understates `Content-Length` (a chunked response, +/// say) could OOM the process despite a declared limit. Reading incrementally +/// enforces the cap while the bytes arrive. Same argument, and same shape, as +/// `tinymemory-sources`' `read_body_capped` -- that guard was written for the +/// web-page reader and simply had not been applied on this path. +async fn read_capped(response: reqwest::Response, path: &str) -> anyhow::Result> { + use futures::StreamExt; + if let Some(len) = response.content_length() { + if len > MAX_RESPONSE_BYTES { + anyhow::bail!( + "memory API {path} response exceeds {MAX_RESPONSE_BYTES}-byte limit \ + (Content-Length={len})" + ); + } + } + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| format!("memory API {path} body read failed"))?; + body.extend_from_slice(&chunk); + if body.len() as u64 > MAX_RESPONSE_BYTES { + anyhow::bail!( + "memory API {path} response exceeds {MAX_RESPONSE_BYTES}-byte limit \ + (read {} bytes)", + body.len() + ); + } + } + Ok(body) +} + impl HttpClient { /// Builds a client that optionally authenticates with a bearer token. pub(crate) fn bearer(endpoint: &str, credential: Option<&str>) -> anyhow::Result { @@ -135,9 +179,8 @@ impl HttpClient { if !status.is_success() { return Err(self.status_error(path, status)); } - response - .json() - .await + let body = read_capped(response, path).await?; + serde_json::from_slice(&body) .with_context(|| format!("memory API {path} returned invalid JSON")) } @@ -153,10 +196,8 @@ impl HttpClient { if !status.is_success() { return Err(self.status_error(path, status)); } - response - .text() - .await - .context("memory API response was unreadable") + let body = read_capped(response, path).await?; + String::from_utf8(body).context("memory API response was not valid UTF-8") } /// Sends a request whose successful response body is not needed. diff --git a/adapters/remote/src/mem0.rs b/adapters/remote/src/mem0.rs index d92546d..747e94b 100644 --- a/adapters/remote/src/mem0.rs +++ b/adapters/remote/src/mem0.rs @@ -119,17 +119,52 @@ struct Mem0Dialect { } impl Mem0Dialect { + /// Largest listing this adapter will request in one call. + /// + /// Every exact-CRUD path here enumerates through [`Self::values`], so this + /// is the ceiling on the whole store, not on one page. + const LISTING_TOP_K: usize = 1000; + /// Fetches Mem0's administrative memory listing. + /// + /// # A hard ceiling, deliberately loud + /// + /// This is a single unpaginated request, and it is the ONLY enumeration + /// path in this adapter -- `get`, `list`, `count` and `export_page` all + /// route through it. Past the ceiling the results are not merely + /// incomplete, they are silently WRONG: `get(ns, key)` for a record beyond + /// the cut-off returns `Ok(None)`, which the contract defines as "no such + /// entry", so a caller reads "deleted" where the truth is "present but + /// past the window". + /// + /// Returning an error instead is the honest failure. A full response is + /// indistinguishable from a truncated one -- both are exactly `top_k` + /// items -- so this cannot detect truncation, only its own boundary, and + /// it refuses at that boundary rather than answering wrongly. Paginating + /// properly needs Mem0's paging parameters verified against a live + /// service; guessing them here would trade a loud failure for a quiet one. async fn values(&self) -> anyhow::Result> { + let top_k = Self::LISTING_TOP_K; let response: Value = self .client - .json(Method::GET, "memories?top_k=1000", None) + .json(Method::GET, &format!("memories?top_k={top_k}"), None) .await?; - Ok(response + let results = response .get("results") .and_then(Value::as_array) .cloned() - .unwrap_or_default()) + .unwrap_or_default(); + if results.len() >= top_k { + anyhow::bail!( + "mem0 returned {} memories, this adapter's unpaginated listing ceiling. \ + Exact reads (get/list/count/export) cannot be answered correctly beyond \ + it -- a record past the window would read as absent -- so the adapter \ + refuses rather than answering wrongly. Recall is unaffected (it queries \ + mem0's search API directly).", + results.len() + ); + } + Ok(results) } /// Decodes a Mem0 result containing TinyMemory-owned metadata. diff --git a/adapters/tinycortex/src/engine/mod.rs b/adapters/tinycortex/src/engine/mod.rs index a9a2c25..9a1ae2f 100644 --- a/adapters/tinycortex/src/engine/mod.rs +++ b/adapters/tinycortex/src/engine/mod.rs @@ -1024,6 +1024,7 @@ impl MemorySourceSink for TinycortexProvider { items: Vec, taint: MemoryTaint, ) -> Result { + let items_len = items.len(); let namespace = format!("source:{source_id}"); let mut outcome = IngestOutcome::default(); for item in items { @@ -1063,8 +1064,20 @@ impl MemorySourceSink for TinycortexProvider { outcome.written = outcome.written.saturating_add(1); outcome.ids.push(id); } - Err(_) => { - outcome.skipped = outcome.skipped.saturating_add(1); + // A write failure is NOT `skipped`. The contract defines that + // field as "units the driver recognised as already present" + // (`IngestOutcome::skipped`), so counting a failed write there + // reports a locked database, a full disk or a dead embedder as + // a successful no-op: the sync caller marks the items done and + // they are never written. Propagate instead — a partial batch + // has no truthful representation in `IngestOutcome`, and a + // caller that wants best-effort ingestion can catch this. + Err(error) => { + return Err(MemoryError::Other(anyhow::anyhow!( + "source ingest failed after {} of {} item(s) were written: {error}", + outcome.written, + items_len + ))); } } } diff --git a/core/src/store/namespace_store/documents.rs b/core/src/store/namespace_store/documents.rs index 9c317f4..db3e6bd 100644 --- a/core/src/store/namespace_store/documents.rs +++ b/core/src/store/namespace_store/documents.rs @@ -46,11 +46,19 @@ impl UnifiedMemory { let document_id = input .document_id .or(existing_document_id) - .unwrap_or_else(|| { - let ts = Self::now_ts() as u64; - let short = &Uuid::new_v4().to_string()[..8]; - format!("{ts}_{short}") - }); + // Derived from (namespace, key), NOT random. The lookup above and + // the write below are separated by `.await`s, so two concurrent + // stores of a not-yet-existing key both miss and both mint an id. + // The ROW is safe -- `ON CONFLICT(namespace, key) DO UPDATE` keeps + // exactly one -- but that clause does not update `document_id`, + // and each writer has already written `vector_chunks` under ITS + // OWN id. The loser's chunks are then unreachable from the row, so + // `forget` (which deletes chunks by the row's document_id) leaves + // them behind and recall keeps returning content the caller + // deleted. A deterministic id makes both writers choose the same + // one, so the second write updates the first's chunks instead of + // orphaning them. + .unwrap_or_else(|| Self::derive_document_id(&namespace, &key)); let now = Self::now_ts(); let created_at = { let conn = self.conn.lock(); @@ -659,8 +667,60 @@ impl UnifiedMemory { } Ok(json!({"deleted": deleted, "namespace": ns, "documentId": document_id })) } + + /// A document id derived from `(namespace, key)`. + /// + /// Deterministic so two concurrent first-writes of one key agree, which is + /// what keeps `vector_chunks` addressable from the row. Hashed rather than + /// concatenated so the id is a fixed-width opaque token whatever the + /// namespace or key contains; the zero byte is a domain separator, so + /// ("a","bc") and ("ab","c") cannot collide. + pub(crate) fn derive_document_id(namespace: &str, key: &str) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(namespace.as_bytes()); + hasher.update([0u8]); + hasher.update(key.as_bytes()); + format!("{:x}", hasher.finalize())[..32].to_string() + } } #[cfg(test)] #[path = "documents_tests.rs"] mod tests; + +#[cfg(test)] +mod document_id_tests { + use super::UnifiedMemory; + + /// Two concurrent first-writes of one key must choose the SAME document + /// id. If they do not, each writes `vector_chunks` under its own id, the + /// `ON CONFLICT(namespace, key)` row keeps only one of them, and the + /// loser's chunks outlive `forget` — deleted content stays recallable. + #[test] + fn the_id_is_derived_from_namespace_and_key_not_random() { + let a = UnifiedMemory::derive_document_id("notes", "q3-plan"); + let b = UnifiedMemory::derive_document_id("notes", "q3-plan"); + assert_eq!(a, b, "the same key must derive the same id"); + assert_ne!( + a, + UnifiedMemory::derive_document_id("notes", "q4-plan"), + "different keys must not collide" + ); + assert_ne!( + a, + UnifiedMemory::derive_document_id("other", "q3-plan"), + "the namespace must participate" + ); + } + + /// The separator matters: without it ("a","bc") and ("ab","c") hash the + /// same bytes and two distinct records share one id. + #[test] + fn the_namespace_key_boundary_cannot_be_shifted() { + assert_ne!( + UnifiedMemory::derive_document_id("a", "bc"), + UnifiedMemory::derive_document_id("ab", "c") + ); + } +} diff --git a/core/src/store/namespace_store/init.rs b/core/src/store/namespace_store/init.rs index a8d0093..c3296f5 100644 --- a/core/src/store/namespace_store/init.rs +++ b/core/src/store/namespace_store/init.rs @@ -405,7 +405,7 @@ impl UnifiedMemory { if trimmed.is_empty() { return GLOBAL_NAMESPACE.to_string(); } - trimmed + let sanitized: String = trimmed .chars() .map(|ch| { if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '/' { @@ -414,7 +414,19 @@ impl UnifiedMemory { '_' } }) - .collect() + .collect(); + // `/` is kept so a namespace can be hierarchical, but a LEADING one + // makes the result an absolute path — and `Path::join` with an + // absolute path discards the base entirely, so + // `memory_dir/namespaces/` vanishes and the namespace addresses + // anywhere on the filesystem. `clear_namespace` calls + // `remove_dir_all` on that path. (`..` is already neutralised above: + // `.` is not in the allow-list, so it becomes `_`.) + let sanitized = sanitized.trim_start_matches('/').to_string(); + if sanitized.is_empty() { + return GLOBAL_NAMESPACE.to_string(); + } + sanitized } /// Resolved memory subdirectory for this store instance (e.g. @@ -474,6 +486,36 @@ mod tests { } } + /// A namespace beginning with `/` must not escape the workspace: + /// `Path::join` with an absolute path DISCARDS the base, so + /// `memory_dir/namespaces/` would vanish and `clear_namespace`'s + /// `remove_dir_all` would run against an arbitrary absolute path. + #[test] + fn a_namespace_cannot_escape_the_workspace() { + for hostile in [ + "/Users/me/Documents", + "//tmp/x", + "///etc", + "/", + "a/../../etc", + "../../etc", + ] { + let sanitized = UnifiedMemory::sanitize_namespace(hostile); + assert!( + !sanitized.starts_with('/'), + "{hostile:?} sanitized to {sanitized:?}, which is absolute" + ); + let dir = std::path::Path::new("/w/memory") + .join("namespaces") + .join(&sanitized); + assert!( + dir.starts_with("/w/memory/namespaces"), + "{hostile:?} escaped to {}", + dir.display() + ); + } + } + #[test] fn namespace_dir_uses_sanitized_namespace() { let tmp = TempDir::new().unwrap(); From 7f605169717c07902a9d9d92dcb428edb16405ef Mon Sep 17 00:00:00 2001 From: Shanu Date: Wed, 19 Aug 2026 17:21:13 +0530 Subject: [PATCH 2/2] Serialise same-key writes; check a body chunk before allocating it Both of CodeRabbit's findings on the previous commit, both correct, both on code that commit added. 1. The body cap appended each chunk and then checked the total, so a single chunk larger than the remaining budget was allocated in full before the limit was noticed -- the exact allocation the cap exists to prevent. The size is checked before the append now, with a `checked_add` so the length arithmetic cannot wrap. 2. The deterministic document id stopped two writers ORPHANING each other's chunks, but it did not ORDER them, and I claimed more than it delivered. The row write and the chunk replacement are separated by embedding, which awaits: writer A can update the row, await the embedder, and have B update the row and replace the chunks in between. A's chunks then land beside B's content, and A's trailing chunks survive if A produced more. Same-key writes now hold a lock for the whole operation. The lock is process-global and keyed by (db path, namespace, key), not per-instance: the same store file can be opened by more than one `UnifiedMemory`, and an instance-level lock would not serialise those. Both write paths take it -- the metadata-only upsert writes the same row, so it must not interleave with a full write either. Keyed per document rather than per store so two different keys never queue behind one slow embedding; a test pins that, and that two workspaces never contend. cargo test --workspace: 1418 passed, 0 failed cargo clippy --workspace --all-targets: clean --- adapters/remote/src/common.rs | 14 +++- core/src/store/namespace_store/documents.rs | 92 +++++++++++++++++++++ 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/adapters/remote/src/common.rs b/adapters/remote/src/common.rs index 0675f2d..c8a8622 100644 --- a/adapters/remote/src/common.rs +++ b/adapters/remote/src/common.rs @@ -72,14 +72,20 @@ async fn read_capped(response: reqwest::Response, path: &str) -> anyhow::Result< let mut stream = response.bytes_stream(); while let Some(chunk) = stream.next().await { let chunk = chunk.with_context(|| format!("memory API {path} body read failed"))?; - body.extend_from_slice(&chunk); - if body.len() as u64 > MAX_RESPONSE_BYTES { + // Check BEFORE appending: one oversized chunk would otherwise be + // allocated in full before the limit is noticed, which is the + // allocation this cap exists to prevent. + let next_len = body + .len() + .checked_add(chunk.len()) + .context("memory API response length overflowed")?; + if next_len as u64 > MAX_RESPONSE_BYTES { anyhow::bail!( "memory API {path} response exceeds {MAX_RESPONSE_BYTES}-byte limit \ - (read {} bytes)", - body.len() + (would reach {next_len} bytes)" ); } + body.extend_from_slice(&chunk); } Ok(body) } diff --git a/core/src/store/namespace_store/documents.rs b/core/src/store/namespace_store/documents.rs index db3e6bd..7c39a6c 100644 --- a/core/src/store/namespace_store/documents.rs +++ b/core/src/store/namespace_store/documents.rs @@ -33,6 +33,19 @@ impl UnifiedMemory { if key.is_empty() { return Err("document key cannot be empty".to_string()); } + // Serialise writers of one key for the WHOLE operation. A deterministic + // document id stops two writers orphaning each other's chunks, but it + // does not ORDER them: the row write and the chunk replacement are + // separated by embedding, which awaits. Without this, writer A can + // update the row, await the embedder, and have B update the row and + // replace the chunks in between -- leaving B's content beside A's + // chunks, plus A's trailing chunks if A had more. The metadata-only + // path below takes the same lock: it writes the same row, so it must + // not interleave with a full write either. Same guard shape as the + // sync path's per-connection lock. + let _write_guard = Self::document_write_lock(&self.db_path, &namespace, &key) + .lock_owned() + .await; let existing_document_id = { let conn = self.conn.lock(); conn.query_row( @@ -229,6 +242,19 @@ impl UnifiedMemory { if key.is_empty() { return Err("document key cannot be empty".to_string()); } + // Serialise writers of one key for the WHOLE operation. A deterministic + // document id stops two writers orphaning each other's chunks, but it + // does not ORDER them: the row write and the chunk replacement are + // separated by embedding, which awaits. Without this, writer A can + // update the row, await the embedder, and have B update the row and + // replace the chunks in between -- leaving B's content beside A's + // chunks, plus A's trailing chunks if A had more. The metadata-only + // path below takes the same lock: it writes the same row, so it must + // not interleave with a full write either. Same guard shape as the + // sync path's per-connection lock. + let _write_guard = Self::document_write_lock(&self.db_path, &namespace, &key) + .lock_owned() + .await; let existing_document_id = { let conn = self.conn.lock(); conn.query_row( @@ -668,6 +694,39 @@ impl UnifiedMemory { Ok(json!({"deleted": deleted, "namespace": ns, "documentId": document_id })) } + /// The write lock for one `(database, namespace, key)`. + /// + /// Process-global rather than per-instance: the same store file can be + /// opened by more than one `UnifiedMemory`, and a lock living on the + /// instance would not serialise those. Keyed by db path so two workspaces + /// never contend. + /// + /// The table only grows, bounded by the number of distinct keys this + /// process has written -- one `Arc` and an unlocked mutex each. + fn document_write_lock( + db_path: &std::path::Path, + namespace: &str, + key: &str, + ) -> std::sync::Arc> { + use std::collections::HashMap; + use std::sync::{Arc, Mutex, OnceLock}; + type Table = Mutex>>>; + static LOCKS: OnceLock = OnceLock::new(); + let table = LOCKS.get_or_init(|| Mutex::new(HashMap::new())); + let id = ( + db_path.to_string_lossy().into_owned(), + namespace.to_owned(), + key.to_owned(), + ); + // Recover from a poisoned table: it holds `Arc`s only, so a panicking + // writer leaves nothing torn, and refusing every later write would be + // a worse failure than continuing. + let mut table = table + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Arc::clone(table.entry(id).or_default()) + } + /// A document id derived from `(namespace, key)`. /// /// Deterministic so two concurrent first-writes of one key agree, which is @@ -714,6 +773,39 @@ mod document_id_tests { ); } + /// The guard must be per key, not global: two different keys writing at + /// once must not serialise, or every concurrent write in the process + /// queues behind one slow embedding. + #[test] + fn the_write_lock_is_per_key_and_shared_per_key() { + let db = std::path::Path::new("/w/memory/memory.db"); + let a1 = UnifiedMemory::document_write_lock(db, "notes", "k1"); + let a2 = UnifiedMemory::document_write_lock(db, "notes", "k1"); + let b = UnifiedMemory::document_write_lock(db, "notes", "k2"); + let other_ns = UnifiedMemory::document_write_lock(db, "other", "k1"); + let other_db = UnifiedMemory::document_write_lock( + std::path::Path::new("/w2/memory/memory.db"), + "notes", + "k1", + ); + assert!( + std::sync::Arc::ptr_eq(&a1, &a2), + "same key must share one lock" + ); + assert!( + !std::sync::Arc::ptr_eq(&a1, &b), + "different keys must not contend" + ); + assert!( + !std::sync::Arc::ptr_eq(&a1, &other_ns), + "the namespace must participate" + ); + assert!( + !std::sync::Arc::ptr_eq(&a1, &other_db), + "two workspaces must not contend" + ); + } + /// The separator matters: without it ("a","bc") and ("ab","c") hash the /// same bytes and two distinct records share one id. #[test]