Reserve resident WAL task capacity beyond legacy merge RPCs

crates/lance-context-master/src/rollout_append.rs

@@ -27,6 +27,14 @@ pub struct AppendConfig {
     /// Uses the ordinary durable task queue; creates no Kubernetes Jobs.
     #[arg(long, env = "ROLLOUT_APPEND_LOCAL", default_value_t = false)]
     pub rollout_append_local: bool,
+    /// Additional task slots reserved for owned local append merges. Legacy RPCs
+    /// cannot claim these slots. All local tasks share the existing Arrow budget.
+    #[arg(
+        long,
+        env = "ROLLOUT_APPEND_LOCAL_TASK_CONCURRENCY",
+        default_value_t = 1
+    )]
+    pub rollout_append_local_task_concurrency: usize,
     /// Shared Arrow-buffer budget across every local append task in this master.
     #[arg(
         long,
@@ -51,6 +59,7 @@ impl Default for AppendConfig {
             rollout_append_max_generations: 64,
             rollout_append_max_bytes: 67_108_864,
             rollout_append_local: false,
+            rollout_append_local_task_concurrency: 1,
             rollout_append_local_memory_bytes: 2_147_483_648,
             rollout_append_reconcile_interval_secs: 30,
         }

crates/lance-context-master/src/scheduler.rs

@@ -1120,6 +1120,21 @@ pub fn spawn_scheduler(state: &Arc<MasterState>) -> JoinHandle<()> {
     if let Some(merge) = merge_sem {
         spawn_pool_poller(state.clone(), merge, TaskKinds::MERGE_WAL, false);
     }
+    if state.config.append.rollout_append_local
+        && state.config.append.rollout_append_local_task_concurrency > 0
+    {
+        // Reserve before claiming. A saturated legacy RPC pool cannot consume
+        // this capacity; buffers still use the process-wide local Arrow budget.
+        spawn_filtered_pool_poller(
+            state.clone(),
+            Arc::new(Semaphore::new(
+                state.config.append.rollout_append_local_task_concurrency,
+            )),
+            TaskKinds::MERGE_WAL,
+            false,
+            true,
+        );
+    }
     general
 }
 
