fix: avoid same-store merge read-lock and memory-budget deadlock

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<tokio::sync::Mutex<()>> {
+        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(

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<tokio::sync::Mutex<()>> {
+        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(

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<Arc<MergeMemoryBudget>>,
+    /// Serializes split preparation/commit on this handle, before taking a store lock.
+    merge_preparation_gate: Arc<tokio::sync::Mutex<()>>,
     /// Timestamp of the last successful [`Self::compact`] on this handle.
     last_compaction: Option<DateTime<Utc>>,
     /// 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<tokio::sync::Mutex<()>> {
+        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?;

crates/lance-context-server/src/routes/generic.rs

@@ -333,6 +333,8 @@ pub(crate) async fn merge_generic_wal_owned(
     Path(name): Path<String>,
 ) -> Result<Json<serde_json::Value>, 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 = {

crates/lance-context-server/src/routes/rollouts.rs

@@ -747,6 +747,8 @@ pub(crate) async fn merge_wal_owned(
     Path(name): Path<String>,
 ) -> Result<Json<MergeWalResponse>, 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

crates/lance-context-server/src/sweeper.rs

@@ -94,6 +94,8 @@ impl Sweepable for Arc<RwLock<RolloutStore>> {
     }
 
     async fn merge_if_due(&self) -> Result<usize, String> {
+        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<RwLock<RolloutStore>> {
     }
 
     async fn merge_wal(&self) -> Result<usize, String> {
+        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<RwLock<GenericStore>> {
     }
 
     async fn merge_if_due(&self) -> Result<usize, String> {
+        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<RwLock<GenericStore>> {
     }
 
     async fn merge_wal(&self) -> Result<usize, String> {
+        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,