diff --git a/crates/lance-context-master/src/executors.rs b/crates/lance-context-master/src/executors.rs index 67c6da75..49bceca1 100644 --- a/crates/lance-context-master/src/executors.rs +++ b/crates/lance-context-master/src/executors.rs @@ -264,6 +264,11 @@ pub(crate) fn spawn(state: &Arc) { }); } +#[cfg(test)] +pub(crate) async fn heartbeat_loop_for_test(state: &Arc) -> Result<(), String> { + heartbeat_loop(state).await +} + async fn heartbeat_loop(state: &Arc) -> Result<(), String> { let keys = Keys::new(&state.config.etcd.etcd_prefix); let mut client = state.task_store.etcd_client().clone(); @@ -580,3 +585,55 @@ mod tests { assert!(h[&id].draining); } } + +#[cfg(test)] +mod shadow_isolation { + /// Spec test 5: shadow assignments are written only by the planner and + /// read only by the planner, `load_headroom` and the operator route. No + /// scheduler, merge, catch-up or executor code path may consult them + /// until P2 deliberately introduces binding. This is a source-level + /// check so a stray read fails CI, not production. + #[test] + fn no_executor_code_path_reads_assignments() { + let allowed = ["planner.rs", "executors.rs", "routes.rs"]; + let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); + let mut offenders = Vec::new(); + for entry in std::fs::read_dir(&dir).unwrap().flatten() { + let path = entry.path(); + let Some(name) = path.file_name().and_then(|n| n.to_str()) else { + continue; + }; + if !name.ends_with(".rs") || allowed.contains(&name) { + continue; + } + let text = std::fs::read_to_string(&path).unwrap(); + if text.contains("assignments_prefix") + || text.contains(".assignment(") + || text.contains("/assignments/") + { + offenders.push(name.to_string()); + } + } + // The catchup/ subdirectory too. + for sub in ["catchup"] { + if let Ok(rd) = std::fs::read_dir(dir.join(sub)) { + for entry in rd.flatten() { + let text = std::fs::read_to_string(entry.path()).unwrap_or_default(); + if text.contains("assignments") { + offenders.push(format!("{sub}/{}", entry.file_name().to_string_lossy())); + } + } + } + } + assert!( + offenders.is_empty(), + "executor paths read assignments: {offenders:?}" + ); + // routes.rs may only read through load_headroom. + let routes = std::fs::read_to_string(dir.join("routes.rs")).unwrap(); + assert!( + !routes.contains("assignments_prefix"), + "routes must go through load_headroom" + ); + } +} diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs index dfe10e7c..a0330986 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/crates/lance-context-master/src/lib.rs @@ -17,6 +17,7 @@ pub mod rollout_append; pub mod routes; pub mod scanner; pub mod scheduler; +pub mod scoring; pub mod state; pub mod stats_store; pub mod task_store; diff --git a/crates/lance-context-master/src/planner.rs b/crates/lance-context-master/src/planner.rs index cbf6bae1..b6bf1a7c 100644 --- a/crates/lance-context-master/src/planner.rs +++ b/crates/lance-context-master/src/planner.rs @@ -281,6 +281,7 @@ pub(crate) async fn reconcile_once( let mut ops = Vec::new(); let mut written = 0usize; let mut total_pending = 0u64; + let folded: BTreeMap = tables.clone(); for (target, mut table) in tables { total_pending += table.pending_generations(); let key = keys.demand(&target); @@ -337,14 +338,27 @@ pub(crate) async fn reconcile_once( return Err("leadership lost during reconcile".into()); } } - // Executor view: heartbeats plus reservations. Read every tick so the - // leader always has it warm; nothing is placed from it yet. + // Executor view: heartbeats plus reservations, then the shadow pass. let executors_keys = crate::executors::Keys::new(&state.config.etcd.etcd_prefix); match crate::executors::load_headroom(&client, &executors_keys).await { Ok(headroom) => { metrics::gauge!("planner_executors").set(headroom.len() as f64); let reserved: u64 = headroom.values().map(|h| h.bytes_reserved).sum(); metrics::gauge!("planner_reserved_bytes_total").set(reserved as f64); + if let Err(error) = shadow_pass( + state, + keys, + &executors_keys, + leader_token, + &folded, + headroom, + now, + ) + .await + { + tracing::warn!(%error, "planner shadow pass failed"); + metrics::counter!("planner_shadow_errors_total").increment(1); + } } Err(error) => { tracing::warn!(%error, "planner could not load executor headroom"); @@ -363,6 +377,178 @@ pub(crate) async fn reconcile_once( Ok(written) } +/// Shadow placement: score every table, order per design §4.2, place +/// against headroom, and record what *would* be dispatched as `Shadow` +/// assignments that no executor reads. Then compare with reality: for each +/// table the planner would place, is a MergeWal task already active or +/// queued? For each table it would not, is one running anyway? Each +/// mismatch is a disagreement; the counter has to approach zero before P2. +/// +/// Shadow assignments are replaced wholesale every pass (they are a view, +/// not a reservation anyone binds), guarded by the leader token. +async fn shadow_pass( + state: &Arc, + keys: &Keys, + executors_keys: &crate::executors::Keys, + leader_token: &str, + tables: &BTreeMap, + mut headroom: BTreeMap, + now: i64, +) -> Result<(), String> { + use crate::executors::{Assignment, AssignmentState}; + use crate::scoring::{planner_order, score_merge, Class, MergePolicy}; + use lance_context_api::TaskKind; + + let policy = MergePolicy { + min_generations: state.config.merge_wal_min_generations.max(1) as u64, + ..MergePolicy::default() + }; + let mut scored: Vec<_> = tables + .iter() + .filter_map(|(target, demand)| score_merge(target, demand, &policy, now)) + .collect(); + scored.sort_by(planner_order); + for s in &scored { + metrics::gauge!("planner_merge_score", "target" => s.target.clone()).set(s.score); + } + let by_class = |c: Class| scored.iter().filter(|s| s.class == c).count(); + metrics::gauge!("planner_units", "class" => "critical").set(by_class(Class::Critical) as f64); + metrics::gauge!("planner_units", "class" => "normal").set(by_class(Class::Normal) as f64); + metrics::gauge!("planner_units", "class" => "tail").set(by_class(Class::Tail) as f64); + + // Only Shadow records from previous passes are replaced; nothing else + // under assignments/ is touched. + let mut client = state.task_store.etcd_client().clone(); + let existing = client + .get( + executors_keys.assignments_prefix(), + Some(GetOptions::new().with_prefix()), + ) + .await + .map_err(|e| e.to_string())?; + let mut ops: Vec = existing + .kvs() + .iter() + .filter(|kv| { + serde_json::from_slice::(kv.value()) + .is_ok_and(|a| a.state == AssignmentState::Shadow) + }) + .map(|kv| TxnOp::delete(kv.key().to_vec(), None)) + .collect(); + // Headroom already counts previous Shadow records; release them in the + // arithmetic too, since they are about to be replaced. + for h in headroom.values_mut() { + h.slots_reserved.clear(); + h.bytes_reserved = 0; + h.assignments = 0; + } + for kv in existing.kvs() { + if let Ok(a) = serde_json::from_slice::(kv.value()) { + if a.state != AssignmentState::Shadow && a.holds_capacity() { + if let Some(h) = headroom.get_mut(&a.executor) { + for (k, n) in &a.reserved_slots { + *h.slots_reserved.entry(*k).or_insert(0) += n; + } + h.bytes_reserved += a.reserved_bytes; + h.assignments += 1; + } + } + } + } + + let mut would_place: Vec<(String, String)> = Vec::new(); + let mut unplaceable = 0usize; + for s in &scored { + if s.class == Class::Tail { + continue; // the real sweeps do not touch tails either + } + // Estimated cost: staging buffers are ~2x decoded bytes, capped at + // the per-pass byte limit the append path enforces. + let expected_bytes = (s.pending_bytes.saturating_mul(2)) + .min(state.config.append.rollout_append_max_bytes as u64 * 2) + .max(1); + let pick = headroom + .values_mut() + .filter(|h| h.fits(TaskKind::MergeWal, expected_bytes)) + .max_by_key(|h| h.bytes_free()); + let Some(h) = pick else { + unplaceable += 1; + continue; + }; + *h.slots_reserved.entry(TaskKind::MergeWal).or_insert(0) += 1; + h.bytes_reserved += expected_bytes; + h.assignments += 1; + let unit_id = generate_id(); + let a = Assignment { + v: crate::executors::SCHEMA_VERSION, + unit_id: unit_id.clone(), + kind: TaskKind::MergeWal, + target: s.target.clone(), + needs_write_turn: true, + executor: h.executor.clone(), + planner_token: leader_token.to_string(), + reserved_slots: [(TaskKind::MergeWal, 1)].into_iter().collect(), + reserved_bytes: expected_bytes, + state: AssignmentState::Shadow, + created_ms: now, + bind_deadline_ms: now + 30_000, + }; + ops.push(TxnOp::put( + executors_keys.assignment(&s.target, &unit_id), + serde_json::to_vec(&a).unwrap(), + None, + )); + would_place.push((s.target.clone(), h.executor.clone())); + } + metrics::gauge!("planner_shadow_placements").set(would_place.len() as f64); + metrics::gauge!("planner_shadow_unplaceable").set(unplaceable as f64); + + for chunk in ops.chunks(100) { + let response = client + .txn( + Txn::new() + .when(vec![Compare::value( + keys.leader_token().as_str(), + CompareOp::Equal, + leader_token, + )]) + .and_then(chunk.to_vec()), + ) + .await + .map_err(|e| e.to_string())?; + if !response.succeeded() { + return Err("leadership lost during shadow pass".into()); + } + } + + // Disagreement: planner says place vs. reality has no task; or planner + // says nothing (no demand scored, or tail) vs. reality is running one. + let placed: std::collections::BTreeSet<&str> = + would_place.iter().map(|(t, _)| t.as_str()).collect(); + let mut disagreements = 0u64; + for s in &scored { + let active = state + .task_store + .get_active_id(TaskKind::MergeWal, &s.target) + .await + .map_err(|e| e.to_string())? + .is_some(); + let planner_wants = placed.contains(s.target.as_str()); + if planner_wants != active { + disagreements += 1; + tracing::debug!( + target = %s.target, class = ?s.class, score = s.score, + planner_wants, reality_active = active, + "planner shadow disagreement" + ); + } + } + metrics::counter!("scheduler_shadow_disagreements_total", "kind" => "merge_wal") + .increment(disagreements); + metrics::gauge!("planner_shadow_disagreements_last_pass").set(disagreements as f64); + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -560,6 +746,71 @@ mod tests { assert_eq!(report.tables[0].pending_generations, 42); assert_eq!(report.tables[0].shards, 2); + // Shadow placement: with this master's heartbeat live, the hot table + // (42 pending >= 8) gets exactly one Shadow assignment against it; + // cold (0 pending) gets none. No real MergeWal task exists, so the + // pass records one disagreement for hot. + let hb_runner = tokio::spawn({ + let first = first.clone(); + async move { crate::executors::heartbeat_loop_for_test(&first).await } + }); + let ekeys = crate::executors::Keys::new(&cfg.etcd.etcd_prefix); + let deadline = std::time::Instant::now() + Duration::from_secs(15); + let shadows: Vec = loop { + let got = client + .get( + ekeys.assignments_prefix(), + Some(GetOptions::new().with_prefix()), + ) + .await + .unwrap(); + let v: Vec = got + .kvs() + .iter() + .filter_map(|kv| serde_json::from_slice(kv.value()).ok()) + .collect(); + if !v.is_empty() { + break v; + } + assert!( + std::time::Instant::now() < deadline, + "no shadow assignment written" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + }; + assert_eq!(shadows.len(), 1, "{shadows:?}"); + assert_eq!(shadows[0].target, "hot"); + assert_eq!(shadows[0].state, crate::executors::AssignmentState::Shadow); + assert_eq!(shadows[0].kind, lance_context_api::TaskKind::MergeWal); + assert_eq!(shadows[0].executor, first.admission.status().executor_id); + assert!(shadows[0].reserved_bytes > 0); + // Headroom reflects the shadow reservation. + let h = crate::executors::load_headroom(&client, &ekeys) + .await + .unwrap(); + assert_eq!(h[&shadows[0].executor].assignments, 1); + // Enqueue a real MergeWal task for hot: on the next pass the planner + // and reality agree, and the shadow row is replaced, not duplicated. + first + .task_store + .enqueue(lance_context_api::TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(1500)).await; + let got = client + .get( + ekeys.assignments_prefix(), + Some(GetOptions::new().with_prefix()), + ) + .await + .unwrap(); + assert_eq!( + got.kvs().len(), + 1, + "shadow rows are replaced, never accumulated" + ); + hb_runner.abort(); + // A deposed leader's reconcile is refused: forge a different token. let err = reconcile_once(&first, &keys, "not-the-leader") .await diff --git a/crates/lance-context-master/src/scoring.rs b/crates/lance-context-master/src/scoring.rs new file mode 100644 index 00000000..c69a7bd8 --- /dev/null +++ b/crates/lance-context-master/src/scoring.rs @@ -0,0 +1,265 @@ +//! Scoring and class assignment for merge demand (P1.7, shadow). +//! +//! Design: `docs/design/unified-maintenance-scheduler.md` §4.2. Pure +//! functions over `TableDemand` plus a clock and a policy. P1 scores only +//! `MergeWal`; compaction and index demand are folded into `TableDemand` +//! but not yet scored, so the shadow planner can be compared against the +//! merge loops it would eventually replace before anything else moves. +//! +//! Ordering rule (review round 2): class 1 is served **oldest first** by +//! wait time, not by score; score orders classes 2–5 only. Promotion to +//! class 1 is therefore a guarantee of service, not a nudge. + +use std::cmp::Ordering; + +use lance_context_merge::demand::TableDemand; +use serde::Serialize; + +/// Per-table merge policy. Defaults mirror today's sweep thresholds. +#[derive(Debug, Clone, Serialize)] +pub struct MergePolicy { + /// `merge_score >= 1` at this many pending generations. + pub min_generations: u64, + /// `merge_score >= 1` at this many pending bytes (0 = ignore bytes). + pub min_bytes: u64, + /// Aging horizon: a backlog this old doubles its score. + pub max_age_secs: u64, + /// At or above this, class 1 regardless of score. + pub critical_generations: u64, + /// A tail that has waited this long is promoted to class 1. + pub max_turn_wait_secs: u64, +} + +impl Default for MergePolicy { + fn default() -> Self { + Self { + min_generations: 8, + min_bytes: 0, + max_age_secs: 3600, + critical_generations: 256, + max_turn_wait_secs: 3600, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum Class { + /// Reserved for recovery units; not produced from demand in P1. + Recovery = 0, + /// Pending ≥ critical, or promoted for age. Served oldest-first. + Critical = 1, + /// Reserved for commit-ready preparations; not produced in P1. + CommitReady = 2, + /// `merge_score >= 1`. + Normal = 3, + /// Reserved for compaction/index demand; not produced in P1. + Maintenance = 4, + /// `0 < merge_score < 1`: low-count tail. + Tail = 5, +} + +/// One scored unit of merge work. +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct Scored { + pub target: String, + pub class: Class, + /// Dimensionless; ≥ 1 means "due". + pub score: f64, + pub pending_generations: u64, + pub pending_bytes: u64, + /// Age of the oldest pending generation, if known. + pub wait_ms: Option, + /// Why this class. + pub reason: &'static str, +} + +/// Score one table's merge demand. `None` when nothing is pending. +pub fn score_merge( + target: &str, + demand: &TableDemand, + policy: &MergePolicy, + now_ms: i64, +) -> Option { + let pending = demand.pending_generations(); + let bytes = demand.pending_bytes(); + if pending == 0 { + return None; + } + let wait_ms = demand.oldest_pending_ms().map(|t| (now_ms - t).max(0)); + let by_gens = pending as f64 / policy.min_generations.max(1) as f64; + let by_bytes = if policy.min_bytes > 0 { + bytes as f64 / policy.min_bytes as f64 + } else { + 0.0 + }; + let age_boost = match wait_ms { + Some(w) if policy.max_age_secs > 0 => { + 1.0 + (w as f64 / 1000.0) / policy.max_age_secs as f64 + } + _ => 1.0, + }; + let score = by_gens.max(by_bytes) * age_boost; + let (class, reason) = if pending >= policy.critical_generations { + (Class::Critical, "pending_at_or_above_critical") + } else if wait_ms.is_some_and(|w| w >= (policy.max_turn_wait_secs as i64).saturating_mul(1000)) + { + (Class::Critical, "promoted_for_age") + } else if score >= 1.0 { + (Class::Normal, "merge_score_due") + } else { + (Class::Tail, "below_threshold") + }; + Some(Scored { + target: target.to_string(), + class, + score, + pending_generations: pending, + pending_bytes: bytes, + wait_ms, + reason, + }) +} + +/// Planner ordering: class ascending; within class 1, oldest wait first +/// (unknown wait sorts last); within other classes, score descending. +/// Target name breaks remaining ties so the order is total and stable. +pub fn planner_order(a: &Scored, b: &Scored) -> Ordering { + a.class.cmp(&b.class).then_with(|| { + if a.class == Class::Critical { + match (a.wait_ms, b.wait_ms) { + (Some(x), Some(y)) => y.cmp(&x), + (Some(_), None) => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (None, None) => Ordering::Equal, + } + } else { + b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal) + } + .then_with(|| a.target.cmp(&b.target)) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use lance_context_merge::demand::{DemandEvent, EventSource, SCHEMA_VERSION}; + + fn demand(pending: u64, oldest_ms: i64) -> TableDemand { + let mut t = TableDemand::default(); + // One shard sealed to `pending`, oldest generation flushed at oldest_ms. + for g in 1..=pending { + let _ = t.fold( + &DemandEvent { + v: SCHEMA_VERSION, + shard: "s".into(), + sealed_through: g, + sealed_bytes_through: g * 1_000_000, + merged_through: None, + flushed_at_ms: oldest_ms + g as i64, + writer_epoch: 1, + source: EventSource::Writer, + merged_epoch: None, + sealed_times: Default::default(), + }, + g as i64, + 0, + ); + } + t + } + + #[test] + fn classes_follow_thresholds() { + let p = MergePolicy::default(); + let now = 10_000_000; + assert!(score_merge("x", &TableDemand::default(), &p, now).is_none()); + let tail = score_merge("x", &demand(3, now - 1000), &p, now).unwrap(); + assert_eq!(tail.class, Class::Tail); + assert!(tail.score < 1.0); + let normal = score_merge("x", &demand(16, now - 1000), &p, now).unwrap(); + assert_eq!(normal.class, Class::Normal); + assert!(normal.score >= 2.0); + let critical = score_merge("x", &demand(300, now - 1000), &p, now).unwrap(); + assert_eq!(critical.class, Class::Critical); + assert_eq!(critical.reason, "pending_at_or_above_critical"); + } + + #[test] + fn age_lifts_score_and_promotes_an_old_tail() { + let p = MergePolicy { + max_age_secs: 100, + max_turn_wait_secs: 200, + ..Default::default() + }; + let now = 1_000_000_000; + let fresh = score_merge("x", &demand(2, now), &p, now).unwrap(); + let aged = score_merge("x", &demand(2, now - 100_000), &p, now).unwrap(); + assert!(aged.score > fresh.score); + assert!( + (aged.score / fresh.score - 2.0).abs() < 0.05, + "doubled at max_age" + ); + assert_eq!(aged.class, Class::Tail, "aged but below promotion"); + let promoted = score_merge("x", &demand(2, now - 250_000), &p, now).unwrap(); + assert_eq!(promoted.class, Class::Critical); + assert_eq!(promoted.reason, "promoted_for_age"); + } + + /// Spec test 6 / review round 2: class 1 is oldest-first. A low-score + /// table promoted for age is placed before a high-score critical merge + /// that has waited less. + #[test] + fn class_one_is_oldest_first_not_score_first() { + let p = MergePolicy { + max_turn_wait_secs: 100, + ..Default::default() + }; + let now = 1_000_000_000; + let hot = score_merge("hot", &demand(4000, now - 10_000), &p, now).unwrap(); + let tail = score_merge("tail", &demand(2, now - 500_000), &p, now).unwrap(); + assert_eq!(hot.class, Class::Critical); + assert_eq!(tail.class, Class::Critical); + assert!(hot.score > tail.score * 100.0); + let mut v = [hot.clone(), tail.clone()]; + v.sort_by(planner_order); + assert_eq!( + v[0].target, "tail", + "older wait wins class 1 despite tiny score" + ); + // Outside class 1, score decides. + let a = score_merge("a", &demand(32, now - 1000), &p, now).unwrap(); + let b = score_merge("b", &demand(16, now - 900_000_000), &p, now).unwrap(); + assert_eq!((a.class, b.class), (Class::Normal, Class::Critical)); + let mut v = [a.clone(), b.clone()]; + v.sort_by(planner_order); + assert_eq!(v[0].target, "b", "class before score"); + } + + #[test] + fn ordering_is_total_and_stable() { + let p = MergePolicy::default(); + let now = 1_000_000_000; + let mut items: Vec = (0..20) + .map(|i| { + let pending = [2u64, 16, 300][i % 3]; + score_merge( + &format!("t{i:02}"), + &demand(pending, now - (i as i64) * 1000), + &p, + now, + ) + .unwrap() + }) + .collect(); + let mut again = items.clone(); + again.reverse(); + items.sort_by(planner_order); + again.sort_by(planner_order); + assert_eq!(items, again); + let classes: Vec<_> = items.iter().map(|s| s.class).collect(); + let mut sorted = classes.clone(); + sorted.sort(); + assert_eq!(classes, sorted, "classes ascend"); + } +} diff --git a/docs/design/scheduler-p1-data-model.md b/docs/design/scheduler-p1-data-model.md index d6b35f75..43c19cf6 100644 --- a/docs/design/scheduler-p1-data-model.md +++ b/docs/design/scheduler-p1-data-model.md @@ -160,9 +160,17 @@ Rules (review point 4): ### 2.3 P1 scope for executors P1 writes heartbeats from every master (all are executors regardless of -leadership) and from catch-up Jobs. The planner computes placements and -writes assignments **only in shadow mode**: state `shadow`, never read by any -executor, garbage-collected after 10 min. No executor binds anything in P1. +leadership). The planner scores `MergeWal` demand only (design §4.2: classes +Critical / Normal / Tail; class 1 oldest‑first), places against headroom, and +writes assignments **only in shadow mode**: state `shadow`, replaced wholesale +every pass, never read by any executor (a source‑level test enforces this), no +binding. Tails are not placed, matching the real sweeps. Each pass then +compares its placements with whether a real `MergeWal` task is active per +table and increments `scheduler_shadow_disagreements_total{kind}`; the +per‑pass count is `planner_shadow_disagreements_last_pass`. Compaction and +index scoring, commit‑ready class 2 and `max_consecutive_turns` come with P2's +preparation integration; catch‑up Job heartbeats come when Jobs become an +executor kind. ## 3. Leader