Enqueue observed WAL backlog before a full stats scan completes

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

@@ -23,12 +23,102 @@ use tokio::task::JoinHandle;
 use crate::scheduler::generic_target;
 use crate::state::MasterState;
 use crate::stats_store::StatRow;
-use lance_context_api::ExperimentSummary;
+use lance_context_api::{ExperimentSummary, TaskKind};
 
 /// Per-experiment open+observe timeout so one wedged dataset cannot stall a
 /// scan round.
 const OBSERVE_TIMEOUT: Duration = Duration::from_secs(30);
 
+type ScanObservation = (String, Option<(StatRow, bool)>);
+
+/// Early admission uses the ordinary durable queue, never a separate executor.
+/// Bound both new tasks and total control-plane time per scan. A degraded queue
+/// must not turn table observation into thousands of serial request timeouts.
+struct ScanMergeAdmission {
+    remaining: usize,
+    time_left: Duration,
+}
+
+impl ScanMergeAdmission {
+    fn new(state: &MasterState) -> Self {
+        Self {
+            remaining: if state.config.merge_wal_interval_secs == 0 {
+                0
+            } else {
+                crate::scheduler::MAX_SWEEP_ENQUEUE
+            },
+            time_left: Duration::from_secs(10),
+        }
+    }
+
+    async fn observe(&mut self, state: &Arc<MasterState>, row: &StatRow) {
+        if self.remaining == 0
+            || self.time_left.is_zero()
+            || row.pending_wal_generations <= 0
+            || row.pending_wal_generations < state.config.merge_wal_min_generations
+        {
+            return;
+        }
+        let start = std::time::Instant::now();
+        let admit = async {
+            if state
+                .task_store
+                .is_cooling_down(TaskKind::MergeWal, &row.name)
+                .await?
+            {
+                return Ok(false);
+            }
+            let before = state
+                .task_store
+                .get_active_id(TaskKind::MergeWal, &row.name)
+                .await?;
+            // Use the ordinary enqueue even when an active pointer exists:
+            // it validates the record and repairs a stale terminal pointer.
+            let task = crate::scheduler::enqueue(state, TaskKind::MergeWal, &row.name).await?;
+            // Another producer may have won the race. Count conservatively:
+            // admission must never exceed its bound because of concurrency.
+            Ok::<_, lance::Error>(before.as_deref() != Some(task.id.as_str()))
+        };
+        let result = tokio::time::timeout(self.time_left.min(Duration::from_secs(2)), admit).await;
+        self.time_left = self.time_left.saturating_sub(start.elapsed());
+        let outcome = match result {
+            Ok(Ok(true)) => {
+                self.remaining -= 1;
+                "enqueued"
+            }
+            Ok(Ok(false)) => "already_active_or_cooling_down",
+            result => {
+                // A timeout/error can follow a committed enqueue. Do not
+                // replay it or modify ownership; leave reconciliation to the
+                // durable queue and stop early admission for this round.
+                self.remaining = 0;
+                tracing::warn!(target = %row.name, ?result, "scan early merge admission stopped; observation continues");
+                "unavailable"
+            }
+        };
+        metrics::counter!("master_scan_merge_admission_total", "outcome" => outcome).increment(1);
+    }
+}
+
+/// Consume each completed observation before waiting for the rest of the scan.
+/// Incremental (base-version-unchanged) observations also recount pending WAL,
+/// so their counts are fresh and eligible. Failed observations are not.
+async fn collect_observations(
+    state: &Arc<MasterState>,
+    observations: impl futures::Stream<Item = ScanObservation>,
+) -> Vec<ScanObservation> {
+    futures::pin_mut!(observations);
+    let mut admission = ScanMergeAdmission::new(state);
+    let mut observed = Vec::new();
+    while let Some(result) = observations.next().await {
+        if let (_, Some((row, _))) = &result {
+            admission.observe(state, row).await;
+        }
+        observed.push(result);
+    }
+    observed
+}
+
 /// Bound on the compaction half of one `_stats` maintenance pass, which runs
 /// under the store lock, so a slow object store cannot wedge the scanner loop.
 ///
