From f0d65de74467b518913165f36164b3d526567c57 Mon Sep 17 00:00:00 2001 From: Beinan Date: Sun, 26 Jul 2026 08:11:56 +0000 Subject: [PATCH] fix(store): retry MemWAL init and shard-claim races in ContextStore::add MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hardening, not a fix for a live single-process bug. Read the scope note below. `ContextStore::add` has two unretried races against a concurrent writer of the same dataset: 1. `initialize_mem_wal` is a CreateIndex transaction. On a fresh dataset two writers both see `has_mem_wal == false` and both try to create it; the loser gets "Retryable commit conflict ... preempted by concurrent transaction CreateIndex" and the entire append fails. 2. Opening a `mem_wal_writer` claims the shard epoch; a concurrent open loses with "Failed to claim shard ... another writer claimed epoch N". `add` opens a fresh writer per group per call, so this recurs in steady state. Both are transient — the loser only has to re-read and reopen — and Lance marks them retryable in the message text itself. Each now gets a bounded retry (5 attempts). The MemWAL path re-checks after a conflict rather than retrying blindly: if the other writer's CreateIndex landed, there is nothing left to do. Rows carry caller-supplied ids and are de-duplicated by id at read time, so a retried append cannot double-count. Note the claim-loss message contains neither "fenced" nor "Fenced", so reusing the rollout store's `is_fenced_error` would not have matched the case this path actually hits. `is_shard_contention_error` covers both wordings. Scope: not reachable in a single process today `AppState::get_or_open_context_store` caches by name with a double-check, so one name maps to one `Arc>`, and every mutating route holds the write lock — appends are serialized. `#[derive(Clone)]` on `ContextStore` has no value-level clone in any non-test code. The Python binding is one store per `Context`. So a single process cannot produce the contending writers. What it does cover is multi-replica: `ContextStore` shards by *data* (`derive_region_id(bot_id, session_id)`), not by writer identity, so two server replicas — each with its own single store — writing the same `(bot_id, session_id)` land on the same MemWAL shard with no cross-process lock. `RolloutStore` avoids this by sharding on writer identity (`shard_id` from the instance hostname); `ContextStore` has no equivalent. There is no server deployment manifest in this repo, so this is a latent architectural gap rather than a defect anyone is currently hitting. The retries are cheap and correct either way, but the durable fix for multi-replica context writes would be per-writer sharding, not retry. Tests use two clones to simulate two independent writers of one dataset, which is the multi-process shape. Verified non-vacuous: forcing both retry predicates to `false` makes both tests fail. Co-Authored-By: Claude --- crates/lance-context-core/src/store.rs | 126 +++++++++-- .../tests/context_store_concurrent_add.rs | 197 ++++++++++++++++++ 2 files changed, 310 insertions(+), 13 deletions(-) create mode 100644 crates/lance-context-core/tests/context_store_concurrent_add.rs diff --git a/crates/lance-context-core/src/store.rs b/crates/lance-context-core/src/store.rs index 5344ade..41a0943 100644 --- a/crates/lance-context-core/src/store.rs +++ b/crates/lance-context-core/src/store.rs @@ -484,12 +484,27 @@ impl ContextStore { groups.entry(key).or_default().push(entry.clone()); } - // Ensure MemWAL is initialized (once for the dataset) + // Ensure MemWAL is initialized (once for the dataset). + // + // This is a `CreateIndex` transaction, so two writers reaching a fresh + // dataset at the same time both see `has_mem_wal == false` and both try + // to create it; the loser gets "Retryable commit conflict ... preempted + // by concurrent transaction CreateIndex". That is a one-time race on the + // very first write, but it failed the whole append, so a concurrent + // cold start could reject every row from one writer. + // + // Re-check after a conflict rather than retrying blindly: if the other + // writer's CreateIndex landed, MemWAL now exists and there is nothing + // left to do. { - let indices = self.dataset.load_indices().await?; - let has_mem_wal = indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME); + const MAX_INIT_ATTEMPTS: usize = 5; + let mut attempt = 0; + loop { + let indices = self.dataset.load_indices().await?; + if indices.iter().any(|i| i.name == MEM_WAL_INDEX_NAME) { + break; + } - if !has_mem_wal { // ZoneMap indices are not supported by MemWAL; exclude them let maintained_indexes: Vec = indices .iter() @@ -498,26 +513,81 @@ impl ContextStore { }) .map(|i| i.name.clone()) .collect(); - self.dataset + match self + .dataset .initialize_mem_wal() .unsharded() .maintained_indexes(maintained_indexes) .execute() - .await?; + .await + { + Ok(()) => break, + Err(err) if is_retryable_commit_conflict(&err) => { + attempt += 1; + if attempt >= MAX_INIT_ATTEMPTS { + return Err(err); + } + // Reload so the next iteration observes the winner's + // committed index rather than our stale manifest. + self.dataset.checkout_latest().await?; + } + Err(err) => return Err(err), + } } } for ((bot_id, session_id), group_entries) in groups { let region_id = Self::derive_region_id(&bot_id, &session_id); let batch = self.records_to_batch(&group_entries)?; - let config = ShardWriterConfig { - shard_id: region_id, - ..Default::default() - }; - let writer = self.dataset.mem_wal_writer(region_id, config).await?; - writer.put(vec![batch]).await?; - writer.close().await?; + // Retry on shard-claim contention. + // + // `ContextStore` is `Clone`, so `add(&mut self)` does not prevent a + // second concurrent writer, and shards are derived from the record's + // own `(bot_id, session_id)` rather than from writer identity — so + // two writers appending to the *same session* necessarily target the + // same shard. Opening a `mem_wal_writer` claims that shard's epoch, + // and a concurrent open loses the claim with "another writer claimed + // epoch N" (or fences an already-open writer). + // + // Both are transient: the losing writer only has to reopen against + // the now-current epoch. Without this retry the append surfaced a + // hard error to the caller — reproducible with two clones appending + // to one session (see tests/context_store_concurrent_add.rs). + // + // Bounded, because a persistent claim failure is a real fault and + // must not spin forever. Rows carry caller-supplied ids and are + // de-duplicated by id at read time, so a retried append cannot + // double-count. + const MAX_CLAIM_ATTEMPTS: usize = 5; + let mut attempt = 0; + loop { + let config = ShardWriterConfig { + shard_id: region_id, + ..Default::default() + }; + let result = async { + let writer = self.dataset.mem_wal_writer(region_id, config).await?; + writer.put(vec![batch.clone()]).await?; + writer.close().await + } + .await; + + match result { + Ok(()) => break, + Err(err) if is_shard_contention_error(&err) => { + attempt += 1; + if attempt >= MAX_CLAIM_ATTEMPTS { + return Err(err); + } + // Brief backoff so the winning writer can finish its + // claim before we re-read the manifest. + tokio::time::sleep(std::time::Duration::from_millis(10 * attempt as u64)) + .await; + } + Err(err) => return Err(err), + } + } } Ok(self.dataset.manifest.version) @@ -3104,6 +3174,36 @@ where .and_then(|col| col.as_ref().as_any().downcast_ref::()) } +/// Whether a Lance commit failed because a concurrent transaction preempted it. +/// +/// Lance signals this as retryable in the message itself ("Please retry"), and +/// the loser only needs to re-read the manifest and re-evaluate. Matching on the +/// text is unfortunate but is how the rest of this crate classifies Lance +/// errors; the typed variant is not exposed. +fn is_retryable_commit_conflict(err: &LanceError) -> bool { + let text = err.to_string(); + text.contains("Retryable commit conflict") || text.contains("was preempted by concurrent") +} + +/// Whether a Lance error means this writer lost a race for a MemWAL shard and +/// should reopen against the current epoch. +/// +/// Two distinct wordings, both transient: +/// - **claim loss** — "Failed to claim shard ... another writer claimed epoch N"; +/// a concurrent open won the claim before ours landed. +/// - **fence** — an already-open writer was superseded by a newer epoch. +/// +/// Note the claim-loss text contains neither "fenced" nor "Fenced", so matching +/// only on those (as the rollout store's `is_fenced_error` does) would miss the +/// case `ContextStore` actually hits, since it opens a fresh writer per append. +fn is_shard_contention_error(err: &LanceError) -> bool { + let text = err.to_string(); + text.contains("fenced") + || text.contains("Fenced") + || text.contains("Failed to claim shard") + || text.contains("another writer claimed epoch") +} + /// Render a SQL `IN (...)` value list as comma-separated quoted string /// literals, escaping any embedded single quotes. Callers ensure the input is /// non-empty before building the surrounding `IN ()` clause. diff --git a/crates/lance-context-core/tests/context_store_concurrent_add.rs b/crates/lance-context-core/tests/context_store_concurrent_add.rs new file mode 100644 index 0000000..b6042e1 --- /dev/null +++ b/crates/lance-context-core/tests/context_store_concurrent_add.rs @@ -0,0 +1,197 @@ +//! Does `ContextStore` lose data when two handles append concurrently? +//! +//! `ContextStore` is `#[derive(Clone)]`, so `add(&mut self)` does not prevent a +//! second concurrent writer — cloning the store yields a second mutable handle +//! onto the same dataset. Each `add` opens a fresh `mem_wal_writer` per +//! `(bot_id, session_id)` group, and opening claims the shard epoch, which +//! fences any other live writer of that shard. Unlike `RolloutStore::add` and +//! `DatagenStore::append`, this path has **no fence retry**. +//! +//! Shards are derived from the record's own `(bot_id, session_id)` +//! (`derive_region_id`), not from writer identity, so two concurrent writers +//! touching the same session necessarily target the same shard. +//! +//! This test exists to establish, empirically, whether that is a real data-loss +//! path or whether something upstream serializes it. + +#![recursion_limit = "256"] + +use lance_context_core::{ContextRecord, ContextStore, ContextStoreOptions}; + +fn rec(id: &str, session: &str) -> ContextRecord { + ContextRecord { + id: id.to_string(), + external_id: None, + run_id: "run".to_string(), + bot_id: Some("bot".to_string()), + session_id: Some(session.to_string()), + tenant: None, + source: None, + created_at: chrono::Utc::now(), + role: "user".to_string(), + state_metadata: None, + metadata: None, + relationships: vec![], + expires_at: None, + retention_policy: None, + lifecycle_status: "active".to_string(), + retired_at: None, + retired_reason: None, + supersedes_id: None, + superseded_by_id: None, + content_type: "text/plain".to_string(), + text_payload: Some("hello".to_string()), + binary_payload: None, + payload_uri: None, + payload_size: None, + payload_checksum: None, + embedding: None, + } +} + +/// Two clones appending to the **same** `(bot_id, session_id)` — therefore the +/// same MemWAL shard — concurrently. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_clones_same_session_do_not_lose_rows() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + + // Deliberately NO serial warm-up: both writers race from a cold dataset, so + // this also covers the one-time `initialize_mem_wal` (CreateIndex) conflict + // on first write, not just steady-state shard-claim contention. + let store = ContextStore::open_with_options(&uri, ContextStoreOptions::default()) + .await + .unwrap(); + + let n = 20; + let mut a = store.clone(); + let mut b = store.clone(); + + let (ra, rb) = tokio::join!( + async move { + let mut errs = Vec::new(); + for i in 0..n { + if let Err(e) = a.add(&[rec(&format!("a-{i}"), "shared")]).await { + errs.push(format!("a-{i}: {e}")); + } + } + errs + }, + async move { + let mut errs = Vec::new(); + for i in 0..n { + if let Err(e) = b.add(&[rec(&format!("b-{i}"), "shared")]).await { + errs.push(format!("b-{i}: {e}")); + } + } + errs + } + ); + + let reader = ContextStore::open_with_options(&uri, ContextStoreOptions::default()) + .await + .unwrap(); + let listed = reader.list(None, None).await.unwrap(); + let ids: std::collections::HashSet = listed.iter().map(|r| r.id.clone()).collect(); + + let missing: Vec = (0..n) + .flat_map(|i| [format!("a-{i}"), format!("b-{i}")]) + .filter(|id| !ids.contains(id)) + .collect(); + + eprintln!( + "errors_a={} errors_b={} listed={} missing={}", + ra.len(), + rb.len(), + listed.len(), + missing.len() + ); + if !ra.is_empty() { + eprintln!("first a error: {}", ra[0]); + } + if !rb.is_empty() { + eprintln!("first b error: {}", rb[0]); + } + + // An append that returned Ok must be readable. Silent loss is the failure + // mode under test: a fenced writer whose `put` succeeded into a dead epoch. + assert!( + missing.is_empty(), + "{} appends returned Ok but are not readable: {missing:?}", + missing.len() + ); + assert!( + ra.is_empty() && rb.is_empty(), + "concurrent appends to the same session errored: a={ra:?} b={rb:?}" + ); +} + +/// Two clones appending to **different** sessions — therefore different shards. +/// Should be entirely contention-free; this isolates whether any failure above +/// is shard contention or something broader. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_clones_distinct_sessions_do_not_lose_rows() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + + // Deliberately NO serial warm-up: both writers race from a cold dataset, so + // this also covers the one-time `initialize_mem_wal` (CreateIndex) conflict + // on first write, not just steady-state shard-claim contention. + let store = ContextStore::open_with_options(&uri, ContextStoreOptions::default()) + .await + .unwrap(); + + let n = 20; + let mut a = store.clone(); + let mut b = store.clone(); + + let (ra, rb) = tokio::join!( + async move { + let mut errs = Vec::new(); + for i in 0..n { + if let Err(e) = a.add(&[rec(&format!("a-{i}"), "session-a")]).await { + errs.push(format!("a-{i}: {e}")); + } + } + errs + }, + async move { + let mut errs = Vec::new(); + for i in 0..n { + if let Err(e) = b.add(&[rec(&format!("b-{i}"), "session-b")]).await { + errs.push(format!("b-{i}: {e}")); + } + } + errs + } + ); + + let reader = ContextStore::open_with_options(&uri, ContextStoreOptions::default()) + .await + .unwrap(); + let listed = reader.list(None, None).await.unwrap(); + let ids: std::collections::HashSet = listed.iter().map(|r| r.id.clone()).collect(); + + let missing: Vec = (0..n) + .flat_map(|i| [format!("a-{i}"), format!("b-{i}")]) + .filter(|id| !ids.contains(id)) + .collect(); + + eprintln!( + "distinct-shards: errors_a={} errors_b={} listed={} missing={}", + ra.len(), + rb.len(), + listed.len(), + missing.len() + ); + + assert!( + missing.is_empty(), + "{} appends to distinct sessions returned Ok but are not readable: {missing:?}", + missing.len() + ); + assert!( + ra.is_empty() && rb.is_empty(), + "concurrent appends to distinct sessions errored: a={ra:?} b={rb:?}" + ); +}