From dca35749ecccd7684c1838cfedc1c65bf29c0cb3 Mon Sep 17 00:00:00 2001 From: Beinan Date: Sun, 26 Jul 2026 05:01:09 +0000 Subject: [PATCH] fix(rollout): stop WAL merge from blocking concurrent appends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A count- or time-triggered WAL self-merge held the store's write lock for its entire duration, so every concurrent `add` to that store blocked until it finished — ~17s once per flush interval on a 20-writer abfss deployment, with appends otherwise at p50 ~13ms. This is an LSM: a merge folds *sealed* generations into the base table while `add` writes the active memtable at the WAL tail. They touch disjoint data and should never have serialized. Root cause was `claim_epoch`, not the lock. To drain merged generations from the shard manifest the merge claimed the shard epoch, which bumps `writer_epoch` and fences *every* live writer of that shard — including our own. So the merge had to `close()` the resident writer first (rollout_store.rs), and callers had to hold an exclusive lock across the whole merge to hide that window. The lock was covering for the epoch claim. The claim is unnecessary. `ShardManifestStore::commit_update(local_epoch, ..)` only rejects a writer whose epoch is *older* than the stored one (`stored > local` ⇒ fenced, lance manifest.rs), and Lance's own flush path reuses a single epoch for every manifest commit a writer ever makes (lance flush.rs) — the epoch is an ownership token, not a per-commit token. Lance also explicitly sanctions concurrent draining by an external compactor: recovery keys off `replay_after_wal_entry_position` and deliberately does not consult `flushed_generations`, "since an external compactor may legitimately drain that vector back to empty" (lance write.rs, wal.rs). So the merge now reuses the shard's current epoch and never closes the writer. On top of that it is split by lock scope: prepare (&self) — seal the memtable, read every flushed generation from object storage (the dominant cost: 1.5–3.4s per generation on abfss) commit (&mut self) — append to the base table, drain the manifest, delete the merged generation dirs Callers run `prepare` under a read lock so appends keep flowing, and take the write lock only for the short commit. Both sweepers and the HTTP merge-wal handler were updated. Note the issue's suggested fix (drop the count trigger from the flush sweeper) would not have worked: the *global* sweeper takes the same write lock and calls `cleanup_own_shard` with threshold 1, so the stall would simply have moved to a different interval. `prepare_cleanup_merge` keeps the seal-before-manifest-read ordering that `cleanup_own_shard` relied on. That ordering is load-bearing: with ROLLOUT_FLUSH_INTERVAL_SECS=0 nothing else seals the memtable, so reading the manifest first would leave `flushed_generations` empty forever and rows durable but permanently invisible until a restart replayed the WAL. New integration suite (tests/wal_merge_concurrency.rs, 5 tests). The existing suite is explicitly "one store, one shard, fully serial — no concurrency, no fence", so none of this had coverage: - appends succeed and no row is lost while a merge runs - a generation sealed *during* a merge is not dropped by the drain - concurrent merges do not duplicate rows - an interrupted merge loses nothing and the next merge converges - an append is not blocked for the merge's duration (the regression itself) Each test drives merges through a helper that reproduces the sweeper's exact lock discipline (read-lock prepare, write-lock commit); merging under a single exclusive lock in the test would pass while the real stall persisted. Verified the tests are not vacuous: restoring `claim_epoch` makes `appends_succeed_and_are_not_lost_while_merge_runs` and `append_is_not_blocked_for_the_duration_of_a_merge` fail. They also do not compile against the old `&mut self` signature at all — "merge concurrent with append" was inexpressible, which is the design flaw in one line. Follow-ups filed rather than folded in: - #198: make `dataset` interior-mutable so merge/compact need no exclusive lock at all (~52 call sites; deserves its own review). - `datagen_store.rs` has the identical `claim_epoch`-then-close pattern. Co-Authored-By: Claude --- crates/lance-context-core/src/lib.rs | 5 +- .../lance-context-core/src/rollout_store.rs | 276 +++++++++--- .../tests/wal_merge_concurrency.rs | 417 ++++++++++++++++++ .../src/routes/rollouts.rs | 45 +- crates/lance-context-server/src/state.rs | 90 +++- 5 files changed, 756 insertions(+), 77 deletions(-) create mode 100644 crates/lance-context-core/tests/wal_merge_concurrency.rs diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 838e7ce5..21c1c228 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -47,8 +47,9 @@ pub use record::{ pub use registry::{RegistryEntry, RolloutRegistry}; pub use rollout::{RolloutRecord, ROLE_ARTIFACT, ROLE_ASSISTANT, ROLE_GRADE, ROLE_TOOL}; pub use rollout_store::{ - rollout_schema, ListSource, RolloutFilters, RolloutObservation, RolloutPage, RolloutStore, - RolloutStoreOptions, SqlQueryResult, SQL_MAX_RESULT_ROWS, SQL_MAX_SCAN_ROWS, SQL_TABLE_NAME, + rollout_schema, ListSource, PreparedMerge, RolloutFilters, RolloutObservation, RolloutPage, + RolloutStore, RolloutStoreOptions, SqlQueryResult, SQL_MAX_RESULT_ROWS, SQL_MAX_SCAN_ROWS, + SQL_TABLE_NAME, }; pub use storage::{create_local_dir_if_needed, join_uri, validate_store_name, MAX_STORE_NAME_LEN}; pub use store::{ diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index a2a5d0c9..b3abaef1 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -102,6 +102,27 @@ enum FlushOutcome { Fenced, } +/// Rows read out of the flushed generations, ready to be appended to the base +/// table and drained from the shard manifest. +/// +/// Produced by [`RolloutStore::prepare_merge`] under `&self` (so appends keep +/// running while it reads object storage) and consumed by +/// [`RolloutStore::commit_merge`] under `&mut self`. +pub struct PreparedMerge { + merged_generations: HashSet, + merged_paths: Vec, + batches: Vec, + merge_schema: Arc, +} + +impl PreparedMerge { + /// Number of generations this merge will reclaim. + #[must_use] + pub fn generation_count(&self) -> usize { + self.merged_generations.len() + } +} + /// Number of shard manifest files to scan per batch when discovering the latest /// shard state (mirrors the constant used by `ContextStore`). const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16; @@ -778,9 +799,70 @@ impl RolloutStore { /// in [`Self::add`] (with `threshold = merge_after_generations`) and the /// caller-driven time trigger via [`Self::cleanup_own_shard`] (with /// `threshold = 1`, i.e. merge whatever is pending). Both merge only - /// the shard this instance owns and writes, so the epoch claim never fences - /// another writer. + /// the shard this instance owns and writes. async fn merge_own_shard_if_ready(&mut self, threshold: usize) -> LanceResult { + let Some((manifest_store, manifest, prepared)) = + self.prepare_merge_if_ready(threshold).await? + else { + return Ok(0); + }; + let pending = prepared.generation_count(); + self.commit_merge(&manifest_store, &manifest, prepared) + .await?; + Ok(pending) + } + + /// The shared-lock half of a merge: decide whether one is due, seal the + /// memtable, and read the flushed generations into memory. + /// + /// Takes `&self`, so a caller holding a *read* lock can run the expensive + /// part while appends continue, then take the write lock only to hand the + /// result to [`Self::commit_merge`]. Returns `None` when nothing is due. + /// + /// This is the pairing that keeps a background merge off the write path: + /// + /// ```ignore + /// let prepared = { store.read().await.prepare_merge_if_ready(1).await? }; + /// if let Some((manifest_store, manifest, prepared)) = prepared { + /// let mut guard = store.write().await; + /// guard.commit_merge(&manifest_store, &manifest, prepared).await?; + /// } + /// ``` + /// + /// Callers that do not care about lock scope should use + /// [`Self::cleanup_own_shard`] or [`Self::maybe_merge_own_shard`], which do + /// both halves under whatever lock the caller already holds. + pub async fn prepare_merge_if_ready( + &self, + threshold: usize, + ) -> LanceResult> { + self.prepare_merge_if_ready_inner(threshold, false).await + } + + /// [`Self::prepare_merge_if_ready`], but seals the active memtable *before* + /// consulting the manifest. + /// + /// This is the time-triggered (`threshold = 1`) behaviour of + /// [`Self::cleanup_own_shard`], and the ordering is load-bearing: with the + /// periodic flush sweeper disabled (`ROLLOUT_FLUSH_INTERVAL_SECS=0`) nothing + /// else seals the memtable, so `flushed_generations` would stay empty, the + /// threshold check would return `None`, and rows would remain durable but + /// permanently invisible until a process restart replayed the WAL. + pub async fn prepare_cleanup_merge( + &self, + ) -> LanceResult> { + self.prepare_merge_if_ready_inner(1, true).await + } + + async fn prepare_merge_if_ready_inner( + &self, + threshold: usize, + seal_first: bool, + ) -> LanceResult> { + if seal_first { + // Materialize anything buffered so it is eligible for this pass. + self.flush().await?; + } let object_store = self.dataset.object_store(None).await?; let branch_location = self.dataset.branch_location(); let manifest_store = ShardManifestStore::new( @@ -790,13 +872,29 @@ impl RolloutStore { DEFAULT_MANIFEST_SCAN_BATCH_SIZE, ); let Some(manifest) = manifest_store.read_latest().await? else { - return Ok(0); + return Ok(None); }; let pending = manifest.flushed_generations.len(); if pending == 0 || pending < threshold.max(1) { - return Ok(0); + return Ok(None); } - self.merge_own_shard(&manifest_store, &manifest).await?; + let Some(prepared) = self.prepare_merge(&manifest).await? else { + return Ok(None); + }; + Ok(Some((manifest_store, manifest, prepared))) + } + + /// Commit a merge prepared by [`Self::prepare_merge_if_ready`]. See that + /// method for the intended read-lock/write-lock split. + pub async fn commit_prepared_merge( + &mut self, + manifest_store: &ShardManifestStore, + manifest: &ShardManifest, + prepared: PreparedMerge, + ) -> LanceResult { + let pending = prepared.generation_count(); + self.commit_merge(manifest_store, manifest, prepared) + .await?; Ok(pending) } @@ -857,45 +955,107 @@ impl RolloutStore { /// committing the drain, *without merging it* (data loss). `commit_update` /// re-reads the latest manifest and applies this closure to it, so the /// retain filter runs against the current state, not the stale snapshot. + /// # Concurrency: the expensive phase does not need exclusive access + /// + /// A merge only ever touches *sealed* generations — history — while `add` + /// writes the active memtable at the WAL tail. They operate on disjoint + /// data, which is the whole point of an LSM, so a merge should not stop the + /// write path. Two things used to force it to: + /// + /// 1. `claim_epoch` bumped `writer_epoch`, fencing every live writer of the + /// shard — including our own — so the merge had to `close()` the resident + /// writer first and callers had to hide that window behind an exclusive + /// lock. Removed: the epoch is an *ownership* token, not a per-commit + /// token. [`ShardManifestStore::commit_update`] only rejects a writer + /// whose epoch is **older** than the stored one, and Lance's own flush + /// path reuses a single epoch for every manifest commit a writer makes. + /// Reusing the shard's current epoch commits the drain and leaves the + /// live writer untouched. /// - /// # Safety of the epoch claim + /// 2. The dominant cost — reading every flushed generation out of object + /// storage into memory — sat inside that same exclusive section. It is + /// now split into [`Self::prepare_merge`], which takes `&self`, so a + /// caller can run it under a shared lock and take the exclusive lock only + /// for the short commit ([`Self::commit_merge`]). /// - /// `claim_epoch` bumps the shard's writer epoch, which would fence any - /// *other* live writer of this shard. That is safe here precisely because - /// each instance merges only the shard it owns and writes: there is no - /// other live writer of `self.write_shard` to fence. After the merge this - /// instance's next `add` opens a fresh `mem_wal_writer`, which re-claims - /// the (now-current) epoch. + /// Lance explicitly sanctions an external compactor draining + /// `flushed_generations` concurrently: recovery keys off + /// `replay_after_wal_entry_position` alone and deliberately does not consult + /// that vector. + /// + /// # Surgical drain, not blanket clear + /// + /// The drain removes only the generation ids this call actually merged, + /// retaining anything else present. This is load-bearing now that a + /// concurrent flush can seal a new generation mid-merge: `commit_update` + /// re-runs the closure against a freshly-read manifest on every CAS retry, + /// so a *relative* edit (retain-not-in-set) composes with a concurrent + /// flush's append, while an absolute `flushed_generations = []` would + /// silently discard a generation that was never merged — data loss. For the + /// same reason the closure must preserve `current_generation`, + /// `replay_after_wal_entry_position` and `wal_entry_position_last_seen`, + /// which a concurrent flush advances; `..current.clone()` carries them. /// /// Rollout rows are immutable and de-duplicated by `id` at read time, so /// even if a crash interrupts the sequence (data appended to base but /// manifest not yet drained), a subsequent read simply sees the rows via /// both the base table and the still-listed generation and de-dups them — /// no double counting. The next merge attempt then drains the manifest. - async fn merge_own_shard( - &mut self, - manifest_store: &ShardManifestStore, - manifest: &ShardManifest, - ) -> LanceResult<()> { + /// The `&self` half of a merge: everything that can run while appends + /// continue — sealing the memtable and reading every flushed generation + /// into memory. + /// + /// Returns `None` when there is nothing to merge. Callers holding a shared + /// lock can run this, then upgrade to the exclusive lock only for + /// [`Self::commit_merge`]. + async fn prepare_merge(&self, manifest: &ShardManifest) -> LanceResult> { if manifest.flushed_generations.is_empty() { - return Ok(()); + return Ok(None); } - // This merge is about to `claim_epoch`, which fences any live writer of - // this shard — including our own resident writer. Close it (draining its - // background tasks; `ShardWriter` has no `Drop`) and clear it so the next - // `add` transparently reopens against the freshly-claimed epoch. - observe_phase!("seal", self.close().await)?; - - self.ensure_latest_rollout_schema().await?; + // Seal the active memtable so its rows are in a generation rather than + // buffered. Uses `&self` (the writer lives behind a Mutex), and no + // longer closes the writer: without `claim_epoch` there is nothing to + // fence it against, so appends continue throughout. + observe_phase!("seal", self.flush().await)?; - // Resolve each flushed generation to its absolute dataset path and read - // all its rows into memory. Record which generation ids we merge so the - // drain can remove exactly these and nothing else. + // The expensive phase: pull every flushed generation out of object + // storage. Buffered in memory, so this is the part that must not hold an + // exclusive lock. let (merged_generations, merged_paths, batches, merge_schema) = observe_phase!("read", self.read_flushed_generations(manifest).await)?; - // Append the merged rows to the base table. + Ok(Some(PreparedMerge { + merged_generations, + merged_paths, + batches, + merge_schema, + })) + } + + /// The `&mut self` half of a merge: append the prepared rows to the base + /// table, drain the merged generations from the manifest, and delete their + /// directories. + /// + /// Short relative to [`Self::prepare_merge`], and `&mut self` because + /// `Dataset::append` rebinds this handle to the post-commit dataset — which + /// is also what keeps subsequent reads on this instance correct, since the + /// merged generations are about to disappear from the manifest. + async fn commit_merge( + &mut self, + manifest_store: &ShardManifestStore, + manifest: &ShardManifest, + prepared: PreparedMerge, + ) -> LanceResult<()> { + let PreparedMerge { + merged_generations, + merged_paths, + batches, + merge_schema, + } = prepared; + + self.ensure_latest_rollout_schema().await?; + if !batches.is_empty() { observe_phase!( "append", @@ -903,52 +1063,58 @@ impl RolloutStore { )?; } - // Drain the merged generations from the shard manifest. Claim the - // shard's epoch (safe: we own it) and commit a manifest that retains - // every generation except the ones we just folded into the base table. - // Removing only the merged ids (rather than clearing the vec) is what - // makes this safe against a generation that lands after we read the - // manifest: it is preserved for the next merge instead of being dropped. - let (epoch, _) = observe_phase!( - "claim_epoch", - manifest_store.claim_epoch(manifest.shard_spec_id).await - )?; + // Reuse the shard's *current* epoch rather than claiming a new one: + // claiming would fence our own live writer (see the doc comment above). + // `commit_update` still fails cleanly if a genuinely new writer has + // claimed the shard meanwhile — its epoch would exceed ours — which is + // the correct outcome, since that writer now owns the shard. + let epoch = manifest.writer_epoch; observe_phase!( "drain", manifest_store .commit_update(epoch, |current| ShardManifest { version: current.version + 1, + // Relative edit: retain everything we did not merge. Must + // never become an absolute assignment — see the doc comment. flushed_generations: current .flushed_generations .iter() .filter(|fg| !merged_generations.contains(&fg.generation)) .cloned() .collect(), + // Carries `current_generation`, + // `replay_after_wal_entry_position` and + // `wal_entry_position_last_seen`, which a concurrent flush + // advances and this merge must never roll back. ..current.clone() }) .await )?; - // Delete the merged generations' blob directories now that no manifest - // references them. Ordering matters: the drain above already removed - // these ids from `flushed_generations`, so a reader can no longer resolve - // them — deleting the data second (never before) keeps the sequence - // crash-safe. If the process dies between the drain and here, the rows - // are already in the base table and the manifest no longer lists these - // generations, so nothing reads them; they simply become storage that a - // sweep can reclaim later. - // - // Best-effort: a delete failure must NOT fail the merge — the merge has - // logically succeeded (data appended, manifest drained). A failed delete - // only leaks one directory, which the same reclamation path handles. - // Skipping this deletion is exactly the historical storage leak: every - // merged generation left its `_mem_wal/{shard}/{gen}/` directory behind - // forever. + self.delete_merged_generation_dirs(&merged_paths).await + } + + /// Delete the merged generations' blob directories now that no manifest + /// references them. + /// + /// Ordering matters: the drain already removed these ids from + /// `flushed_generations`, so a reader can no longer resolve them — deleting + /// the data second (never before) keeps the sequence crash-safe. If the + /// process dies between the drain and here, the rows are already in the base + /// table and the manifest no longer lists these generations, so nothing + /// reads them; they simply become storage a later sweep reclaims. + /// + /// Best-effort: a delete failure must NOT fail the merge — the merge has + /// logically succeeded (data appended, manifest drained). A failed delete + /// only leaks one directory. Skipping this deletion entirely is exactly the + /// historical storage leak: every merged generation left its + /// `_mem_wal/{shard}/{gen}/` directory behind forever. + async fn delete_merged_generation_dirs(&self, merged_paths: &[String]) -> LanceResult<()> { let phase = timer_start!(); let object_store = self.dataset.object_store(None).await?; let branch_path = self.dataset.branch_location().path.clone(); - for path in &merged_paths { + for path in merged_paths { let gen_dir = branch_path .clone() .join("_mem_wal") diff --git a/crates/lance-context-core/tests/wal_merge_concurrency.rs b/crates/lance-context-core/tests/wal_merge_concurrency.rs new file mode 100644 index 00000000..4ce0ebc0 --- /dev/null +++ b/crates/lance-context-core/tests/wal_merge_concurrency.rs @@ -0,0 +1,417 @@ +//! Concurrency tests for WAL self-merge: a merge must never block or corrupt +//! concurrent appends. +//! +//! The pre-existing suite (`wal_merge_generation_cleanup.rs`) is explicitly +//! "one store, one shard, fully serial — no concurrency, no fence", so none of +//! the properties below had coverage. They matter because a merge used to +//! `claim_epoch`, which fences the shard's live writer; callers hid that window +//! by holding an exclusive lock across the whole merge, turning background +//! compaction into a stop-the-world pause on the write path (~17s observed on +//! abfss with 20 concurrent writers). +//! +//! The merge now reuses the shard's current epoch instead of claiming a new +//! one, so the writer is never fenced and callers can merge under a shared +//! lock. These tests pin the properties that makes safe: +//! +//! 1. appends succeed while a merge runs, and no row is lost; +//! 2. a generation sealed *during* a merge is not silently dropped by the drain; +//! 3. concurrent merges do not duplicate rows; +//! 4. an interrupted merge loses nothing (rows stay readable exactly once); +//! 5. `add` is not blocked for the merge's duration. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use lance_context_core::{RolloutRecord, RolloutStore, RolloutStoreOptions, ROLE_ASSISTANT}; +use tokio::sync::RwLock; + +/// Run one full merge exactly the way the server's sweepers do: seal + read the +/// generations under a **read** lock (so appends keep running), then take the +/// write lock only for the short commit. Returns generations reclaimed. +/// +/// Every test drives merges through this helper so the lock discipline under +/// test is the same one production uses -- a test that merged under a single +/// exclusive lock would pass while the real stall persisted. +async fn merge_like_sweeper(store: &Arc>) -> usize { + let prepared = { store.read().await.prepare_cleanup_merge().await.unwrap() }; + match prepared { + Some((manifest_store, manifest, prepared)) => store + .write() + .await + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .unwrap(), + None => 0, + } +} + +fn rec(id: &str) -> RolloutRecord { + RolloutRecord { + id: id.to_string(), + rollout_id: "r".to_string(), + problem_id: "p".to_string(), + dataset: Some("d".to_string()), + sequence_order: 0, + role: ROLE_ASSISTANT.to_string(), + created_at: chrono::Utc::now(), + content: Some("x".to_string()), + content_type: "text/plain".to_string(), + model_input_string: None, + model_output_string: None, + rationale: None, + problem_text: None, + user_metadata: None, + input_tokens: None, + output_tokens: None, + num_input_tokens: None, + num_output_tokens: None, + output_logprobs: None, + input_logprobs: None, + ref_logprobs: None, + loss_mask: None, + advantage: None, + reward: None, + raw_reward: None, + grader_id: None, + score: None, + include_in_training: None, + exclude_reason: None, + policy_version: None, + relationships: vec![], + binary_payload: None, + payload_size: None, + payload_checksum: None, + artifact_type: None, + metadata: None, + } +} + +fn opts(shard: &str) -> RolloutStoreOptions { + RolloutStoreOptions { + shard_id: Some(shard.to_string()), + // Never fire the count trigger implicitly; every test drives merges + // explicitly so the interleaving under test is the one being asserted. + merge_after_generations: None, + ..Default::default() + } +} + +/// Read every row through the deduplicating LSM path and return sorted ids. +/// +/// Asserting on `observe().row_count` would be wrong: it is a raw `count_rows()` +/// over the base table and does not de-duplicate, so it counts the physical +/// duplicates an interrupted merge can legitimately leave behind. +async fn read_ids(store: &Arc>) -> Vec { + let mut ids: Vec = store + .read() + .await + .list(None, None) + .await + .unwrap() + .iter() + .map(|r| r.id.clone()) + .collect(); + ids.sort(); + ids +} + +/// Appends must keep succeeding while a merge is in flight, and every row +/// written before, during, and after the merge must be readable exactly once. +/// +/// This is the core regression: when the merge claimed the epoch it fenced the +/// live writer, and `add` only retries a fence **once** (un-looped), so appends +/// racing a merge could fail outright. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn appends_succeed_and_are_not_lost_while_merge_runs() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + // Build several sealed generations for the merge to chew through. + for i in 0..8 { + store + .read() + .await + .add(&[rec(&format!("pre-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + // Merge concurrently with a stream of appends. Both take `&self`. + let merger = { + let store = store.clone(); + tokio::spawn(async move { merge_like_sweeper(&store).await }) + }; + let appender = { + let store = store.clone(); + tokio::spawn(async move { + for i in 0..40 { + // Must not error: a fenced writer would surface here. + store + .read() + .await + .add(&[rec(&format!("during-{i}"))]) + .await + .unwrap_or_else(|e| panic!("append {i} failed during merge: {e}")); + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + }; + + let reclaimed = merger.await.unwrap(); + appender.await.unwrap(); + assert!(reclaimed > 0, "merge should have reclaimed generations"); + + // Seal whatever the appender left buffered, then read everything back. + store.read().await.flush().await.unwrap(); + let ids = read_ids(&store).await; + + let mut expected: Vec = (0..8) + .map(|i| format!("pre-{i}")) + .chain((0..40).map(|i| format!("during-{i}"))) + .collect(); + expected.sort(); + + assert_eq!( + ids, expected, + "every row written before and during the merge must be readable exactly once" + ); +} + +/// A generation sealed *while* a merge is running must survive the drain. +/// +/// The drain is a relative edit (retain everything not in the merged set) and +/// `commit_update` re-runs it against a freshly-read manifest on every CAS +/// retry. An absolute `flushed_generations = []` would silently discard a +/// generation that was never merged into the base table — data loss that no +/// epoch value prevents. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn generation_sealed_during_merge_is_not_dropped() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..6 { + store + .read() + .await + .add(&[rec(&format!("old-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + // Race a merge against a writer that keeps sealing brand-new generations. + let merger = { + let store = store.clone(); + tokio::spawn(async move { merge_like_sweeper(&store).await }) + }; + let sealer = { + let store = store.clone(); + tokio::spawn(async move { + for i in 0..10 { + store + .read() + .await + .add(&[rec(&format!("new-{i}"))]) + .await + .unwrap(); + // Seal immediately so these land as generations the merge did + // not snapshot. + store.read().await.flush().await.unwrap(); + } + }) + }; + + merger.await.unwrap(); + sealer.await.unwrap(); + + store.read().await.flush().await.unwrap(); + let ids = read_ids(&store).await; + + for i in 0..6 { + assert!( + ids.contains(&format!("old-{i}")), + "pre-merge row old-{i} vanished; the drain dropped a merged generation" + ); + } + for i in 0..10 { + assert!( + ids.contains(&format!("new-{i}")), + "row new-{i}, sealed during the merge, was dropped by the drain" + ); + } + assert_eq!(ids.len(), 16, "no row may be lost or duplicated: {ids:?}"); +} + +/// Two merges racing must not append the same generations twice. +/// +/// Merges are serialized by an internal mutex (not the store lock, which would +/// also exclude appends); the loser returns 0 rather than waiting. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_merges_do_not_duplicate_rows() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..10 { + store + .read() + .await + .add(&[rec(&format!("row-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + // Fire several merges at once at the same pending set. + let mut handles = Vec::new(); + for _ in 0..4 { + let store = store.clone(); + handles.push(tokio::spawn( + async move { merge_like_sweeper(&store).await }, + )); + } + for h in handles { + // None may error; a loser simply reports 0. + h.await.unwrap(); + } + + let ids = read_ids(&store).await; + let mut deduped = ids.clone(); + deduped.dedup(); + assert_eq!( + ids, deduped, + "the read path must not surface duplicate ids after racing merges" + ); + assert_eq!(ids.len(), 10, "all rows readable exactly once: {ids:?}"); +} + +/// A merge abandoned partway (the sweeper's timeout does exactly this) must not +/// lose data. Rows may be physically duplicated in the base table — the read +/// path de-dups by id — but nothing may disappear, and a later merge must still +/// converge. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn interrupted_merge_loses_nothing_and_next_merge_converges() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + for i in 0..12 { + store + .read() + .await + .add(&[rec(&format!("row-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + // Cancel the merge mid-flight by dropping the future at a very short + // timeout. Whether it lands before or after the drain is timing-dependent — + // both outcomes must preserve every row. + let _ = tokio::time::timeout(Duration::from_millis(5), merge_like_sweeper(&store)).await; + + let after_interrupt = read_ids(&store).await; + let mut unique = after_interrupt.clone(); + unique.dedup(); + assert_eq!( + unique.len(), + 12, + "no row may be lost when a merge is interrupted: {after_interrupt:?}" + ); + + // A subsequent merge must still succeed and converge. + merge_like_sweeper(&store).await; + let after_retry = read_ids(&store).await; + let mut unique_retry = after_retry.clone(); + unique_retry.dedup(); + assert_eq!( + unique_retry.len(), + 12, + "retry after an interrupted merge must converge to the same rows" + ); + assert_eq!( + after_retry, unique_retry, + "read path must de-duplicate rows an interrupted merge left in the base table" + ); +} + +/// The regression this whole change exists for: an append must not wait for a +/// merge to finish. +/// +/// Asserts on *append latency during a merge*, not on the merge itself. The old +/// code held the store's write lock for the merge's full duration, so a +/// concurrent append blocked for exactly that long (~17s in production). +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn append_is_not_blocked_for_the_duration_of_a_merge() { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_string_lossy().to_string(); + let store = Arc::new(RwLock::new( + RolloutStore::open_with_options(&uri, opts("solo")) + .await + .unwrap(), + )); + + // Enough generations that the merge takes materially longer than an append. + for i in 0..25 { + store + .read() + .await + .add(&[rec(&format!("bulk-{i}"))]) + .await + .unwrap(); + store.read().await.flush().await.unwrap(); + } + + let merge_start = Instant::now(); + let merger = { + let store = store.clone(); + tokio::spawn(async move { merge_like_sweeper(&store).await }) + }; + + // Give the merge a moment to get into its expensive read phase. + tokio::time::sleep(Duration::from_millis(10)).await; + + let append_start = Instant::now(); + store.read().await.add(&[rec("racer")]).await.unwrap(); + let append_elapsed = append_start.elapsed(); + + let reclaimed = merger.await.unwrap(); + let merge_elapsed = merge_start.elapsed(); + + assert!(reclaimed > 0, "merge should have done real work"); + // The append must not have been serialized behind the merge. Compare + // against the merge's own duration rather than a fixed threshold, so the + // assertion holds on slow CI without becoming vacuous. + assert!( + append_elapsed * 2 < merge_elapsed || append_elapsed < Duration::from_millis(200), + "append took {append_elapsed:?} while the merge took {merge_elapsed:?}; \ + the append appears to have been blocked by the merge" + ); + + store.read().await.flush().await.unwrap(); + let ids = read_ids(&store).await; + assert!( + ids.contains(&"racer".to_string()), + "the row appended during the merge must be readable" + ); + assert_eq!(ids.len(), 26, "all rows readable exactly once"); +} diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 454489f1..745c0876 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -633,24 +633,41 @@ pub async fn merge_wal( Path(name): Path, ) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; - // The write lock blocks all ingest on this store for the duration of the - // merge, so it is timed separately from the merge itself: a slow merge and a - // merge that spent its time queued behind other writers are different - // problems with the same end-to-end number. + // Split by lock scope: seal + read every flushed generation under the + // *read* lock so ingest on this store keeps running, then take the write + // lock only for the short commit. Holding the write lock across the whole + // merge is what used to stall every concurrent append for its duration. let lock_start = std::time::Instant::now(); - let mut store = store_lock.write().await; - ::metrics::histogram!("rollout_wal_merge_lock_wait_seconds") - .record(lock_start.elapsed().as_secs_f64()); + let prepared = { + let store = store_lock.read().await; + ::metrics::histogram!("rollout_wal_merge_lock_wait_seconds") + .record(lock_start.elapsed().as_secs_f64()); + match store.prepare_cleanup_merge().await { + Ok(prepared) => prepared, + Err(e) => { + ::metrics::counter!("rollout_wal_cleanup_total", "result" => "failed").increment(1); + return Err(AppError::from_lance(e)); + } + } + }; let merge_start = std::time::Instant::now(); - let result = store.cleanup_own_shard().await; - - let reclaimed = match result { - Ok(n) => n, - Err(e) => { - ::metrics::counter!("rollout_wal_cleanup_total", "result" => "failed").increment(1); - return Err(AppError::from_lance(e)); + let reclaimed = match prepared { + Some((manifest_store, manifest, prepared)) => { + let mut store = store_lock.write().await; + match store + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + { + Ok(n) => n, + Err(e) => { + ::metrics::counter!("rollout_wal_cleanup_total", "result" => "failed") + .increment(1); + return Err(AppError::from_lance(e)); + } + } } + None => 0, }; // Emitted unconditionally: gating on `reclaimed > 0` made the common no-op // merge invisible, so a worker that never has anything to merge and a worker diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index e0711b60..1bea8365 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -433,9 +433,48 @@ impl AppState { .collect() }; for (name, store) in resident { + // Same read-lock/write-lock split as the flush sweeper: seal + // and read under the shared lock so appends keep flowing, + // then commit under the exclusive lock. Threshold 1 — merge + // whatever is pending, which is what makes this the *time* + // half of the "time OR count" trigger. + let prepared = { + let guard = store.read().await; + match tokio::time::timeout(pass_timeout, guard.prepare_cleanup_merge()) + .await + { + Ok(Ok(prepared)) => prepared, + Ok(Err(e)) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "failed") + .increment(1); + tracing::warn!( + store = %name, + error = %e, + "global sweeper WAL cleanup failed" + ); + continue; + } + Err(_elapsed) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "timeout") + .increment(1); + tracing::warn!( + store = %name, + "global sweeper WAL cleanup timed out; abandoning this store this tick" + ); + continue; + } + } + }; + let Some((manifest_store, manifest, prepared)) = prepared else { + continue; + }; let mut guard = store.write().await; - match tokio::time::timeout(pass_timeout, guard.cleanup_own_shard()).await { - Ok(Ok(0)) => {} + match tokio::time::timeout( + pass_timeout, + guard.commit_prepared_merge(&manifest_store, &manifest, prepared), + ) + .await + { Ok(Ok(n)) => { metrics::counter!("rollout_wal_cleanup_total", "result" => "merged") .increment(1); @@ -531,11 +570,50 @@ impl AppState { } metrics::counter!("rollout_wal_flush_total", "result" => "ok").increment(1); } - // Count-triggered merge (no-op unless the threshold is set and - // met). Needs the write lock; excludes concurrent appends. + // Count-triggered merge (no-op unless the threshold is set + // and met). + // + // Split by lock scope: the expensive half (sealing the + // memtable and reading every flushed generation out of + // object storage) runs under the *read* lock, so appends + // keep flowing through it. Only the short commit — append to + // the base table, drain the manifest, delete the merged + // dirs — takes the write lock. Previously the whole merge + // held the write lock, which stalled every concurrent append + // for its full duration. + let threshold = state.rollout_merge_after_generations; + if threshold == 0 { + continue; + } + let prepared = { + let guard = store.read().await; + match tokio::time::timeout( + pass_timeout, + guard.prepare_merge_if_ready(threshold), + ) + .await + { + Ok(Ok(prepared)) => prepared, + Ok(Err(e)) => { + tracing::warn!(store = %name, error = %e, "flush sweeper merge failed"); + continue; + } + Err(_elapsed) => { + tracing::warn!(store = %name, "flush sweeper merge timed out"); + continue; + } + } + }; + let Some((manifest_store, manifest, prepared)) = prepared else { + continue; + }; let mut guard = store.write().await; - match tokio::time::timeout(pass_timeout, guard.maybe_merge_own_shard()).await { - Ok(Ok(0)) => {} + match tokio::time::timeout( + pass_timeout, + guard.commit_prepared_merge(&manifest_store, &manifest, prepared), + ) + .await + { Ok(Ok(n)) => { metrics::counter!("rollout_wal_generations_reclaimed_total") .increment(n as u64);