@@ -1133,6 +1148,16 @@ fn spawn_pool_poller(
     pool: Arc<Semaphore>,
     kinds: TaskKinds,
     report_depth: bool,
+) -> tokio::task::JoinHandle<()> {
+    spawn_filtered_pool_poller(state, pool, kinds, report_depth, false)
+}
+
+fn spawn_filtered_pool_poller(
+    state: Arc<MasterState>,
+    pool: Arc<Semaphore>,
+    kinds: TaskKinds,
+    report_depth: bool,
+    resident_only: bool,
 ) -> tokio::task::JoinHandle<()> {
     tokio::spawn(async move {
         loop {
@@ -1158,7 +1183,15 @@ fn spawn_pool_poller(
                     break;
                 };
                 let claim_start = std::time::Instant::now();
-                match state.task_store.claim_next_of_kinds(kinds).await {
+                let claim = if resident_only {
+                    state
+                        .task_store
+                        .claim_resident_merge(&state.config.append.rollout_append_targets)
+                        .await
+                } else {
+                    state.task_store.claim_next_of_kinds(kinds).await
+                };
+                match claim {
                     Ok(Some(claim)) => {
                         let claim_elapsed = claim_start.elapsed();
                         // Local capacity was reserved before the claim. No
@@ -2409,6 +2442,116 @@ mod tests {
         server.abort();
     }
 
+    #[tokio::test]
+    #[ignore = "requires ETCD_TEST_ENDPOINTS"]
+    async fn reserved_resident_merge_finishes_while_legacy_pool_is_saturated() {
+        use axum::{routing::post, Json, Router};
+        use lance_context_core::{RolloutStore, RolloutStoreOptions};
+        let started = Arc::new(tokio::sync::Notify::new());
+        let release = Arc::new(Semaphore::new(0));
+        let app = Router::new().route(
+            "/api/v1/internal/merge-wal/{name}",
+            post({
+                let started = started.clone();
+                let release = release.clone();
+                move || {
+                    let started = started.clone();
+                    let release = release.clone();
+                    async move {
+                        started.notify_one();
+                        release.acquire().await.unwrap().forget();
+                        Json(serde_json::json!({"reclaimed": 1}))
+                    }
+                }
+            }),
+        );
+        let dir = TempDir::new().unwrap();
+        let mut cfg = config(&dir);
+        cfg.merge_wal_concurrency = 1;
+        cfg.append.rollout_append_local = true;
+        cfg.append.rollout_append_targets = vec!["exp".into()];
+        cfg.catchup.shards = vec!["writer".into()];
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        cfg.worker_endpoints = vec![format!("http://{}", listener.local_addr().unwrap())];
+        let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
+        let state = MasterState::new(cfg).await.unwrap();
+        let writer = RolloutStore::open_with_options(
+            &state.rollout_uri("exp"),
+            RolloutStoreOptions {
+                shard_id: Some("writer".into()),
+                merge_after_generations: Some(0),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+        let dto = serde_json::from_value(serde_json::json!({
+            "id":"one", "rollout_id":"r", "content":"test"
+        }))
+        .unwrap();
+        writer
+            .add(&[lance_context_core::rollout_record_from_add_request(&dto)])
+            .await
+            .unwrap();
+        writer.flush().await.unwrap();
+        let legacy = enqueue(&state, TaskKind::MergeWal, "legacy-blocked")
+            .await
+            .unwrap();
+        let scheduler = spawn_scheduler(&state);
+        tokio::time::timeout(Duration::from_secs(10), started.notified())
+            .await
+            .unwrap();
+        let queued = enqueue(&state, TaskKind::MergeWal, "legacy-queued")
+            .await
+            .unwrap();
+        let resident = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap();
+        let result = await_terminal(&state, &resident.id).await;
+        assert_eq!(result.state, TaskState::Done, "{:?}", result.error);
+        assert_eq!(writer.pending_wal_generations().await.unwrap(), 0);
+        let reader = lance::Dataset::open(&state.rollout_uri("exp"))
+            .await
+            .unwrap();
+        assert_eq!(reader.count_rows(None).await.unwrap(), 1);
+        assert_eq!(
+            state
+                .task_store
+                .get(&legacy.id)
+                .await
+                .unwrap()
+                .unwrap()
+                .state,
+            TaskState::Running
+        );
+        assert_eq!(
+            state
+                .task_store
+                .get(&queued.id)
+                .await
+                .unwrap()
+                .unwrap()
+                .state,
+            TaskState::Queued
+        );
+        // The reserved poller obeys the same joined admission drain.
+        state.admission.begin_drain();
+        release.add_permits(1);
+        tokio::time::timeout(Duration::from_secs(10), state.admission.wait_drained())
+            .await
+            .unwrap();
+        assert_eq!(
+            state
+                .task_store
+                .get(&queued.id)
+                .await
+                .unwrap()
+                .unwrap()
+                .state,
+            TaskState::Queued
+        );
+        scheduler.abort();
+        server.abort();
+    }
+
     /// MergeWal fans out to every configured worker endpoint and sums the
     /// reclaimed counts. Uses a tiny in-process stub server per "worker".
     #[tokio::test]

crates/lance-context-master/src/task_store.rs

@@ -460,7 +460,26 @@ impl TaskStore {
     /// saturated one.
     pub async fn claim_next_of_kinds(&self, kinds: TaskKinds) -> lance::Result<Option<TaskClaim>> {
         self.inner.recover_orphaned().await?;
-        let (claim, dependency_failed) = self.inner.claim_next(kinds, None).await?;
+        let (claim, dependency_failed) = self.inner.claim_next(kinds, None, None).await?;
+        if dependency_failed {
+            if let Err(error) = self.prune_terminal_history().await {
+                tracing::warn!(error = %error, "failed to prune task history");
+            }
+        }
+        Ok(claim)
+    }
+
+    /// Reserved resident capacity filters before any dependency/ownership RPC.
+    /// Keep canonical claims and all external-owner/fencing checks unchanged.
+    pub(crate) async fn claim_resident_merge(
+        &self,
+        targets: &[String],
+    ) -> lance::Result<Option<TaskClaim>> {
+        self.inner.recover_orphaned().await?;
+        let (claim, dependency_failed) = self
+            .inner
+            .claim_next(TaskKinds::MERGE_WAL, None, Some(targets))
+            .await?;
         if dependency_failed {
             if let Err(error) = self.prune_terminal_history().await {
                 tracing::warn!(error = %error, "failed to prune task history");
@@ -485,7 +504,7 @@ impl TaskStore {
         }
         Ok(self
             .inner
-            .claim_next(TaskKinds::MERGE_WAL, Some((&id, job)))
+            .claim_next(TaskKinds::MERGE_WAL, Some((&id, job)), None)
             .await?
             .0)
     }
@@ -790,6 +809,7 @@ impl EtcdTaskStore {
         &self,
         kinds: TaskKinds,
         only_id: Option<(&str, &str)>,
+        resident_targets: Option<&[String]>,
     ) -> lance::Result<(Option<TaskClaim>, bool)> {
         let prefix = only_id.map_or_else(|| self.queue_prefix(), |(id, _)| self.queue_key(id));
         let range_end = if only_id.is_some() {
@@ -834,6 +854,15 @@ impl EtcdTaskStore {
                 if !kinds.contains(task.kind) {
                     continue;
                 }
+                if resident_targets.is_some_and(|targets| {
+                    task.kind != TaskKind::MergeWal
+                        || task.target.starts_with("generic:")
+                        || !self.rollout.owned(&task.target)
+                        || self.rollout.draining(&task.target)
+                        || !targets.iter().any(|t| t == "*" || t == &task.target)
+                }) {
+                    continue;
+                }
                 match self.dependency_status(&task).await? {
                     DependencyStatus::Ready => {}
                     DependencyStatus::Waiting => continue,
@@ -2497,6 +2526,84 @@ mod tests {
         (dir, cfg, store)
     }
 
+    #[tokio::test]
+    #[ignore = "requires ETCD_TEST_ENDPOINTS"]
+    async fn resident_claim_filter_preserves_legacy_draining_and_external_owners() {
+        let (_dir, mut cfg, _) = preparation_test_store().await;
+        cfg.merge_rollout.owned_targets = ["hot", "draining", "excluded", "external", "generic:gs"]
+            .map(str::to_owned)
+            .to_vec();
+        cfg.merge_rollout.drain_targets = vec!["draining".into()];
+        let store = TaskStore::open(&cfg).await.unwrap();
+        let mut queued = Vec::new();
+        for target in [
+            "legacy",
+            "draining",
+            "excluded",
+            "external",
+            "generic:gs",
+            "hot",
+        ] {
+            queued.push(
+                store
+                    .enqueue(TaskKind::MergeWal, target, vec![])
+                    .await
+                    .unwrap(),
+            );
+        }
+        let key = crate::catchup::store::active_key(&cfg.etcd.etcd_prefix, "external");
+        store
+            .inner
+            .client
+            .clone()
+            .put(key.clone(), "external-owner", None)
+            .await
+            .unwrap();
+        let targets = ["hot", "draining", "external", "generic:gs"].map(str::to_owned);
+        let claim = store.claim_resident_merge(&targets).await.unwrap().unwrap();
+        assert_eq!(claim.task.target, "hot");
+        // A second replica cannot take the running canonical task.
+        let other = TaskStore::open(&cfg).await.unwrap();
+        assert!(other
+            .claim_resident_merge(&targets)
+            .await
+            .unwrap()
+            .is_none());
+        for task in &queued[..5] {
+            assert_eq!(
+                store.get(&task.id).await.unwrap().unwrap().state,
+                TaskState::Queued
+            );
+        }
+        assert_eq!(
+            store.inner.get_text(&key).await.unwrap().as_deref(),
+            Some("external-owner")
+        );
+        store.finish(claim, Ok("done".into())).await.unwrap();
+        let wildcard = store
+            .claim_resident_merge(&["*".into()])
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(wildcard.task.target, "excluded");
+        store.finish(wildcard, Ok("done".into())).await.unwrap();
+        assert!(store
+            .claim_resident_merge(&["*".into()])
+            .await
+            .unwrap()
+            .is_none());
+        store
+            .inner
+            .client
+            .clone()
+            .delete(
+                cfg.etcd.etcd_prefix,
+                Some(etcd_client::DeleteOptions::new().with_prefix()),
+            )
+            .await
+            .unwrap();
+    }
+
     #[tokio::test]
     #[ignore = "requires ETCD_TEST_ENDPOINTS"]
     async fn index_preparation_allows_merge_and_shares_compact_commit_turn() {

docs/resident-wal-recovery.md

@@ -50,3 +50,14 @@ This removes external refill for already-owned resident tables. It does not
 certify cancellation of legacy unfenced worker handlers, fix serial worker
 flush scheduling, or replay an unsealed historical raw WAL tail. Those cases
 must retain their existing recovery protection until separately validated.
+
+Resident masters also reserve `ROLLOUT_APPEND_LOCAL_TASK_CONCURRENCY` task slots
+(default 1, 0 disables the reservation) in addition to `MERGE_WAL_CONCURRENCY`.
+Only locally enabled, owned, non-draining rollout merges can use these slots.
+Eligibility is checked before dependency reads or claim/lock mutation. The
+ordinary merge pool remains available to all merge tasks, including legacy
+worker RPCs; a legacy request cannot consume the reserved capacity. Local
+Arrow buffers still share `ROLLOUT_APPEND_LOCAL_MEMORY_BYTES` across both pools.
+Canonical claims, external catch-up ownership, progress watchdogs, and process
+admission drain apply equally to the reserved pool. This does not cancel an
+existing legacy owner of the same table.