diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index e29afdff..14cb482c 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/crates/lance-context-master/src/scanner.rs @@ -13,10 +13,13 @@ use std::time::Duration; use chrono::Utc; use futures::stream::{self, StreamExt}; -use lance_context_core::{CompactionConfig, RolloutStore, RolloutStoreOptions}; +use lance_context_core::{ + CompactionConfig, GenericStore, GenericStoreOptions, RolloutStore, RolloutStoreOptions, +}; use tokio::sync::Semaphore; use tokio::task::JoinHandle; +use crate::scheduler::generic_target; use crate::state::MasterState; use crate::stats_store::StatRow; use lance_context_api::ExperimentSummary; @@ -176,7 +179,37 @@ async fn try_scan_once(state: &Arc, maintain: bool) -> lance::Resul async fn scan_once_inner(state: &Arc) -> lance::Result { let scan_start = std::time::Instant::now(); - let entries = state.registry.write().await.list().await?; + // Rollout experiments under their bare name; generic stores under the + // `generic:` task-target prefix, so they occupy distinct rows and the + // WAL-merge sweep enqueues them as generic tasks. Their MemWAL merges + // depend on this: the sweep reads nothing but this table. + let mut entries: Vec = state + .registry + .write() + .await + .list() + .await? + .into_iter() + .map(|e| ScanEntry { + name: e.name, + uri: e.uri, + kind: ScanKind::Rollout, + }) + .collect(); + entries.extend( + state + .generic_registry + .write() + .await + .list() + .await? + .into_iter() + .map(|e| ScanEntry { + name: generic_target(&e.name), + uri: e.uri, + kind: ScanKind::Generic, + }), + ); let live: HashSet = entries.iter().map(|e| e.name.clone()).collect(); let concurrency = state.config.scan_concurrency.max(1); @@ -196,6 +229,7 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { }; let previous = Arc::new(previous); let rollout_options = state.rollout_store_options(); + let generic_options = state.generic_store_options(); // Observe experiments concurrently (bounded). `None` means this round could // not observe that experiment; its previous row is carried over below @@ -204,9 +238,18 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { .map(|entry| { let previous = previous.clone(); let rollout_options = rollout_options.clone(); + let generic_options = generic_options.clone(); async move { let prev = previous.get(&entry.name); - match observe_one(&entry.name, &entry.uri, prev, rollout_options).await { + let observed = match entry.kind { + ScanKind::Rollout => { + observe_one(&entry.name, &entry.uri, prev, rollout_options).await + } + ScanKind::Generic => { + observe_generic(&entry.name, &entry.uri, prev, generic_options).await + } + }; + match observed { Ok(result) => (entry.name, Some(result)), Err(e) => { tracing::warn!(store = %entry.name, error = %e, "scan: observe failed"); @@ -271,8 +314,15 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { // is equivalent and keeps the critical section to one commit. let retire_after = Duration::from_secs(state.config.stats_cold_retire_secs); let retire_start = std::time::Instant::now(); + // Retirement opens each row through the rollout store; generic rows are + // observed only, never retired by the master. + let rollout_rows: Vec = snapshot + .iter() + .filter(|row| !row.name.starts_with("generic:")) + .cloned() + .collect(); let retired = retire_cold_experiments( - &snapshot, + &rollout_rows, retire_after, Utc::now().timestamp_millis(), state.rollout_store_options(), @@ -373,6 +423,90 @@ async fn scan_once_inner(state: &Arc) -> lance::Result { /// rows, and a `count_rows` over the base table. At tens of thousands of /// experiments, almost all of them cold, that is the difference between a scan /// that finishes inside its interval and one that never does. +/// Which registry a scan entry came from, which decides how it is opened. +#[derive(Debug, Clone, Copy)] +enum ScanKind { + Rollout, + Generic, +} + +/// One store to observe this round. `name` is the stats-row / task-target +/// name (prefixed for generic stores); `uri` is the physical dataset. +struct ScanEntry { + name: String, + uri: String, + kind: ScanKind, +} + +/// Observe one generic store. Only the fields the WAL-merge sweep and the UI +/// need: version (for the skip-if-unchanged shortcut), base row count, +/// fragment count and pending MemWAL generations. Compaction counters stay at +/// their carried-over values because the master does not compact generic +/// stores. +async fn observe_generic( + name: &str, + uri: &str, + previous: Option<&StatRow>, + options: GenericStoreOptions, +) -> lance::Result<(StatRow, bool)> { + let open = GenericStore::open_existing(uri, options); + let store = match tokio::time::timeout(OBSERVE_TIMEOUT, open).await { + Ok(Ok(store)) => store, + Ok(Err(e)) => return Err(e), + Err(_) => { + return Err(lance::Error::io(format!( + "open timed out for store '{name}'" + ))) + } + }; + + let current_version = store.version() as i64; + if let Some(prev) = previous { + if prev.version != StatRow::UNKNOWN_VERSION && prev.version == current_version { + let mut row = prev.clone(); + row.uri = uri.to_string(); + row.scanned_at = Utc::now().timestamp_millis(); + return Ok((row, true)); + } + } + + let observe = async { + let pending = store.pending_wal_generations().await?; + let rows = store.count_base_rows().await?; + Ok::<_, lance::Error>((pending, rows)) + }; + let (pending, rows) = match tokio::time::timeout(OBSERVE_TIMEOUT, observe).await { + Ok(Ok(v)) => v, + Ok(Err(e)) => return Err(e), + Err(_) => { + return Err(lance::Error::io(format!( + "observe timed out for store '{name}'" + ))); + } + }; + let compaction = store.compaction_stats(); + let (last_compaction, total_compactions) = previous + .map_or((StatRow::NO_COMPACTION, 0), |prev| { + (prev.last_compaction, prev.total_compactions) + }); + + Ok(( + StatRow { + name: name.to_string(), + uri: uri.to_string(), + row_count: rows as i64, + fragment_count: compaction.total_fragments as i64, + last_updated: Utc::now().timestamp_millis(), + pending_wal_generations: pending as i64, + last_compaction, + total_compactions, + scanned_at: Utc::now().timestamp_millis(), + version: current_version, + }, + false, + )) +} + async fn observe_one( name: &str, uri: &str, @@ -793,6 +927,65 @@ mod incremental_scan_tests { ); } + /// A generic store's pending MemWAL generations are observed and written + /// under its `generic:` target name. This is the row the WAL-merge sweep + /// reads; without it the master never schedules a generic store's merge + /// and 20 workers race to commit into its base table on their own timers. + #[tokio::test] + async fn generic_store_is_observed_with_pending_wal() { + use lance_context_api::{ColumnSpec, ColumnType, SchemaSpec, ID_COLUMN}; + use lance_context_core::{GenericStore, GenericStoreOptions}; + + let dir = TempDir::new().unwrap(); + let uri = dir + .path() + .join("g.generic.lance") + .to_string_lossy() + .to_string(); + let spec = SchemaSpec::new(vec![( + ID_COLUMN.to_string(), + ColumnSpec::required(ColumnType::String { large: false }), + )]); + { + let store = GenericStore::open( + &uri, + spec, + GenericStoreOptions { + seal_on_add: true, + ..Default::default() + }, + ) + .await + .unwrap(); + for i in 0..3 { + store + .add(&[serde_json::json!({"id": format!("r{i}")}) + .as_object() + .unwrap() + .clone()]) + .await + .unwrap(); + } + } + + let name = generic_target("g"); + let (row, skipped) = observe_generic(&name, &uri, None, GenericStoreOptions::default()) + .await + .unwrap(); + assert!(!skipped); + assert_eq!(row.name, "generic:g"); + assert_eq!(row.pending_wal_generations, 3); + assert_ne!(row.version, StatRow::UNKNOWN_VERSION); + + // Unchanged version: skipped like rollout, row reused. + let (again, skipped) = + observe_generic(&name, &uri, Some(&row), GenericStoreOptions::default()) + .await + .unwrap(); + assert!(skipped); + assert_eq!(again.pending_wal_generations, 3); + } + /// A write moves the base version, so the next scan must observe fully and /// pick up the new counts. #[tokio::test] diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index dc4f02e6..67a58c38 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -44,6 +44,35 @@ use crate::task_store::{TaskClaim, TaskKinds}; /// anything still over the threshold is picked up by the next tick. const MAX_SWEEP_ENQUEUE: usize = 256; +/// Task-target prefix marking a generic store. Store names match +/// `[A-Za-z0-9_][A-Za-z0-9._-]*`, so a `:` can never appear in a bare name +/// and the prefix is unambiguous. Carrying the kind in the target string -- +/// rather than adding a field to `TaskRecord` -- keeps the etcd task schema, +/// dedupe keys and per-target locks exactly as they are: a generic store's +/// MergeWal is locked and de-duped under `generic:`, so it can never +/// collide with a rollout experiment of the same name. +const GENERIC_TARGET_PREFIX: &str = "generic:"; + +/// Which store kind a task target names. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StoreKind { + Rollout, + Generic, +} + +/// Split a task target into its store kind and bare name. +fn parse_target(target: &str) -> (StoreKind, &str) { + match target.strip_prefix(GENERIC_TARGET_PREFIX) { + Some(name) => (StoreKind::Generic, name), + None => (StoreKind::Rollout, target), + } +} + +/// Build the task target for a generic store. +pub(crate) fn generic_target(name: &str) -> String { + format!("{GENERIC_TARGET_PREFIX}{name}") +} + fn kind_label(kind: TaskKind) -> &'static str { match kind { TaskKind::Compact => "compact", @@ -141,7 +170,14 @@ async fn run_task(state: &Arc, claim: TaskClaim, timing: TaskClaimT /// Compact one experiment. The task-store claim owns the per-experiment write /// lock for the full execution. -async fn run_compaction(state: &Arc, name: &str) -> Result { +async fn run_compaction(state: &Arc, target: &str) -> Result { + let (kind, name) = parse_target(target); + if kind != StoreKind::Rollout { + // Base-table compaction of generic stores is not scheduled by the + // master yet; only WAL merges are. Refuse rather than open the store + // through the rollout code path with the wrong URI and schema. + return Err(format!("compaction is not scheduled for {kind:?} stores")); + } compact_inner(state, name).await } @@ -149,7 +185,11 @@ async fn run_compaction(state: &Arc, name: &str) -> Result, name: &str) -> Result { +async fn run_index_id(state: &Arc, target: &str) -> Result { + let (kind, name) = parse_target(target); + if kind != StoreKind::Rollout { + return Err(format!("id indexing is not scheduled for {kind:?} stores")); + } index_id_inner(state, name).await } @@ -199,19 +239,25 @@ struct MergeWalReply { /// its own shard; a worker that owns no data for `name` reports 0 (or 404, which /// we tolerate). Succeeds if at least one endpoint responded; fails only when /// there are no endpoints or every one errored. -async fn run_merge_wal(state: &Arc, name: &str) -> Result { +async fn run_merge_wal(state: &Arc, target: &str) -> Result { let endpoints = &state.config.worker_endpoints; if endpoints.is_empty() { return Err("no worker endpoints configured (--worker-endpoints)".to_string()); } + let (kind, name) = parse_target(target); + let route = match kind { + StoreKind::Rollout => "api/v1/internal/merge-wal", + StoreKind::Generic => "api/v1/generic", + }; let calls = endpoints.iter().map(|ep| { let http = state.http.clone(); - let url = format!( - "{}/api/v1/internal/merge-wal/{}", - ep.trim_end_matches('/'), - name - ); + let url = match kind { + StoreKind::Rollout => format!("{}/{}/{}", ep.trim_end_matches('/'), route, name), + StoreKind::Generic => { + format!("{}/{}/{}/merge-wal", ep.trim_end_matches('/'), route, name) + } + }; async move { // Per-worker timing: `join_all` means the slowest worker sets the // whole task's latency, so without this one straggler is @@ -236,7 +282,18 @@ async fn run_merge_wal(state: &Arc, name: &str) -> Result) -> lance::Result| async move { + hits.fetch_add(1, Ordering::SeqCst); + Json(serde_json::json!({ "reclaimed": 1 })) + }) + }) + .route("/api/v1/generic/{name}/merge-wal", { + let hits = generic_hits.clone(); + post(move |Path(name): Path| async move { + assert_eq!( + name, "gs", + "the bare name reaches the worker, not the prefix" + ); + hits.fetch_add(1, Ordering::SeqCst); + Json(serde_json::json!({ "reclaimed": 4 })) + }) + }); + 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, app).await.unwrap() }); + + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.worker_endpoints = vec![format!("http://{addr}")]; + let state = MasterState::new(cfg).await.unwrap(); + let worker = spawn_scheduler(&state); + + let rec = enqueue(&state, TaskKind::MergeWal, &generic_target("gs")) + .await + .unwrap(); + let status = await_terminal(&state, &rec.id).await; + assert_eq!(status.state, TaskState::Done, "got {status:?}"); + assert!(status.detail.unwrap().contains("merged 4 generations")); + assert_eq!(generic_hits.load(Ordering::SeqCst), 1); + assert_eq!(rollout_hits.load(Ordering::SeqCst), 0); + worker.abort(); + } + + /// Workers are called one at a time. Every worker's merge commits to the + /// same base table and Lance's commit-conflict retry gives up after 30s, so + /// concurrent fan-out is N writers racing one commit point -- the failure + /// the master exists to prevent. + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn merge_wal_calls_workers_serially() { + use axum::{routing::post, Json, Router}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + // Each stub records the max number of in-flight calls it observed + // across the fleet through a shared counter. + let inflight = Arc::new(AtomicUsize::new(0)); + let max_inflight = Arc::new(AtomicUsize::new(0)); + async fn spawn_stub(inflight: Arc, max_inflight: Arc) -> String { + let app = Router::new().route( + "/api/v1/internal/merge-wal/{name}", + post(move || { + let inflight = inflight.clone(); + let max_inflight = max_inflight.clone(); + async move { + let now = inflight.fetch_add(1, Ordering::SeqCst) + 1; + max_inflight.fetch_max(now, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(150)).await; + inflight.fetch_sub(1, Ordering::SeqCst); + Json(serde_json::json!({ "reclaimed": 1 })) + } + }), + ); + 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, app).await.unwrap() }); + format!("http://{addr}") + } + + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.worker_endpoints = vec![ + spawn_stub(inflight.clone(), max_inflight.clone()).await, + spawn_stub(inflight.clone(), max_inflight.clone()).await, + spawn_stub(inflight.clone(), max_inflight.clone()).await, + ]; + let state = MasterState::new(cfg).await.unwrap(); + let worker = spawn_scheduler(&state); + + let rec = enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); + let status = await_terminal(&state, &rec.id).await; + assert_eq!(status.state, TaskState::Done, "got {status:?}"); + assert!(status.detail.unwrap().contains("3/3 workers")); + assert_eq!( + max_inflight.load(Ordering::SeqCst), + 1, + "fan-out must never have more than one worker merging at a time" + ); + worker.abort(); + } + + /// The compaction sweep must not enqueue generic rows: the master does not + /// compact generic stores, and `run_compaction` would otherwise open the + /// wrong URI through the rollout code path. + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn compaction_sweep_skips_generic_rows() { + use crate::stats_store::StatRow; + + let dir = TempDir::new().unwrap(); + let state = MasterState::new(config(&dir)).await.unwrap(); + let seed = |name: String, uri: String| StatRow { + version: StatRow::UNKNOWN_VERSION, + name, + uri, + row_count: 0, + fragment_count: 50, + last_updated: 0, + pending_wal_generations: 0, + last_compaction: StatRow::NO_COMPACTION, + total_compactions: 0, + scanned_at: 0, + }; + { + let mut stats = state.stats.lock().await; + stats + .upsert(&seed("r".to_string(), state.rollout_uri("r"))) + .await + .unwrap(); + stats + .upsert(&seed(generic_target("g"), state.generic_uri("g"))) + .await + .unwrap(); + } + let queued = sweep_candidates(&state).await.unwrap(); + assert_eq!(queued, 1, "only the rollout row is compacted"); + let tasks = state.task_store.list().await.unwrap(); + assert!(tasks.iter().all(|t| t.target == "r"), "{tasks:?}"); + } + /// A `Compact` runs even while every `MergeWal` slot is occupied by slow /// fan-outs and the queue is dominated by MergeWal tasks. /// @@ -1007,18 +1217,27 @@ mod tests { let mut stats = state.stats.lock().await; stats.upsert(&seed("hot", 5)).await.unwrap(); // >= threshold stats.upsert(&seed("cold", 1)).await.unwrap(); // < threshold + // A generic store over threshold is swept too, as a generic task. + let mut g = seed(&generic_target("ghot"), 7); + g.uri = state.generic_uri("ghot"); + stats.upsert(&g).await.unwrap(); } let queued = sweep_merge_wal_candidates(&state).await.unwrap(); - assert_eq!(queued, 1, "only the over-threshold experiment is swept"); + assert_eq!( + queued, 2, + "both over-threshold stores are swept, rollout and generic" + ); let tasks = state.task_store.list().await.unwrap(); let merge_tasks: Vec<_> = tasks .iter() .filter(|t| t.kind == TaskKind::MergeWal) .collect(); - assert_eq!(merge_tasks.len(), 1); - assert_eq!(merge_tasks[0].target, "hot"); + assert_eq!(merge_tasks.len(), 2); + let mut targets: Vec<&str> = merge_tasks.iter().map(|t| t.target.as_str()).collect(); + targets.sort(); + assert_eq!(targets, vec!["generic:ghot", "hot"]); // Second sweep must de-dupe against the still-queued MergeWal. sweep_merge_wal_candidates(&state).await.unwrap(); @@ -1030,7 +1249,7 @@ mod tests { .into_iter() .filter(|t| t.kind == TaskKind::MergeWal) .count(); - assert_eq!(merge_after, 1, "duplicate MergeWal is de-duped"); + assert_eq!(merge_after, 2, "duplicate MergeWal is de-duped"); } /// Minimal rollout record builder for tests (the core struct has no /// `Default`). diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index c21330fe..722256ea 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -4,7 +4,8 @@ use std::num::NonZeroUsize; use std::sync::Arc; use lance_context_core::{ - join_uri, CompactionConfig, RolloutRegistry, RolloutStore, RolloutStoreOptions, Session, + join_uri, CompactionConfig, GenericStoreOptions, RolloutRegistry, RolloutStore, + RolloutStoreOptions, Session, }; use lru::LruCache; use tokio::sync::{Mutex, RwLock, Semaphore}; @@ -60,6 +61,15 @@ fn build_compaction_config(config: &MasterConfig) -> CompactionConfig { pub struct MasterState { /// Durable directory of which rollout stores exist. pub registry: RwLock, + /// Durable directory of which generic stores exist. Same registry format + /// as rollout (the data plane writes `_registry.generic.lance` with the + /// same `RolloutRegistry` type); read here so the stats scan and the + /// WAL-merge sweep cover generic stores. Without this, generic stores' + /// MemWAL merges were left to each worker's own timer, so 20 shards raced + /// to commit into one base table and lost to `Too many concurrent writers` + /// while the master -- whose etcd-locked MergeWal task exists to serialize + /// exactly that -- never heard of them. + pub generic_registry: RwLock, /// Periodically-refreshed per-experiment metrics (master-owned). pub stats: Mutex, /// Last snapshot written to the stats table, kept in memory so @@ -130,8 +140,10 @@ impl MasterState { let rollout_session = build_rollout_session(config.rollout_cache_bytes); let compaction_concurrency = config.compaction_concurrency.max(1); let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); + let generic_registry_uri = join_uri(&base_uri, "_registry.generic.lance"); let stats_uri = join_uri(&base_uri, "_stats.rollout.lance"); let mut registry = RolloutRegistry::open_or_create(®istry_uri, None).await?; + let generic_registry = RolloutRegistry::open_or_create(&generic_registry_uri, None).await?; let backfilled = discovery::backfill_registry(&config.data_dir, &mut registry).await?; if backfilled > 0 { tracing::info!( @@ -143,6 +155,7 @@ impl MasterState { task_store.release_coordination_lock(init_guard).await?; let state = Arc::new(Self { registry: RwLock::new(registry), + generic_registry: RwLock::new(generic_registry), stats: Mutex::new(stats), stats_cache: RwLock::new(Arc::new(Vec::new())), record_stores: Mutex::new(LruCache::new( @@ -167,6 +180,21 @@ impl MasterState { join_uri(&self.base_uri, &format!("{}.rollout.lance", name)) } + /// Physical generic dataset URI for `name`, matching the data-plane's + /// `generic_uri` convention (`{name}.generic.lance`). + pub fn generic_uri(&self, name: &str) -> String { + join_uri(&self.base_uri, &format!("{}.generic.lance", name)) + } + + /// Options for every generic store opened by the master (read-only + /// observation; the master never writes a generic store's WAL). + pub(crate) fn generic_store_options(&self) -> GenericStoreOptions { + GenericStoreOptions { + session: self.rollout_session.clone(), + ..Default::default() + } + } + /// Options for every rollout store opened by the master. pub(crate) fn rollout_store_options(&self) -> RolloutStoreOptions { RolloutStoreOptions {