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

[beinan]

Hardening, not a fix for a live single-process bug. Rescoped after @beinan correctly pushed back — see the scope section.

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. Shard epoch claim. 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 ("Please retry"). 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's 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.

⚠️ 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

I originally framed this as a live bug with a "20/20 appends failed" figure. That was overclaiming, and the correction matters:

• AppState::get_or_open_context_store caches by name with a double-check, so one name maps to one Arc<RwLock<ContextStore>>.
• Every mutating route holds store_lock.write().await — appends are serialized.
• #[derive(Clone)] on ContextStore has no value-level clone in any non-test code (verified).
• The Python binding is one store per Context.

A single process cannot produce the contending writers. My test manufactures them.

What it does cover: 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, each satisfying the single-store guarantee — 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), documented as "each instance owns exactly one shard and no two instances ever contend". ContextStore has no equivalent.

There is no server deployment manifest in this repo (deploy/kubernetes/ has only master.yaml, and master does not write context stores), so this is a latent architectural gap, not 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.

Testing

Tests use two clones to simulate two independent writers of one dataset — the multi-process shape. Verified non-vacuous: forcing both retry predicates to false makes both tests fail.

Now rebased onto main; the CI change that was bundled here shipped separately as #202.

Open question

Will a context store ever be written by more than one server replica? If the answer is a firm no, this can be closed as unnecessary complexity — I'd rather drop it than carry speculative hardening.

🤖 Generated with Claude Code

[beinan]

Correction — I overstated this

@beinan asked whether this is a real bug given a guarantee of not initializing two stores. Having checked the actual call sites, the answer is largely no, and my framing was wrong.

Within a single process, that guarantee already holds:

• AppState::get_or_open_context_store (state.rs:368) caches by name with a double-check, so one name maps to one Arc<RwLock<ContextStore>>.
• Every mutating route takes store_lock.write().await — appends are serialized.
• #[derive(Clone)] on ContextStore has no value-level clone in any non-test code (verified by grep). My test constructs a writer pair that production never creates.
• The Python binding is one store per Context object.

So the "20/20 appends failed" figure is from a scenario I manufactured, not one the server can reach. Presenting it as a live production bug was overclaiming, and the PR description should not have omitted the precondition.

What remains

One case the single-process guarantee cannot cover: 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, each satisfying the guarantee — writing the same (bot_id, session_id) land on the same MemWAL shard, with no cross-process lock. Contrast RolloutStore, which shards by writer identity (shard_id from hostname/StatefulSet ordinal) precisely so instances never contend.

That said, I found no evidence in this repo that the server is deployed multi-replica: deploy/kubernetes/ contains only master.yaml, and master does not write context stores. So this is a latent architectural gap, not a defect anyone is currently hitting.

Plan

• The CI fix — which is unambiguous and blocks nothing else — is split out into fix(ci): run Rust integration tests, not just unit tests #202.
• This PR will be rescoped to the retry hardening alone, with honest framing (multi-replica hardening, not a live bug), or closed if we decide ContextStore will never be written by multiple replicas.

The retries themselves are still defensible: both errors are ones Lance itself marks retryable ("Please retry" is in the message text), and the claim-loss message contains neither "fenced" nor "Fenced", so the rollout store's is_fenced_error would not have matched it. But that is hardening, not firefighting.

Question that decides it: will a context store ever be written by more than one server replica? If not, this should just be closed.

[beinan]

Closing this — my fix makes things worse

test (stable) failed on the rebased branch, and it is not a CI flake. My own test hung for 19 minutes until the 30-minute job timeout killed it:

test concurrent_clones_distinct_sessions_do_not_lose_rows ... ok
test concurrent_clones_same_session_do_not_lose_rows has been running for over 60 seconds
##[error]The operation was canceled.
Terminate orphan process: pid (6932) (context_store_concurrent_add-...)

Reproduced locally: 2 of 5 runs hang indefinitely on the same-session test. Distinct sessions never hangs (3/3 clean). So it is specifically shard contention — and it is my change that causes it.

Root cause

ShardWriter::wait_for_flush_drain (lance write.rs:1832) is an unbounded loop waiting on flush watchers. A writer that gets fenced while inside close() waits on a watcher nobody will ever complete.

My retry loop actively creates that situation: every retry opens a new writer, which claims the shard epoch, which fences the previous writer that may still be inside close(). Lance also retries the claim 16× internally (manifest.rs:423), so my 5 attempts nest into up to 80 contending claims. It is a livelock I introduced, not one I found.

Assessment

Putting this together with the earlier correction:

1. The scenario is not reachable in a single process — the server caches one Arc<RwLock<ContextStore>> per name and serializes writes.
2. There is no evidence the server is deployed multi-replica (deploy/kubernetes/ has only master.yaml).
3. The retry does not actually work under real contention — it converts a clean error into a hang, which is strictly worse. An error is recoverable; a hung writer holding a flush watcher is not.

Retrying at this layer is the wrong mechanism. The claim is not a lock you can politely re-take: reopening fences whoever is mid-close(). The durable fix for multi-replica context writes is per-writer sharding — what RolloutStore already does with shard_id from the instance hostname — so writers never target the same shard in the first place.

Filing that as an issue instead. Closing this PR.

Salvage from the work: the CI fix shipped separately in #202 and is the reason this was caught at all — under the old --lib these integration tests never ran, so a hang-inducing change would have merged clean.