feat: let masters provision bounded WAL catch-up jobs

crates/lance-context-core/src/merge_write_scope.rs

@@ -141,11 +141,22 @@ impl Drop for Active {
     }
 }
 
+#[cfg(test)]
 pub(crate) async fn shield<F, T, E>(write: F) -> std::result::Result<T, E>
 where
     F: Future<Output = std::result::Result<T, E>> + Send + 'static,
     T: Send + 'static,
     E: From<Error> + Send + 'static,
+{
+    shield_observed(write, |result| result.is_ok()).await
+}
+
+async fn shield_observed<F, T, E, P>(write: F, progressed: P) -> std::result::Result<T, E>
+where
+    F: Future<Output = std::result::Result<T, E>> + Send + 'static,
+    T: Send + 'static,
+    E: From<Error> + Send + 'static,
+    P: FnOnce(&std::result::Result<T, E>) -> bool + Send + 'static,
 {
     let Ok(scope) = CURRENT.try_with(Arc::clone) else {
         return write.await;
@@ -166,7 +177,7 @@ where
     let leaf = tokio::spawn(async move {
         let mut active = active;
         let result = write.await;
-        if result.is_ok() {
+        if progressed(&result) {
             active.scope.completed_steps.fetch_add(1, Ordering::Relaxed);
         }
         if result.is_err() {
@@ -209,18 +220,21 @@ pub(crate) async fn drain_generations(
         authorize(&resource, next.version).await?;
         let writer = store.clone();
         let scope = CURRENT.try_with(Arc::clone).ok();
-        let (next, written) = shield(async move {
-            let written = writer.write(&next).await;
-            if written
-                .as_ref()
-                .is_err_and(|e| !e.to_string().contains("already exists"))
-            {
-                if let Some(scope) = scope {
-                    scope.mark_uncertain();
+        let (next, written) = shield_observed(
+            async move {
+                let written = writer.write(&next).await;
+                if written
+                    .as_ref()
+                    .is_err_and(|e| !e.to_string().contains("already exists"))
+                {
+                    if let Some(scope) = scope {
+                        scope.mark_uncertain();
+                    }
                 }
-            }
-            Ok::<_, Error>((next, written))
-        })
+                Ok::<_, Error>((next, written))
+            },
+            |result| result.as_ref().is_ok_and(|(_, written)| written.is_ok()),
+        )
         .await?;
         match written {
             Ok(_) => return Ok(next),
@@ -384,28 +398,31 @@ impl CommitHandler for GuardedCommit {
             let path = base_path.clone();
             let store = object_store.clone();
             let scope = CURRENT.try_with(Arc::clone).ok();
-            let (location, committed) = shield(async move {
-                let location = inner
-                    .commit(
-                        &mut owned_manifest,
-                        indices,
-                        &path,
-                        &store,
-                        manifest_writer,
-                        naming_scheme,
-                        transaction,
-                    )
-                    .await;
-                // Only a definite conflict proves this conditional commit did
-                // not happen. A transport/storage error can arrive after the
-                // remote service accepted it, even though the Rust future ended.
-                if matches!(&location, Err(CommitError::OtherError(_))) {
-                    if let Some(scope) = scope {
-                        scope.mark_uncertain();
+            let (location, committed) = shield_observed(
+                async move {
+                    let location = inner
+                        .commit(
+                            &mut owned_manifest,
+                            indices,
+                            &path,
+                            &store,
+                            manifest_writer,
+                            naming_scheme,
+                            transaction,
+                        )
+                        .await;
+                    // Only a definite conflict proves this conditional commit did
+                    // not happen. A transport/storage error can arrive after the
+                    // remote service accepted it, even though the Rust future ended.
+                    if matches!(&location, Err(CommitError::OtherError(_))) {
+                        if let Some(scope) = scope {
+                            scope.mark_uncertain();
+                        }
                     }
-                }
-                Ok::<_, CommitError>((location, owned_manifest))
-            })
+                    Ok::<_, CommitError>((location, owned_manifest))
+                },
+                |result| result.as_ref().is_ok_and(|(location, _)| location.is_ok()),
+            )
             .await?;
             *manifest = committed;
             location
@@ -503,6 +520,21 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn conditional_conflicts_and_errors_do_not_count_as_progress() {
+        let scope = MergeWriteScope::new();
+        for committed in [false, false, true] {
+            scope
+                .run(shield_observed(
+                    async move { Ok::<_, Error>(committed) },
+                    |result| matches!(result, Ok(true)),
+                ))
+                .await
+                .unwrap();
+        }
+        assert_eq!(scope.completed_steps(), 1);
+    }
+
     #[tokio::test]
     async fn confirmed_fence_can_abort_a_permanently_stalled_leaf() {
         struct Writing(Arc<AtomicBool>);

crates/lance-context-core/src/store_base.rs

@@ -1209,7 +1209,6 @@ impl StorageBase {
             let mut current_batches = Vec::new();
             let mut stream = gen_dataset.scan().try_into_stream().await?;
             while let Some(batch) = stream.try_next().await? {
-                crate::merge_write_scope::checkpoint();
                 if batch.num_rows() > 0 {
                     let batch = align_batch_to_schema(batch, merge_schema.clone())?;
                     buffered_bytes = buffered_bytes.saturating_add(batch.get_array_memory_size());
@@ -1229,6 +1228,7 @@ impl StorageBase {
                         }
                     }
                     current_batches.push(batch);
+                    crate::merge_write_scope::checkpoint();
                 }
             }
             generation_batches.push((flushed.generation, current_batches));

crates/lance-context-master/src/catchup/executor.rs

@@ -0,0 +1,162 @@
+use super::Result;
+use crate::{config::MasterConfig, maintenance_execution, state::MasterState};
+use lance_context_api::TaskKind;
+use lance_context_core::{
+    GenericStore, GenericStoreOptions, MergeMemoryBudget, RolloutStore, RolloutStoreOptions,
+};
+use lance_context_merge::MaintenanceKind;
+use std::{sync::Arc, time::Duration};
+
+/// One dedicated process, one table, serial shards, bounded passes. No server,
+/// scanner, ingestion writer, or fleet task consumer starts in this mode.
+pub async fn execute(mut config: MasterConfig, target: &str) -> Result<()> {
+    config.catchup.enabled = false;
+    config.catchup.validate()?;
+    if !config.merge_rollout.owned(target) || config.merge_rollout.draining(target) {
+        return Err("catch-up executor requires a non-draining owned target".into());
+    }
+    let job = config
+        .catchup
+        .job_name
+        .clone()
+        .ok_or("missing catch-up Job identity")?;
+    let state = MasterState::new(config).await.map_err(|e| e.to_string())?;
+    // Enqueue through the normal durable queue, then claim only this table.
+    state
+        .task_store
+        .enqueue(TaskKind::MergeWal, target, Vec::new())
+        .await
+        .map_err(|e| e.to_string())?;
+    let claim = tokio::time::timeout(Duration::from_secs(30), async {
+        loop {
+            if let Some(claim) = state
+                .task_store
+                .claim_merge_target(target, &job)
+                .await
+                .map_err(|e| e.to_string())?
+            {
+                return Ok::<_, String>(claim);
+            }
+            tokio::time::sleep(Duration::from_secs(1)).await;
+        }
+    })
+    .await
+    .map_err(|_| "catch-up admission remained busy for 30 seconds")??;
+    let result = async {
+        // The claim CAS has excluded any live reconciler. A previous process
+        // may still have a storage PUT in flight: fence before opening writers.
+        let coordinator = state.task_store.merge_coordinator();
+        if let Some(old) = coordinator.get(target).await? {
+            crate::merge_execution::ensure_recovery_due(&coordinator, &old).await?;
+            crate::merge_execution::recover_execution(
+                &state,
+                &coordinator,
+                &state.task_store.merge_claim(&claim),
+                old,
+            )
+            .await?;
+        }
+        maintenance_execution::run_as(
+            &state,
+            &claim,
+            MaintenanceKind::Catchup,
+            merge_passes(&state, target),
+        )
+        .await
+    }
+    .await;
+    let error = result.as_ref().err().cloned();
+    state
+        .task_store
+        .finish(claim, result)
+        .await
+        .map_err(|e| e.to_string())?;
+    if let Some(error) = error {
+        return Err(error);
+    }
+    Ok(())
+}
+
+async fn merge_passes(state: &Arc<MasterState>, target: &str) -> Result<String> {
+    let config = &state.config.catchup;
+    let session = RolloutStore::build_session(96 * 1024 * 1024, 32 * 1024 * 1024);
+    let budget = MergeMemoryBudget::new(config.merge_memory_bytes);
+    let mut total = 0usize;
+    let started = tokio::time::Instant::now();
+    // This is only an admission time slice. A progressing operation already
+    // admitted may run beyond it; cancellation is based on lack of progress.
+    let admit_for = Duration::from_secs(config.slice_secs);
+    // Bound each slice even under continuous ingestion; later fresh stats may
+    // request another Job. Visit every shard before repeating any hot shard.
+    for _ in 0..16 {
+        let before = total;
+        for shard in &config.shards {
+            if started.elapsed() >= admit_for {
+                return Ok(format!(
+                    "dedicated catch-up slice merged {total} generations"
+                ));
+            }
+            let reclaimed = if let Some(name) = target.strip_prefix("generic:") {
+                let options = GenericStoreOptions {
+                    shard_id: Some(shard.clone()),
+                    merge_after_generations: Some(0),
+                    merge_max_generations: Some(8),
+                    merge_max_bytes: Some(config.merge_max_bytes),
+                    merge_budget: Some(budget.clone()),
+                    session: Some(session.clone()),
+                    pending_generations_max: Some(0),
+                    ..Default::default()
+                };
+                let mut store = GenericStore::open_existing(&state.generic_uri(name), options)
+                    .await
+                    .map_err(|e| e.to_string())?;
+                // No resident writer was opened, so cleanup only sees sealed WAL.
+                match store
+                    .prepare_cleanup_merge()
+                    .await
+                    .map_err(|e| e.to_string())?
+                {
+                    Some((manifest_store, manifest, prepared)) => store
+                        .commit_prepared_merge(&manifest_store, &manifest, prepared)
+                        .await
+                        .map_err(|e| e.to_string())?,
+                    None => 0,
+                }
+            } else {
+                let options = RolloutStoreOptions {
+                    shard_id: Some(shard.clone()),
+                    merge_after_generations: Some(0),
+                    merge_max_generations: Some(8),
+                    merge_max_bytes: Some(config.merge_max_bytes),
+                    merge_budget: Some(budget.clone()),
+                    session: Some(session.clone()),
+                    pending_generations_max: Some(0),
+                    ..Default::default()
+                };
+                let mut store =
+                    RolloutStore::open_existing_with_options(&state.rollout_uri(target), options)
+                        .await
+                        .map_err(|e| e.to_string())?;
+                match store
+                    .prepare_merge_if_ready(1)
+                    .await
+                    .map_err(|e| e.to_string())?
+                {
+                    Some((manifest_store, manifest, prepared)) => store
+                        .commit_prepared_merge(&manifest_store, &manifest, prepared)
+                        .await
+                        .map_err(|e| e.to_string())?,
+                    None => 0,
+                }
+            };
+            total += reclaimed;
+            metrics::counter!("master_catchup_generations_reclaimed_total")
+                .increment(reclaimed as u64);
+            tracing::info!(%target, %shard, reclaimed, total, "dedicated catch-up committed");
+        }
+        if total == before {
+            break;
+        }
+    }
+    Ok(format!("dedicated catch-up merged {total} generations"))
+}