From 711989950009f462631d52a0b55b7504b1b9bed5 Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 16 Jul 2026 01:48:28 +0000 Subject: [PATCH] feat(master): add etcd scheduler backend --- .github/workflows/rust-test.yml | 25 + Cargo.lock | 127 ++ README.md | 2 +- crates/lance-context-master/Cargo.toml | 1 + crates/lance-context-master/src/config.rs | 58 +- crates/lance-context-master/src/main.rs | 2 +- crates/lance-context-master/src/routes.rs | 67 +- crates/lance-context-master/src/scanner.rs | 46 +- crates/lance-context-master/src/scheduler.rs | 293 ++--- crates/lance-context-master/src/state.rs | 210 +--- .../lance-context-master/src/stats_store.rs | 16 +- crates/lance-context-master/src/task_store.rs | 1108 ++++++++++++++++- deploy/kubernetes/README.md | 37 +- deploy/kubernetes/master-etcd.yaml | 97 ++ deploy/kubernetes/master.yaml | 2 + 15 files changed, 1626 insertions(+), 465 deletions(-) create mode 100644 deploy/kubernetes/master-etcd.yaml diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index 7430d2f..6e92b02 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -48,5 +48,30 @@ jobs: run: cargo test -p lance-context-core -p lance-context-master --no-run - name: Run unit tests run: cargo test -p lance-context-core -p lance-context-master --lib + - name: Start etcd for HA scheduler test + run: | + ETCD_VERSION=3.7.0 + curl -fsSL -o /tmp/etcd.tar.gz \ + "https://github.com/etcd-io/etcd/releases/download/v${ETCD_VERSION}/etcd-v${ETCD_VERSION}-linux-amd64.tar.gz" + mkdir -p /tmp/etcd + tar -xzf /tmp/etcd.tar.gz -C /tmp/etcd --strip-components=1 + nohup /tmp/etcd/etcd \ + --data-dir /tmp/etcd-data \ + --listen-client-urls http://127.0.0.1:2379 \ + --advertise-client-urls http://127.0.0.1:2379 \ + >/tmp/etcd.log 2>&1 & + for _ in $(seq 1 30); do + curl -fsS http://127.0.0.1:2379/health && exit 0 + sleep 1 + done + cat /tmp/etcd.log + exit 1 + - name: Run etcd scheduler integration test + env: + ETCD_TEST_ENDPOINTS: http://127.0.0.1:2379 + run: | + cargo test -p lance-context-master \ + task_store::tests::etcd_coordinates_dedupe_claims_and_target_locks \ + -- --ignored --exact - name: Run doc tests run: cargo test -p lance-context-core --doc diff --git a/Cargo.lock b/Cargo.lock index a9145c0..1f6e1a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3753,6 +3753,24 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "etcd-client" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef5da6e9a6ae89f4a91f80ba1caae45a5a924397a19947e18f5121a43285e9bc" +dependencies = [ + "http 1.4.1", + "prost 0.14.3", + "tokio", + "tokio-stream", + "tonic", + "tonic-build", + "tonic-prost", + "tonic-prost-build", + "tower", + "tower-service", +] + [[package]] name = "ethnum" version = "1.5.3" @@ -4676,6 +4694,19 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.10.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.6.0" @@ -5476,6 +5507,7 @@ dependencies = [ "axum", "chrono", "clap", + "etcd-client", "futures", "lance 7.0.0", "lance-context-api", @@ -8209,6 +8241,8 @@ dependencies = [ "prettyplease", "prost 0.14.3", "prost-types 0.14.3", + "pulldown-cmark", + "pulldown-cmark-to-cmark", "regex", "syn 2.0.117", "tempfile", @@ -8288,6 +8322,26 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "pulldown-cmark" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f068eba8e7071c5f9511831b44f32c740d5adf574e990f946ddb53db2f314e" +dependencies = [ + "bitflags 2.11.1", + "memchr", + "unicase", +] + +[[package]] +name = "pulldown-cmark-to-cmark" +version = "22.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50793def1b900256624a709439404384204a5dc3a6ec580281bfaac35e882e90" +dependencies = [ + "pulldown-cmark", +] + [[package]] name = "pyo3" version = "0.25.1" @@ -9243,6 +9297,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -10553,6 +10608,75 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "axum", + "base64", + "bytes", + "h2 0.4.14", + "http 1.4.1", + "http-body 1.0.1", + "http-body-util", + "hyper 1.10.1", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2 0.6.4", + "sync_wrapper", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-build" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322" +dependencies = [ + "prettyplease", + "proc-macro2", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost 0.14.3", + "tonic", +] + +[[package]] +name = "tonic-prost-build" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27" +dependencies = [ + "prettyplease", + "proc-macro2", + "prost-build 0.14.3", + "prost-types 0.14.3", + "quote", + "syn 2.0.117", + "tempfile", + "tonic-build", +] + [[package]] name = "tower" version = "0.5.3" @@ -10561,9 +10685,12 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.14.0", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", diff --git a/README.md b/README.md index e2a33e9..f5f13ef 100644 --- a/README.md +++ b/README.md @@ -138,7 +138,7 @@ crates/lance-context-client # HTTP client for the server crates/lance-context-master # Control plane, scheduler, and admin UI crates/lance-context # Re-export crate for downstream clients python/ # Python bindings (PyO3) + tests -deploy/kubernetes/ # Kubernetes master + RocksDB PVC example +deploy/kubernetes/ # Kubernetes master examples: RocksDB/PVC or etcd HA examples/ # Runnable example projects ``` diff --git a/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml index 4d340ed..ec93c5b 100644 --- a/crates/lance-context-master/Cargo.toml +++ b/crates/lance-context-master/Cargo.toml @@ -21,6 +21,7 @@ arrow-schema = "58" axum = { version = "0.8", features = ["json"] } chrono = { version = "0.4", default-features = false, features = ["clock"] } clap = { version = "4", features = ["derive", "env"] } +etcd-client = { version = "0.19", features = ["tls"] } futures = "0.3" lance = "7.0.0" metrics = "0.24" diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index 4028dfe..f8fa1dc 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -1,6 +1,15 @@ //! Command-line / environment configuration for the master control-plane. -use clap::Parser; +use clap::{Parser, ValueEnum}; + +/// Durable scheduler state backend. +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum TaskStoreBackend { + /// Single-master RocksDB on a dedicated local PVC. + Rocksdb, + /// Shared etcd state with lease-based task claims for HA masters. + Etcd, +} /// Control-plane (master) process for lance-context rollout stores. #[derive(Debug, Clone, Parser)] @@ -56,9 +65,20 @@ pub struct MasterConfig { #[arg(long, env = "TASK_CONCURRENCY", default_value_t = 4)] pub task_concurrency: usize, + /// Scheduler state backend. `rocksdb` requires exactly one master and a + /// dedicated PVC; `etcd` supports multiple stateless master replicas. + #[arg( + long, + env = "TASK_STORE_BACKEND", + value_enum, + default_value_t = TaskStoreBackend::Rocksdb + )] + pub task_store_backend: TaskStoreBackend, + /// Local RocksDB directory for durable scheduler task state. This must be /// backed by persistent local storage in production; it must not be an - /// object-store URI or a volume shared by multiple master processes. + /// object-store URI or a volume shared by multiple master processes. Used + /// only when `TASK_STORE_BACKEND=rocksdb`. #[arg( long, env = "TASK_DB_PATH", @@ -66,6 +86,40 @@ pub struct MasterConfig { )] pub task_db_path: String, + /// Comma-separated etcd v3 endpoints. Required when + /// `TASK_STORE_BACKEND=etcd`. + #[arg(long, env = "ETCD_ENDPOINTS", value_delimiter = ',')] + pub etcd_endpoints: Vec, + + /// Namespace for all lance-context master keys in etcd. + #[arg(long, env = "ETCD_PREFIX", default_value = "/lance-context/master")] + pub etcd_prefix: String, + + /// Optional etcd username. `ETCD_PASSWORD` must also be set. + #[arg(long, env = "ETCD_USERNAME")] + pub etcd_username: Option, + + /// Optional etcd password. `ETCD_USERNAME` must also be set. + #[arg(long, env = "ETCD_PASSWORD")] + pub etcd_password: Option, + + /// Optional PEM CA certificate path for etcd TLS. + #[arg(long, env = "ETCD_CA_CERT")] + pub etcd_ca_cert: Option, + + /// Optional PEM client certificate path for etcd mutual TLS. + #[arg(long, env = "ETCD_CLIENT_CERT")] + pub etcd_client_cert: Option, + + /// Optional PEM client private-key path for etcd mutual TLS. + #[arg(long, env = "ETCD_CLIENT_KEY")] + pub etcd_client_key: Option, + + /// TTL for etcd task claims and distributed locks. The master renews leases + /// while work is running; orphaned tasks are requeued after expiry. + #[arg(long, env = "ETCD_LEASE_TTL_SECS", default_value_t = 30)] + pub etcd_lease_ttl_secs: i64, + /// Maximum number of terminal scheduler tasks retained in RocksDB and the /// queue UI. Queued and running tasks are never pruned, and at least one /// terminal task is retained for status polling. diff --git a/crates/lance-context-master/src/main.rs b/crates/lance-context-master/src/main.rs index a13f37e..a46a6c0 100644 --- a/crates/lance-context-master/src/main.rs +++ b/crates/lance-context-master/src/main.rs @@ -48,7 +48,7 @@ async fn main() { // Background stats scanner (detached; stops when the last Arc drops). let _scanner = scanner::spawn_scanner(&state); - // Single serial compaction worker (+ optional periodic auto-sweep). + // Durable scheduler poller (+ optional coordinated auto-sweep). let _scheduler = scheduler::spawn_scheduler(&state); let mut app = Router::new().nest("/api/v1", routes::api_router()); diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 338f8cb..2314661 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -85,7 +85,7 @@ pub async fn list_experiments( Query(params): Query, ) -> Result, MasterError> { let search = params.search.as_deref().filter(|s| !s.is_empty()); - let stats = state.stats.lock().await; + let mut stats = state.stats.lock().await; let total = stats.count(search).await.map_err(MasterError::from_lance)?; let rows = stats .list(search, params.limit, params.offset) @@ -118,7 +118,7 @@ pub async fn get_experiment( .await .map_err(MasterError::from_lance)?; } - let stats = state.stats.lock().await; + let mut stats = state.stats.lock().await; match stats.get(&name).await.map_err(MasterError::from_lance)? { Some(row) => Ok(Json(ExperimentDetail { summary: row.into_summary(), @@ -245,15 +245,17 @@ fn task_to_compact_status(task: &TaskRecord) -> CompactJobStatus { } /// The most recent `Compact` task for `name`, if any (by `enqueued_at`). -async fn latest_compact_task(state: &Arc, name: &str) -> Option { - state - .tasks - .lock() - .await - .values() +async fn latest_compact_task( + state: &Arc, + name: &str, +) -> lance::Result> { + Ok(state + .task_store + .list() + .await? + .into_iter() .filter(|t| t.kind == TaskKind::Compact && t.target == name) - .max_by_key(|t| t.enqueued_at) - .cloned() + .max_by_key(|t| t.enqueued_at)) } /// `POST /api/v1/experiments/{name}/compact` — enqueue a manual compaction. @@ -288,11 +290,16 @@ pub async fn compact_experiment( pub async fn compact_status( State(state): State>, Path(name): Path, -) -> Json { - match latest_compact_task(&state, &name).await { - Some(task) => Json(task_to_compact_status(&task)), - None => Json(CompactJobStatus::None), - } +) -> Result, MasterError> { + Ok( + match latest_compact_task(&state, &name) + .await + .map_err(MasterError::from_lance)? + { + Some(task) => Json(task_to_compact_status(&task)), + None => Json(CompactJobStatus::None), + }, + ) } /// `POST /api/v1/tasks` — enqueue a task of any kind. Returns 202 Accepted with @@ -321,10 +328,16 @@ pub async fn enqueue_task( } /// `GET /api/v1/tasks` — all tasks (queue + recent history), newest first. -pub async fn list_tasks(State(state): State>) -> Json { - let mut tasks: Vec = state.tasks.lock().await.values().cloned().collect(); +pub async fn list_tasks( + State(state): State>, +) -> Result, MasterError> { + let mut tasks = state + .task_store + .list() + .await + .map_err(MasterError::from_lance)?; tasks.sort_by_key(|t| std::cmp::Reverse(t.enqueued_at)); - Json(TaskListResponse { tasks }) + Ok(Json(TaskListResponse { tasks })) } /// `GET /api/v1/tasks/{id}` — a single task by id. @@ -332,7 +345,12 @@ pub async fn get_task( State(state): State>, Path(id): Path, ) -> Result, MasterError> { - match state.tasks.lock().await.get(&id).cloned() { + match state + .task_store + .get(&id) + .await + .map_err(MasterError::from_lance)? + { Some(task) => Ok(Json(task)), None => Err(MasterError::NotFound(format!("task '{}' not found", id))), } @@ -404,7 +422,7 @@ fn blob_filename(record: &lance_context_core::RolloutRecord) -> String { #[cfg(test)] mod tests { use super::*; - use crate::config::MasterConfig; + use crate::config::{MasterConfig, TaskStoreBackend}; use chrono::Utc; use lance_context_core::{RolloutRecord, RolloutRegistry, RolloutStore}; use serde_json::json; @@ -422,11 +440,20 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, + task_store_backend: TaskStoreBackend::Rocksdb, task_db_path: dir .path() .join("master-tasks") .to_string_lossy() .to_string(), + etcd_endpoints: vec![], + etcd_prefix: "/test".to_string(), + etcd_username: None, + etcd_password: None, + etcd_ca_cert: None, + etcd_client_cert: None, + etcd_client_key: None, + etcd_lease_ttl_secs: 30, task_history_limit: 1_000, ui_dir: None, } diff --git a/crates/lance-context-master/src/scanner.rs b/crates/lance-context-master/src/scanner.rs index adf8bc8..024faf4 100644 --- a/crates/lance-context-master/src/scanner.rs +++ b/crates/lance-context-master/src/scanner.rs @@ -26,6 +26,34 @@ const OBSERVE_TIMEOUT: Duration = Duration::from_secs(30); /// for experiments no longer in the registry. Returns the number of /// experiments successfully observed. pub async fn scan_once(state: &Arc) -> lance::Result { + let guard = state.task_store.coordination_lock("stats-writer").await?; + let result = scan_once_inner(state).await; + let release = state.task_store.release_coordination_lock(guard).await; + match (result, release) { + (Ok(count), Ok(())) => Ok(count), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error), + } +} + +async fn try_scan_once(state: &Arc) -> lance::Result> { + let Some(guard) = state + .task_store + .try_coordination_lock("stats-writer") + .await? + else { + return Ok(None); + }; + let result = scan_once_inner(state).await; + let release = state.task_store.release_coordination_lock(guard).await; + match (result, release) { + (Ok(count), Ok(())) => Ok(Some(count)), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error), + } +} + +async fn scan_once_inner(state: &Arc) -> lance::Result { let scan_start = std::time::Instant::now(); let entries = state.registry.write().await.list().await?; let live: HashSet = entries.iter().map(|e| e.name.clone()).collect(); @@ -111,7 +139,7 @@ async fn observe_one(state: &Arc, name: &str, uri: &str) -> lance:: // Carry compaction counters forward across scans. let (last_compaction, total_compactions) = { - let stats = state.stats.lock().await; + let mut stats = state.stats.lock().await; match stats.get(name).await { Ok(Some(prev)) => (prev.last_compaction, prev.total_compactions), _ => (StatRow::NO_COMPACTION, 0), @@ -133,8 +161,13 @@ async fn observe_one(state: &Arc, name: &str, uri: &str) -> lance:: /// Refresh one experiment immediately and persist its new stats row. pub async fn refresh_one(state: &Arc, name: &str, uri: &str) -> lance::Result<()> { - let row = observe_one(state, name, uri).await?; - state.stats.lock().await.upsert(&row).await + let guard = state.task_store.coordination_lock("stats-writer").await?; + let result = match observe_one(state, name, uri).await { + Ok(row) => state.stats.lock().await.upsert(&row).await, + Err(error) => Err(error), + }; + let release = state.task_store.release_coordination_lock(guard).await; + result.and(release) } /// Spawn the periodic scanner. Returns `None` when the interval is `0`. @@ -148,8 +181,11 @@ pub fn spawn_scanner(state: &Arc) -> Option> { let mut ticker = tokio::time::interval(Duration::from_secs(interval_secs)); loop { ticker.tick().await; - match scan_once(&state).await { - Ok(n) => tracing::info!(experiments = n, "stats scan complete"), + match try_scan_once(&state).await { + Ok(Some(n)) => tracing::info!(experiments = n, "stats scan complete"), + Ok(None) => { + tracing::debug!("stats scan skipped; another master owns the writer lock") + } Err(e) => tracing::warn!(error = %e, "stats scan round failed"), } } diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index f155a40..effb28b 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -1,17 +1,16 @@ //! Unified task scheduler. //! -//! The master runs a single scheduler that drains one durable queue with -//! **bounded concurrency**. Task records are written to RocksDB before they -//! enter the in-memory dispatch channel; queued and interrupted tasks are -//! reloaded after a master restart. It executes three kinds of task +//! Each master polls one durable queue with **bounded concurrency**. RocksDB +//! supports one master; etcd uses atomic lease-backed claims so multiple +//! stateless masters can drain the same queue. It executes three kinds of task //! ([`TaskKind`]): //! //! - **Compact** — rewrites an experiment's base-table fragments. Two `Rewrite`s //! on the *same* dataset conflict in Lance's conflict matrix, so compaction of //! one experiment is serialized against itself (and against `IndexId`) via a -//! per-name in-flight gate ([`MasterState::inflight_dataset_writes`]). Distinct -//! experiments compact concurrently. `Rewrite` vs the data-plane's `Append` is -//! non-conflicting, so this runs safely alongside live ingest. +//! per-name task-store lock. Distinct experiments compact concurrently. +//! `Rewrite` vs the data-plane's `Append` is non-conflicting, so this runs +//! safely alongside live ingest. //! - **MergeWal** — folds flushed MemWAL generations back into the base table. //! The master cannot do this itself without fencing the live shard writer, so //! it fans out to every configured worker endpoint and each worker merges its @@ -19,7 +18,7 @@ //! - **IndexId** — builds a ZoneMap scalar index on the base table's `id` column //! (runs locally on the master). It commits a `CreateIndex`, which can conflict //! with a concurrent `Compact` `Rewrite` on the same dataset, so it shares the -//! per-name in-flight gate with `Compact`. +//! per-name task-store lock with `Compact`. //! //! Each task runs in its own `tokio::spawn`, so one task's failure never affects //! another. A global [`Semaphore`] bounds how many run at once. @@ -28,13 +27,14 @@ use std::sync::Arc; use std::time::Duration; use chrono::Utc; -use lance_context_api::{TaskKind, TaskRecord, TaskState}; -use lance_context_core::{generate_id, CompactionConfig, RolloutStore, RolloutStoreOptions}; +use lance_context_api::{TaskKind, TaskRecord}; +use lance_context_core::{CompactionConfig, RolloutStore, RolloutStoreOptions}; use tokio::sync::Semaphore; use tokio::task::JoinHandle; -use crate::state::{prune_terminal_history, MasterState}; +use crate::state::MasterState; use crate::stats_store::StatRow; +use crate::task_store::TaskClaim; /// Build the [`CompactionConfig`] the scheduler applies, from master config. pub fn compaction_config(state: &MasterState) -> CompactionConfig { @@ -46,11 +46,6 @@ pub fn compaction_config(state: &MasterState) -> CompactionConfig { } } -/// Current Unix-millisecond timestamp. -fn now_ms() -> i64 { - Utc::now().timestamp_millis() -} - fn kind_label(kind: TaskKind) -> &'static str { match kind { TaskKind::Compact => "compact", @@ -69,97 +64,14 @@ pub async fn enqueue( kind: TaskKind, target: &str, ) -> lance::Result { - let mut tasks = state.tasks.lock().await; - if matches!(kind, TaskKind::Compact | TaskKind::IndexId) { - if let Some(existing) = tasks.values().find(|t| { - t.kind == kind - && t.target == target - && matches!(t.state, TaskState::Queued | TaskState::Running) - }) { - return Ok(existing.clone()); - } - } - - let record = TaskRecord { - id: generate_id(), - kind, - target: target.to_string(), - state: TaskState::Queued, - error: None, - detail: None, - enqueued_at: now_ms(), - started_at: None, - finished_at: None, - }; - state.task_store.put(&record)?; - tasks.insert(record.id.clone(), record.clone()); - drop(tasks); - // Unbounded send only fails if the receiver was dropped (scheduler gone). - let _ = state.task_tx.send(record.id.clone()); + let record = state.task_store.enqueue(kind, target).await?; metrics::counter!("master_task_enqueued_total", "kind" => kind_label(kind)).increment(1); - metrics::gauge!("master_task_queue_depth").increment(1.0); Ok(record) } -/// Persist a task mutation before publishing it to the in-memory cache. -async fn update_task( - state: &Arc, - id: &str, - f: impl FnOnce(&mut TaskRecord), -) -> lance::Result<()> { - let mut tasks = state.tasks.lock().await; - let Some(current) = tasks.get(id) else { - return Ok(()); - }; - let mut updated = current.clone(); - f(&mut updated); - state.task_store.put(&updated)?; - let terminal = matches!(updated.state, TaskState::Done | TaskState::Failed); - tasks.insert(id.to_string(), updated); - if terminal { - prune_terminal_history( - &state.task_store, - &mut tasks, - state.config.task_history_limit, - )?; - } - Ok(()) -} - -/// Execute one task by id: dispatch on its kind, recording lifecycle timestamps -/// and the terminal state. -async fn run_task(state: &Arc, id: String) { - metrics::gauge!("master_task_queue_depth").decrement(1.0); - - let Some(task) = state.tasks.lock().await.get(&id).cloned() else { - return; // task vanished (should not happen) - }; - update_task(state, &id, |t| { - t.state = TaskState::Running; - t.started_at = Some(now_ms()); - }) - .await - .unwrap_or_else(|error| { - tracing::error!(task = %id, error = %error, "failed to persist task start"); - }); - - let task_started = state - .tasks - .lock() - .await - .get(&id) - .is_some_and(|task| task.state == TaskState::Running); - if !task_started { - metrics::gauge!("master_task_queue_depth").increment(1.0); - let retry_tx = state.task_tx.clone(); - let retry_id = id.clone(); - tokio::spawn(async move { - tokio::time::sleep(Duration::from_secs(5)).await; - let _ = retry_tx.send(retry_id); - }); - return; - } - +/// Execute one claimed task and atomically publish its terminal state. +async fn run_task(state: &Arc, claim: TaskClaim) { + let task = claim.task.clone(); let started = std::time::Instant::now(); let outcome = match task.kind { TaskKind::Compact => run_compaction(state, &task.target).await, @@ -173,46 +85,18 @@ async fn run_task(state: &Arc, id: String) { metrics::counter!("master_tasks_total", "kind" => kind_label(task.kind), "result" => result) .increment(1); - update_task(state, &id, |t| { - t.finished_at = Some(now_ms()); - match outcome { - Ok(detail) => { - t.state = TaskState::Done; - t.detail = Some(detail); - } - Err(error) => { - tracing::warn!(task = %id, target = %t.target, error = %error, "task failed"); - t.state = TaskState::Failed; - t.error = Some(error); - } - } - }) - .await - .unwrap_or_else(|error| { - tracing::error!(task = %id, error = %error, "failed to persist task completion"); - }); + if let Err(error) = &outcome { + tracing::warn!(task = %task.id, target = %task.target, error, "task failed"); + } + if let Err(error) = state.task_store.finish(claim, outcome).await { + tracing::error!(task = %task.id, error = %error, "failed to persist task completion"); + } } -/// Compact one experiment. Serialized per experiment via the shared base-table -/// write gate so two `Rewrite`s (or a `Rewrite` racing a `CreateIndex`) on the -/// same dataset never conflict. Returns a human-readable outcome summary on -/// success. +/// 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 { - // Per-name serial gate: refuse if another base-table write for this - // experiment is already running. (De-dup at enqueue makes this rare, but - // the gate is the real guarantee against conflicting commits.) - { - let mut inflight = state.inflight_dataset_writes.lock().await; - if !inflight.insert(name.to_string()) { - return Err(format!( - "a base-table write for '{name}' is already in progress" - )); - } - } - // Ensure the gate is released no matter how the inner call exits. - let result = compact_inner(state, name).await; - state.inflight_dataset_writes.lock().await.remove(name); - result + compact_inner(state, name).await } /// Build a ZoneMap scalar index on one experiment's `id` column. Shares the @@ -220,17 +104,7 @@ async fn run_compaction(state: &Arc, name: &str) -> Result, name: &str) -> Result { - { - let mut inflight = state.inflight_dataset_writes.lock().await; - if !inflight.insert(name.to_string()) { - return Err(format!( - "a base-table write for '{name}' is already in progress" - )); - } - } - let result = index_id_inner(state, name).await; - state.inflight_dataset_writes.lock().await.remove(name); - result + index_id_inner(state, name).await } async fn index_id_inner(state: &Arc, name: &str) -> Result { @@ -329,10 +203,18 @@ async fn run_merge_wal(state: &Arc, name: &str) -> Result, name: &str, store: &RolloutStore) { + let guard = match state.task_store.coordination_lock("stats-writer").await { + Ok(guard) => guard, + Err(e) => { + tracing::warn!(store = %name, error = %e, "stats writer lock failed"); + return; + } + }; let obs = match store.observe().await { Ok(obs) => obs, Err(e) => { tracing::warn!(store = %name, error = %e, "post-compaction observe failed"); + let _ = state.task_store.release_coordination_lock(guard).await; return; } }; @@ -355,12 +237,33 @@ async fn update_stats_after_compaction(state: &Arc, name: &str, sto if let Err(e) = stats.upsert(&row).await { tracing::warn!(store = %name, error = %e, "post-compaction stats upsert failed"); } + drop(stats); + if let Err(e) = state.task_store.release_coordination_lock(guard).await { + tracing::warn!(store = %name, error = %e, "stats writer unlock failed"); + } } /// Enqueue a `Compact` task for every experiment whose fragment count is at or /// above the configured threshold, honoring quiet hours. Reads candidates from /// the stats table. pub async fn sweep_candidates(state: &Arc) -> lance::Result { + let Some(guard) = state + .task_store + .try_coordination_lock("compaction-sweep") + .await? + else { + return Ok(0); + }; + let result = sweep_candidates_inner(state).await; + let release = state.task_store.release_coordination_lock(guard).await; + match (result, release) { + (Ok(count), Ok(())) => Ok(count), + (Err(error), _) => Err(error), + (Ok(_), Err(error)) => Err(error), + } +} + +async fn sweep_candidates_inner(state: &Arc) -> lance::Result { let config = compaction_config(state); // Quiet-hours gate applies to the whole sweep. if in_quiet_hours(&config) { @@ -389,18 +292,8 @@ fn in_quiet_hours(config: &CompactionConfig) -> bool { .any(|(start, end)| hour >= *start && hour < *end) } -/// Spawn the scheduler dispatcher plus (optionally) the periodic auto-sweep. The -/// dispatcher owns the queue receiver and runs up to `task_concurrency` tasks at -/// once, each in its own detached task. Returns the dispatcher handle. Panics if -/// called twice (receiver already taken). +/// Spawn the scheduler poller plus the optional periodic auto-sweep. pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { - let mut rx = state - .task_rx - .try_lock() - .expect("scheduler: state lock") - .take() - .expect("spawn_scheduler called more than once"); - // Optional periodic auto-sweep feeds the same queue. let interval_secs = state.config.compaction_interval_secs; if interval_secs > 0 { @@ -420,32 +313,35 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { } let concurrency = state.config.task_concurrency.max(1); - let queue_depth = state - .tasks - .try_lock() - .map(|tasks| { - tasks - .values() - .filter(|task| task.state == TaskState::Queued) - .count() - }) - .unwrap_or(0); - metrics::gauge!("master_task_queue_depth").set(queue_depth as f64); let sem = Arc::new(Semaphore::new(concurrency)); let dispatch_state = state.clone(); tokio::spawn(async move { - while let Some(id) = rx.recv().await { - // Acquire a slot before spawning so at most `concurrency` run at once. - let permit = sem - .clone() - .acquire_owned() - .await - .expect("semaphore never closed"); - let st = dispatch_state.clone(); - tokio::spawn(async move { - run_task(&st, id).await; - drop(permit); - }); + loop { + if let Ok(queued) = dispatch_state.task_store.queue_depth().await { + metrics::gauge!("master_task_queue_depth").set(queued as f64); + } + while sem.available_permits() > 0 { + match dispatch_state.task_store.claim_next().await { + Ok(Some(claim)) => { + let permit = sem + .clone() + .acquire_owned() + .await + .expect("semaphore never closed"); + let st = dispatch_state.clone(); + tokio::spawn(async move { + run_task(&st, claim).await; + drop(permit); + }); + } + Ok(None) => break, + Err(error) => { + tracing::warn!(error = %error, "scheduler queue poll failed"); + break; + } + } + } + tokio::time::sleep(Duration::from_millis(500)).await; } }) } @@ -453,7 +349,8 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { #[cfg(test)] mod tests { use super::*; - use crate::config::MasterConfig; + use crate::config::{MasterConfig, TaskStoreBackend}; + use lance_context_api::TaskState; use tempfile::TempDir; fn config(dir: &TempDir) -> MasterConfig { @@ -469,11 +366,20 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, + task_store_backend: TaskStoreBackend::Rocksdb, task_db_path: dir .path() .join("master-tasks") .to_string_lossy() .to_string(), + etcd_endpoints: vec![], + etcd_prefix: "/test".to_string(), + etcd_username: None, + etcd_password: None, + etcd_ca_cert: None, + etcd_client_cert: None, + etcd_client_key: None, + etcd_lease_ttl_secs: 30, task_history_limit: 1_000, ui_dir: None, } @@ -483,9 +389,9 @@ mod tests { async fn await_terminal(state: &Arc, id: &str) -> TaskRecord { for _ in 0..100 { tokio::time::sleep(Duration::from_millis(50)).await; - if let Some(t) = state.tasks.lock().await.get(id) { + if let Some(t) = state.task_store.get(id).await.unwrap() { if matches!(t.state, TaskState::Done | TaskState::Failed) { - return t.clone(); + return t; } } } @@ -528,6 +434,7 @@ mod tests { state .task_store .list() + .await .unwrap() .into_iter() .find(|task| task.id == rec.id) @@ -586,15 +493,7 @@ mod tests { let a = enqueue(&state, TaskKind::Compact, "x").await.unwrap(); let b = enqueue(&state, TaskKind::Compact, "x").await.unwrap(); assert_eq!(a.id, b.id, "second enqueue returns the same task"); - // Only one task in the map, one message on the queue. - assert_eq!(state.tasks.lock().await.len(), 1); - assert_eq!(state.task_store.list().unwrap().len(), 1); - let mut rx = state.task_rx.lock().await.take().unwrap(); - let mut n = 0; - while rx.try_recv().is_ok() { - n += 1; - } - assert_eq!(n, 1); + assert_eq!(state.task_store.list().await.unwrap().len(), 1); } /// A MergeWal task with no configured endpoints fails fast with a clear msg. diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index af96a97..456f66d 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -1,11 +1,9 @@ //! Shared master state: durable registry + stats table. -use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use lance_context_api::{TaskRecord, TaskState}; use lance_context_core::{join_uri, RolloutRegistry}; -use tokio::sync::{mpsc, Mutex, RwLock}; +use tokio::sync::{Mutex, RwLock}; use crate::config::MasterConfig; use crate::discovery; @@ -16,9 +14,9 @@ use crate::task_store::TaskStore; /// /// The data-plane owns steady-state registry writes (store create/delete); the /// master only adds missing legacy rows during startup discovery, then reads it -/// to enumerate experiments. The `stats_store` is owned and written exclusively -/// by the master. Both are wrapped in locks because their mutating methods take -/// `&mut self`. +/// to enumerate experiments. Master replicas coordinate stats-table writes +/// through the task store. The Lance handles are wrapped in locks because their +/// mutating methods take `&mut self`. pub struct MasterState { /// Durable directory of which rollout stores exist. pub registry: RwLock, @@ -28,31 +26,25 @@ pub struct MasterState { pub base_uri: String, /// Effective configuration. pub config: MasterConfig, - /// Enqueues a task id for the scheduler. Both automatic sweeps and manual - /// API triggers push here; the scheduler drains it with bounded concurrency - /// (see [`crate::scheduler`]). - pub task_tx: mpsc::UnboundedSender, - /// Receiver half, taken once by [`crate::scheduler::spawn_scheduler`]. - pub task_rx: Mutex>>, - /// Durable task records on the master-local RocksDB PVC. + /// Durable scheduler queue. RocksDB is single-master; etcd provides shared + /// CAS/lease semantics for stateless HA master replicas. pub task_store: TaskStore, - /// Every task (queued / running / terminal) keyed by task id, for the queue - /// API, de-duplication, and UI. This is a cache of `task_store`, rebuilt at - /// startup. Terminal history is bounded by `task_history_limit`. - pub tasks: Mutex>, - /// Experiment names with a base-table write currently executing — the - /// per-name serial gate. A `Compact` or `IndexId` task must not run while - /// another base-table-mutating task on the same experiment is in flight: - /// two `Rewrite`s (or a `Rewrite` and a `CreateIndex`) on one dataset can - /// conflict in Lance's commit matrix, forcing wasteful retries. - pub inflight_dataset_writes: Mutex>, /// Shared HTTP client for fanning `MergeWal` tasks out to worker endpoints. pub http: reqwest::Client, } impl MasterState { - /// Open the registry, stats dataset, and local durable task store. + /// Open the registry, stats dataset, and configured durable task store. pub async fn new(config: MasterConfig) -> lance::Result> { + let task_store = TaskStore::open(&config).await?; + // Serialize first-time registry/stats creation and legacy backfill in + // etcd mode. Followers wait briefly rather than racing Lance creates. + let init_guard = loop { + if let Some(guard) = task_store.try_coordination_lock("state-init").await? { + break guard; + } + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + }; let base_uri = config.data_dir.clone(); let registry_uri = join_uri(&base_uri, "_registry.rollout.lance"); let stats_uri = join_uri(&base_uri, "_stats.rollout.lance"); @@ -65,58 +57,15 @@ impl MasterState { ); } let stats = StatsStore::open_or_create(&stats_uri, None).await?; - let task_store = TaskStore::open(&config.task_db_path)?; - let mut tasks = task_store - .list()? - .into_iter() - .map(|task| (task.id.clone(), task)) - .collect::>(); - - let mut recovered_running = 0usize; - for task in tasks.values_mut() { - if task.state == TaskState::Running { - task.state = TaskState::Queued; - task.started_at = None; - task.finished_at = None; - task.error = None; - task.detail = None; - task_store.put(task)?; - recovered_running += 1; - } - } - let pruned = prune_terminal_history(&task_store, &mut tasks, config.task_history_limit)?; - let mut pending = tasks - .values() - .filter(|task| task.state == TaskState::Queued) - .map(|task| (task.enqueued_at, task.id.clone())) - .collect::>(); - pending.sort(); - let pending_ids = pending.into_iter().map(|(_, id)| id).collect::>(); - - let (task_tx, task_rx) = mpsc::unbounded_channel(); + task_store.release_coordination_lock(init_guard).await?; let state = Arc::new(Self { registry: RwLock::new(registry), stats: Mutex::new(stats), base_uri, config, - task_tx, - task_rx: Mutex::new(Some(task_rx)), task_store, - tasks: Mutex::new(tasks), - inflight_dataset_writes: Mutex::new(HashSet::new()), http: reqwest::Client::new(), }); - for id in &pending_ids { - let _ = state.task_tx.send(id.clone()); - } - if !pending_ids.is_empty() || recovered_running > 0 || pruned > 0 { - tracing::info!( - queued = pending_ids.len(), - recovered_running, - pruned, - "loaded durable scheduler tasks" - ); - } Ok(state) } @@ -127,41 +76,11 @@ impl MasterState { } } -pub(crate) fn prune_terminal_history( - task_store: &TaskStore, - tasks: &mut HashMap, - limit: usize, -) -> lance::Result { - // Keep at least the latest terminal task so legacy per-experiment status - // polling can observe completion even if the configured limit is zero. - let limit = limit.max(1); - let mut terminal = tasks - .values() - .filter(|task| matches!(task.state, TaskState::Done | TaskState::Failed)) - .map(|task| { - ( - task.finished_at.unwrap_or(task.enqueued_at), - task.id.clone(), - ) - }) - .collect::>(); - terminal.sort_by_key(|(finished_at, _)| std::cmp::Reverse(*finished_at)); - let ids = terminal - .into_iter() - .skip(limit) - .map(|(_, id)| id) - .collect::>(); - task_store.delete_many(ids.iter().map(String::as_str))?; - for id in &ids { - tasks.remove(id); - } - Ok(ids.len()) -} - #[cfg(test)] mod tests { use super::*; - use lance_context_api::TaskKind; + use crate::config::TaskStoreBackend; + use lance_context_api::{TaskKind, TaskState}; use lance_context_core::RolloutStore; use tempfile::TempDir; @@ -177,11 +96,20 @@ mod tests { target_rows_per_fragment: 1_048_576, worker_endpoints: vec![], task_concurrency: 4, + task_store_backend: TaskStoreBackend::Rocksdb, task_db_path: dir .path() .join("master-tasks") .to_string_lossy() .to_string(), + etcd_endpoints: vec![], + etcd_prefix: "/test".to_string(), + etcd_username: None, + etcd_password: None, + etcd_ca_cert: None, + etcd_client_cert: None, + etcd_client_key: None, + etcd_lease_ttl_secs: 30, task_history_limit: 1_000, ui_dir: None, } @@ -202,77 +130,23 @@ mod tests { } #[tokio::test] - async fn startup_requeues_queued_and_interrupted_tasks() { + async fn startup_requeues_interrupted_local_tasks() { let dir = TempDir::new().unwrap(); let config = test_config(&dir); - { - let state = MasterState::new(config.clone()).await.unwrap(); - for (id, task_state) in [ - ("queued-task", TaskState::Queued), - ("running-task", TaskState::Running), - ("done-task", TaskState::Done), - ] { - let task = TaskRecord { - id: id.to_string(), - kind: TaskKind::Compact, - target: "legacy".to_string(), - state: task_state, - error: None, - detail: None, - enqueued_at: 1, - started_at: (task_state == TaskState::Running).then_some(2), - finished_at: (task_state == TaskState::Done).then_some(3), - }; - state.task_store.put(&task).unwrap(); - } - } + let first = MasterState::new(config.clone()).await.unwrap(); + let task = first + .task_store + .enqueue(TaskKind::Compact, "legacy") + .await + .unwrap(); + let claim = first.task_store.claim_next().await.unwrap().unwrap(); + assert_eq!(claim.task.state, TaskState::Running); + drop(claim); + drop(first); let state = MasterState::new(config).await.unwrap(); - let tasks = state.tasks.lock().await; - assert_eq!(tasks["queued-task"].state, TaskState::Queued); - assert_eq!(tasks["running-task"].state, TaskState::Queued); - assert_eq!(tasks["running-task"].started_at, None); - assert_eq!(tasks["done-task"].state, TaskState::Done); - drop(tasks); - - let mut rx = state.task_rx.lock().await.take().unwrap(); - let mut recovered = Vec::new(); - while let Ok(id) = rx.try_recv() { - recovered.push(id); - } - recovered.sort(); - assert_eq!(recovered, ["queued-task", "running-task"]); - } - - #[test] - fn terminal_history_pruning_keeps_active_tasks() { - let dir = TempDir::new().unwrap(); - let store = TaskStore::open(dir.path().join("tasks")).unwrap(); - let mut tasks = HashMap::new(); - for (id, state, timestamp) in [ - ("queued", TaskState::Queued, 1), - ("old", TaskState::Done, 2), - ("new", TaskState::Failed, 3), - ] { - let task = TaskRecord { - id: id.to_string(), - kind: TaskKind::Compact, - target: "experiment".to_string(), - state, - error: None, - detail: None, - enqueued_at: timestamp, - started_at: None, - finished_at: matches!(state, TaskState::Done | TaskState::Failed) - .then_some(timestamp), - }; - store.put(&task).unwrap(); - tasks.insert(id.to_string(), task); - } - - assert_eq!(prune_terminal_history(&store, &mut tasks, 1).unwrap(), 1); - assert!(tasks.contains_key("queued")); - assert!(tasks.contains_key("new")); - assert!(!tasks.contains_key("old")); + let recovered = state.task_store.get(&task.id).await.unwrap().unwrap(); + assert_eq!(recovered.state, TaskState::Queued); + assert_eq!(recovered.started_at, None); } } diff --git a/crates/lance-context-master/src/stats_store.rs b/crates/lance-context-master/src/stats_store.rs index a1238bf..a8b27e9 100644 --- a/crates/lance-context-master/src/stats_store.rs +++ b/crates/lance-context-master/src/stats_store.rs @@ -7,7 +7,8 @@ //! the UI list endpoint can render row/fragment counts without opening every //! dataset on each request. //! -//! The master is the sole writer of this table; the data-plane never touches it. +//! Master replicas coordinate a single writer through the configured task-store +//! backend; the data-plane never touches this table. use std::collections::HashMap; use std::sync::Arc; @@ -170,6 +171,7 @@ impl StatsStore { /// Insert or replace the stats row for `row.name` (delete-then-append, /// idempotent). Callers must serialize mutations. pub async fn upsert(&mut self, row: &StatRow) -> LanceResult<()> { + self.dataset.checkout_latest().await?; self.delete_row(&row.name).await?; let batch = Self::row_to_batch(row)?; let schema = Arc::new(stats_schema()); @@ -181,6 +183,7 @@ impl StatsStore { /// Remove the stats row for `name`, if present. No-op when absent. pub async fn remove(&mut self, name: &str) -> LanceResult<()> { + self.dataset.checkout_latest().await?; self.delete_row(name).await } @@ -193,7 +196,8 @@ impl StatsStore { } /// Fetch a single experiment's stats row by name. - pub async fn get(&self, name: &str) -> LanceResult> { + pub async fn get(&mut self, name: &str) -> LanceResult> { + self.dataset.checkout_latest().await?; let escaped = name.replace('\'', "''"); let mut scanner = self.dataset.scan(); scanner.filter(&format!("name = '{}'", escaped))?; @@ -210,7 +214,8 @@ impl StatsStore { /// Total number of rows matching an optional case-sensitive substring on /// `name`, ignoring pagination. - pub async fn count(&self, search: Option<&str>) -> LanceResult { + pub async fn count(&mut self, search: Option<&str>) -> LanceResult { + self.dataset.checkout_latest().await?; let mut scanner = self.dataset.scan(); scanner.project(&["name"])?; if let Some(q) = search { @@ -227,11 +232,12 @@ impl StatsStore { /// List stats rows, optionally filtered by a `name` substring, ordered by /// name, with `limit`/`offset` pagination applied in memory. pub async fn list( - &self, + &mut self, search: Option<&str>, limit: usize, offset: usize, ) -> LanceResult> { + self.dataset.checkout_latest().await?; let mut scanner = self.dataset.scan(); if let Some(q) = search { scanner.filter(&Self::like_filter(q))?; @@ -384,7 +390,7 @@ mod tests { let mut s = new_store(&dir).await; s.upsert(&sample("persist", 7)).await.unwrap(); } - let s = new_store(&dir).await; + let mut s = new_store(&dir).await; assert_eq!(s.get("persist").await.unwrap().unwrap().row_count, 7); } diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index b60926b..0b99e3b 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -1,23 +1,266 @@ -//! Durable scheduler task storage backed by RocksDB. +//! Durable scheduler task storage. +//! +//! RocksDB is the single-master backend. etcd adds compare-and-swap enqueue, +//! lease-backed task claims, and distributed target locks so several stateless +//! masters can safely drain one shared queue. +use std::collections::HashSet; use std::path::Path; +use std::sync::Arc; +use std::time::Duration; -use lance_context_api::TaskRecord; +use chrono::Utc; +use etcd_client::{ + Certificate, Client, Compare, CompareOp, ConnectOptions, GetOptions, Identity, PutOptions, + TlsOptions, Txn, TxnOp, +}; +use lance_context_api::{TaskKind, TaskRecord, TaskState}; +use lance_context_core::generate_id; use rocksdb::{Direction, IteratorMode, Options, WriteBatch, WriteOptions, DB}; +use tokio::sync::{oneshot, Mutex}; +use tokio::task::JoinHandle; + +use crate::config::{MasterConfig, TaskStoreBackend}; const TASK_KEY_PREFIX: &[u8] = b"task/"; +const TASK_POLL_BATCH: usize = 256; -/// Single-process RocksDB store for scheduler task records. -/// -/// Every write uses RocksDB's WAL with `sync=true`: the scheduler never -/// acknowledges an enqueue or lifecycle transition before it is durable on -/// the master PVC. +#[derive(Clone)] pub struct TaskStore { + inner: Arc, + history_limit: usize, +} + +enum TaskStoreInner { + Rocks(Box), + Etcd(Box), +} + +struct RocksTaskStore { db: DB, + operation_lock: Mutex<()>, + claimed_tasks: Mutex>, + claimed_targets: Mutex>, +} + +struct EtcdTaskStore { + client: Client, + prefix: String, + lease_ttl: i64, +} + +/// Ownership of one running task. Dropping an etcd claim stops lease renewal; +/// etcd then removes its claim and target-lock keys, allowing recovery. +pub struct TaskClaim { + pub task: TaskRecord, + backend: ClaimBackend, +} + +enum ClaimBackend { + Rocks { + target_locked: bool, + }, + Etcd { + token: String, + lease_id: i64, + claim_key: String, + target_key: Option, + keepalive: LeaseKeepalive, + }, +} + +/// Short-lived distributed guard used to serialize stats-table writers. +pub struct CoordinationGuard { + backend: GuardBackend, +} + +enum GuardBackend { + Rocks, + Etcd { + token: String, + lease_id: i64, + key: String, + keepalive: LeaseKeepalive, + }, +} + +struct LeaseKeepalive { + stop: Option>, + task: JoinHandle<()>, +} + +impl Drop for LeaseKeepalive { + fn drop(&mut self) { + if let Some(stop) = self.stop.take() { + let _ = stop.send(()); + } + self.task.abort(); + } } impl TaskStore { - pub fn open(path: impl AsRef) -> lance::Result { + pub async fn open(config: &MasterConfig) -> lance::Result { + let inner = match config.task_store_backend { + TaskStoreBackend::Rocksdb => { + TaskStoreInner::Rocks(Box::new(RocksTaskStore::open(&config.task_db_path)?)) + } + TaskStoreBackend::Etcd => { + TaskStoreInner::Etcd(Box::new(EtcdTaskStore::connect(config).await?)) + } + }; + let store = Self { + inner: Arc::new(inner), + history_limit: config.task_history_limit.max(1), + }; + store.recover_orphaned().await?; + store.prune_terminal_history().await?; + Ok(store) + } + + #[must_use] + pub fn is_distributed(&self) -> bool { + matches!(self.inner.as_ref(), TaskStoreInner::Etcd(_)) + } + + /// Atomically enqueue a task. Compact and IndexId use an etcd/local dedupe + /// key while queued or running; MergeWal always creates a new task. + pub async fn enqueue(&self, kind: TaskKind, target: &str) -> lance::Result { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.enqueue(kind, target).await, + TaskStoreInner::Etcd(store) => store.enqueue(kind, target).await, + } + } + + pub async fn list(&self) -> lance::Result> { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.list(), + TaskStoreInner::Etcd(store) => store.list().await, + } + } + + pub async fn get(&self, id: &str) -> lance::Result> { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.get(id), + TaskStoreInner::Etcd(store) => store.get(id).await, + } + } + + pub async fn queue_depth(&self) -> lance::Result { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => Ok(store + .list()? + .into_iter() + .filter(|task| task.state == TaskState::Queued) + .count()), + TaskStoreInner::Etcd(store) => store.queue_depth().await, + } + } + + /// Claim the oldest runnable task. In etcd mode the queued->running update, + /// claim lease, and per-experiment write lock are one transaction. + pub async fn claim_next(&self) -> lance::Result> { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.claim_next().await, + TaskStoreInner::Etcd(store) => { + store.recover_orphaned().await?; + store.claim_next().await + } + } + } + + /// Finish a claimed task and release its claim/target lock. + pub async fn finish( + &self, + claim: TaskClaim, + outcome: Result, + ) -> lance::Result<()> { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.finish(claim, outcome).await?, + TaskStoreInner::Etcd(store) => store.finish(claim, outcome).await?, + } + self.prune_terminal_history().await.map(|_| ()) + } + + /// Try to acquire a named coordination lock without waiting. This is used + /// around the shared Lance stats table, whose mutations must have one writer + /// across all master replicas. + pub async fn try_coordination_lock( + &self, + name: &str, + ) -> lance::Result> { + match self.inner.as_ref() { + TaskStoreInner::Rocks(_) => Ok(Some(CoordinationGuard { + backend: GuardBackend::Rocks, + })), + TaskStoreInner::Etcd(store) => store.try_coordination_lock(name).await, + } + } + + pub async fn coordination_lock(&self, name: &str) -> lance::Result { + loop { + if let Some(guard) = self.try_coordination_lock(name).await? { + return Ok(guard); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + pub async fn release_coordination_lock(&self, guard: CoordinationGuard) -> lance::Result<()> { + match (self.inner.as_ref(), guard.backend) { + (TaskStoreInner::Rocks(_), GuardBackend::Rocks) => Ok(()), + ( + TaskStoreInner::Etcd(store), + GuardBackend::Etcd { + token, + lease_id, + key, + keepalive, + }, + ) => { + drop(keepalive); + store.delete_owned_key(&key, &token).await?; + store.revoke_lease(lease_id).await + } + _ => Err(lance::Error::io( + "coordination guard belongs to a different task backend", + )), + } + } + + async fn recover_orphaned(&self) -> lance::Result { + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.recover_running().await, + TaskStoreInner::Etcd(store) => store.recover_orphaned().await, + } + } + + async fn prune_terminal_history(&self) -> lance::Result { + let mut terminal = self + .list() + .await? + .into_iter() + .filter(|task| matches!(task.state, TaskState::Done | TaskState::Failed)) + .collect::>(); + terminal + .sort_by_key(|task| std::cmp::Reverse(task.finished_at.unwrap_or(task.enqueued_at))); + let ids = terminal + .into_iter() + .skip(self.history_limit) + .map(|task| task.id) + .collect::>(); + if ids.is_empty() { + return Ok(0); + } + match self.inner.as_ref() { + TaskStoreInner::Rocks(store) => store.delete_many(ids.iter().map(String::as_str))?, + TaskStoreInner::Etcd(store) => store.delete_many(&ids).await?, + } + Ok(ids.len()) + } +} + +impl RocksTaskStore { + fn open(path: impl AsRef) -> lance::Result { let path = path.as_ref(); let text = path.to_string_lossy(); if text.contains("://") { @@ -36,7 +279,6 @@ impl TaskStore { )) })?; } - let mut options = Options::default(); options.create_if_missing(true); options.set_max_open_files(64); @@ -47,19 +289,46 @@ impl TaskStore { path.display() )) })?; - Ok(Self { db }) + Ok(Self { + db, + operation_lock: Mutex::new(()), + claimed_tasks: Mutex::new(HashSet::new()), + claimed_targets: Mutex::new(HashSet::new()), + }) } - pub fn put(&self, task: &TaskRecord) -> lance::Result<()> { - let value = serde_json::to_vec(task).map_err(|err| { - lance::Error::io(format!("failed to encode task '{}': {err}", task.id)) - })?; + async fn enqueue(&self, kind: TaskKind, target: &str) -> lance::Result { + let _operation = self.operation_lock.lock().await; + if requires_dedupe(kind) { + if let Some(existing) = self.list()?.into_iter().find(|task| { + task.kind == kind + && task.target == target + && matches!(task.state, TaskState::Queued | TaskState::Running) + }) { + return Ok(existing); + } + } + let task = new_task(kind, target); + self.put(&task)?; + Ok(task) + } + + fn put(&self, task: &TaskRecord) -> lance::Result<()> { + let value = encode_task(task)?; self.db .put_opt(task_key(&task.id), value, &sync_write_options()) .map_err(|err| lance::Error::io(format!("failed to persist task '{}': {err}", task.id))) } - pub fn list(&self) -> lance::Result> { + fn get(&self, id: &str) -> lance::Result> { + self.db + .get(task_key(id)) + .map_err(|err| lance::Error::io(format!("failed to read task '{id}': {err}")))? + .map(|value| decode_task(&value, id)) + .transpose() + } + + fn list(&self) -> lance::Result> { let mut tasks = Vec::new(); for item in self .db @@ -70,20 +339,73 @@ impl TaskStore { if !key.starts_with(TASK_KEY_PREFIX) { break; } - let task = serde_json::from_slice(&value).map_err(|err| { - lance::Error::io(format!( - "failed to decode task record '{}': {err}", - String::from_utf8_lossy(&key) - )) - })?; - tasks.push(task); + tasks.push(decode_task(&value, &String::from_utf8_lossy(&key))?); } Ok(tasks) } - pub fn delete_many<'a>(&self, ids: impl IntoIterator) -> lance::Result<()> { + async fn claim_next(&self) -> lance::Result> { + let _operation = self.operation_lock.lock().await; + let claimed_tasks = self.claimed_tasks.lock().await; + let mut claimed_targets = self.claimed_targets.lock().await; + let mut queued = self + .list()? + .into_iter() + .filter(|task| task.state == TaskState::Queued && !claimed_tasks.contains(&task.id)) + .collect::>(); + queued.sort_by_key(|task| task.enqueued_at); + let Some(mut task) = queued.into_iter().find(|task| { + !requires_target_lock(task.kind) || !claimed_targets.contains(&task.target) + }) else { + return Ok(None); + }; + let target_locked = requires_target_lock(task.kind); + if target_locked { + claimed_targets.insert(task.target.clone()); + } + drop(claimed_targets); + drop(claimed_tasks); + self.claimed_tasks.lock().await.insert(task.id.clone()); + task.state = TaskState::Running; + task.started_at = Some(now_ms()); + self.put(&task)?; + Ok(Some(TaskClaim { + task, + backend: ClaimBackend::Rocks { target_locked }, + })) + } + + async fn finish(&self, claim: TaskClaim, outcome: Result) -> lance::Result<()> { + let ClaimBackend::Rocks { target_locked } = claim.backend else { + return Err(lance::Error::io("task claim belongs to etcd")); + }; + let _operation = self.operation_lock.lock().await; + let mut task = claim.task; + apply_outcome(&mut task, outcome); + self.put(&task)?; + self.claimed_tasks.lock().await.remove(&task.id); + if target_locked { + self.claimed_targets.lock().await.remove(&task.target); + } + Ok(()) + } + + async fn recover_running(&self) -> lance::Result { + let _operation = self.operation_lock.lock().await; + let mut recovered = 0; + for mut task in self.list()? { + if task.state == TaskState::Running { + requeue(&mut task); + self.put(&task)?; + recovered += 1; + } + } + Ok(recovered) + } + + fn delete_many<'a>(&self, ids: impl IntoIterator) -> lance::Result<()> { let mut batch = WriteBatch::default(); - let mut count = 0usize; + let mut count = 0; for id in ids { batch.delete(task_key(id)); count += 1; @@ -97,6 +419,566 @@ impl TaskStore { } } +impl EtcdTaskStore { + async fn connect(config: &MasterConfig) -> lance::Result { + if config.etcd_endpoints.is_empty() { + return Err(lance::Error::io( + "ETCD_ENDPOINTS is required when TASK_STORE_BACKEND=etcd", + )); + } + if config.etcd_lease_ttl_secs < 5 { + return Err(lance::Error::io("ETCD_LEASE_TTL_SECS must be at least 5")); + } + let mut options = ConnectOptions::new() + .with_connect_timeout(Duration::from_secs(5)) + .with_timeout(Duration::from_secs(10)) + .with_keep_alive(Duration::from_secs(10), Duration::from_secs(3)) + .with_require_leader(true); + match (&config.etcd_username, &config.etcd_password) { + (Some(username), Some(password)) => { + options = options.with_user(username, password); + } + (None, None) => {} + _ => { + return Err(lance::Error::io( + "ETCD_USERNAME and ETCD_PASSWORD must be configured together", + )) + } + } + if let Some(path) = &config.etcd_ca_cert { + let pem = std::fs::read(path).map_err(|err| { + lance::Error::io(format!("failed to read ETCD_CA_CERT '{path}': {err}")) + })?; + let mut tls = TlsOptions::new().ca_certificate(Certificate::from_pem(pem)); + match (&config.etcd_client_cert, &config.etcd_client_key) { + (Some(cert), Some(key)) => { + let cert_pem = std::fs::read(cert).map_err(|err| { + lance::Error::io(format!("failed to read ETCD_CLIENT_CERT '{cert}': {err}")) + })?; + let key_pem = std::fs::read(key).map_err(|err| { + lance::Error::io(format!("failed to read ETCD_CLIENT_KEY '{key}': {err}")) + })?; + tls = tls.identity(Identity::from_pem(cert_pem, key_pem)); + } + (None, None) => {} + _ => { + return Err(lance::Error::io( + "ETCD_CLIENT_CERT and ETCD_CLIENT_KEY must be configured together", + )) + } + } + options = options.with_tls(tls); + } else if config.etcd_client_cert.is_some() || config.etcd_client_key.is_some() { + return Err(lance::Error::io( + "ETCD_CA_CERT is required when configuring an etcd client certificate", + )); + } + let client = Client::connect(config.etcd_endpoints.clone(), Some(options)) + .await + .map_err(etcd_error("connect to etcd"))?; + Ok(Self { + client, + prefix: config.etcd_prefix.trim_end_matches('/').to_string(), + lease_ttl: config.etcd_lease_ttl_secs, + }) + } + + async fn enqueue(&self, kind: TaskKind, target: &str) -> lance::Result { + let task = new_task(kind, target); + let task_key = self.task_key(&task.id); + let queue_key = self.queue_key(&task.id); + let value = encode_task(&task)?; + let mut client = self.client.clone(); + if let Some(dedupe_key) = self.dedupe_key(kind, target) { + for _ in 0..4 { + let txn = Txn::new() + .when([Compare::version(dedupe_key.as_str(), CompareOp::Equal, 0)]) + .and_then([ + TxnOp::put(task_key.as_str(), value.clone(), None), + TxnOp::put(queue_key.as_str(), value.clone(), None), + TxnOp::put(dedupe_key.as_str(), task.id.as_bytes(), None), + ]); + if client + .txn(txn) + .await + .map_err(etcd_error("enqueue task"))? + .succeeded() + { + return Ok(task); + } + if let Some(existing_id) = self.get_text(&dedupe_key).await? { + if let Some(existing) = self.get(&existing_id).await? { + if matches!(existing.state, TaskState::Queued | TaskState::Running) { + return Ok(existing); + } + } + self.delete_owned_key(&dedupe_key, &existing_id).await?; + } + } + return Err(lance::Error::io( + "failed to resolve concurrent etcd task enqueue", + )); + } + client + .txn(Txn::new().and_then([ + TxnOp::put(task_key, value.clone(), None), + TxnOp::put(queue_key, value, None), + ])) + .await + .map_err(etcd_error("enqueue task"))?; + Ok(task) + } + + async fn get(&self, id: &str) -> lance::Result> { + let mut client = self.client.clone(); + let response = client + .get(self.task_key(id), None) + .await + .map_err(etcd_error("read task"))?; + response + .kvs() + .first() + .map(|kv| decode_task(kv.value(), id)) + .transpose() + } + + async fn list(&self) -> lance::Result> { + let mut client = self.client.clone(); + let response = client + .get(self.tasks_prefix(), Some(GetOptions::new().with_prefix())) + .await + .map_err(etcd_error("list tasks"))?; + response + .kvs() + .iter() + .map(|kv| decode_task(kv.value(), &String::from_utf8_lossy(kv.key()))) + .collect() + } + + async fn queue_depth(&self) -> lance::Result { + let mut client = self.client.clone(); + let response = client + .get( + self.queue_prefix(), + Some(GetOptions::new().with_prefix().with_count_only()), + ) + .await + .map_err(etcd_error("count queued tasks"))?; + usize::try_from(response.count()) + .map_err(|_| lance::Error::io("etcd returned an invalid queue count")) + } + + async fn claim_next(&self) -> lance::Result> { + let mut client = self.client.clone(); + let response = client + .get( + self.queue_prefix(), + Some( + GetOptions::new() + .with_prefix() + .with_limit(TASK_POLL_BATCH as i64), + ), + ) + .await + .map_err(etcd_error("list queued tasks"))?; + let queued = response + .kvs() + .iter() + .map(|kv| decode_task(kv.value(), &String::from_utf8_lossy(kv.key()))) + .collect::>>()?; + + for mut task in queued { + let token = generate_id(); + let lease_id = self.grant_lease().await?; + let claim_key = self.claim_key(&task.id); + let target_key = + requires_target_lock(task.kind).then(|| self.target_lock_key(&task.target)); + let queue_key = self.queue_key(&task.id); + let running_key = self.running_key(&task.id); + let queued_value = encode_task(&task)?; + task.state = TaskState::Running; + task.started_at = Some(now_ms()); + let running_value = encode_task(&task)?; + let mut compares = vec![ + Compare::value(self.task_key(&task.id), CompareOp::Equal, queued_value), + Compare::version(claim_key.as_str(), CompareOp::Equal, 0), + Compare::version(queue_key.as_str(), CompareOp::Greater, 0), + ]; + if let Some(key) = &target_key { + compares.push(Compare::version(key.as_str(), CompareOp::Equal, 0)); + } + let lease_options = Some(PutOptions::new().with_lease(lease_id)); + let mut operations = vec![ + TxnOp::put(self.task_key(&task.id), running_value.clone(), None), + TxnOp::delete(queue_key, None), + TxnOp::put(running_key, running_value, None), + TxnOp::put(claim_key.as_str(), token.as_bytes(), lease_options.clone()), + ]; + if let Some(key) = &target_key { + operations.push(TxnOp::put( + key.as_str(), + token.as_bytes(), + lease_options.clone(), + )); + } + let mut client = self.client.clone(); + let claimed = client + .txn(Txn::new().when(compares).and_then(operations)) + .await + .map_err(etcd_error("claim task"))? + .succeeded(); + if claimed { + let keepalive = self.start_keepalive(lease_id).await?; + return Ok(Some(TaskClaim { + task, + backend: ClaimBackend::Etcd { + token, + lease_id, + claim_key, + target_key, + keepalive, + }, + })); + } + self.revoke_lease(lease_id).await?; + } + Ok(None) + } + + async fn finish(&self, claim: TaskClaim, outcome: Result) -> lance::Result<()> { + let ClaimBackend::Etcd { + token, + lease_id, + claim_key, + target_key, + keepalive, + } = claim.backend + else { + return Err(lance::Error::io("task claim belongs to RocksDB")); + }; + let mut task = claim.task; + apply_outcome(&mut task, outcome); + let mut operations = vec![ + TxnOp::put(self.task_key(&task.id), encode_task(&task)?, None), + TxnOp::delete(claim_key.as_str(), None), + TxnOp::delete(self.running_key(&task.id), None), + ]; + if let Some(key) = &target_key { + operations.push(TxnOp::delete(key.as_str(), None)); + } + if let Some(key) = self.dedupe_key(task.kind, &task.target) { + operations.push(TxnOp::delete(key, None)); + } + let mut client = self.client.clone(); + let completed = client + .txn( + Txn::new() + .when([Compare::value( + claim_key.as_str(), + CompareOp::Equal, + token.as_bytes(), + )]) + .and_then(operations), + ) + .await + .map_err(etcd_error("complete task"))? + .succeeded(); + drop(keepalive); + self.revoke_lease(lease_id).await?; + if !completed { + return Err(lance::Error::io(format!( + "task '{}' lost its etcd claim before completion", + task.id + ))); + } + Ok(()) + } + + async fn recover_orphaned(&self) -> lance::Result { + let mut client = self.client.clone(); + let response = client + .get(self.running_prefix(), Some(GetOptions::new().with_prefix())) + .await + .map_err(etcd_error("list running tasks"))?; + let running = response + .kvs() + .iter() + .map(|kv| decode_task(kv.value(), &String::from_utf8_lossy(kv.key()))) + .collect::>>()?; + let mut recovered = 0; + for mut task in running { + let running_value = encode_task(&task)?; + let claim_key = self.claim_key(&task.id); + let running_key = self.running_key(&task.id); + requeue(&mut task); + let txn = Txn::new() + .when([ + Compare::value(self.task_key(&task.id), CompareOp::Equal, running_value), + Compare::version(claim_key.as_str(), CompareOp::Equal, 0), + Compare::version(running_key.as_str(), CompareOp::Greater, 0), + ]) + .and_then([ + TxnOp::put(self.task_key(&task.id), encode_task(&task)?, None), + TxnOp::put(self.queue_key(&task.id), encode_task(&task)?, None), + TxnOp::delete(running_key, None), + ]); + let mut client = self.client.clone(); + if client + .txn(txn) + .await + .map_err(etcd_error("recover orphaned task"))? + .succeeded() + { + recovered += 1; + } + } + Ok(recovered) + } + + async fn try_coordination_lock(&self, name: &str) -> lance::Result> { + let lease_id = self.grant_lease().await?; + let token = generate_id(); + let key = format!("{}/coordination/{}", self.prefix, encode_segment(name)); + let txn = Txn::new() + .when([Compare::version(key.as_str(), CompareOp::Equal, 0)]) + .and_then([TxnOp::put( + key.as_str(), + token.as_bytes(), + Some(PutOptions::new().with_lease(lease_id)), + )]); + let mut client = self.client.clone(); + if !client + .txn(txn) + .await + .map_err(etcd_error("acquire coordination lock"))? + .succeeded() + { + self.revoke_lease(lease_id).await?; + return Ok(None); + } + let keepalive = self.start_keepalive(lease_id).await?; + Ok(Some(CoordinationGuard { + backend: GuardBackend::Etcd { + token, + lease_id, + key, + keepalive, + }, + })) + } + + async fn grant_lease(&self) -> lance::Result { + let mut client = self.client.clone(); + client + .lease_grant(self.lease_ttl, None) + .await + .map(|response| response.id()) + .map_err(etcd_error("grant lease")) + } + + async fn start_keepalive(&self, lease_id: i64) -> lance::Result { + let mut client = self.client.clone(); + let (mut keeper, mut stream) = client + .lease_keep_alive(lease_id) + .await + .map_err(etcd_error("start lease keepalive"))?; + let interval = Duration::from_secs((self.lease_ttl / 3).max(1) as u64); + let (stop_tx, mut stop_rx) = oneshot::channel(); + let task = tokio::spawn(async move { + loop { + tokio::select! { + _ = &mut stop_rx => break, + _ = tokio::time::sleep(interval) => { + if keeper.keep_alive().await.is_err() { + break; + } + match tokio::time::timeout(interval, stream.message()).await { + Ok(Ok(Some(response))) if response.ttl() > 0 => {} + _ => break, + } + } + } + } + }); + Ok(LeaseKeepalive { + stop: Some(stop_tx), + task, + }) + } + + async fn revoke_lease(&self, lease_id: i64) -> lance::Result<()> { + let mut client = self.client.clone(); + client + .lease_revoke(lease_id) + .await + .map(|_| ()) + .map_err(etcd_error("revoke lease")) + } + + async fn delete_owned_key(&self, key: &str, owner: &str) -> lance::Result<()> { + let txn = Txn::new() + .when([Compare::value(key, CompareOp::Equal, owner.as_bytes())]) + .and_then([TxnOp::delete(key, None)]); + let mut client = self.client.clone(); + client + .txn(txn) + .await + .map(|_| ()) + .map_err(etcd_error("release owned key")) + } + + async fn get_text(&self, key: &str) -> lance::Result> { + let mut client = self.client.clone(); + let response = client + .get(key, None) + .await + .map_err(etcd_error("read etcd key"))?; + response + .kvs() + .first() + .map(|kv| { + std::str::from_utf8(kv.value()) + .map(str::to_string) + .map_err(|err| lance::Error::io(format!("invalid UTF-8 in etcd key: {err}"))) + }) + .transpose() + } + + async fn delete_many(&self, ids: &[String]) -> lance::Result<()> { + for chunk in ids.chunks(100) { + let operations = chunk + .iter() + .map(|id| TxnOp::delete(self.task_key(id), None)) + .collect::>(); + let mut client = self.client.clone(); + client + .txn(Txn::new().and_then(operations)) + .await + .map_err(etcd_error("prune task history"))?; + } + Ok(()) + } + + fn tasks_prefix(&self) -> String { + format!("{}/tasks/", self.prefix) + } + + fn queue_prefix(&self) -> String { + format!("{}/queue/", self.prefix) + } + + fn running_prefix(&self) -> String { + format!("{}/running/", self.prefix) + } + + fn task_key(&self, id: &str) -> String { + format!("{}{id}", self.tasks_prefix()) + } + + fn queue_key(&self, id: &str) -> String { + format!("{}{id}", self.queue_prefix()) + } + + fn running_key(&self, id: &str) -> String { + format!("{}{id}", self.running_prefix()) + } + + fn claim_key(&self, id: &str) -> String { + format!("{}/claims/{id}", self.prefix) + } + + fn target_lock_key(&self, target: &str) -> String { + format!("{}/target-locks/{}", self.prefix, encode_segment(target)) + } + + fn dedupe_key(&self, kind: TaskKind, target: &str) -> Option { + requires_dedupe(kind).then(|| { + format!( + "{}/dedupe/{}/{}", + self.prefix, + kind_label(kind), + encode_segment(target) + ) + }) + } +} + +fn new_task(kind: TaskKind, target: &str) -> TaskRecord { + TaskRecord { + id: generate_id(), + kind, + target: target.to_string(), + state: TaskState::Queued, + error: None, + detail: None, + enqueued_at: now_ms(), + started_at: None, + finished_at: None, + } +} + +fn apply_outcome(task: &mut TaskRecord, outcome: Result) { + task.finished_at = Some(now_ms()); + match outcome { + Ok(detail) => { + task.state = TaskState::Done; + task.detail = Some(detail); + task.error = None; + } + Err(error) => { + task.state = TaskState::Failed; + task.error = Some(error); + task.detail = None; + } + } +} + +fn requeue(task: &mut TaskRecord) { + task.state = TaskState::Queued; + task.started_at = None; + task.finished_at = None; + task.error = None; + task.detail = None; +} + +fn requires_dedupe(kind: TaskKind) -> bool { + matches!(kind, TaskKind::Compact | TaskKind::IndexId) +} + +fn requires_target_lock(kind: TaskKind) -> bool { + matches!(kind, TaskKind::Compact | TaskKind::IndexId) +} + +fn kind_label(kind: TaskKind) -> &'static str { + match kind { + TaskKind::Compact => "compact", + TaskKind::MergeWal => "merge-wal", + TaskKind::IndexId => "index-id", + } +} + +fn now_ms() -> i64 { + Utc::now().timestamp_millis() +} + +fn encode_segment(value: &str) -> String { + value + .as_bytes() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +fn encode_task(task: &TaskRecord) -> lance::Result> { + serde_json::to_vec(task) + .map_err(|err| lance::Error::io(format!("failed to encode task '{}': {err}", task.id))) +} + +fn decode_task(value: &[u8], key: &str) -> lance::Result { + serde_json::from_slice(value) + .map_err(|err| lance::Error::io(format!("failed to decode task '{key}': {err}"))) +} + fn task_key(id: &str) -> Vec { let mut key = Vec::with_capacity(TASK_KEY_PREFIX.len() + id.len()); key.extend_from_slice(TASK_KEY_PREFIX); @@ -110,59 +992,163 @@ fn sync_write_options() -> WriteOptions { options } +fn etcd_error(action: &'static str) -> impl FnOnce(etcd_client::Error) -> lance::Error { + move |err| lance::Error::io(format!("failed to {action}: {err}")) +} + #[cfg(test)] mod tests { use super::*; - use lance_context_api::{TaskKind, TaskState}; use tempfile::TempDir; - fn task(id: &str, state: TaskState, enqueued_at: i64) -> TaskRecord { - TaskRecord { - id: id.to_string(), - kind: TaskKind::Compact, - target: "experiment".to_string(), - state, - error: None, - detail: None, - enqueued_at, - started_at: None, - finished_at: None, + fn config(dir: &TempDir) -> MasterConfig { + MasterConfig { + data_dir: dir.path().to_string_lossy().to_string(), + host: "127.0.0.1".to_string(), + port: 0, + stats_scan_interval_secs: 0, + scan_concurrency: 4, + compaction_interval_secs: 0, + min_fragments: 16, + target_rows_per_fragment: 1_048_576, + worker_endpoints: vec![], + task_concurrency: 4, + task_store_backend: TaskStoreBackend::Rocksdb, + task_db_path: dir.path().join("tasks").to_string_lossy().to_string(), + etcd_endpoints: vec![], + etcd_prefix: "/test".to_string(), + etcd_username: None, + etcd_password: None, + etcd_ca_cert: None, + etcd_client_cert: None, + etcd_client_key: None, + etcd_lease_ttl_secs: 30, + task_history_limit: 1_000, + ui_dir: None, } } - #[test] - fn survives_reopen_and_deletes_in_batch() { + #[tokio::test] + async fn rocksdb_survives_reopen_and_dedupes() { let dir = TempDir::new().unwrap(); - let path = dir.path().join("tasks"); - { - let store = TaskStore::open(&path).unwrap(); - store.put(&task("a", TaskState::Queued, 1)).unwrap(); - store.put(&task("b", TaskState::Done, 2)).unwrap(); - } + let cfg = config(&dir); + let first = TaskStore::open(&cfg).await.unwrap(); + let a = first + .enqueue(TaskKind::Compact, "experiment") + .await + .unwrap(); + let b = first + .enqueue(TaskKind::Compact, "experiment") + .await + .unwrap(); + assert_eq!(a.id, b.id); + drop(first); - let store = TaskStore::open(&path).unwrap(); - let mut tasks = store.list().unwrap(); - tasks.sort_by_key(|task| task.enqueued_at); - assert_eq!( - tasks - .iter() - .map(|task| task.id.as_str()) - .collect::>(), - ["a", "b"] - ); + let reopened = TaskStore::open(&cfg).await.unwrap(); + assert_eq!(reopened.get(&a.id).await.unwrap().unwrap(), a); + } - store.delete_many(["a"]).unwrap(); - let tasks = store.list().unwrap(); - assert_eq!(tasks.len(), 1); - assert_eq!(tasks[0].id, "b"); + #[tokio::test] + async fn rocksdb_claims_and_finishes() { + let dir = TempDir::new().unwrap(); + let store = TaskStore::open(&config(&dir)).await.unwrap(); + let task = store + .enqueue(TaskKind::IndexId, "experiment") + .await + .unwrap(); + let claim = store.claim_next().await.unwrap().unwrap(); + assert_eq!(claim.task.id, task.id); + assert_eq!(claim.task.state, TaskState::Running); + store + .finish(claim, Ok("indexed".to_string())) + .await + .unwrap(); + let finished = store.get(&task.id).await.unwrap().unwrap(); + assert_eq!(finished.state, TaskState::Done); + assert_eq!(finished.detail.as_deref(), Some("indexed")); } - #[test] - fn rejects_object_store_uri() { - let error = match TaskStore::open("s3://bucket/tasks") { + #[tokio::test] + async fn rejects_object_store_path_for_rocksdb() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.task_db_path = "s3://bucket/tasks".to_string(); + let error = match TaskStore::open(&cfg).await { Ok(_) => panic!("object-store URI should be rejected"), Err(error) => error, }; assert!(error.to_string().contains("must be local")); } + + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn etcd_coordinates_dedupe_claims_and_target_locks() { + 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.task_store_backend = TaskStoreBackend::Etcd; + 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 first = TaskStore::open(&cfg).await.unwrap(); + let second = TaskStore::open(&cfg).await.unwrap(); + let (a, b) = tokio::join!( + first.enqueue(TaskKind::Compact, "experiment"), + second.enqueue(TaskKind::Compact, "experiment") + ); + let a = a.unwrap(); + let b = b.unwrap(); + assert_eq!(a.id, b.id, "concurrent enqueue must dedupe"); + + let (claim_a, claim_b) = tokio::join!(first.claim_next(), second.claim_next()); + let mut claims = [claim_a.unwrap(), claim_b.unwrap()] + .into_iter() + .flatten() + .collect::>(); + assert_eq!(claims.len(), 1, "only one master may claim a task"); + let compact_claim = claims.pop().unwrap(); + + first + .enqueue(TaskKind::IndexId, "experiment") + .await + .unwrap(); + assert!( + second.claim_next().await.unwrap().is_none(), + "Compact and IndexId must share the experiment write lock" + ); + + first + .finish(compact_claim, Ok("compacted".to_string())) + .await + .unwrap(); + let index_claim = second.claim_next().await.unwrap().unwrap(); + assert_eq!(index_claim.task.kind, TaskKind::IndexId); + second + .finish(index_claim, Ok("indexed".to_string())) + .await + .unwrap(); + + let orphan = first.enqueue(TaskKind::MergeWal, "orphan").await.unwrap(); + let abandoned = first.claim_next().await.unwrap().unwrap(); + assert_eq!(abandoned.task.id, orphan.id); + drop(abandoned); + tokio::time::sleep(Duration::from_secs(6)).await; + let recovered = second.claim_next().await.unwrap().unwrap(); + assert_eq!(recovered.task.id, orphan.id); + second + .finish(recovered, Ok("recovered".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(); + } } diff --git a/deploy/kubernetes/README.md b/deploy/kubernetes/README.md index 01450a2..d9f350e 100644 --- a/deploy/kubernetes/README.md +++ b/deploy/kubernetes/README.md @@ -1,7 +1,11 @@ # Kubernetes deployment -`master.yaml` deploys one `lance-context-master` process with a dedicated -ReadWriteOnce PVC for its RocksDB scheduler state. +Two scheduler persistence modes are supported: + +- `master.yaml`: one master with scheduler state in RocksDB on a dedicated + ReadWriteOnce PVC. +- `master-etcd.yaml`: three stateless master replicas sharing scheduler state + through a separate etcd cluster. Before applying it: @@ -10,8 +14,31 @@ Before applying it: 2. Set object-store credentials and `WORKER_ENDPOINTS` for your environment. 3. Replace the example image name and `UI_DIR` with paths from the published master image. -4. Set `storageClassName` on the PVC when the cluster has no default class. +4. For RocksDB, set `storageClassName` on the PVC when the cluster has no + default class. +5. For etcd, replace `ETCD_ENDPOINTS` and create the referenced auth/TLS + Secrets. Remove the auth variables or TLS volume/configuration when the + application etcd cluster does not use them. + +Do not apply both manifests at once. -Keep the master at one replica. RocksDB is an embedded single-process database; -the PVC must not be mounted by multiple active master pods. The `Recreate` +## RocksDB mode + +Keep the master at one replica. RocksDB is an embedded single-process database, +and its PVC must not be mounted by multiple active master pods. The `Recreate` strategy ensures an old pod releases the volume before its replacement starts. + +## etcd mode + +Use a separate application etcd cluster, not the Kubernetes control-plane +etcd. Each etcd member owns its own persistent volume; master pods have no PVC. + +Task enqueue de-duplication, queued-to-running claims, and per-experiment +Compact/IndexId locks use etcd transactions. Claims and locks are attached to +renewed leases. If a master disappears, another replica requeues the task after +the lease expires. Execution is at-least-once, so task implementations must +remain idempotent across crash recovery. + +The `_stats.rollout.lance` table remains in `DATA_DIR`. etcd coordinates its +single-writer sections, while readers reload the latest Lance manifest so every +master replica sees current stats. diff --git a/deploy/kubernetes/master-etcd.yaml b/deploy/kubernetes/master-etcd.yaml new file mode 100644 index 0000000..4c1e52c --- /dev/null +++ b/deploy/kubernetes/master-etcd.yaml @@ -0,0 +1,97 @@ +apiVersion: v1 +kind: Service +metadata: + name: lance-context-master +spec: + selector: + app: lance-context-master + ports: + - name: http + port: 8090 + targetPort: http +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: lance-context-master +spec: + replicas: 3 + selector: + matchLabels: + app: lance-context-master + template: + metadata: + labels: + app: lance-context-master + spec: + containers: + - name: master + image: lance-context-master:latest + imagePullPolicy: IfNotPresent + env: + # Replace with the same shared object-store prefix used by workers. + - name: DATA_DIR + value: s3://replace-me/rollouts + - name: TASK_STORE_BACKEND + value: etcd + # Point this at a separate application etcd cluster, not the + # Kubernetes control-plane etcd. + - name: ETCD_ENDPOINTS + value: https://etcd-0.etcd:2379,https://etcd-1.etcd:2379,https://etcd-2.etcd:2379 + - name: ETCD_PREFIX + value: /lance-context/master + - name: ETCD_LEASE_TTL_SECS + value: "30" + - name: ETCD_USERNAME + valueFrom: + secretKeyRef: + name: lance-context-etcd + key: username + optional: true + - name: ETCD_PASSWORD + valueFrom: + secretKeyRef: + name: lance-context-etcd + key: password + optional: true + - name: ETCD_CA_CERT + value: /etc/lance-context/etcd/ca.crt + - name: ETCD_CLIENT_CERT + value: /etc/lance-context/etcd/tls.crt + - name: ETCD_CLIENT_KEY + value: /etc/lance-context/etcd/tls.key + - name: TASK_HISTORY_LIMIT + value: "1000" + - name: MASTER_PORT + value: "8090" + - name: UI_DIR + value: /app/ui + ports: + - name: http + containerPort: 8090 + readinessProbe: + httpGet: + path: /metrics + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /metrics + port: http + initialDelaySeconds: 15 + periodSeconds: 20 + resources: + requests: + cpu: 250m + memory: 512Mi + limits: + memory: 2Gi + volumeMounts: + - name: etcd-tls + mountPath: /etc/lance-context/etcd + readOnly: true + volumes: + - name: etcd-tls + secret: + secretName: lance-context-etcd-client-tls diff --git a/deploy/kubernetes/master.yaml b/deploy/kubernetes/master.yaml index 165ca10..d3ea7eb 100644 --- a/deploy/kubernetes/master.yaml +++ b/deploy/kubernetes/master.yaml @@ -48,6 +48,8 @@ spec: # Replace with the same shared object-store prefix used by workers. - name: DATA_DIR value: s3://replace-me/rollouts + - name: TASK_STORE_BACKEND + value: rocksdb - name: TASK_DB_PATH value: /var/lib/lance-context-master/tasks.rocksdb - name: TASK_HISTORY_LIMIT