diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index f2b9ccc..316c4f3 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -152,8 +152,17 @@ pub struct MasterConfig { /// both backlogs are growing and both need to drain. Giving merge its own /// budget decouples them. /// + /// The budget alone was not sufficient: it bounds concurrency, not + /// admission order. Because `MergeWal` is numerically dominant (the + /// `merge_wal_interval_secs` sweep enqueues one per over-threshold + /// experiment), a single unfiltered poller kept claiming MergeWal tasks and + /// a queued `Compact` was never reached. The scheduler therefore runs one + /// poller per pool, each claiming only the kinds it can run, so this budget + /// now also determines admission. + /// /// `0` disables the separate budget and falls back to sharing the general - /// `task_concurrency` pool. + /// `task_concurrency` pool -- which also reinstates the single-poller + /// behavior, so prefer a non-zero value when both kinds are in play. #[arg(long, env = "MERGE_WAL_CONCURRENCY", default_value_t = 4)] pub merge_wal_concurrency: usize, diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 34783b7..dc4f02e 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -34,7 +34,7 @@ use tokio::task::JoinHandle; use crate::state::MasterState; use crate::stats_store::StatRow; -use crate::task_store::TaskClaim; +use crate::task_store::{TaskClaim, TaskKinds}; /// Maximum tasks a single auto-sweep may enqueue. /// @@ -447,7 +447,14 @@ async fn sweep_merge_wal_inner(state: &Arc) -> lance::Result Ok(queued) } -/// Spawn the scheduler poller plus the optional periodic auto-sweep. +/// Spawn the scheduler pollers plus the optional periodic auto-sweep. +/// +/// Returns the handle of the *general* poller only. When a separate WAL-merge +/// budget is configured there is also a second poller (and the auto-sweep task) +/// whose handles are dropped: like the sweep, they are meant to live as long as +/// the process. Aborting the returned handle therefore stops general dispatch, +/// not every scheduler task -- adequate for tests, which drop the whole +/// `MasterState` immediately after. pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { // Optional periodic compaction auto-sweep feeds the same queue. let interval_secs = state.config.compaction_interval_secs; @@ -494,38 +501,84 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { // into the shared pool. let merge_sem = (state.config.merge_wal_concurrency > 0) .then(|| Arc::new(Semaphore::new(state.config.merge_wal_concurrency))); - let dispatch_state = state.clone(); + + // One poller per pool, each claiming only the kinds its pool runs. + // + // A single poller drawing from one unfiltered `claim_next` could not keep + // the pools independent, because claiming is destructive: it deletes the + // queue key, grants a lease, and takes the per-experiment target lock. The + // claimed task's kind then decided which pool it drew from, so a claim made + // on behalf of an idle pool could return a task belonging to a saturated + // one and park it on `acquire_owned()` -- holding its target lock while + // idle. With `MergeWal` numerically dominant (the 600s sweep enqueues one + // per over-threshold experiment), nearly every claim returned a MergeWal, + // and the loop's own `while` guard stayed true only because the *general* + // pool had permits. The result was a dispatcher that spun claiming MergeWal + // tasks it could not run while a lone `Compact` sat queued behind them -- + // starvation that adding replicas or raising `task_concurrency` could not + // fix, since every added slot was filled the same way. + // + // Splitting the pollers makes each one claim only what it can run, so a + // free compaction slot reaches past any number of queued MergeWal tasks. + // FIFO order within a kind is unchanged. + let general = spawn_pool_poller( + state.clone(), + sem, + if merge_sem.is_some() { + TaskKinds::GENERAL + } else { + // No separate merge budget: the general pool runs everything, so it + // must still be allowed to claim MergeWal. + TaskKinds::ANY + }, + true, + ); + if let Some(merge) = merge_sem { + spawn_pool_poller(state.clone(), merge, TaskKinds::MERGE_WAL, false); + } + general +} + +/// Spawn one dispatch loop bound to a single execution pool. +/// +/// `kinds` restricts what this loop will claim; `report_depth` designates the +/// one loop that publishes the shared queue-depth gauge, so running several +/// pollers does not multiply that metric. +fn spawn_pool_poller( + state: Arc, + pool: Arc, + kinds: TaskKinds, + report_depth: bool, +) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { loop { - if let Ok(queued) = dispatch_state.task_store.queue_depth().await { - metrics::gauge!("master_task_queue_depth").set(queued as f64); + if report_depth { + if let Ok(queued) = state.task_store.queue_depth().await { + metrics::gauge!("master_task_queue_depth").set(queued as f64); + } } - // Poll while *either* pool can accept work; the claimed task's kind - // decides which one it draws from below. - while sem.available_permits() > 0 - || merge_sem - .as_ref() - .is_some_and(|s| s.available_permits() > 0) - { + while pool.available_permits() > 0 { let claim_start = std::time::Instant::now(); - match dispatch_state.task_store.claim_next().await { + match state.task_store.claim_next_of_kinds(kinds).await { Ok(Some(claim)) => { let claim_elapsed = claim_start.elapsed(); // The task is already claimed at this point (queue key // deleted, lease granted, target lock held), so time // spent here is a claimed-but-idle task holding its - // per-experiment lock — worth seeing separately. + // per-experiment lock — worth seeing separately. This + // loop only claims kinds its own pool runs, so the wait + // is now bounded by that pool's own occupancy. let permit_start = std::time::Instant::now(); - let pool = match (claim.task.kind, merge_sem.as_ref()) { - (TaskKind::MergeWal, Some(merge)) => merge.clone(), - _ => sem.clone(), - }; - let permit = pool.acquire_owned().await.expect("semaphore never closed"); + let permit = pool + .clone() + .acquire_owned() + .await + .expect("semaphore never closed"); let timing = TaskClaimTiming { claim: claim_elapsed, permit_wait: permit_start.elapsed(), }; - let st = dispatch_state.clone(); + let st = state.clone(); tokio::spawn(async move { run_task(&st, claim, timing).await; drop(permit); @@ -756,6 +809,98 @@ mod tests { worker.abort(); } + /// A `Compact` runs even while every `MergeWal` slot is occupied by slow + /// fan-outs and the queue is dominated by MergeWal tasks. + /// + /// This is the end-to-end form of the starvation bug: the merge pool is + /// saturated by stub workers that never return, and 20 MergeWal tasks are + /// enqueued *ahead* of the Compact. Before the split-poller fix the single + /// dispatch loop kept claiming MergeWal tasks (parking them on the merge + /// semaphore while they held their locks) and the trailing Compact was + /// never reached, so this test would hang to its timeout. + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn compact_runs_while_merge_wal_pool_is_saturated() { + use axum::{routing::post, Router}; + + // A worker that accepts the merge call and then never responds, so the + // MergeWal task occupies its slot for the duration of the test. + let hang = Router::new().route( + "/api/v1/internal/merge-wal/{name}", + post(|| async { + std::future::pending::<()>().await; + String::new() + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, hang).await.unwrap() }); + + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.worker_endpoints = vec![format!("http://{addr}")]; + cfg.merge_wal_concurrency = 2; + cfg.task_concurrency = 2; + let state = MasterState::new(cfg).await.unwrap(); + + // Build a compactable store before the dispatcher starts. + let name = "starved"; + let uri = state.rollout_uri(name); + { + let mut store = RolloutStore::open(&uri).await.unwrap(); + for i in 0..4 { + let rec = rollout_record(&format!("r{i}")); + store.add(&[rec]).await.unwrap(); + store.cleanup_own_shard().await.unwrap(); + } + } + state + .registry + .write() + .await + .upsert(name, &uri) + .await + .unwrap(); + crate::scanner::scan_once(&state).await.unwrap(); + + // Saturate and over-fill the merge queue *before* the Compact. + for i in 0..20 { + enqueue(&state, TaskKind::MergeWal, &format!("exp-{i}")) + .await + .unwrap(); + } + let compact = enqueue(&state, TaskKind::Compact, name).await.unwrap(); + + let worker = spawn_scheduler(&state); + + let status = await_terminal(&state, &compact.id).await; + assert_eq!( + status.state, + TaskState::Done, + "Compact must drain while MergeWal saturates its own pool: {status:?}" + ); + + // Sanity: the merge backlog really is still stuck, i.e. the Compact did + // not simply win because the MergeWal tasks all completed. + let stuck = state + .task_store + .list() + .await + .unwrap() + .into_iter() + .filter(|t| { + t.kind == TaskKind::MergeWal + && matches!(t.state, TaskState::Queued | TaskState::Running) + }) + .count(); + assert!( + stuck > 0, + "test must leave MergeWal work outstanding, else it proves nothing" + ); + + worker.abort(); + } + /// A dependent task runs only after its dependency reaches `Done`: an /// `index_id` depending on a `compact` must start after compaction finishes, /// so the two never contend for the shared per-experiment base-table gate. diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 203de90..96411b4 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -25,6 +25,48 @@ use crate::config::MasterConfig; const TASK_POLL_BATCH: usize = 256; +/// A set of [`TaskKind`]s a claim is allowed to return. +/// +/// Small enough to copy; used to let each scheduler execution pool claim only +/// the kinds it can actually run, so a saturated pool's backlog never hides a +/// runnable task belonging to an idle pool. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TaskKinds { + compact: bool, + merge_wal: bool, + index_id: bool, +} + +impl TaskKinds { + /// Every kind -- the unfiltered claim used outside the scheduler. + pub const ANY: Self = Self { + compact: true, + merge_wal: true, + index_id: true, + }; + /// Only `MergeWal`, the kind with its own dedicated budget. + pub const MERGE_WAL: Self = Self { + compact: false, + merge_wal: true, + index_id: false, + }; + /// Everything the general pool runs: `Compact` and `IndexId`. + pub const GENERAL: Self = Self { + compact: true, + merge_wal: false, + index_id: true, + }; + + #[must_use] + pub fn contains(self, kind: TaskKind) -> bool { + match kind { + TaskKind::Compact => self.compact, + TaskKind::MergeWal => self.merge_wal, + TaskKind::IndexId => self.index_id, + } + } +} + #[derive(Clone)] pub struct TaskStore { inner: Arc, @@ -116,8 +158,25 @@ impl TaskStore { /// Claim the oldest runnable task. The queued->running update, claim lease, /// and per-experiment write lock are one etcd transaction. pub async fn claim_next(&self) -> lance::Result> { + self.claim_next_of_kinds(TaskKinds::ANY).await + } + + /// Claim the oldest runnable task whose kind is in `kinds`, skipping over + /// queued tasks of other kinds instead of stopping at them. + /// + /// The scheduler needs this because it runs several execution pools with + /// separate budgets. Claiming is destructive -- it deletes the queue key, + /// grants a lease, and takes the per-experiment target lock -- so claiming a + /// task the caller has no free slot for does not merely waste a poll: it + /// parks a claimed task on a semaphore while it holds its target lock. With + /// one kind numerically dominant, an unfiltered claim returns that kind + /// nearly every time, so a scarce kind can wait behind it indefinitely even + /// though its own pool is idle. Filtering at claim time keeps the queue's + /// FIFO order within each kind while letting an idle pool reach past a + /// saturated one. + pub async fn claim_next_of_kinds(&self, kinds: TaskKinds) -> lance::Result> { self.inner.recover_orphaned().await?; - let (claim, dependency_failed) = self.inner.claim_next().await?; + let (claim, dependency_failed) = self.inner.claim_next(kinds).await?; if dependency_failed { if let Err(error) = self.prune_terminal_history().await { tracing::warn!(error = %error, "failed to prune task history"); @@ -395,7 +454,7 @@ impl EtcdTaskStore { .map_err(|_| lance::Error::io("etcd returned an invalid queue count")) } - async fn claim_next(&self) -> lance::Result<(Option, bool)> { + async fn claim_next(&self, kinds: TaskKinds) -> lance::Result<(Option, bool)> { let prefix = self.queue_prefix(); let range_end = prefix_range_end(prefix.as_bytes()); let mut start_key = prefix.into_bytes(); @@ -427,6 +486,12 @@ impl EtcdTaskStore { .collect::>>()?; for mut task in queued { + // Skip kinds this caller cannot run *before* the dependency + // probe: an unrunnable kind should cost nothing, and this is + // what lets a scarce kind be found behind a dominant one. + if !kinds.contains(task.kind) { + continue; + } match self.dependency_status(&task).await? { DependencyStatus::Ready => {} DependencyStatus::Waiting => continue, @@ -1135,4 +1200,86 @@ mod tests { .await .unwrap(); } + + /// A kind-filtered claim reaches past a large backlog of another kind. + /// + /// This is the starvation fix in miniature: many `MergeWal` tasks are + /// enqueued *before* a single `Compact`, so the `Compact` is last in the + /// queue's FIFO order. An unfiltered claim returns MergeWal every time -- + /// asserted below, so the test still describes the old behavior -- while a + /// `GENERAL`-filtered claim must skip all of them and find the Compact. + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn filtered_claim_reaches_compact_behind_merge_wal_backlog() { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS") + .expect("ETCD_TEST_ENDPOINTS must point to a test etcd"); + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.etcd_endpoints = endpoint.split(',').map(str::to_string).collect(); + cfg.etcd_prefix = format!("/lance-context/test/{}", generate_id()); + cfg.etcd_lease_ttl_secs = 5; + + let store = TaskStore::open(&cfg).await.unwrap(); + + // 30 MergeWal tasks (distinct targets, so dedupe keeps them all), then + // the Compact last -- worst case for FIFO order. + for i in 0..30 { + store + .enqueue(TaskKind::MergeWal, &format!("exp-{i}"), Vec::new()) + .await + .unwrap(); + } + let compact = store + .enqueue(TaskKind::Compact, "starved", Vec::new()) + .await + .unwrap(); + + // Unfiltered: FIFO hands back a MergeWal, never the trailing Compact. + let any = store.claim_next().await.unwrap().unwrap(); + assert_eq!( + any.task.kind, + TaskKind::MergeWal, + "the backlog head must still be MergeWal -- otherwise this test is \ + not exercising the starvation scenario" + ); + store.finish(any, Ok("done".to_string())).await.unwrap(); + + // Filtered to what the general pool runs: must skip the whole backlog. + let claimed = store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .expect("a Compact behind a MergeWal backlog must still be claimable"); + assert_eq!(claimed.task.kind, TaskKind::Compact); + assert_eq!(claimed.task.id, compact.id); + store.finish(claimed, Ok("done".to_string())).await.unwrap(); + + // And with no Compact left, the general filter reports empty rather + // than falling back to the still-large MergeWal backlog. + assert!( + store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .is_none(), + "GENERAL must not claim MergeWal" + ); + // The merge poller still sees its own backlog. + let merge = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .expect("MergeWal backlog must remain claimable by its own pool"); + assert_eq!(merge.task.kind, TaskKind::MergeWal); + store.finish(merge, Ok("done".to_string())).await.unwrap(); + + let mut client = Client::connect(cfg.etcd_endpoints, None).await.unwrap(); + client + .delete( + cfg.etcd_prefix, + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } }