From 8acbc2ced05e13e35205fc5ed3cde80e63e258d6 Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Wed, 21 Oct 2026 11:41:05 +0000 Subject: [PATCH] fix: admit same-store merges before acquiring the read lock --- .../lance-context-core/src/generic_store.rs | 8 ++ .../lance-context-core/src/rollout_store.rs | 8 ++ crates/lance-context-core/src/store_base.rs | 13 +++ .../src/routes/generic.rs | 2 + .../src/routes/rollouts.rs | 2 + crates/lance-context-server/src/sweeper.rs | 99 +++++++++++++++++++ 6 files changed, 132 insertions(+) diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index b966499..27602d4 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -663,6 +663,14 @@ impl GenericStore { self.base.cleanup_own_shard().await } + /// Admission for a split merge on this resident store. Clone under a short + /// store read lock, release that lock, then await the gate; retain the guard + /// until commit finishes. Reads, ingestion and other tables do not take it. + #[must_use] + pub fn merge_preparation_gate(&self) -> Arc> { + self.base.merge_preparation_gate() + } + /// The shared-lock half of [`Self::maybe_merge_wal`]: `None` unless the /// count trigger is configured and met. pub async fn prepare_count_merge( diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index ee126bc..9c94bdf 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -667,6 +667,14 @@ impl RolloutStore { self.base.prepare_merge_if_ready(threshold).await } + /// Admission for a split merge on this resident store. Clone under a short + /// store read lock, release that lock, then await the gate; retain the guard + /// until commit finishes. Reads, ingestion and other tables do not take it. + #[must_use] + pub fn merge_preparation_gate(&self) -> Arc> { + self.base.merge_preparation_gate() + } + /// 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( diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 1bf4de0..36a6480 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -403,6 +403,8 @@ pub(crate) struct StorageBase { pending_generations_max: usize, /// Process-wide merge byte budget; `None` means unbounded. merge_budget: Option>, + /// Serializes split preparation/commit on this handle, before taking a store lock. + merge_preparation_gate: Arc>, /// Timestamp of the last successful [`Self::compact`] on this handle. last_compaction: Option>, /// Number of successful compactions performed by this handle. @@ -556,6 +558,7 @@ impl StorageBase { pending_generations_max: pending_generations_max .unwrap_or(DEFAULT_PENDING_GENERATIONS_MAX), merge_budget, + merge_preparation_gate: Arc::new(tokio::sync::Mutex::new(())), last_compaction: None, total_compactions: 0, last_compaction_error: None, @@ -940,6 +943,14 @@ impl StorageBase { Ok(if committed { pending } else { 0 }) } + /// Clone this gate under a short read lock, then drop that read guard before + /// awaiting the gate. Hold the permit through preparation and commit. + /// Otherwise a second preparation can hold a read lock while waiting for + /// memory owned by the first preparation, which needs the write lock to commit. + pub fn merge_preparation_gate(&self) -> Arc> { + self.merge_preparation_gate.clone() + } + /// The shared-lock half of a merge: decide whether one is due and read the /// flushed generations into memory. /// @@ -949,6 +960,8 @@ impl StorageBase { /// due. /// /// ```ignore + /// let merge_gate = { store.read().await.merge_preparation_gate() }; + /// let _merge = merge_gate.lock().await; /// let prepared = { store.read().await.prepare_merge_if_ready(1).await? }; /// if let Some((manifest_store, manifest, prepared)) = prepared { /// store.write().await.commit_prepared_merge(&manifest_store, &manifest, prepared).await?; diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index eaece28..b5b64e2 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -333,6 +333,8 @@ pub(crate) async fn merge_generic_wal_owned( Path(name): Path, ) -> Result, AppError> { let store = state.get_or_open_generic_store(&name).await?; + let merge_gate = { store.read().await.merge_preparation_gate() }; + let _merge = merge_gate.lock().await; // 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 = { diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 6e01e72..3a2dcf4 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -747,6 +747,8 @@ pub(crate) async fn merge_wal_owned( Path(name): Path, ) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; + let merge_gate = { store_lock.read().await.merge_preparation_gate() }; + let _merge = merge_gate.lock().await; // 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 diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index 3ca5690..aa7ec10 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -94,6 +94,8 @@ impl Sweepable for Arc> { } async fn merge_if_due(&self) -> Result { + let merge_gate = { self.read().await.merge_preparation_gate() }; + let _merge = merge_gate.lock().await; let prepared = { let guard = self.read().await; guard @@ -113,6 +115,8 @@ impl Sweepable for Arc> { } async fn merge_wal(&self) -> Result { + let merge_gate = { self.read().await.merge_preparation_gate() }; + let _merge = merge_gate.lock().await; let prepared = { let guard = self.read().await; guard @@ -179,6 +183,8 @@ impl Sweepable for Arc> { } async fn merge_if_due(&self) -> Result { + let merge_gate = { self.read().await.merge_preparation_gate() }; + let _merge = merge_gate.lock().await; // 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 @@ -202,6 +208,8 @@ impl Sweepable for Arc> { } async fn merge_wal(&self) -> Result { + let merge_gate = { self.read().await.merge_preparation_gate() }; + let _merge = merge_gate.lock().await; let prepared = { let guard = self.read().await; guard @@ -860,6 +868,97 @@ mod tests { assert_eq!(report["snapshot"]["takeover_authorized"], false); } + /// A prepared merge retains its memory reservation until the write-lock + /// commit. A competing prepare must not hold a read lock while waiting + /// for those same bytes, or neither merge can finish. + #[tokio::test] + async fn concurrent_same_store_merges_do_not_pin_budget_behind_read_lock() { + let dir = TempDir::new().unwrap(); + let budget = lance_context_core::merge_budget::MergeMemoryBudget::new(1024 * 1024); + let store = GenericStore::open( + &dir.path().to_string_lossy(), + spec(), + GenericStoreOptions { + merge_after_generations: Some(1), + merge_max_bytes: Some(1024 * 1024), + merge_budget: Some(budget.clone()), + seal_on_add: true, + ..Default::default() + }, + ) + .await + .unwrap(); + for i in 0..3 { + store + .add(&[json!({"id": format!("r{i}")}).as_object().unwrap().clone()]) + .await + .unwrap(); + } + let store = Arc::new(RwLock::new(store)); + let (count, cleanup) = tokio::time::timeout(Duration::from_secs(20), async { + tokio::join!(store.merge_if_due(), store.merge_wal()) + }) + .await + .expect("same-store preparation must not hold a read lock waiting for the committing merge's budget"); + assert_eq!(count.unwrap() + cleanup.unwrap(), 3); + assert_eq!( + store.read().await.pending_wal_generations().await.unwrap(), + 0 + ); + assert_eq!(store.read().await.count_base_rows().await.unwrap(), 3); + assert_eq!(budget.reserved(), 0); + } + + #[tokio::test] + async fn merge_admission_wait_releases_store_lock_and_other_tables_keep_merging() { + let first_dir = TempDir::new().unwrap(); + let second_dir = TempDir::new().unwrap(); + let first = generic_with_pending(&first_dir, 3, 3).await; + let second = generic_with_pending(&second_dir, 3, 3).await; + let gate = { first.read().await.merge_preparation_gate() }; + let held = gate.lock().await; + let waiting = first.merge_if_due(); + tokio::pin!(waiting); + assert!(futures::poll!(&mut waiting).is_pending()); + // Admission must await the gate outside even a shared store guard. + assert!(first.try_write().is_ok()); + tokio::time::timeout(Duration::from_secs(30), async { + first + .read() + .await + .add(&[json!({"id": "concurrent-ingest"}) + .as_object() + .unwrap() + .clone()]) + .await + .unwrap(); + assert_eq!( + first.read().await.pending_wal_generations().await.unwrap(), + 4 + ); + assert_eq!(second.merge_if_due().await.unwrap(), 3); + assert_eq!( + second.read().await.pending_wal_generations().await.unwrap(), + 0 + ); + }) + .await + .expect("reads, ingestion and another table must progress while one merge waits"); + drop(held); + assert_eq!( + tokio::time::timeout(Duration::from_secs(30), waiting) + .await + .unwrap() + .unwrap(), + 4 + ); + assert_eq!( + first.read().await.pending_wal_generations().await.unwrap(), + 0 + ); + assert_eq!(first.read().await.count_base_rows().await.unwrap(), 4); + } + async fn generic_with_pending( dir: &TempDir, merge_after_generations: usize,