fix(store): retry MemWAL init and shard-claim races in ContextStore::add (multi-replica hardening)

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<String> = 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::<A>())
 }
 
+/// 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.

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<String> = listed.iter().map(|r| r.id.clone()).collect();
+
+    let missing: Vec<String> = (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<String> = listed.iter().map(|r| r.id.clone()).collect();
+
+    let missing: Vec<String> = (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:?}"
+    );
+}