From ed94d19d276a292684a58400c234eba29c1d051e Mon Sep 17 00:00:00 2001 From: Beinan Date: Wed, 2 Sep 2026 00:26:13 +0000 Subject: [PATCH] fix(master): stop MergeWal starving Compact in the scheduler Compact tasks could sit queued forever whenever MergeWal load was present. In production the 600s auto-sweep kept ~50-80 MergeWal tasks queued/running, fragments grew to 1,085,065 (1-10 rows each), and since every worker's open Dataset handle keeps the full manifest resident, 18 of 20 workers were OOMKilled. Scaling the master 3->8 replicas did not help -- every added slot was consumed the same way. Root cause is admission, not concurrency. MERGE_WAL_CONCURRENCY already gave merge its own semaphore, but a single dispatch loop called one unfiltered claim_next() and only decided which pool to draw from *after* the claim. Claiming is destructive -- it deletes the queue key, grants a lease, and takes the per-experiment target lock -- so a claim made while the general pool had permits could return a MergeWal, then park it on the saturated merge semaphore holding its target lock. With MergeWal numerically dominant, nearly every claim returned one, and the loop's `while` guard stayed true only because the *general* pool was idle. The dispatcher spun claiming work it could not run while a lone Compact sat behind it. Note this is narrower than "no fairness in claim_next": task ids are UUIDv7, so the queue is genuinely FIFO and ordering was never arbitrary. Plain FIFO is enough once each pool claims only what it can run. Fix: add TaskKinds, a small kind set threaded into claim_next, and run one poller per pool -- GENERAL (Compact + IndexId) and MERGE_WAL -- each claiming only its own kinds. The filter is applied before the dependency probe, so skipped kinds cost nothing. FIFO order within a kind is unchanged, and no new config is introduced. Setting MERGE_WAL_CONCURRENCY=0 still shares one pool, which necessarily reinstates single-poller behavior; the config docs now say so. Tests (both etcd-backed, both verified to fail without the fix): - filtered_claim_reaches_compact_behind_merge_wal_backlog: 30 MergeWal enqueued ahead of one Compact; asserts an unfiltered claim still returns MergeWal (so the scenario is real), that GENERAL reaches the trailing Compact, that GENERAL never returns MergeWal, and that the merge pool still drains its own backlog. - compact_runs_while_merge_wal_pool_is_saturated: end-to-end through the dispatcher against a stub worker that never responds; asserts the Compact reaches Done *and* that MergeWal work is still outstanding, so it cannot pass by the backlog draining. Co-Authored-By: Claude Opus 5 --- crates/lance-context-master/src/config.rs | 11 +- crates/lance-context-master/src/scheduler.rs | 185 ++++++++++++++++-- crates/lance-context-master/src/task_store.rs | 151 +++++++++++++- 3 files changed, 324 insertions(+), 23 deletions(-) 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(); + } }