From db8a8f6858311910e92817fae323399495f6542f Mon Sep 17 00:00:00 2001 From: Beinan Wang Date: Thu, 22 Oct 2026 00:43:41 +0000 Subject: [PATCH] Schedule observed WAL backlog before full scan completion --- crates/lance-context-master/src/scanner.rs | 293 ++++++++++++++++++- crates/lance-context-master/src/scheduler.rs | 2 +- 2 files changed, 289 insertions(+), 6 deletions(-) diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index d4bd9fe8..985fb2c0 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/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, 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, + observations: impl futures::Stream, +) -> Vec { + 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) -> Option> { })) } +#[cfg(test)] +mod early_merge_tests { + use super::*; + use clap::Parser; + + async fn fixture(enabled: bool) -> (tempfile::TempDir, Arc) { + 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}; diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 0dbda13f..8d0e3dc4 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/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.