From 9851d260c70bb1d388157daed27a97c924b8853a Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Tue, 22 Sep 2026 21:10:56 +0000 Subject: [PATCH] fix(server): run the count-triggered WAL merge under the prepare/commit lock split (#257) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sweepable::flush for rollout and generic sealed under the read lock, then took the write lock and called maybe_merge_own_shard / maybe_merge_wal, which reads up to MERGE_MAX_GENERATIONS generations from object storage before committing. The exclusive lock was held across that whole read (seconds to tens of seconds on ABFS), stalling every add/get/list on the store and the next 30s flush tick behind it. In production this surfaced as 'flush sweeper timed out' 52 times in 15 minutes and request handlers failing with 'Too many concurrent writers'. The cleanup path already used the prepare (read lock) / commit (write lock) split. This applies it to the count trigger too: - StorageBase::prepare_count_merge(): the shared-lock half of maybe_merge_own_shard, None when the trigger is disabled or not met. Passthroughs on RolloutStore and GenericStore. - Sweepable gains merge_if_due(); flush() now only seals. flush_pass runs merge_if_due after a successful seal and reports it under the cleanup counters, so a failing merge no longer shows up as a failing flush. - The flush sweeper's three kind passes run with tokio::join!, like the cleanup sweeper since #256, so the merge it now carries for generic is not queued behind a slow rollout walk. - POST /generic/{name}/merge-wal uses the same split; the rollout route already did. Tests: the count trigger merges at threshold, is a no-op below it, treats merge_after_generations=0 as disabled (not threshold 0), and completes while a reader holds the shared lock for the whole merge — which deadlocks under the old exclusive-lock-across-prepare shape. Co-Authored-By: Claude Fable 5 --- .../lance-context-core/src/generic_store.rs | 8 + .../lance-context-core/src/rollout_store.rs | 8 + crates/lance-context-core/src/store_base.rs | 16 + .../src/routes/generic.rs | 21 +- crates/lance-context-server/src/state.rs | 33 +- crates/lance-context-server/src/sweeper.rs | 282 ++++++++++++++---- 6 files changed, 299 insertions(+), 69 deletions(-) diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index 873df0e..6db6aab 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -382,6 +382,14 @@ impl GenericStore { self.base.cleanup_own_shard().await } + /// The shared-lock half of [`Self::maybe_merge_wal`]: `None` unless the + /// count trigger is configured and met. + pub async fn prepare_count_merge( + &self, + ) -> LanceResult> { + self.base.prepare_count_merge().await + } + /// The shared-lock half of [`Self::cleanup_wal`]: seal, then read a /// budgeted prefix of flushed generations into memory. Callers holding a /// read lock run this while appends continue, then take the write lock diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index cf5e620..d6d890d 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -629,6 +629,14 @@ impl RolloutStore { self.base.prepare_merge_if_ready(threshold).await } + /// The shared-lock half of [`Self::maybe_merge_own_shard`]: `None` unless + /// the count trigger is configured and met. + pub async fn prepare_count_merge( + &self, + ) -> LanceResult> { + self.base.prepare_count_merge().await + } + /// [`Self::prepare_merge_if_ready`], but seals the active memtable first — /// the time-triggered behavior of [`Self::cleanup_own_shard`]. pub async fn prepare_cleanup_merge( diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 0e75933..a886506 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -783,6 +783,22 @@ impl StorageBase { self.prepare_merge_if_ready_inner(threshold, false).await } + /// The shared-lock half of the *count*-triggered merge: prepare only if the + /// shard has at least `merge_after_generations` flushed generations pending + /// (`0` disables the trigger, so this returns `None`). Same threshold as + /// [`Self::maybe_merge_own_shard`], but split so the expensive generation + /// read runs under a read lock and only [`Self::commit_prepared_merge`] + /// needs the write lock. + pub async fn prepare_count_merge( + &self, + ) -> LanceResult> { + if self.merge_after_generations == 0 { + return Ok(None); + } + self.prepare_merge_if_ready_inner(self.merge_after_generations, false) + .await + } + /// [`Self::prepare_merge_if_ready`], but seals the active memtable *before* /// consulting the manifest — the time-triggered (`threshold = 1`) behavior /// of [`Self::cleanup_own_shard`]. See that method for why the ordering is diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index 69ccda2..2a30a05 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -271,8 +271,25 @@ pub async fn merge_generic_wal( Path(name): Path, ) -> Result, AppError> { let store = state.get_or_open_generic_store(&name).await?; - let mut guard = store.write().await; - let reclaimed = guard.cleanup_wal().await.map_err(AppError::from_lance)?; + // Same prepare/commit split as the sweeper: the object-storage read of the + // generations runs under the shared lock so the store keeps serving. + let prepared = { + let guard = store.read().await; + guard + .prepare_cleanup_merge() + .await + .map_err(AppError::from_lance)? + }; + let reclaimed = match prepared { + Some((manifest_store, manifest, prepared)) => { + let mut guard = store.write().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .map_err(AppError::from_lance)? + } + None => 0, + }; Ok(Json(serde_json::json!({ "reclaimed": reclaimed }))) } diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index ef6d15f..e5150ff 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -945,8 +945,8 @@ impl AppState { /// deferred seal and genuinely depend on this; datagen seals on each append, /// so its pass is a no-op in steady state and is kept only for symmetry. /// - /// For rollout it also runs the count-triggered merge - /// ([`RolloutStore::maybe_merge_own_shard`]): the read-amplification bound + /// For rollout and generic it also runs the count-triggered merge + /// ([`sweeper::Sweepable::merge_if_due`]): the read-amplification bound /// that formerly lived on the append path now rides this timer. The heavier /// time-based cleanup/merge remains on [`Self::spawn_global_sweeper`]. /// @@ -969,15 +969,26 @@ impl AppState { let Some(state) = weak.upgrade() else { return; }; - sweeper::flush_pass(sweeper::resident(&state.rollout_stores).await, pass_timeout) - .await; - // Datagen seals on every append, so its pass is a no-op in - // steady state; generic stores default to a deferred seal and - // genuinely depend on this. - sweeper::flush_pass(sweeper::resident(&state.datagen_stores).await, pass_timeout) - .await; - sweeper::flush_pass(sweeper::resident(&state.generic_stores).await, pass_timeout) - .await; + // Kinds run concurrently for the same reason the cleanup sweeper + // does: this pass now carries the count-triggered merge, so a + // slow walk over hundreds of rollout stores must not delay + // generic's tick. Datagen seals on every append, so its pass is + // a no-op in steady state; generic stores default to a deferred + // seal and genuinely depend on this. + tokio::join!( + sweeper::flush_pass( + sweeper::resident(&state.rollout_stores).await, + pass_timeout + ), + sweeper::flush_pass( + sweeper::resident(&state.datagen_stores).await, + pass_timeout + ), + sweeper::flush_pass( + sweeper::resident(&state.generic_stores).await, + pass_timeout + ), + ); } })) } diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index b629cfa..c13cc07 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -37,8 +37,16 @@ pub(crate) trait Sweepable: Send + Sync + 'static { /// Seal the active memtable. A no-op for a store that seals on write. fn flush(&self) -> impl std::future::Future> + Send; - /// Fold pending flushed generations into the base table; returns how many - /// were reclaimed. + /// Fold this instance's pending generations into the base table **if** the + /// count trigger is configured and met; returns how many were reclaimed. + /// Rides the flush timer so read amplification is bounded between the + /// slower time-triggered passes. Default: no count trigger. + fn merge_if_due(&self) -> impl std::future::Future> + Send { + async { Ok(0) } + } + + /// Fold **every** pending flushed generation into the base table; returns + /// how many were reclaimed. fn merge_wal(&self) -> impl std::future::Future> + Send; } @@ -50,24 +58,29 @@ impl Sweepable for Arc> { async fn flush(&self) -> Result<(), String> { // Read lock: `flush` is `&self`, so concurrent appends are not blocked. let guard = self.read().await; - let result = guard.flush().await.map_err(|e| e.to_string()); - if result.is_ok() { - // The count-triggered merge rides this timer; it is a no-op unless - // the threshold is configured and met. - drop(guard); - let mut guard = self.write().await; + guard.flush().await.map_err(|e| e.to_string()) + } + + async fn merge_if_due(&self) -> Result { + let prepared = { + let guard = self.read().await; guard - .maybe_merge_own_shard() + .prepare_count_merge() .await - .map_err(|e| e.to_string())?; - } - result + .map_err(|e| e.to_string())? + }; + // Only now take the write lock, and only for the short commit. + let Some((manifest_store, manifest, prepared)) = prepared else { + return Ok(0); + }; + let mut guard = self.write().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .map_err(|e| e.to_string()) } async fn merge_wal(&self) -> Result { - // The prepare/commit split: seal and read every flushed generation - // under the *shared* lock so appends keep running, then take the - // exclusive lock only for the short commit. let prepared = { let guard = self.read().await; guard @@ -75,6 +88,7 @@ impl Sweepable for Arc> { .await .map_err(|e| e.to_string())? }; + // Only now take the write lock, and only for the short commit. let Some((manifest_store, manifest, prepared)) = prepared else { return Ok(0); }; @@ -110,23 +124,34 @@ impl Sweepable for Arc> { } async fn flush(&self) -> Result<(), String> { - // Same shape as rollout: generic stores default to a deferred seal and - // take the same one-row-per-append traffic, so the count-triggered - // merge rides this timer too. Without it a hot generic store depends - // entirely on the slower cleanup sweeper reaching it. let guard = self.read().await; - let result = guard.flush().await.map_err(|e| e.to_string()); - if result.is_ok() { - drop(guard); - let mut guard = self.write().await; - guard.maybe_merge_wal().await.map_err(|e| e.to_string())?; - } - result + guard.flush().await.map_err(|e| e.to_string()) + } + + async fn merge_if_due(&self) -> Result { + // Generic stores default to a deferred seal and take the same + // one-row-per-append traffic as rollout, so the count trigger rides + // this timer too. Without it a hot generic store depends entirely on + // the slower cleanup sweeper reaching it. + let prepared = { + let guard = self.read().await; + guard + .prepare_count_merge() + .await + .map_err(|e| e.to_string())? + }; + // Only now take the write lock, and only for the short commit. + let Some((manifest_store, manifest, prepared)) = prepared else { + return Ok(0); + }; + let mut guard = self.write().await; + guard + .commit_prepared_merge(&manifest_store, &manifest, prepared) + .await + .map_err(|e| e.to_string()) } async fn merge_wal(&self) -> Result { - // Prepare under the shared lock so appends and reads keep running while - // generations are read; hold the exclusive lock only for the commit. let prepared = { let guard = self.read().await; guard @@ -134,6 +159,7 @@ impl Sweepable for Arc> { .await .map_err(|e| e.to_string())? }; + // Only now take the write lock, and only for the short commit. let Some((manifest_store, manifest, prepared)) = prepared else { return Ok(0); }; @@ -170,6 +196,15 @@ pub(crate) async fn flush_pass(stores: Vec<(String, S)>, pass_time Ok(Ok(())) => { metrics::counter!("rollout_wal_flush_total", "result" => "ok", "kind" => kind) .increment(1); + // The count-triggered merge rides this timer, but it is a merge, + // not a flush: its outcome is reported under the cleanup + // counters so a failing merge cannot masquerade as a failing + // flush on the dashboards. + report_merge( + &name, + kind, + tokio::time::timeout(pass_timeout, store.merge_if_due()).await, + ); } Ok(Err(error)) => { metrics::counter!("rollout_wal_flush_total", "result" => "failed", "kind" => kind) @@ -189,34 +224,169 @@ pub(crate) async fn flush_pass(stores: Vec<(String, S)>, pass_time pub(crate) async fn merge_pass(stores: Vec<(String, S)>, pass_timeout: Duration) { let kind = S::kind(); for (name, store) in stores { - match tokio::time::timeout(pass_timeout, store.merge_wal()).await { - Ok(Ok(0)) => {} - Ok(Ok(reclaimed)) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "merged", "kind" => kind) - .increment(1); - metrics::counter!("rollout_wal_generations_reclaimed_total", "kind" => kind) - .increment(reclaimed as u64); - tracing::info!( - store = %name, - kind, - reclaimed, - "sweeper merged flushed generations" - ); - } - Ok(Err(error)) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "failed", "kind" => kind) - .increment(1); - tracing::warn!(store = %name, kind, %error, "sweeper WAL cleanup failed"); - } - Err(_elapsed) => { - metrics::counter!("rollout_wal_cleanup_total", "result" => "timeout", "kind" => kind) - .increment(1); - tracing::warn!( - store = %name, - kind, - "sweeper WAL cleanup timed out; abandoning this store this tick" - ); - } + report_merge( + &name, + kind, + tokio::time::timeout(pass_timeout, store.merge_wal()).await, + ); + } +} + +/// Record one merge attempt's outcome on the cleanup counters and log. +fn report_merge( + name: &str, + kind: &'static str, + outcome: Result, tokio::time::error::Elapsed>, +) { + match outcome { + Ok(Ok(0)) => {} + Ok(Ok(reclaimed)) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "merged", "kind" => kind) + .increment(1); + metrics::counter!("rollout_wal_generations_reclaimed_total", "kind" => kind) + .increment(reclaimed as u64); + tracing::info!(store = %name, kind, reclaimed, "sweeper merged flushed generations"); + } + Ok(Err(error)) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "failed", "kind" => kind) + .increment(1); + tracing::warn!(store = %name, kind, %error, "sweeper WAL cleanup failed"); + } + Err(_elapsed) => { + metrics::counter!("rollout_wal_cleanup_total", "result" => "timeout", "kind" => kind) + .increment(1); + tracing::warn!( + store = %name, + kind, + "sweeper WAL cleanup timed out; abandoning this store this tick" + ); } } } + +#[cfg(test)] +mod tests { + use super::*; + use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; + use lance_context_core::GenericStoreOptions; + use serde_json::json; + use tempfile::TempDir; + + fn spec() -> SchemaSpec { + SchemaSpec::new(vec![( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + )]) + } + + async fn generic_with_pending( + dir: &TempDir, + merge_after_generations: usize, + pending: usize, + ) -> Arc> { + let uri = dir.path().to_string_lossy().to_string(); + let store = GenericStore::open( + &uri, + spec(), + GenericStoreOptions { + merge_after_generations: Some(merge_after_generations), + seal_on_add: true, + ..Default::default() + }, + ) + .await + .unwrap(); + for i in 0..pending { + store + .add(&[json!({"id": format!("r{i}")}).as_object().unwrap().clone()]) + .await + .unwrap(); + } + assert_eq!(store.pending_wal_generations().await.unwrap(), pending); + Arc::new(RwLock::new(store)) + } + + /// The count trigger is what bounds read amplification between the slower + /// time-triggered passes; a generic store at the threshold must merge on a + /// flush tick. This is the path #256 added and #257 fixed the locking of. + #[tokio::test] + async fn generic_merge_if_due_folds_pending_generations_at_threshold() { + let dir = TempDir::new().unwrap(); + let store = generic_with_pending(&dir, 3, 3).await; + + let reclaimed = store.merge_if_due().await.unwrap(); + + assert_eq!(reclaimed, 3); + assert_eq!( + store.read().await.pending_wal_generations().await.unwrap(), + 0 + ); + assert_eq!(store.read().await.count_base_rows().await.unwrap(), 3); + } + + #[tokio::test] + async fn generic_merge_if_due_is_a_noop_below_threshold() { + let dir = TempDir::new().unwrap(); + let store = generic_with_pending(&dir, 5, 2).await; + + assert_eq!(store.merge_if_due().await.unwrap(), 0); + assert_eq!( + store.read().await.pending_wal_generations().await.unwrap(), + 2 + ); + } + + /// `merge_after_generations = 0` means "count trigger off". It must not be + /// read as "threshold 0, merge every tick" — which is what the underlying + /// `threshold.max(1)` would do if the zero were passed straight through. + #[tokio::test] + async fn generic_merge_if_due_respects_disabled_count_trigger() { + let dir = TempDir::new().unwrap(); + let store = generic_with_pending(&dir, 0, 4).await; + + assert_eq!(store.merge_if_due().await.unwrap(), 0); + assert_eq!( + store.read().await.pending_wal_generations().await.unwrap(), + 4 + ); + // The time trigger is unaffected and still drains everything. + assert_eq!(store.merge_wal().await.unwrap(), 4); + assert_eq!( + store.read().await.pending_wal_generations().await.unwrap(), + 0 + ); + } + + /// The expensive half of a merge must not hold the exclusive lock: a + /// concurrent reader takes the shared lock while `prepare` is in flight. + /// If `merge_if_due` held the write lock across the read, this would + /// deadlock (the reader waits on the merge, the merge holds the lock). + #[tokio::test] + async fn generic_merge_if_due_does_not_block_readers_during_prepare() { + let dir = TempDir::new().unwrap(); + let store = generic_with_pending(&dir, 3, 3).await; + + // Hold a read lock for the whole merge. With the prepare/commit split, + // prepare proceeds under a second shared lock and commit waits only for + // this guard to drop; without the split the merge would need the write + // lock up front and this test would hang. + let reader = store.read().await; + let merge = tokio::spawn({ + let store = store.clone(); + async move { store.merge_if_due().await } + }); + // Give the merge time to reach the commit phase (it needs the write + // lock there), proving prepare completed under the shared lock. + tokio::time::sleep(Duration::from_millis(200)).await; + assert!(!merge.is_finished(), "commit must wait for the live reader"); + assert_eq!(reader.pending_wal_generations().await.unwrap(), 3); + drop(reader); + + let reclaimed = tokio::time::timeout(Duration::from_secs(30), merge) + .await + .expect("merge must finish once the reader releases the lock") + .unwrap() + .unwrap(); + assert_eq!(reclaimed, 3); + } +}