@@ -336,7 +426,7 @@ async fn scan_once_inner(
     // not observe that experiment; its previous row is carried over below
     // rather than dropped.
     operation.phase("observe_tables");
-    let observed: Vec<(String, Option<(StatRow, bool)>)> = stream::iter(entries)
+    let observations = stream::iter(entries)
         .map(|entry| {
             let previous = previous.clone();
             let rollout_options = rollout_options.clone();
@@ -379,9 +469,8 @@ async fn scan_once_inner(
                 }
             }
         })
-        .buffer_unordered(concurrency)
-        .collect()
-        .await;
+        .buffer_unordered(concurrency);
+    let observed = collect_observations(state, observations).await;
 
     let count = observed.iter().filter(|(_, row)| row.is_some()).count();
     let failed = observed.len() - count;
@@ -978,6 +1067,200 @@ pub fn spawn_scanner(state: &Arc<MasterState>) -> Option<JoinHandle<()>> {
     }))
 }
 
+#[cfg(test)]
+mod early_merge_tests {
+    use super::*;
+    use clap::Parser;
+
+    async fn fixture(enabled: bool) -> (tempfile::TempDir, Arc<MasterState>) {
+        let dir = tempfile::tempdir().unwrap();
+        let mut config = crate::config::MasterConfig::parse_from(["test"]);
+        config.data_dir = dir.path().to_string_lossy().into_owned();
+        config.etcd.etcd_endpoints = std::env::var("ETCD_TEST_ENDPOINTS")
+            .expect("test requires isolated etcd")
+            .split(',')
+            .map(str::to_owned)
+            .collect();
+        config.etcd.etcd_prefix = format!("/early-scan-test/{}", lance_context_core::generate_id());
+        config.merge_wal_interval_secs = if enabled { 600 } else { 0 };
+        config.merge_wal_min_generations = 8;
+        config.stats_scan_interval_secs = 0;
+        config.worker_endpoints.clear();
+        (dir, MasterState::new(config).await.unwrap())
+    }
+
+    fn row(name: &str, pending: i64) -> StatRow {
+        StatRow {
+            name: name.into(),
+            uri: format!("/unused/{name}"),
+            row_count: 0,
+            fragment_count: 0,
+            last_updated: 0,
+            pending_wal_generations: pending,
+            last_compaction: StatRow::NO_COMPACTION,
+            total_compactions: 0,
+            scanned_at: Utc::now().timestamp_millis(),
+            version: 1,
+        }
+    }
+
+    #[tokio::test]
+    #[ignore = "requires ETCD_TEST_ENDPOINTS"]
+    async fn hot_observation_enqueues_before_slow_scan_finishes() {
+        let (_dir, state) = fixture(true).await;
+        let (release, blocked) = tokio::sync::oneshot::channel();
+        // `true` is the unchanged-base fast path: its WAL recount is fresh.
+        let ready = stream::iter(vec![(
+            "generic:hot".into(),
+            Some((row("generic:hot", 80), true)),
+        )]);
+        let slow = stream::once(async move {
+            blocked.await.unwrap();
+            ("slow".into(), None)
+        });
+        let scan_state = state.clone();
+        let scan =
+            tokio::spawn(async move { collect_observations(&scan_state, ready.chain(slow)).await });
+        let id = tokio::time::timeout(Duration::from_secs(5), async {
+            loop {
+                if let Some(id) = state
+                    .task_store
+                    .get_active_id(TaskKind::MergeWal, "generic:hot")
+                    .await
+                    .unwrap()
+                {
+                    break id;
+                }
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .expect("hot target waited for the rest of the scan");
+        assert!(!scan.is_finished());
+        // The ordinary queue sees exactly the same durable task.
+        assert_eq!(
+            crate::scheduler::enqueue(&state, TaskKind::MergeWal, "generic:hot")
+                .await
+                .unwrap()
+                .id,
+            id
+        );
+        release.send(()).unwrap();
+        assert_eq!(scan.await.unwrap().len(), 2);
+        assert!(state
+            .task_store
+            .get_active_id(TaskKind::MergeWal, "slow")
+            .await
+            .unwrap()
+            .is_none());
+    }
+
+    #[tokio::test]
+    #[ignore = "requires ETCD_TEST_ENDPOINTS"]
+    async fn early_admission_respects_disable_threshold_and_new_task_budget() {
+        let (_disabled_dir, disabled) = fixture(false).await;
+        collect_observations(
+            &disabled,
+            stream::iter(vec![("hot".into(), Some((row("hot", 80), false)))]),
+        )
+        .await;
+        assert!(disabled
+            .task_store
+            .get_active_id(TaskKind::MergeWal, "hot")
+            .await
+            .unwrap()
+            .is_none());
+
+        let (_dir, state) = fixture(true).await;
+        let existing = crate::scheduler::enqueue(&state, TaskKind::MergeWal, "active")
+            .await
+            .unwrap();
+        let mut admission = ScanMergeAdmission::new(&state);
+        admission.remaining = 2;
+        for target in [
+            "active",
+            "active",
+            "small",
+            "new-a",
+            "active",
+            "generic:new-b",
+            "over-budget",
+        ] {
+            admission
+                .observe(&state, &row(target, if target == "small" { 7 } else { 80 }))
+                .await;
+        }
+        assert_eq!(
+            state
+                .task_store
+                .get_active_id(TaskKind::MergeWal, "active")
+                .await
+                .unwrap(),
+            Some(existing.id)
+        );
+        for target in ["new-a", "generic:new-b"] {
+            assert!(state
+                .task_store
+                .get_active_id(TaskKind::MergeWal, target)
+                .await
+                .unwrap()
+                .is_some());
+        }
+        for target in ["small", "over-budget"] {
+            assert!(state
+                .task_store
+                .get_active_id(TaskKind::MergeWal, target)
+                .await
+                .unwrap()
+                .is_none());
+        }
+        // An exhausted time budget also leaves scanning intact without writes.
+        let mut admission = ScanMergeAdmission::new(&state);
+        admission.time_left = Duration::ZERO;
+        admission.observe(&state, &row("time-budget", 80)).await;
+        assert!(state
+            .task_store
+            .get_active_id(TaskKind::MergeWal, "time-budget")
+            .await
+            .unwrap()
+            .is_none());
+    }
+
+    #[tokio::test]
+    #[ignore = "requires ETCD_TEST_ENDPOINTS"]
+    async fn queue_error_preserves_scan_results_and_stops_early_admission() {
+        let (_dir, state) = fixture(true).await;
+        let task = crate::scheduler::enqueue(&state, TaskKind::MergeWal, "broken")
+            .await
+            .unwrap();
+        // Fault injection is confined to this test's unique local-etcd prefix.
+        // The active pointer now references a malformed queue record.
+        state
+            .task_store
+            .etcd_client()
+            .clone()
+            .put(
+                format!("{}/tasks/{}", state.config.etcd.etcd_prefix, task.id),
+                "malformed",
+                None,
+            )
+            .await
+            .unwrap();
+        let input = vec![
+            ("broken".into(), Some((row("broken", 80), false))),
+            ("later".into(), Some((row("later", 90), false))),
+        ];
+        let observed = collect_observations(&state, stream::iter(input.clone())).await;
+        assert_eq!(observed, input, "admission failure must not lose stats");
+        assert!(state
+            .task_store
+            .get_active_id(TaskKind::MergeWal, "later")
+            .await
+            .unwrap()
+            .is_none());
+    }
+}
+
 #[cfg(test)]
 mod maintenance_alerting_tests {
     use std::sync::atomic::{AtomicU64, Ordering};

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

@@ -41,7 +41,7 @@ use crate::task_store::{TaskClaim, TaskKinds};
 /// task per row over the threshold, so at tens of thousands of experiments a
 /// single tick could flood the queue. Capping keeps each tick's work bounded;
 /// anything still over the threshold is picked up by the next tick.
-const MAX_SWEEP_ENQUEUE: usize = 256;
+pub(crate) const MAX_SWEEP_ENQUEUE: usize = 256;
 
 /// How many candidate rows a merge sweep reads per new task it may enqueue.
 /// Candidates already in flight are deduped and do not count against the cap.