From c504de30b473ac41b308ae04744a15f7f95b50bd Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Thu, 1 Oct 2026 17:37:52 +0000 Subject: [PATCH 1/9] fix: own merge execution and persist bounded shard recovery --- .github/workflows/rust-test.yml | 4 +- Cargo.lock | 20 + Cargo.toml | 1 + crates/lance-context-core/Cargo.toml | 4 + crates/lance-context-core/src/lib.rs | 1 + .../src/merge_write_scope.rs | 238 +++++++ .../lance-context-core/src/rollout_store.rs | 120 ++++ crates/lance-context-core/src/store_base.rs | 54 +- crates/lance-context-master/Cargo.toml | 2 + crates/lance-context-master/src/config.rs | 6 +- crates/lance-context-master/src/lib.rs | 1 + .../src/merge_execution.rs | 563 ++++++++++++++++ crates/lance-context-master/src/routes.rs | 21 + crates/lance-context-master/src/scheduler.rs | 424 ++++++------- crates/lance-context-master/src/task_store.rs | 259 +++++++- crates/lance-context-merge/Cargo.toml | 17 + crates/lance-context-merge/src/failure.rs | 277 ++++++++ crates/lance-context-merge/src/lib.rs | 600 ++++++++++++++++++ crates/lance-context-server/Cargo.toml | 2 + crates/lance-context-server/src/config.rs | 9 + crates/lance-context-server/src/main.rs | 2 + .../src/merge_execution.rs | 362 +++++++++++ .../src/routes/generic.rs | 12 + crates/lance-context-server/src/routes/mod.rs | 12 + .../src/routes/rollouts.rs | 12 + crates/lance-context-server/src/state.rs | 30 + docs/merge-recovery.md | 104 +++ 27 files changed, 2880 insertions(+), 277 deletions(-) create mode 100644 crates/lance-context-core/src/merge_write_scope.rs create mode 100644 crates/lance-context-master/src/merge_execution.rs create mode 100644 crates/lance-context-merge/Cargo.toml create mode 100644 crates/lance-context-merge/src/failure.rs create mode 100644 crates/lance-context-merge/src/lib.rs create mode 100644 crates/lance-context-server/src/merge_execution.rs create mode 100644 docs/merge-recovery.md diff --git a/.github/workflows/rust-test.yml b/.github/workflows/rust-test.yml index 696fd81..d25bb6a 100644 --- a/.github/workflows/rust-test.yml +++ b/.github/workflows/rust-test.yml @@ -81,7 +81,7 @@ jobs: - name: Run etcd-backed HA test suite env: ETCD_TEST_ENDPOINTS: http://127.0.0.1:2379 - run: cargo test -p lance-context-master --lib -- --ignored + run: cargo test -p lance-context-master -p lance-context-server -p lance-context-merge -- --ignored coverage: runs-on: ubuntu-24.04 @@ -131,7 +131,7 @@ jobs: ETCD_TEST_ENDPOINTS: http://127.0.0.1:2379 run: | cargo llvm-cov --no-report --all-features \ - -p lance-context-master --lib -- --ignored + -p lance-context-master -p lance-context-server -p lance-context-merge -- --ignored - name: Merge coverage reports run: cargo llvm-cov report --lcov --output-path lcov.info - name: Upload coverage to Codecov diff --git a/Cargo.lock b/Cargo.lock index ae51867..0ed20d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6349,6 +6349,7 @@ dependencies = [ "arrow-json 58.3.0", "arrow-schema 58.3.0", "arrow-select 58.3.0", + "async-trait", "base64", "chrono", "datafusion 54.1.0", @@ -6357,10 +6358,13 @@ dependencies = [ "lance-context-api", "lance-graph", "lance-index 9.0.0", + "lance-io 9.0.0", "lance-namespace 9.0.0", + "lance-table 9.0.0", "lancedb", "metrics", "metrics-util", + "object_store 0.13.2", "serde", "serde_json", "tempfile", @@ -6384,6 +6388,7 @@ dependencies = [ "lance 9.0.0", "lance-context-api", "lance-context-core", + "lance-context-merge", "lance-context-metrics", "lru", "metrics", @@ -6392,11 +6397,24 @@ dependencies = [ "serde_json", "tempfile", "tokio", + "tower", "tower-http", "tracing", "tracing-subscriber", ] +[[package]] +name = "lance-context-merge" +version = "0.1.0" +dependencies = [ + "clap", + "etcd-client", + "serde", + "serde_json", + "tokio", + "uuid", +] + [[package]] name = "lance-context-metrics" version = "0.1.0" @@ -6433,9 +6451,11 @@ dependencies = [ "bytes", "chrono", "clap", + "etcd-client", "futures", "lance-context-api", "lance-context-core", + "lance-context-merge", "lance-context-metrics", "lru", "metrics", diff --git a/Cargo.toml b/Cargo.toml index 2efab19..2cb7893 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = [ + "crates/lance-context-merge", "crates/lance-context-core", "crates/lance-context-metrics", "crates/lance-context", diff --git a/crates/lance-context-core/Cargo.toml b/crates/lance-context-core/Cargo.toml index a5d0f25..ef6be65 100644 --- a/crates/lance-context-core/Cargo.toml +++ b/crates/lance-context-core/Cargo.toml @@ -18,6 +18,10 @@ default = ["metrics"] metrics = ["dep:metrics"] [dependencies] +async-trait = "0.1" +lance-table = "9.0.0" +lance-io = "9.0.0" +object_store = "0.13.2" base64 = "0.22" arrow-array = "58" arrow-ipc = "58" diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index a46587b..c9666d4 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -11,6 +11,7 @@ pub mod generic_codec; mod generic_store; mod id; pub mod merge_budget; +pub mod merge_write_scope; pub mod metrics; mod namespace; mod record; diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs new file mode 100644 index 0000000..061ce0e --- /dev/null +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -0,0 +1,238 @@ +//! Cancel merge preparation without abandoning a manifest commit already sent +//! to storage. The executor must drain this scope before acknowledging terminal. +use futures::stream::BoxStream; +use lance::{Error, Result}; +use lance_io::object_store::ObjectStore; +use lance_table::{ + format::{IndexMetadata, Manifest, Transaction}, + io::commit::{ + CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, ManifestWriter, + }, +}; +use object_store::path::Path; +use std::{ + future::Future, + sync::{Arc, Mutex}, +}; +use tokio::sync::{oneshot, Notify}; + +tokio::task_local! { static CURRENT: Arc; } + +#[derive(Debug, Default)] +struct Progress { + closed: bool, + active: usize, +} + +#[derive(Debug, Default)] +pub struct MergeWriteScope { + progress: Mutex, + changed: Notify, +} + +impl MergeWriteScope { + pub fn new() -> Arc { + Arc::new(Self::default()) + } + + pub async fn run(self: &Arc, future: F) -> F::Output { + CURRENT.scope(self.clone(), future).await + } + + /// Call only after dropping/joining the merge future. Prevent new commits + /// and join all manifest writes that started before cancellation. + pub async fn drain(&self) { + self.progress.lock().unwrap().closed = true; + loop { + let changed = self.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if self.progress.lock().unwrap().active == 0 { + return; + } + changed.await; + } + } +} + +struct Active(Arc); +impl Drop for Active { + fn drop(&mut self) { + self.0.progress.lock().unwrap().active -= 1; + self.0.changed.notify_waiters(); + } +} + +pub(crate) async fn shield(write: F) -> std::result::Result +where + F: Future> + Send + 'static, + T: Send + 'static, + E: From + Send + 'static, +{ + let Ok(scope) = CURRENT.try_with(Arc::clone) else { + return write.await; + }; + { + let mut progress = scope.progress.lock().unwrap(); + if progress.closed { + return Err(Error::io("merge cancelled before manifest commit").into()); + } + progress.active += 1; + } + let active = Active(scope); + let (send, receive) = oneshot::channel(); + tokio::spawn(async move { + let _active = active; + let result = write.await; + let _ = send.send(result); + }); + receive + .await + .map_err(|_| E::from(Error::io("manifest commit executor panicked")))? +} + +#[derive(Debug)] +pub(crate) struct GuardedCommit(pub Arc); + +#[async_trait::async_trait] +#[allow(clippy::too_many_arguments)] +impl CommitHandler for GuardedCommit { + async fn commit( + &self, + manifest: &mut Manifest, + indices: Option>, + base_path: &Path, + object_store: &ObjectStore, + manifest_writer: ManifestWriter, + naming_scheme: ManifestNamingScheme, + transaction: Option, + ) -> std::result::Result { + let inner = self.0.clone(); + let mut owned_manifest = manifest.clone(); + let path = base_path.clone(); + let store = object_store.clone(); + let (location, committed) = shield(async move { + let location = inner + .commit( + &mut owned_manifest, + indices, + &path, + &store, + manifest_writer, + naming_scheme, + transaction, + ) + .await; + Ok::<_, CommitError>((location, owned_manifest)) + }) + .await?; + *manifest = committed; + location + } + + async fn resolve_latest_location( + &self, + base_path: &Path, + object_store: &ObjectStore, + ) -> Result { + self.0 + .resolve_latest_location(base_path, object_store) + .await + } + async fn resolve_version_location( + &self, + base_path: &Path, + version: u64, + object_store: &dyn object_store::ObjectStore, + ) -> Result { + self.0 + .resolve_version_location(base_path, version, object_store) + .await + } + async fn version_exists( + &self, + base_path: &Path, + version: u64, + object_store: &dyn object_store::ObjectStore, + naming: ManifestNamingScheme, + ) -> Result { + self.0 + .version_exists(base_path, version, object_store, naming) + .await + } + fn list_detached_manifest_locations<'a>( + &self, + base_path: &Path, + object_store: &'a ObjectStore, + ) -> BoxStream<'a, Result> { + self.0 + .list_detached_manifest_locations(base_path, object_store) + } + fn list_manifest_locations<'a>( + &self, + base_path: &Path, + object_store: &'a ObjectStore, + sorted: bool, + ) -> BoxStream<'a, Result> { + self.0 + .list_manifest_locations(base_path, object_store, sorted) + } + fn list_manifest_locations_since<'a>( + &self, + base_path: &Path, + object_store: &'a ObjectStore, + since: u64, + ) -> BoxStream<'a, Result> { + self.0 + .list_manifest_locations_since(base_path, object_store, since) + } + async fn delete(&self, base_path: &Path) -> Result<()> { + self.0.delete(base_path).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[tokio::test] + async fn dropping_merge_waits_for_the_surviving_manifest_write() { + let scope = MergeWriteScope::new(); + let (entered, began) = oneshot::channel(); + let (release, blocked) = oneshot::channel(); + let committed = Arc::new(AtomicBool::new(false)); + let flag = committed.clone(); + let task_scope = scope.clone(); + let merge = tokio::spawn(async move { + task_scope + .run(shield(async move { + entered.send(()).unwrap(); + blocked.await.unwrap(); + flag.store(true, Ordering::SeqCst); + Ok::<_, Error>(()) + })) + .await + }); + began.await.unwrap(); + merge.abort(); + assert!(merge.await.unwrap_err().is_cancelled()); + let draining_scope = scope.clone(); + let drain = tokio::spawn(async move { draining_scope.drain().await }); + tokio::task::yield_now().await; + assert!( + !drain.is_finished(), + "cannot hand off while storage write survives" + ); + release.send(()).unwrap(); + drain.await.unwrap(); + assert!(committed.load(Ordering::SeqCst)); + assert!( + scope + .run(shield(async { Ok::<_, Error>(()) })) + .await + .is_err(), + "cancelled scope cannot issue another commit" + ); + } +} diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index b15c0db..28bc2bf 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -2410,6 +2410,126 @@ mod tests { ); } + #[tokio::test] + async fn cancelled_merge_joins_inflight_commit_then_retries_without_losing_blob_or_new_wal() { + use crate::merge_write_scope::{GuardedCommit, MergeWriteScope}; + use lance_table::format::{IndexMetadata, Manifest, Transaction}; + use lance_table::io::commit::{ + CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, ManifestWriter, + }; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[derive(Debug)] + struct PausedCommit { + inner: Arc, + once: AtomicBool, + entered: Arc, + release: Arc, + } + #[async_trait::async_trait] + #[allow(clippy::too_many_arguments)] + impl CommitHandler for PausedCommit { + async fn commit( + &self, + manifest: &mut Manifest, + indices: Option>, + path: &object_store::path::Path, + store: &lance_io::object_store::ObjectStore, + writer: ManifestWriter, + naming: ManifestNamingScheme, + transaction: Option, + ) -> std::result::Result { + let appending = transaction.as_ref().is_some_and(|tx| { + matches!( + tx.as_pb().operation, + Some(lance_table::format::pb::transaction::Operation::Append(_)) + ) + }); + if appending && !self.once.swap(true, Ordering::SeqCst) { + self.entered.notify_one(); + self.release.notified().await; + } + self.inner + .commit(manifest, indices, path, store, writer, naming, transaction) + .await + } + } + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut store = RolloutStore::open(uri).await.unwrap(); + let bytes = vec![73u8; 2 * 1024 * 1024]; + store + .add(&[artifact_record("large", &bytes)]) + .await + .unwrap(); + store.flush().await.unwrap(); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let handler = Arc::new(GuardedCommit(Arc::new(PausedCommit { + inner: lance_table::io::commit::commit_handler_from_url(uri, &None) + .await + .unwrap(), + once: AtomicBool::new(false), + entered: entered.clone(), + release: release.clone(), + }))); + store.base.dataset = lance::dataset::builder::DatasetBuilder::from_uri(uri) + .with_commit_handler(handler) + .load() + .await + .unwrap(); + let store = Arc::new(tokio::sync::Mutex::new(store)); + let scope = MergeWriteScope::new(); + let worker_scope = scope.clone(); + let writer = store.clone(); + let merge = tokio::spawn(async move { + worker_scope + .run(async { writer.lock().await.cleanup_own_shard().await }) + .await + }); + tokio::time::timeout(std::time::Duration::from_secs(30), entered.notified()) + .await + .unwrap(); + merge.abort(); + assert!(merge.await.unwrap_err().is_cancelled()); + let drain_scope = scope.clone(); + let drain = tokio::spawn(async move { drain_scope.drain().await }); + tokio::task::yield_now().await; + assert!(!drain.is_finished()); + // Writes arriving after cancellation must survive the eventual retry. + { + let store = store.lock().await; + store + .add(&[assistant_record("arrived-during-recovery")]) + .await + .unwrap(); + store.flush().await.unwrap(); + } + release.notify_one(); + tokio::time::timeout(std::time::Duration::from_secs(30), drain) + .await + .unwrap() + .unwrap(); + let mut store = store.lock().await; + assert!(flushed_generation_count(&store).await > 0); + store.base.refresh_latest().await.unwrap(); + assert_eq!( + store.base.dataset.count_rows(None).await.unwrap(), + 1, + "cancelled caller's leaf append must have committed before retry" + ); + store.cleanup_own_shard().await.unwrap(); + assert_eq!(flushed_generation_count(&store).await, 0); + assert_eq!(store.base.dataset.count_rows(None).await.unwrap(), 2); + assert_eq!(store.get_blob("large").await.unwrap().unwrap(), bytes); + assert!(store + .get_by_id("arrived-during-recovery") + .await + .unwrap() + .is_some()); + store.close().await.unwrap(); + } + fn assistant_record(id: &str) -> RolloutRecord { RolloutRecord { id: id.to_string(), diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index d9a4230..6461015 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -1015,6 +1015,10 @@ impl StorageBase { return Ok(false); } + // A prior cancelled merge may have completed a shielded manifest + // commit after this handle's future was dropped. Always reopen latest + // before retrying, including generic stores without schema evolution. + self.refresh_latest().await?; self.ensure_latest_schema().await?; if !batches.is_empty() { @@ -1035,22 +1039,31 @@ impl StorageBase { // writer now owns the shard. let epoch = manifest.writer_epoch; + let drain_store = ShardManifestStore::new( + self.dataset.object_store(None).await?, + &self.dataset.branch_location().path, + self.write_shard, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); observe_phase!( "drain", - manifest_store - .commit_update(epoch, |current| ShardManifest { - version: current.version + 1, - // Relative edit: retain everything we did not merge. Must - // never become an absolute assignment — see the doc comment. - flushed_generations: current - .flushed_generations - .iter() - .filter(|fg| !merged_generations.contains(&fg.generation)) - .cloned() - .collect(), - ..current.clone() - }) - .await + crate::merge_write_scope::shield(async move { + drain_store + .commit_update(epoch, |current| ShardManifest { + version: current.version + 1, + // Relative edit: retain everything we did not merge. Must + // never become an absolute assignment — see the doc comment. + flushed_generations: current + .flushed_generations + .iter() + .filter(|fg| !merged_generations.contains(&fg.generation)) + .cloned() + .collect(), + ..current.clone() + }) + .await + }) + .await )?; self.delete_merged_generation_dirs(&merged_paths).await?; @@ -2016,7 +2029,15 @@ impl StorageBase { storage_options: Option>, session: Option>, ) -> LanceResult { - let mut builder = DatasetBuilder::from_uri(uri); + let store_params = storage_options.clone().map(|options| ObjectStoreParams { + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + options, + ))), + ..Default::default() + }); + let handler = lance_table::io::commit::commit_handler_from_url(uri, &store_params).await?; + let mut builder = DatasetBuilder::from_uri(uri) + .with_commit_handler(Arc::new(crate::merge_write_scope::GuardedCommit(handler))); if let Some(options) = storage_options { builder = builder.with_storage_options(options); } @@ -2052,6 +2073,9 @@ impl StorageBase { }); } params.session = session; + let handler = + lance_table::io::commit::commit_handler_from_url(uri, ¶ms.store_params).await?; + params.commit_handler = Some(Arc::new(crate::merge_write_scope::GuardedCommit(handler))); Dataset::write(batches, uri, Some(params)).await } diff --git a/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml index 0a0fbf5..3ac0027 100644 --- a/crates/lance-context-master/Cargo.toml +++ b/crates/lance-context-master/Cargo.toml @@ -13,6 +13,7 @@ name = "lance-context-master" path = "src/main.rs" [dependencies] +lance-context-merge = { path = "../lance-context-merge" } lance-context-core = { version = "0.6.3", path = "../lance-context-core" } lance-context-api = { version = "0.6.3", path = "../lance-context-api" } lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" } @@ -36,4 +37,5 @@ tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } [dev-dependencies] +tower = { version = "0.5", features = ["util"] } tempfile = "3" diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index 081dd57..e9b0d5a 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -106,9 +106,9 @@ pub struct MasterConfig { #[arg(long, env = "INDEX_AFTER_COMPACTION", default_value_t = true, action = clap::ArgAction::Set)] pub index_after_compaction: bool, - /// Before fanning a merge-wal out to the workers, build the `id` BTree - /// index on the target's base table if it is missing. Without it Lance's - /// `merge_insert` full-scans the base table on every merge. + /// Deprecated compatibility flag. Owned merge executors now ensure the id + /// index under the same execution fence as the merge; the master does not + /// mutate the base table before admitting a worker. #[arg(long, env = "INDEX_BEFORE_MERGE", default_value_t = true, action = clap::ArgAction::Set)] pub index_before_merge: bool, diff --git a/crates/lance-context-master/src/lib.rs b/crates/lance-context-master/src/lib.rs index 1744518..982179a 100644 --- a/crates/lance-context-master/src/lib.rs +++ b/crates/lance-context-master/src/lib.rs @@ -3,6 +3,7 @@ pub mod config; pub mod discovery; pub mod error; +mod merge_execution; pub mod routes; pub mod scanner; pub mod scheduler; diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs new file mode 100644 index 0000000..1b06226 --- /dev/null +++ b/crates/lance-context-master/src/merge_execution.rs @@ -0,0 +1,563 @@ +//! Serial worker merges with durable execution ownership and targeted retries. +use crate::{state::MasterState, task_store::TaskClaim}; +use lance_context_merge::{ClaimProof, Coordinator, Execution, Phase}; +use std::{sync::Arc, time::Duration}; + +const RPC_TIMEOUT: Duration = Duration::from_secs(10); +const POLL_DELAY: Duration = Duration::from_millis(250); +const RETRY_DELAY: Duration = Duration::from_secs(2); +const ATTEMPTS: usize = 3; + +#[derive(serde::Deserialize)] +struct Capabilities { + protocol: u32, + instance: String, + timeout_secs: u64, +} + +pub(crate) async fn run_merge_wal( + state: &Arc, + claim: &TaskClaim, +) -> Result { + if state.config.worker_endpoints.is_empty() { + return Err("no worker endpoints configured".into()); + } + let coordinator = state.task_store.merge_coordinator(); + let proof = state.task_store.merge_claim(claim); + // Reconcile first, before any index/base-table mutation or new fan-out. + if let Some(old) = coordinator.get(&claim.task.target).await? { + let endpoint = old.endpoint.clone(); + if old.phase == Phase::Running { + if let Some(failure) = coordinator.failure(&claim.task.target, &endpoint).await? { + if failure.class == lance_context_merge::failure::FailureClass::OwnershipUnresolved + && failure.next_retry_ms > lance_context_merge::failure::now_ms() + { + return Err(format!( + "merge ownership unresolved; recovery probe at {}; {}", + failure.next_retry_ms, failure.last_error + )); + } + } + } + match reconcile(&state.http, &coordinator, &proof, old, true).await { + Ok(reclaimed) => { + metrics::counter!("master_merge_wal_generations_reclaimed_total") + .increment(reclaimed as u64); + } + Err(error) => { + if coordinator.get(&claim.task.target).await?.is_some() { + return Err(error); + } + tracing::info!(target = %claim.task.target, %error, "previous merge terminated; resuming shards"); + } + } + } + run_workers( + &state.http, + &coordinator, + &proof, + &claim.task.target, + &state.config.worker_endpoints, + ) + .await +} + +async fn run_workers( + http: &reqwest::Client, + coordinator: &Coordinator, + proof: &ClaimProof, + target: &str, + endpoints: &[String], +) -> Result { + let mut pending = Vec::new(); + let mut deferred = Vec::new(); + for endpoint in endpoints { + match coordinator.failure(target, endpoint).await? { + Some(failure) if failure.next_retry_ms > lance_context_merge::failure::now_ms() => { + deferred.push(format!( + "{endpoint}: retry at {}; {:?}: {}", + failure.next_retry_ms, failure.class, failure.last_error + )) + } + _ => pending.push(endpoint.clone()), + } + } + + let mut reclaimed = 0; + let mut errors = Vec::new(); + for attempt in 0..ATTEMPTS { + if attempt > 0 { + tokio::time::sleep(RETRY_DELAY).await; + } + let mut retry = Vec::new(); + errors.clear(); + for endpoint in pending { + let failures_before = coordinator + .failure(target, &endpoint) + .await? + .map_or(0, |failure| failure.consecutive_attempts); + let started = std::time::Instant::now(); + let outcome = one(http, coordinator, proof, target, &endpoint).await; + metrics::histogram!("master_merge_wal_worker_duration_seconds") + .record(started.elapsed().as_secs_f64()); + metrics::counter!("master_merge_wal_workers_total", "result" => if outcome.is_ok() { "ok" } else { "failed" }).increment(1); + match outcome { + Ok(n) => { + reclaimed += n; + } + Err(error) => { + // An execution still in storage is not a failed endpoint we + // may skip. Ownership must first be reconciled to terminal. + if coordinator.get(target).await?.is_some() { + return Err(error); + } + // Executed failures are recorded atomically with release. + // Capability/admission failures have no execution result. + let failure = match coordinator.failure(target, &endpoint).await? { + Some(failure) if failure.consecutive_attempts > failures_before => failure, + _ => { + coordinator + .record_failure(proof, target, &endpoint, &error) + .await? + } + }; + let message = format!( + "{endpoint}: {error} (attempt {}; retry at {}; attention={})", + failure.consecutive_attempts, + failure.next_retry_ms, + failure.needs_attention + ); + if failure.consecutive_attempts < ATTEMPTS as u32 && !failure.needs_attention { + errors.push(message); + retry.push(endpoint); + } else { + deferred.push(message); + } + } + } + } + pending = retry; + if pending.is_empty() { + break; + } + } + metrics::counter!("master_merge_wal_generations_reclaimed_total").increment(reclaimed as u64); + errors.extend(deferred); + if !errors.is_empty() { + return Err(format!("merged {reclaimed} generations; {} worker(s) failed or cooling down (up to {ATTEMPTS} targeted attempts): {}", errors.len(), errors.join("; "))); + } + Ok(format!( + "merged {reclaimed} generations across {}/{} workers", + endpoints.len(), + endpoints.len() + )) +} + +async fn one( + http: &reqwest::Client, + coordinator: &Coordinator, + proof: &ClaimProof, + target: &str, + endpoint: &str, +) -> Result { + let base = format!( + "{}/api/v1/internal/merge-executor", + endpoint.trim_end_matches('/') + ); + let capabilities: Capabilities = http + .get(&base) + .timeout(RPC_TIMEOUT) + .send() + .await + .map_err(|e| e.to_string())? + .error_for_status() + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + if capabilities.protocol != 1 || capabilities.timeout_secs == 0 { + return Err("worker lacks bounded owned-merge protocol".into()); + } + let execution = Execution::new( + target, + endpoint, + &capabilities.instance, + capabilities.timeout_secs.min(600), + ); + // Even a lost reserve response is ambiguous. Recover the fence before + // issuing any other storage operation; never infer absence from an error. + let admission = coordinator.reserve(proof, &execution).await; + match admission { + Ok(true) => {} + _ => { + if let Some(current) = coordinator.get(target).await? { + if current.id == execution.id { + return reconcile(http, coordinator, proof, current, true).await; + } + } + return Err("merge admission failed or task claim lost".into()); + } + } + let started = http + .post(format!("{base}/start")) + .timeout(RPC_TIMEOUT) + .json(&execution) + .send() + .await; + let cancel = !matches!(started, Ok(ref response) if response.status().is_success()); + // A transport error can mean the worker is already writing. It is NOT an + // outcome, and cannot let this task advance to another endpoint. + reconcile(http, coordinator, proof, execution, cancel).await +} + +async fn reconcile( + http: &reqwest::Client, + coordinator: &Coordinator, + proof: &ClaimProof, + initial: Execution, + cancel_immediately: bool, +) -> Result { + reconcile_with_grace( + http, + coordinator, + proof, + initial, + cancel_immediately, + Duration::from_secs(60), + ) + .await +} + +async fn reconcile_with_grace( + http: &reqwest::Client, + coordinator: &Coordinator, + proof: &ClaimProof, + initial: Execution, + cancel_immediately: bool, + grace: Duration, +) -> Result { + let deadline = + tokio::time::Instant::now() + Duration::from_secs(initial.timeout_secs) + RPC_TIMEOUT; + let handoff_deadline = if cancel_immediately { + tokio::time::Instant::now() + } else { + deadline + } + grace; + let mut next_cancel = tokio::time::Instant::now(); + loop { + if tokio::time::Instant::now() >= handoff_deadline { + let error = "merge ownership unresolved: executor did not acknowledge termination; fence retained, recovery requires old writer termination evidence"; + coordinator + .record_failure(proof, &initial.target, &initial.endpoint, error) + .await?; + return Err(error.into()); + } + let current = match coordinator.get(&initial.target).await { + Ok(Some(current)) if current.id == initial.id => current, + Ok(_) => return Err("merge execution ownership changed during reconciliation".into()), + Err(error) => { + tracing::warn!(target = %initial.target, %error, "cannot confirm merge completion; retaining execution fence"); + tokio::time::sleep(RETRY_DELAY).await; + continue; + } + }; + if current.phase == Phase::Finished { + if !coordinator.release(proof, ¤t).await? { + return Err("task claim lost while releasing terminal merge execution".into()); + } + return current.error.map_or(Ok(current.reclaimed), Err); + } + if (cancel_immediately || tokio::time::Instant::now() >= deadline) + && tokio::time::Instant::now() >= next_cancel + { + next_cancel = tokio::time::Instant::now() + RETRY_DELAY; + // CAS reserved work to terminal, or ask its owner to cancel the + // actual storage future. HTTP success alone is not acknowledgement. + if current.phase == Phase::Reserved { + let _ = coordinator.cancel_reserved(¤t).await; + } + let _ = http + .post(format!( + "{}/api/v1/internal/merge-executor/cancel", + current.endpoint.trim_end_matches('/') + )) + .timeout(RPC_TIMEOUT) + .json(¤t) + .send() + .await; + } + tokio::time::sleep(POLL_DELAY).await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::{ + extract::State, + routing::{get, post}, + Json, Router, + }; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Mutex, + }; + use tokio::sync::watch; + + #[derive(Clone)] + struct Worker { + coordinator: Coordinator, + calls: Arc, + fail: bool, + stall_first: bool, + name: &'static str, + events: Arc>>, + } + + async fn worker(worker: Worker) -> (String, tokio::task::JoinHandle<()>) { + let app = Router::new() + .route( + "/api/v1/internal/merge-executor", + get(|| async { + Json(serde_json::json!({"protocol":1,"instance":"test","timeout_secs":1})) + }), + ) + .route( + "/api/v1/internal/merge-executor/start", + post( + |State(w): State, Json(e): Json| async move { + let running = w.coordinator.start(&e).await.unwrap().unwrap(); + tokio::spawn(async move { + let call = w.calls.fetch_add(1, Ordering::SeqCst); + let (_cancel, rx) = watch::channel(false); + let result = lance_context_merge::execute_scoped( + async { + if w.stall_first && call == 0 { + return std::future::pending().await; + } + if w.fail { + Err("injected shard failure".into()) + } else { + Ok(7) + } + }, + Duration::from_secs(1), + rx, + ) + .await; + w.events.lock().unwrap().push(format!( + "{}:{}", + w.name, + if result.is_ok() { "ok" } else { "failed" } + )); + assert!(w.coordinator.finish(&running, result).await.unwrap()); + }); + axum::http::StatusCode::ACCEPTED + }, + ), + ) + .with_state(worker); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = format!("http://{}", listener.local_addr().unwrap()); + ( + address, + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }), + ) + } + + async fn fixture() -> (Coordinator, ClaimProof) { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").expect("isolated local etcd required"); + let mut client = etcd_client::Client::connect([endpoint], None) + .await + .unwrap(); + let prefix = format!("/merge-fanout-tests/{}", Execution::new("", "", "", 1).id); + let lease = client.lease_grant(120, None).await.unwrap().id(); + let proof = ClaimProof { + key: format!("{prefix}/claim"), + token: "owner".into(), + lease_id: lease, + }; + client + .put( + lance_context_merge::target_lock_key(&prefix, "table"), + proof.token.clone(), + Some(etcd_client::PutOptions::new().with_lease(lease)), + ) + .await + .unwrap(); + client + .put( + proof.key.clone(), + proof.token.clone(), + Some(etcd_client::PutOptions::new().with_lease(lease)), + ) + .await + .unwrap(); + (Coordinator::new(client, prefix), proof) + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn missing_executor_returns_attention_without_unlocking_surviving_write() { + let (coordinator, proof) = fixture().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = + tokio::spawn(async move { axum::serve(listener, Router::new()).await.unwrap() }); + let execution = Execution::new("table", &endpoint, "lost-process", 1); + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + let running = coordinator.start(&execution).await.unwrap().unwrap(); + let error = tokio::time::timeout( + Duration::from_secs(3), + reconcile_with_grace( + &reqwest::Client::new(), + &coordinator, + &proof, + running.clone(), + true, + Duration::from_millis(50), + ), + ) + .await + .unwrap() + .unwrap_err(); + assert!(error.contains("ownership unresolved")); + assert_eq!( + coordinator.get("table").await.unwrap(), + Some(running.clone()) + ); + assert!(!coordinator + .reserve(&proof, &Execution::new("table", "replacement", "new", 1)) + .await + .unwrap()); + assert!( + coordinator + .failure("table", &endpoint) + .await + .unwrap() + .unwrap() + .needs_attention + ); + // A late positive acknowledgement still permits normal recovery. + assert!(coordinator + .finish(&running, Err("cancelled".into())) + .await + .unwrap()); + let terminal = coordinator.get("table").await.unwrap().unwrap(); + assert!(coordinator.release(&proof, &terminal).await.unwrap()); + server.abort(); + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn stalled_shard_is_terminated_healthy_shards_advance_then_only_failure_retries() { + let (coordinator, proof) = fixture().await; + let events = Arc::new(Mutex::new(Vec::new())); + let bad = Worker { + coordinator: coordinator.clone(), + calls: Arc::new(AtomicUsize::new(0)), + fail: false, + stall_first: true, + name: "stalled", + events: events.clone(), + }; + let good = Worker { + stall_first: false, + name: "healthy", + calls: Arc::new(AtomicUsize::new(0)), + ..bad.clone() + }; + let (first, first_server) = worker(bad.clone()).await; + let (second, second_server) = worker(good.clone()).await; + let result = tokio::time::timeout( + Duration::from_secs(20), + run_workers( + &reqwest::Client::new(), + &coordinator, + &proof, + "table", + &[first, second], + ), + ) + .await + .unwrap() + .unwrap(); + assert!(result.contains("merged 14")); + assert_eq!(bad.calls.load(Ordering::SeqCst), 2); + assert_eq!(good.calls.load(Ordering::SeqCst), 1); + assert_eq!( + *events.lock().unwrap(), + ["stalled:failed", "healthy:ok", "stalled:ok"] + ); + assert!(coordinator.get("table").await.unwrap().is_none()); + first_server.abort(); + second_server.abort(); + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn partial_success_does_not_hide_exhausted_shard_failure() { + let (coordinator, proof) = fixture().await; + let bad = Worker { + coordinator: coordinator.clone(), + calls: Arc::new(AtomicUsize::new(0)), + fail: true, + stall_first: false, + name: "bad", + events: Arc::new(Mutex::new(Vec::new())), + }; + let good = Worker { + fail: false, + name: "good", + calls: Arc::new(AtomicUsize::new(0)), + ..bad.clone() + }; + let (first, first_server) = worker(bad.clone()).await; + let (second, second_server) = worker(good.clone()).await; + let result = tokio::time::timeout( + Duration::from_secs(25), + run_workers( + &reqwest::Client::new(), + &coordinator, + &proof, + "table", + &[first.clone(), second.clone()], + ), + ) + .await + .unwrap() + .unwrap_err(); + assert!(result.contains("merged 7 generations")); + assert!(result.contains("1 worker(s) failed or cooling down")); + assert_eq!(bad.calls.load(Ordering::SeqCst), 3); + assert_eq!(good.calls.load(Ordering::SeqCst), 1); + // A fresh task must retain the failed shard's backoff, while still + // admitting the healthy shard for newly arrived generations. + let next = run_workers( + &reqwest::Client::new(), + &coordinator, + &proof, + "table", + &[first.clone(), second], + ) + .await + .unwrap_err(); + assert!(next.contains("retry at")); + assert_eq!(bad.calls.load(Ordering::SeqCst), 3); + assert_eq!(good.calls.load(Ordering::SeqCst), 2); + assert_eq!( + coordinator + .failure("table", &first) + .await + .unwrap() + .unwrap() + .consecutive_attempts, + 3 + ); + assert!(coordinator.get("table").await.unwrap().is_none()); + first_server.abort(); + second_server.abort(); + } +} diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index 859eff1..a575367 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -573,6 +573,26 @@ pub async fn enqueue_task( /// lapses, and the last error. A store in this list is broken in a way that /// retrying will not fix (a manifest naming a missing fragment, for example); /// a manual `POST /tasks` still enqueues it. +#[derive(Debug, Deserialize)] +pub struct MergeFailureParams { + pub after: Option, +} + +pub async fn list_merge_failures( + State(state): State>, + Query(params): Query, +) -> Result, MasterError> { + let (failures, next) = state + .task_store + .merge_coordinator() + .failure_page(params.after.as_deref(), 256) + .await + .map_err(MasterError::Internal)?; + Ok(Json( + serde_json::json!({"failures": failures, "next": next}), + )) +} + pub async fn list_cooldowns( State(state): State>, ) -> Result>, MasterError> { @@ -652,6 +672,7 @@ pub fn api_router() -> Router> { .route("/experiments/{name}/rescan", post(rescan_experiment)) .route("/tasks", post(enqueue_task).get(list_tasks)) .route("/scheduler/cooldowns", get(list_cooldowns)) + .route("/scheduler/merge-failures", get(list_merge_failures)) .route("/scheduler/repairs", get(list_repairs)) .route("/tasks/{id}", get(get_task)) .route("/rescan", post(rescan)) diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 09d5fa0..27db7f6 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -149,7 +149,7 @@ async fn run_task(state: &Arc, claim: TaskClaim, timing: TaskClaimT let started = std::time::Instant::now(); let outcome = match task.kind { TaskKind::Compact => run_compaction(state, &task).await, - TaskKind::MergeWal => run_merge_wal(state, &task.target).await, + TaskKind::MergeWal => crate::merge_execution::run_merge_wal(state, &claim).await, TaskKind::IndexId => run_index_id(state, &task.target).await, TaskKind::Repair => run_repair(state, &task).await, }; @@ -168,7 +168,9 @@ async fn run_task(state: &Arc, claim: TaskClaim, timing: TaskClaimT if let Err(error) = &outcome { tracing::warn!(task = %task.id, target = %task.target, error, "task failed"); - if task.kind != TaskKind::Repair && is_missing_fragment_error(error) { + if !matches!(task.kind, TaskKind::Repair | TaskKind::MergeWal) + && is_missing_fragment_error(error) + { // The manifest names a file storage does not have. No retry and // no cooldown changes that; a repair does, so enqueue one now // and re-run this task behind it. @@ -370,214 +372,6 @@ async fn compact_inner( Ok(metrics) } -/// Shape of the worker's merge-wal response (`{ "reclaimed": n }`). -#[derive(serde::Deserialize)] -struct MergeWalReply { - reclaimed: usize, -} - -/// Fan a WAL-merge out to every configured worker endpoint. Each worker merges -/// 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. -/// Make sure the target's base table has a BTree index on `id` before the -/// workers merge into it. -/// -/// `merge_insert` without an exact-answer index on the join key reads the -/// whole base table into a hash join, so a merge's memory and time scale with -/// the base table rather than with the WAL rows being merged. Building the -/// index once here turns every later merge into an indexed probe. This is a -/// base-table write like compaction; the merge task already holds nothing on -/// the table, and `replace(true)` makes a race with a concurrent `IndexId` -/// harmless. Skipped when the index is already present (one manifest read). -async fn ensure_id_btree_index(state: &Arc, name: &str) -> Result<(), String> { - let uri = state.rollout_uri(name); - let mut store = RolloutStore::open_existing_with_options(&uri, state.rollout_store_options()) - .await - .map_err(|e| e.to_string())?; - if store - .has_id_btree_index() - .await - .map_err(|e| e.to_string())? - { - // The BTree exists but every merge and compaction since it was built - // added fragments it does not cover, and `merge_insert` full-scans - // those. Append an index delta over them so the probe stays a probe. - let started = std::time::Instant::now(); - let covered = store - .extend_id_btree_index() - .await - .map_err(|e| format!("extending id index before merge: {e}"))?; - if covered > 0 { - metrics::counter!("master_merge_wal_index_extended_total").increment(1); - tracing::info!( - target = %name, - fragments = covered, - elapsed_secs = started.elapsed().as_secs(), - "extended id BTree index over fragments added since it was built" - ); - } - return Ok(()); - } - let started = std::time::Instant::now(); - store - .create_id_btree_index() - .await - .map_err(|e| format!("building id index before merge: {e}"))?; - metrics::counter!("master_merge_wal_index_built_total").increment(1); - tracing::info!( - target = %name, - elapsed_secs = started.elapsed().as_secs(), - "built id BTree index before merge-wal so merge_insert probes instead of scanning" - ); - Ok(()) -} - -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); - if kind == StoreKind::Rollout && state.config.index_before_merge { - ensure_id_btree_index(state, name).await?; - } - 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 = 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 - // indistinguishable from every worker being slow. Unlabelled -- - // outcome is carried by the counter below, which costs one series - // per value instead of one per bucket per value. - let started = std::time::Instant::now(); - let outcome = merge_wal_one(&http, &url).await; - let result = match &outcome { - Ok(WorkerMerge::Reclaimed(_)) => "ok", - Ok(WorkerMerge::NotFound) => "not_found", - Err(WorkerMergeError::Http(_)) => "http_error", - Err(WorkerMergeError::Transport(_)) => "transport_error", - }; - metrics::histogram!("master_merge_wal_worker_duration_seconds") - .record(started.elapsed().as_secs_f64()); - // Counted per worker per attempt: a 404 is tolerated as "owns no - // shard" and N-1 failures still report task success, so this counter - // is the only place partial failure is visible at all. - metrics::counter!("master_merge_wal_workers_total", "result" => result).increment(1); - outcome - } - }); - - // Serial, not `join_all`: every worker's merge commits a new version of - // the *same* base table, and Lance's commit-conflict retry gives up after - // 30s of wall clock. Fanning out to 20 workers at once is 20 writers - // racing one commit point -- on a throttled object store the retries - // cannot complete in time and most of them fail with "Too many concurrent - // writers". The etcd target lock already guarantees one MergeWal task per - // store; this makes the task itself one writer at a time, which is the - // whole point of routing merges through the master. - let mut results = Vec::with_capacity(endpoints.len()); - for call in calls { - results.push(call.await); - } - let total_workers = results.len(); - let mut reclaimed = 0usize; - let mut ok_workers = 0usize; - let mut failed_workers = 0usize; - let mut last_err = None; - for r in results { - match r { - Ok(WorkerMerge::Reclaimed(n)) => { - reclaimed += n; - ok_workers += 1; - } - Ok(WorkerMerge::NotFound) => { - ok_workers += 1; - } - Err(e) => { - failed_workers += 1; - last_err = Some(e.to_string()); - } - } - } - - metrics::counter!("master_merge_wal_generations_reclaimed_total").increment(reclaimed as u64); - - // A worker that owns no shard for this store answers 404 or reclaims 0, - // and that is fine -- but it must not launder a real failure elsewhere - // into success. If any worker errored and the task as a whole made no - // progress, it failed: this is what a store with a broken base table - // looks like (the one worker holding its data 500s, the rest have - // nothing), and it is what the failure cooldown needs to see. Partial - // progress with some errors is still success; the next sweep retries - // the stragglers. - if ok_workers == 0 || (failed_workers > 0 && reclaimed == 0) { - return Err(format!( - "{}/{total_workers} workers failed and nothing was merged: {}", - failed_workers, - last_err.unwrap_or_else(|| "all workers failed".to_string()) - )); - } - Ok(format!( - "merged {reclaimed} generations across {ok_workers}/{total_workers} workers" - )) -} - -/// One worker's response to a WAL-merge fan-out. -enum WorkerMerge { - Reclaimed(usize), - /// The worker owns no shard for this experiment; tolerated as success. - NotFound, -} - -enum WorkerMergeError { - /// A non-success HTTP status. - Http(String), - /// Connection/timeout/decode failure. - Transport(String), -} - -impl std::fmt::Display for WorkerMergeError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Http(m) | Self::Transport(m) => f.write_str(m), - } - } -} - -/// Issue the merge call to one worker, classifying the failure mode so the -/// caller can label its metrics. -async fn merge_wal_one(http: &reqwest::Client, url: &str) -> Result { - let resp = http - .post(url) - .send() - .await - .map_err(|e| WorkerMergeError::Transport(e.to_string()))?; - let status = resp.status(); - if status == reqwest::StatusCode::NOT_FOUND { - return Ok(WorkerMerge::NotFound); - } - if !status.is_success() { - return Err(WorkerMergeError::Http(format!("{url}: HTTP {status}"))); - } - let body: MergeWalReply = resp - .json() - .await - .map_err(|e| WorkerMergeError::Transport(e.to_string()))?; - Ok(WorkerMerge::Reclaimed(body.reclaimed)) -} - /// Refresh the stats row for `name` after a successful compaction: re-observe /// fragment/row counts and bump `last_compaction`/`total_compactions`. /// @@ -783,6 +577,40 @@ async fn sweep_merge_wal_inner(state: &Arc) -> lance::Result /// not every scheduler task -- adequate for tests, which drop the whole /// `MasterState` immediately after. pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { + // Retry metadata independently of coarse stats sweeps. Every replica may + // enqueue; the existing task dedupe/claim transaction elects one executor. + let retry_state = state.clone(); + tokio::spawn(async move { + let coordinator = retry_state.task_store.merge_coordinator(); + let mut cursor = None; + let mut ticker = tokio::time::interval(Duration::from_secs(15)); + loop { + ticker.tick().await; + match coordinator.failure_page(cursor.as_deref(), 256).await { + Ok((rows, next)) => { + cursor = next; + let mut targets = std::collections::HashSet::new(); + for failure in rows { + if failure.next_retry_ms <= lance_context_merge::failure::now_ms() + && retry_state + .config + .worker_endpoints + .contains(&failure.endpoint) + && targets.insert(failure.target.clone()) + { + if let Err(error) = + enqueue(&retry_state, TaskKind::MergeWal, &failure.target).await + { + tracing::warn!(target = %failure.target, %error, "merge recovery enqueue failed"); + } + } + } + } + Err(error) => tracing::warn!(%error, "merge recovery metadata scan failed"), + } + } + }); + // Optional periodic compaction auto-sweep feeds the same queue. let interval_secs = state.config.compaction_interval_secs; if interval_secs > 0 { @@ -1144,20 +972,35 @@ mod tests { /// merge finds it present and builds nothing. #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] - async fn merge_wal_builds_the_id_btree_first() { + async fn owned_merge_delegates_index_safety_to_worker() { use axum::{routing::post, Json, Router}; - let app = Router::new().route( - "/api/v1/internal/merge-wal/{name}", - post(|| async { Json(serde_json::json!({ "reclaimed": 0 })) }), - ); 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}")]; cfg.index_before_merge = true; + let base = cfg.data_dir.clone(); + let app = Router::new().route( + "/api/v1/internal/merge-wal/{name}", + post( + move |axum::extract::Path(name): axum::extract::Path| { + let base = base.clone(); + async move { + let uri = + lance_context_core::join_uri(&base, &format!("{name}.rollout.lance")); + let mut store = RolloutStore::open(&uri).await.unwrap(); + if !store.has_id_btree_index().await.unwrap() { + store.create_id_btree_index().await.unwrap(); + } + Json(serde_json::json!({"reclaimed":0})) + } + }, + ), + ); + let app = owned_stub(app, &cfg).await; + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); @@ -1347,6 +1190,68 @@ mod tests { worker.abort(); } + /// Wrap test workers in the production ownership wire protocol, preserving + /// their individual route assertions and injected delays/errors. + async fn owned_stub(app: axum::Router, cfg: &MasterConfig) -> axum::Router { + use axum::{ + routing::{get, post}, + Json, + }; + use lance_context_merge::{Coordinator, Execution}; + use tower::ServiceExt; + let client = etcd_client::Client::connect(cfg.etcd_endpoints.clone(), None) + .await + .unwrap(); + let coordinator = Coordinator::new(client, cfg.etcd_prefix.clone()); + axum::Router::new() + .route( + "/api/v1/internal/merge-executor", + get(|| async { + Json(serde_json::json!({"protocol":1,"instance":"stub","timeout_secs":600})) + }), + ) + .route( + "/api/v1/internal/merge-executor/start", + post(move |Json(execution): Json| { + let coordinator = coordinator.clone(); + let app = app.clone(); + async move { + let running = coordinator.start(&execution).await.unwrap().unwrap(); + tokio::spawn(async move { + let path = match execution.target.strip_prefix("generic:") { + Some(name) => format!("/api/v1/generic/{name}/merge-wal"), + None => format!("/api/v1/internal/merge-wal/{}", execution.target), + }; + let response = app + .oneshot( + axum::http::Request::post(path) + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = axum::body::to_bytes(response.into_body(), 65536) + .await + .unwrap(); + let outcome = if status.is_success() { + Ok(serde_json::from_slice::(&body).unwrap() + ["reclaimed"] + .as_u64() + .unwrap() as usize) + } else if status == axum::http::StatusCode::NOT_FOUND { + Ok(0) + } else { + Err(format!("HTTP {status}: {}", String::from_utf8_lossy(&body))) + }; + assert!(coordinator.finish(&running, outcome).await.unwrap()); + }); + axum::http::StatusCode::ACCEPTED + } + }), + ) + } + /// MergeWal fans out to every configured worker endpoint and sums the /// reclaimed counts. Uses a tiny in-process stub server per "worker". #[tokio::test] @@ -1355,11 +1260,12 @@ mod tests { use axum::{routing::post, Json, Router}; // A stub worker that always reports `reclaimed` for any merge call. - async fn spawn_stub(reclaimed: usize) -> String { + async fn spawn_stub(reclaimed: usize, cfg: &MasterConfig) -> String { let app = Router::new().route( "/api/v1/internal/merge-wal/{name}", post(move || async move { Json(serde_json::json!({ "reclaimed": reclaimed })) }), ); + let app = owned_stub(app, cfg).await; let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { @@ -1370,7 +1276,7 @@ mod tests { let dir = TempDir::new().unwrap(); let mut cfg = config(&dir); - cfg.worker_endpoints = vec![spawn_stub(3).await, spawn_stub(2).await]; + cfg.worker_endpoints = vec![spawn_stub(3, &cfg).await, spawn_stub(2, &cfg).await]; let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); @@ -1415,11 +1321,12 @@ mod tests { }); 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 app = owned_stub(app, &cfg).await; + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); @@ -1444,7 +1351,7 @@ mod tests { async fn merge_wal_with_an_erroring_worker_and_no_progress_fails() { use axum::{http::StatusCode, routing::post, Json, Router}; - async fn spawn(ok: bool) -> String { + async fn spawn(ok: bool, cfg: &MasterConfig) -> String { let app = Router::new().route( "/api/v1/internal/merge-wal/{name}", post(move || async move { @@ -1458,6 +1365,7 @@ mod tests { } }), ); + let app = owned_stub(app, cfg).await; 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() }); @@ -1466,14 +1374,40 @@ mod tests { let dir = TempDir::new().unwrap(); let mut cfg = config(&dir); - cfg.worker_endpoints = vec![spawn(true).await, spawn(false).await]; + cfg.worker_endpoints = vec![spawn(true, &cfg).await, spawn(false, &cfg).await]; let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); let rec = enqueue(&state, TaskKind::MergeWal, "broken").await.unwrap(); let status = await_terminal(&state, &rec.id).await; assert_eq!(status.state, TaskState::Failed, "got {status:?}"); - assert!(status.error.unwrap().contains("nothing was merged")); + assert!(status.error.unwrap().contains("merged 0 generations")); + assert!( + state + .task_store + .list() + .await + .unwrap() + .iter() + .all(|task| task.kind != TaskKind::Repair), + "a merge failure must not schedule destructive fragment removal" + ); + assert!( + !state + .task_store + .is_cooling_down(TaskKind::MergeWal, "broken") + .await + .unwrap(), + "one broken shard must not suppress healthy shard scheduling" + ); + let (failures, _) = state + .task_store + .merge_coordinator() + .failure_page(None, 256) + .await + .unwrap(); + assert_eq!(failures.len(), 1); + assert!(failures[0].needs_attention); worker.abort(); } @@ -1491,7 +1425,11 @@ mod tests { // 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 { + async fn spawn_stub( + inflight: Arc, + max_inflight: Arc, + cfg: &MasterConfig, + ) -> String { let app = Router::new().route( "/api/v1/internal/merge-wal/{name}", post(move || { @@ -1506,6 +1444,7 @@ mod tests { } }), ); + let app = owned_stub(app, cfg).await; 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() }); @@ -1515,9 +1454,9 @@ mod tests { 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, + spawn_stub(inflight.clone(), max_inflight.clone(), &cfg).await, + spawn_stub(inflight.clone(), max_inflight.clone(), &cfg).await, + spawn_stub(inflight.clone(), max_inflight.clone(), &cfg).await, ]; let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); @@ -1598,13 +1537,14 @@ mod tests { ); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { axum::serve(listener, hang).await.unwrap() }); let dir = TempDir::new().unwrap(); let mut cfg = config(&dir); cfg.worker_endpoints = vec![format!("http://{addr}")]; cfg.merge_wal_concurrency = 2; cfg.task_concurrency = 2; + let hang = owned_stub(hang, &cfg).await; + tokio::spawn(async move { axum::serve(listener, hang).await.unwrap() }); let state = MasterState::new(cfg).await.unwrap(); // Build a compactable store before the dispatcher starts. @@ -1742,23 +1682,19 @@ mod tests { worker.abort(); } - /// The WAL-merge sweep enqueues a `MergeWal` only for experiments whose - /// pending generation count is at or above the threshold, and de-dupes so a - /// A target that keeps failing is skipped by the sweep after the cooldown - /// threshold, so a permanently broken store stops consuming task slots - /// and blocking healthy stores behind it. Manual enqueue bypasses the - /// cooldown, and a later success clears it. + /// Non-merge maintenance retains target-wide cooldown. Merge failures now + /// use the independent per-shard circuit tested in merge_execution. #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] - async fn repeated_failures_cool_the_target_down_and_sweeps_skip_it() { + async fn repeated_compaction_failures_cool_the_target_down_and_sweeps_skip_it() { use crate::stats_store::StatRow; let dir = TempDir::new().unwrap(); let mut cfg = config(&dir); - cfg.merge_wal_min_generations = 1; + cfg.min_fragments = 1; cfg.task_cooldown_after_failures = 2; cfg.task_cooldown_base_secs = 3600; - // No worker endpoints: every MergeWal fails immediately. + // Missing base dataset: every compaction fails immediately. cfg.worker_endpoints = vec![]; let state = MasterState::new(cfg).await.unwrap(); let worker = spawn_scheduler(&state); @@ -1771,7 +1707,7 @@ mod tests { name: "broken".to_string(), uri: state.rollout_uri("broken"), row_count: 0, - fragment_count: 0, + fragment_count: 50, last_updated: 0, pending_wal_generations: 50, last_compaction: StatRow::NO_COMPACTION, @@ -1787,19 +1723,19 @@ mod tests { assert!( !state .task_store - .is_cooling_down(TaskKind::MergeWal, "broken") + .is_cooling_down(TaskKind::Compact, "broken") .await .unwrap(), "round {round}: must not cool down below the failure threshold" ); - assert_eq!(sweep_merge_wal_candidates(&state).await.unwrap(), 1); + assert_eq!(sweep_candidates(&state).await.unwrap(), 1); let id = state .task_store .list() .await .unwrap() .into_iter() - .filter(|t| t.kind == TaskKind::MergeWal && t.target == "broken") + .filter(|t| t.kind == TaskKind::Compact && t.target == "broken") .max_by_key(|t| t.enqueued_at) .unwrap() .id; @@ -1812,7 +1748,7 @@ mod tests { for _ in 0..40 { if state .task_store - .is_cooling_down(TaskKind::MergeWal, "broken") + .is_cooling_down(TaskKind::Compact, "broken") .await .unwrap() { @@ -1830,13 +1766,13 @@ mod tests { // The sweep now skips it. assert_eq!( - sweep_merge_wal_candidates(&state).await.unwrap(), + sweep_candidates(&state).await.unwrap(), 0, "a cooled-down target must not be re-enqueued by the sweep" ); // A manual enqueue is not gated. - let manual = enqueue(&state, TaskKind::MergeWal, "broken").await.unwrap(); + let manual = enqueue(&state, TaskKind::Compact, "broken").await.unwrap(); assert_eq!(manual.target, "broken"); assert_eq!( await_terminal(&state, &manual.id).await.state, @@ -1868,7 +1804,7 @@ mod tests { worker.abort(); } - /// second sweep does not pile up a duplicate for the same target. + /// WAL sweeps enqueue over-threshold targets and dedupe a second sweep. #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn sweep_merge_wal_enqueues_over_threshold_and_dedupes() { diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 91567da..66d7a64 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -150,6 +150,18 @@ impl Drop for LeaseKeepalive { } impl TaskStore { + pub(crate) fn merge_coordinator(&self) -> lance_context_merge::Coordinator { + lance_context_merge::Coordinator::new(self.inner.client.clone(), self.inner.prefix.clone()) + } + + pub(crate) fn merge_claim(&self, claim: &TaskClaim) -> lance_context_merge::ClaimProof { + lance_context_merge::ClaimProof { + key: self.inner.claim_key(&claim.task.id), + token: claim.backend.token.clone(), + lease_id: claim.backend.lease_id, + } + } + pub async fn open(config: &MasterConfig) -> lance::Result { let store = Self { inner: Arc::new(EtcdTaskStore::connect(config).await?), @@ -233,7 +245,7 @@ impl TaskStore { self.inner.finish(claim, outcome).await?; // Cooldown bookkeeping is best-effort and never fails the completion: // the task's terminal state is already committed above. - if self.cooldown.after_failures > 0 { + if kind != TaskKind::MergeWal && self.cooldown.after_failures > 0 { let result = match failed { Some(error) => { self.inner @@ -262,7 +274,7 @@ impl TaskStore { } pub async fn is_cooling_down(&self, kind: TaskKind, target: &str) -> lance::Result { - if self.cooldown.after_failures == 0 { + if kind == TaskKind::MergeWal || self.cooldown.after_failures == 0 { return Ok(false); } // Below the threshold the record only carries the failure count; it is @@ -603,6 +615,14 @@ impl EtcdTaskStore { } } + let merge_execution = if task.kind == TaskKind::MergeWal { + lance_context_merge::Coordinator::new(self.client.clone(), self.prefix.clone()) + .get(&task.target) + .await + .map_err(lance::Error::io)? + } else { + None + }; let token = generate_id(); let lease_id = self.grant_lease().await?; let claim_key = self.claim_key(&task.id); @@ -620,7 +640,37 @@ impl EtcdTaskStore { 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)); + if let Some(execution) = &merge_execution { + compares.push(Compare::value( + key.as_str(), + CompareOp::Equal, + lance_context_merge::execution_owner(execution), + )); + } else { + compares.push(Compare::version(key.as_str(), CompareOp::Equal, 0)); + } + } + // A lost scheduler lease does not terminate remote storage work. + // MergeWal may claim to reconcile it; every other writer waits. + if task.kind != TaskKind::MergeWal && target_key.is_some() { + compares.push(Compare::version( + lance_context_merge::execution_key(&self.prefix, &task.target), + CompareOp::Equal, + 0, + )); + } + // Remote execution ownership persists, but its reconciler is + // still exclusive and leased. Dependency-chain merge tasks + // must not cancel another live scheduler's execution. + let merge_claim_key = + lance_context_merge::execution_key(&self.prefix, &task.target) + .replace("/merge-executions/", "/merge-claims/"); + if task.kind == TaskKind::MergeWal { + compares.push(Compare::version( + merge_claim_key.as_str(), + CompareOp::Equal, + 0, + )); } let lease_options = Some(PutOptions::new().with_lease(lease_id)); let mut operations = vec![ @@ -629,13 +679,22 @@ impl EtcdTaskStore { 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 { + if task.kind == TaskKind::MergeWal { operations.push(TxnOp::put( - key.as_str(), + merge_claim_key, token.as_bytes(), lease_options.clone(), )); } + if merge_execution.is_none() { + 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)) @@ -680,14 +739,31 @@ impl EtcdTaskStore { keepalive, } = claim.backend; let mut task = claim.task; + let unresolved = if task.kind == TaskKind::MergeWal && outcome.is_err() { + lance_context_merge::Coordinator::new(self.client.clone(), self.prefix.clone()) + .get(&task.target) + .await + .map_err(lance::Error::io)? + } else { + None + }; 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 task.kind == TaskKind::MergeWal { + operations.push(TxnOp::delete( + lance_context_merge::execution_key(&self.prefix, &task.target) + .replace("/merge-executions/", "/merge-claims/"), + None, + )); + } + if unresolved.is_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, &task.depends_on) { operations.push(TxnOp::delete(key, None)); @@ -696,11 +772,22 @@ impl EtcdTaskStore { let completed = client .txn( Txn::new() - .when([Compare::value( - claim_key.as_str(), - CompareOp::Equal, - token.as_bytes(), - )]) + .when([ + Compare::value(claim_key.as_str(), CompareOp::Equal, token.as_bytes()), + match &unresolved { + Some(execution) => Compare::value( + lance_context_merge::execution_key(&self.prefix, &task.target), + CompareOp::Equal, + serde_json::to_vec(execution) + .map_err(|e| lance::Error::io(e.to_string()))?, + ), + None => Compare::version( + lance_context_merge::execution_key(&self.prefix, &task.target), + CompareOp::Equal, + 0, + ), + }, + ]) .and_then(operations), ) .await @@ -1265,7 +1352,7 @@ fn should_dedupe(kind: TaskKind, depends_on: &[String]) -> bool { fn requires_target_lock(kind: TaskKind) -> bool { matches!( kind, - TaskKind::Compact | TaskKind::IndexId | TaskKind::Repair + TaskKind::Compact | TaskKind::IndexId | TaskKind::Repair | TaskKind::MergeWal ) } @@ -1364,6 +1451,152 @@ mod tests { } } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn merge_serializes_other_writers_and_survives_claim_loss() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.etcd_endpoints = std::env::var("ETCD_TEST_ENDPOINTS") + .unwrap() + .split(',') + .map(str::to_string) + .collect(); + cfg.etcd_prefix = format!("/merge-lock-test/{}", generate_id()); + let store = TaskStore::open(&cfg).await.unwrap(); + store + .enqueue(TaskKind::MergeWal, "shared", Vec::new()) + .await + .unwrap(); + let merge = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + store + .enqueue(TaskKind::Compact, "shared", Vec::new()) + .await + .unwrap(); + store + .enqueue(TaskKind::IndexId, "shared", Vec::new()) + .await + .unwrap(); + assert!( + store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .is_none(), + "merge must lock against index and compact" + ); + store + .enqueue(TaskKind::Compact, "unrelated", Vec::new()) + .await + .unwrap(); + let other = store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .unwrap(); + assert_eq!(other.task.target, "unrelated"); + let dependency = other.task.id.clone(); + store.finish(other, Ok("done".into())).await.unwrap(); + // A dependency-chain task has a distinct id and bypasses dedupe. + store + .enqueue(TaskKind::MergeWal, "shared", vec![dependency]) + .await + .unwrap(); + let coordinator = store.merge_coordinator(); + let execution = lance_context_merge::Execution::new("shared", "worker", "boot", 600); + assert!(coordinator + .reserve(&store.merge_claim(&merge), &execution) + .await + .unwrap()); + let running = coordinator.start(&execution).await.unwrap().unwrap(); + assert!( + store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .is_none(), + "only one scheduler may own reconciliation, even with a persistent execution lock" + ); + store.abandon_claim_for_test(merge).await.unwrap(); + store.inner.recover_orphaned().await.unwrap(); + assert!( + store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .is_none(), + "expired claim is not permission to write past live worker execution" + ); + let recovered = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + store + .finish(recovered, Err("merge ownership unresolved".into())) + .await + .unwrap(); + assert_eq!( + coordinator.get("shared").await.unwrap(), + Some(running.clone()) + ); + assert!(store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .is_none()); + let recovered = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + assert!(coordinator + .finish(&running, Err("cancelled".into())) + .await + .unwrap()); + let done = coordinator.get("shared").await.unwrap().unwrap(); + assert!(coordinator + .release(&store.merge_claim(&recovered), &done) + .await + .unwrap()); + store + .finish(recovered, Ok("recovered".into())) + .await + .unwrap(); + let compact = store + .claim_next_of_kinds(TaskKinds::GENERAL) + .await + .unwrap() + .unwrap(); + assert_eq!(compact.task.kind, TaskKind::Compact); + store + .enqueue(TaskKind::MergeWal, "shared", Vec::new()) + .await + .unwrap(); + assert!( + store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .is_none(), + "compact must also block a new merge" + ); + store.finish(compact, Ok("done".into())).await.unwrap(); + store + .inner + .client + .clone() + .delete( + cfg.etcd_prefix, + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + #[tokio::test] async fn connect_requires_etcd_endpoints() { let dir = TempDir::new().unwrap(); diff --git a/crates/lance-context-merge/Cargo.toml b/crates/lance-context-merge/Cargo.toml new file mode 100644 index 0000000..efeda50 --- /dev/null +++ b/crates/lance-context-merge/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "lance-context-merge" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +description = "Durable ownership for cancellable WAL merge execution" + +[dependencies] +clap = { version = "4", features = ["derive", "env"] } +etcd-client = { version = "0.19", features = ["tls"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] } +uuid = { version = "1", features = ["v4"] } + +[dev-dependencies] +tokio = { version = "1", features = ["test-util"] } diff --git a/crates/lance-context-merge/src/failure.rs b/crates/lance-context-merge/src/failure.rs new file mode 100644 index 0000000..3fdea0b --- /dev/null +++ b/crates/lance-context-merge/src/failure.rs @@ -0,0 +1,277 @@ +//! Durable shard failures: task recreation must not reset a retry budget. +use crate::{ClaimProof, Coordinator, Execution, Result}; +use etcd_client::{Compare, CompareOp, GetOptions, TxnOp}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum FailureClass { + Retryable, + Deadline, + DataOrConfiguration, + OwnershipUnresolved, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ShardFailure { + pub target: String, + pub endpoint: String, + pub consecutive_attempts: u32, + pub class: FailureClass, + pub last_error: String, + pub last_failure_ms: u64, + pub next_retry_ms: u64, + pub needs_attention: bool, +} + +pub fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} + +pub fn classify(error: &str) -> FailureClass { + let error = error.to_ascii_lowercase(); + if error.contains("ownership unresolved") { + FailureClass::OwnershipUnresolved + } else if [ + "not found", + "notfound", + "corrupt", + "schema", + "invalid", + "403", + "401", + "protocol", + ] + .iter() + .any(|s| error.contains(s)) + { + FailureClass::DataOrConfiguration + } else if error.contains("deadline") || error.contains("timeout") || error.contains("timed out") + { + FailureClass::Deadline + } else { + FailureClass::Retryable + } +} + +impl ShardFailure { + fn advance( + target: &str, + endpoint: &str, + previous: Option, + error: &str, + now: u64, + ) -> Self { + let attempts = previous.map_or(1, |f| f.consecutive_attempts.saturating_add(1)); + let class = classify(error); + let needs_attention = attempts >= 15 + || matches!( + class, + FailureClass::DataOrConfiguration | FailureClass::OwnershipUnresolved + ); + // First task gets three attempts. Later tasks get one half-open probe, + // never a fresh budget of three. Cap probing at once per hour. + let delay_secs = if needs_attention { + 3600 + } else if attempts < 3 { + 2 + } else { + (30u64.saturating_mul(1 << (attempts - 3).min(5))).min(900) + }; + Self { + target: target.into(), + endpoint: endpoint.into(), + consecutive_attempts: attempts, + class, + last_error: error.chars().take(4096).collect(), + last_failure_ms: now, + next_retry_ms: now.saturating_add(delay_secs * 1000), + needs_attention, + } + } +} + +impl Coordinator { + fn failures_prefix(&self) -> String { + format!("{}/merge-failures/", self.prefix.trim_end_matches('/')) + } + fn failure_key(&self, target: &str, endpoint: &str) -> String { + let hex = |s: &str| { + s.as_bytes() + .iter() + .map(|b| format!("{b:02x}")) + .collect::() + }; + format!( + "{}{}/{}", + self.failures_prefix(), + hex(target), + hex(endpoint) + ) + } + pub async fn failure(&self, target: &str, endpoint: &str) -> Result> { + let response = self + .client + .clone() + .get(self.failure_key(target, endpoint), None) + .await + .map_err(|e| e.to_string())?; + response + .kvs() + .first() + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose() + } + /// Commit the terminal outcome and its retry budget in the same etcd + /// transaction that releases execution ownership. A master crash between + /// release and bookkeeping cannot reset failed attempts. + pub(crate) async fn completion_changes( + &self, + execution: &Execution, + ) -> Result<(Vec, Vec)> { + let key = self.failure_key(&execution.target, &execution.endpoint); + let Some(error) = &execution.error else { + return Ok((Vec::new(), vec![TxnOp::delete(key, None)])); + }; + let response = self + .client + .clone() + .get(key.clone(), None) + .await + .map_err(|e| e.to_string())?; + let previous = response.kvs().first(); + let old = previous + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose()?; + let failure = + ShardFailure::advance(&execution.target, &execution.endpoint, old, error, now_ms()); + Ok(( + vec![Compare::mod_revision( + key.clone(), + CompareOp::Equal, + previous.map_or(0, |kv| kv.mod_revision()), + )], + vec![TxnOp::put(key, serde_json::to_vec(&failure).unwrap(), None)], + )) + } + + pub async fn record_failure( + &self, + proof: &ClaimProof, + target: &str, + endpoint: &str, + error: &str, + ) -> Result { + let key = self.failure_key(target, endpoint); + let response = self + .client + .clone() + .get(key.clone(), None) + .await + .map_err(|e| e.to_string())?; + let previous = response.kvs().first(); + let old = previous + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose()?; + let failure = ShardFailure::advance(target, endpoint, old, error, now_ms()); + let revision = previous.map_or(0, |kv| kv.mod_revision()); + if !self + .transact( + vec![ + Compare::value(proof.key.as_str(), CompareOp::Equal, proof.token.as_bytes()), + Compare::mod_revision(key.clone(), CompareOp::Equal, revision), + ], + vec![TxnOp::put(key, serde_json::to_vec(&failure).unwrap(), None)], + ) + .await? + { + return Err("claim or failure record changed while recording merge failure".into()); + } + Ok(failure) + } + pub async fn clear_failure( + &self, + proof: &ClaimProof, + target: &str, + endpoint: &str, + ) -> Result<()> { + if !self + .transact( + vec![Compare::value( + proof.key.as_str(), + CompareOp::Equal, + proof.token.as_bytes(), + )], + vec![TxnOp::delete(self.failure_key(target, endpoint), None)], + ) + .await? + { + return Err("claim lost while recording merge recovery".into()); + } + Ok(()) + } + /// Bounded, cursor-based metadata scan; never opens table data or WAL. + pub async fn failure_page( + &self, + after: Option<&str>, + limit: i64, + ) -> Result<(Vec, Option)> { + let prefix = self.failures_prefix(); + let (start, options) = match after { + Some(key) if key.starts_with(&prefix) => { + let mut end = prefix.as_bytes().to_vec(); + *end.last_mut().unwrap() += 1; + (format!("{key}\0"), GetOptions::new().with_range(end)) + } + Some(_) => return Err("invalid merge failure cursor".into()), + None => (prefix, GetOptions::new().with_prefix()), + }; + let response = self + .client + .clone() + .get(start, Some(options.with_limit(limit.clamp(1, 256)))) + .await + .map_err(|e| e.to_string())?; + let rows = response + .kvs() + .iter() + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .collect::>>()?; + let next = if response.more() { + response + .kvs() + .last() + .map(|kv| String::from_utf8_lossy(kv.key()).into_owned()) + } else { + None + }; + Ok((rows, next)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn persistent_budget_backoff_and_attention() { + let mut old = None; + for attempt in 1..=20 { + let next = + ShardFailure::advance("table", "worker", old, "temporary storage failure", 1000); + assert_eq!(next.consecutive_attempts, attempt); + assert_eq!(next.needs_attention, attempt >= 15); + assert!(next.next_retry_ms <= 3_601_000); + if attempt == 3 { + assert_eq!(next.next_retry_ms, 31_000); + } + old = Some(next); + } + let broken = + ShardFailure::advance("table", "worker", None, "Not found: /data/1.lance", 1000); + assert!(broken.needs_attention); + assert_eq!(broken.next_retry_ms, 3_601_000); + } +} diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs new file mode 100644 index 0000000..855a938 --- /dev/null +++ b/crates/lance-context-merge/src/lib.rs @@ -0,0 +1,600 @@ +//! Execution fences outlive scheduler leases and HTTP connections. +//! +//! A claim authorizes admission, not cancellation of an existing storage write. +//! Only the executor can publish its terminal outcome. Recovery first cancels +//! and reconciles the old execution; it never clears a running fence on timeout. + +pub mod failure; + +use etcd_client::{Client, Compare, CompareOp, Txn, TxnOp}; +use serde::{Deserialize, Serialize}; + +pub type Result = std::result::Result; + +#[derive(Clone, Debug)] +pub struct ClaimProof { + pub key: String, + pub token: String, + pub lease_id: i64, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum Phase { + Reserved, + Running, + Finished, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub struct Execution { + pub id: String, + pub target: String, + pub endpoint: String, + pub instance: String, + pub timeout_secs: u64, + pub phase: Phase, + pub reclaimed: usize, + pub error: Option, +} + +impl Execution { + pub fn new(target: &str, endpoint: &str, instance: &str, timeout_secs: u64) -> Self { + Self { + id: uuid::Uuid::new_v4().to_string(), + target: target.into(), + endpoint: endpoint.into(), + instance: instance.into(), + timeout_secs, + phase: Phase::Reserved, + reclaimed: 0, + error: None, + } + } + + pub fn finished(&self, outcome: Result) -> Self { + let mut next = self.clone(); + next.phase = Phase::Finished; + match outcome { + Ok(n) => next.reclaimed = n, + Err(e) => next.error = Some(e), + } + next + } +} + +/// Deliberately no lease on these keys. An expired task claim must not let a +/// second endpoint commit while the first endpoint's HTTP handler is still alive. +#[derive(Clone)] +pub struct Coordinator { + client: Client, + prefix: String, +} + +pub fn execution_key(prefix: &str, target: &str) -> String { + let encoded: String = target + .as_bytes() + .iter() + .map(|b| format!("{b:02x}")) + .collect(); + format!( + "{}/merge-executions/{encoded}", + prefix.trim_end_matches('/') + ) +} + +pub fn target_lock_key(prefix: &str, target: &str) -> String { + execution_key(prefix, target).replace("/merge-executions/", "/target-locks/") +} + +pub fn execution_owner(execution: &Execution) -> String { + format!("merge-execution:{}", execution.id) +} + +fn encode(execution: &Execution) -> Vec { + serde_json::to_vec(execution).expect("execution is JSON serializable") +} + +impl Coordinator { + pub fn new(client: Client, prefix: impl Into) -> Self { + Self { + client, + prefix: prefix.into(), + } + } + + pub async fn get(&self, target: &str) -> Result> { + let response = self + .client + .clone() + .get(execution_key(&self.prefix, target), None) + .await + .map_err(|e| e.to_string())?; + response + .kvs() + .first() + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose() + } + + /// Atomic claim check and admission close the delayed-request/lease-loss race. + pub async fn reserve(&self, claim: &ClaimProof, execution: &Execution) -> Result { + if execution.phase != Phase::Reserved || execution.timeout_secs == 0 { + return Err("invalid merge execution reservation".into()); + } + let key = execution_key(&self.prefix, &execution.target); + self.transact( + vec![ + Compare::value(claim.key.as_str(), CompareOp::Equal, claim.token.as_bytes()), + Compare::version(key.as_str(), CompareOp::Equal, 0), + Compare::value( + target_lock_key(&self.prefix, &execution.target), + CompareOp::Equal, + claim.token.as_bytes(), + ), + ], + vec![ + TxnOp::put(key, encode(execution), None), + TxnOp::put( + target_lock_key(&self.prefix, &execution.target), + execution_owner(execution), + None, + ), + ], + ) + .await + } + + pub async fn start(&self, execution: &Execution) -> Result> { + if execution.phase != Phase::Reserved { + return Err("execution is not reserved".into()); + } + let mut running = execution.clone(); + running.phase = Phase::Running; + Ok(self.replace(execution, &running).await?.then_some(running)) + } + + pub async fn finish(&self, running: &Execution, outcome: Result) -> Result { + if running.phase != Phase::Running { + return Err("execution is not running".into()); + } + self.replace(running, &running.finished(outcome)).await + } + + /// Cancelling an unstarted request is a CAS. A late POST cannot start it. + pub async fn cancel_reserved(&self, execution: &Execution) -> Result { + if execution.phase != Phase::Reserved { + return Ok(false); + } + self.replace( + execution, + &execution.finished(Err("cancelled before admission".into())), + ) + .await + } + + /// Never delete on elapsed time, an HTTP error, or claim expiry. Both the + /// terminal result and the current task claim must still match atomically. + pub async fn release(&self, claim: &ClaimProof, execution: &Execution) -> Result { + if execution.phase != Phase::Finished { + return Err("cannot release a live execution".into()); + } + let key = execution_key(&self.prefix, &execution.target); + let (mut compares, mut operations) = self.completion_changes(execution).await?; + compares.extend([ + Compare::value(claim.key.as_str(), CompareOp::Equal, claim.token.as_bytes()), + Compare::value(key.as_str(), CompareOp::Equal, encode(execution)), + Compare::value( + target_lock_key(&self.prefix, &execution.target), + CompareOp::Equal, + execution_owner(execution), + ), + ]); + operations.extend([ + TxnOp::delete(key, None), + TxnOp::put( + target_lock_key(&self.prefix, &execution.target), + claim.token.clone(), + Some(etcd_client::PutOptions::new().with_lease(claim.lease_id)), + ), + ]); + self.transact(compares, operations).await + } + + async fn replace(&self, old: &Execution, new: &Execution) -> Result { + let key = execution_key(&self.prefix, &old.target); + self.transact( + vec![Compare::value(key.as_str(), CompareOp::Equal, encode(old))], + vec![TxnOp::put(key, encode(new), None)], + ) + .await + } + + async fn transact(&self, compares: Vec, operations: Vec) -> Result { + self.client + .clone() + .txn(Txn::new().when(compares).and_then(operations)) + .await + .map(|r| r.succeeded()) + .map_err(|e| e.to_string()) + } +} + +/// The executor owns this future independently of the request handler. On +/// cancellation, drop its scoped storage future *before* publishing Finished. +/// Callers must not pass a detached JoinHandle: dropping one does not stop it. +pub async fn execute_scoped( + work: F, + timeout: std::time::Duration, + mut cancel: tokio::sync::watch::Receiver, +) -> Result +where + F: std::future::Future>, +{ + // Keep the work in this inner scope: select! only drops branch borrows when + // the future was pinned outside it, which would publish completion too soon. + tokio::select! { + biased; + _ = async { + loop { + if *cancel.borrow_and_update() { return; } + if cancel.changed().await.is_err() { std::future::pending::<()>().await; } + } + } => Err("merge execution cancelled".into()), + result = tokio::time::timeout(timeout, work) => + result.unwrap_or_else(|_| Err("merge execution deadline exceeded".into())), + } +} + +#[derive(Clone, Debug, clap::Args)] +pub struct EtcdConfig { + /// Comma-separated etcd v3 endpoints. Scheduler state (task queue, + /// lease-based claims, per-experiment write locks) lives in etcd so several + /// stateless master replicas can share one queue. Required. + #[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, +} + +impl EtcdConfig { + pub async fn connect(&self) -> Result { + use etcd_client::{Certificate, ConnectOptions, Identity, TlsOptions}; + use std::time::Duration; + let config = self; + 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(String::from( + "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| 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| { + format!("failed to read ETCD_CLIENT_CERT '{cert}': {err}") + })?; + let key_pem = std::fs::read(key) + .map_err(|err| format!("failed to read ETCD_CLIENT_KEY '{key}': {err}"))?; + tls = tls.identity(Identity::from_pem(cert_pem, key_pem)); + } + (None, None) => {} + _ => { + return Err(String::from( + "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(String::from( + "ETCD_CA_CERT is required when configuring an etcd client certificate", + )); + } + let client = Client::connect(config.etcd_endpoints.clone(), Some(options)) + .await + .map_err(|e| e.to_string())?; + + Ok(Coordinator::new(client, self.etcd_prefix.clone())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + time::Duration, + }; + use tokio::sync::{oneshot, watch}; + + struct InFlight(Arc); + impl Drop for InFlight { + fn drop(&mut self) { + self.0.store(false, Ordering::SeqCst); + } + } + + #[tokio::test(start_paused = true)] + async fn deadline_drops_the_writer_before_terminal_acknowledgement() { + let writing = Arc::new(AtomicBool::new(false)); + let flag = writing.clone(); + let (_cancel, rx) = watch::channel(false); + let result = execute_scoped( + async move { + flag.store(true, Ordering::SeqCst); + let _guard = InFlight(flag); + std::future::pending::>().await + }, + Duration::from_secs(600), + rx, + ) + .await; + assert!(result.unwrap_err().contains("deadline")); + assert!(!writing.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn disconnected_caller_does_not_detach_uncontrolled_work() { + let writing = Arc::new(AtomicBool::new(false)); + let flag = writing.clone(); + let (cancel, rx) = watch::channel(false); + let (entered, began) = oneshot::channel(); + // The HTTP handler owns only admission, never this executor's lifetime. + let executor = tokio::spawn(execute_scoped( + async move { + flag.store(true, Ordering::SeqCst); + let _guard = InFlight(flag); + entered.send(()).unwrap(); + std::future::pending::>().await + }, + Duration::from_secs(600), + rx, + )); + began.await.unwrap(); + assert!(writing.load(Ordering::SeqCst)); + cancel.send(true).unwrap(); + assert!(executor.await.unwrap().unwrap_err().contains("cancelled")); + assert!(!writing.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn pre_cancelled_request_never_polls_storage() { + let (cancel, rx) = watch::channel(false); + cancel.send(true).unwrap(); + let result = execute_scoped( + async { panic!("must not execute storage") }, + Duration::from_secs(1), + rx, + ) + .await; + assert!(result.is_err()); + } + + // Isolated prefix on a LOCAL test etcd. These tests never use production. + async fn fixture() -> (Coordinator, Client, ClaimProof, i64) { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS") + .expect("ETCD_TEST_ENDPOINTS is required for ignored etcd tests"); + let mut client = Client::connect([endpoint], None).await.unwrap(); + let prefix = format!("/merge-execution-tests/{}", uuid::Uuid::new_v4()); + let lease = client.lease_grant(30, None).await.unwrap().id(); + let proof = ClaimProof { + key: format!("{prefix}/claims/test"), + token: "original".into(), + lease_id: lease, + }; + client + .put( + target_lock_key(&prefix, "table"), + proof.token.clone(), + Some(etcd_client::PutOptions::new().with_lease(lease)), + ) + .await + .unwrap(); + client + .put( + proof.key.clone(), + proof.token.clone(), + Some(etcd_client::PutOptions::new().with_lease(lease)), + ) + .await + .unwrap(); + ( + Coordinator::new(client.clone(), prefix), + client, + proof, + lease, + ) + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn failure_budget_survives_reconnect_and_lost_claim_cannot_clear_it() { + let (coordinator, mut client, proof, lease) = fixture().await; + for count in 1..=3 { + let failure = coordinator + .record_failure(&proof, "table", "worker", "temporary failure") + .await + .unwrap(); + assert_eq!(failure.consecutive_attempts, count); + } + let reconnected = Coordinator::new(client.clone(), coordinator.prefix.clone()); + let failure = reconnected + .failure("table", "worker") + .await + .unwrap() + .unwrap(); + assert_eq!(failure.consecutive_attempts, 3); + assert!(failure.next_retry_ms > failure.last_failure_ms + 20_000); + reconnected + .record_failure(&proof, "table2", "worker", "schema mismatch") + .await + .unwrap(); + let (page, next) = reconnected.failure_page(None, 1).await.unwrap(); + assert_eq!(page.len(), 1); + let (second, end) = reconnected.failure_page(next.as_deref(), 1).await.unwrap(); + assert_eq!(second.len(), 1); + assert_ne!(page[0].target, second[0].target); + assert!(end.is_none()); + client.lease_revoke(lease).await.unwrap(); + assert!(reconnected + .clear_failure(&proof, "table", "worker") + .await + .is_err()); + assert!(reconnected + .record_failure(&proof, "table", "worker", "stale writer") + .await + .is_err()); + assert_eq!( + reconnected + .failure("table", "worker") + .await + .unwrap() + .unwrap() + .consecutive_attempts, + 3 + ); + client + .put(proof.key.clone(), proof.token.clone(), None) + .await + .unwrap(); + reconnected + .clear_failure(&proof, "table", "worker") + .await + .unwrap(); + assert!(reconnected + .failure("table", "worker") + .await + .unwrap() + .is_none()); + client + .delete( + coordinator.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn revoked_claim_cannot_release_surviving_server_work() { + let (coordinator, mut client, proof, lease) = fixture().await; + let operation = Execution::new("table", "worker", "boot", 600); + assert!(coordinator.reserve(&proof, &operation).await.unwrap()); + let running = coordinator.start(&operation).await.unwrap().unwrap(); + client.lease_revoke(lease).await.unwrap(); + assert_eq!( + coordinator.get("table").await.unwrap(), + Some(running.clone()) + ); + let next = ClaimProof { + key: proof.key.clone(), + token: "replacement".into(), + lease_id: 0, + }; + client + .put(next.key.clone(), next.token.clone(), None) + .await + .unwrap(); + assert!(!coordinator + .reserve(&next, &Execution::new("table", "worker2", "boot2", 600)) + .await + .unwrap()); + assert!(coordinator.release(&next, &running).await.is_err()); + assert!(coordinator + .finish(&running, Err("cancelled".into())) + .await + .unwrap()); + let done = coordinator.get("table").await.unwrap().unwrap(); + assert!(!coordinator.release(&proof, &done).await.unwrap()); + assert!(coordinator + .failure("table", "worker") + .await + .unwrap() + .is_none()); + assert!(coordinator.release(&next, &done).await.unwrap()); + assert_eq!( + coordinator + .failure("table", "worker") + .await + .unwrap() + .unwrap() + .consecutive_attempts, + 1, + "releasing a failed execution must atomically persist the retry budget" + ); + assert!(coordinator + .reserve(&next, &Execution::new("table", "worker2", "boot2", 600)) + .await + .unwrap()); + client + .delete( + coordinator.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn delayed_http_start_cannot_resurrect_cancelled_operation() { + let (coordinator, mut client, proof, lease) = fixture().await; + let old = Execution::new("table", "worker", "boot", 600); + assert!(coordinator.reserve(&proof, &old).await.unwrap()); + assert!(coordinator.cancel_reserved(&old).await.unwrap()); + let done = coordinator.get("table").await.unwrap().unwrap(); + assert!(coordinator.release(&proof, &done).await.unwrap()); + let new = Execution::new("table", "worker2", "boot2", 600); + assert!(coordinator.reserve(&proof, &new).await.unwrap()); + assert!(coordinator.start(&old).await.unwrap().is_none()); + assert_eq!(coordinator.get("table").await.unwrap(), Some(new)); + client.lease_revoke(lease).await.unwrap(); + client + .delete( + coordinator.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } +} diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index bebeed0..743051b 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -13,6 +13,7 @@ name = "lance-context-server" path = "src/main.rs" [dependencies] +lance-context-merge = { path = "../lance-context-merge" } lance-context-core = { version = "0.6.7", path = "../lance-context-core" } lance-context-api = { version = "0.6.7", path = "../lance-context-api" } lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" } @@ -32,6 +33,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } uuid = { version = "1", features = ["v4"] } [dev-dependencies] +etcd-client = { version = "0.19", features = ["tls"] } # Snapshotting recorder so tests can assert which metric series an operation # emitted, without installing a process-global Prometheus exporter. metrics-util = { version = "0.19", default-features = false, features = [ diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 5ef4fd9..6ff12b7 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -4,6 +4,15 @@ use clap::Parser; #[command(name = "lance-context-server")] #[command(about = "REST API server for lance-context")] pub struct ServerConfig { + /// Configure etcd to enable the owned merge execution protocol. + #[command(flatten)] + pub merge_etcd: lance_context_merge::EtcdConfig, + + /// Deadline for scoped merge work, including slot acquisition. Already + /// started manifest writes must drain before ownership can be released. + #[arg(long, env = "MERGE_EXECUTION_TIMEOUT_SECS", default_value_t = 600)] + pub merge_execution_timeout_secs: u64, + #[arg(long, default_value = "0.0.0.0")] pub host: String, diff --git a/crates/lance-context-server/src/main.rs b/crates/lance-context-server/src/main.rs index 5a9862a..ce0be21 100644 --- a/crates/lance-context-server/src/main.rs +++ b/crates/lance-context-server/src/main.rs @@ -1,5 +1,6 @@ mod config; mod error; +mod merge_execution; mod routes; mod state; mod sweeper; @@ -87,6 +88,7 @@ async fn main() { // Connections have drained. Deterministically close every resident writer // before the runtime tears down. + state.merge_executions.shutdown().await; state.shutdown().await; } diff --git a/crates/lance-context-server/src/merge_execution.rs b/crates/lance-context-server/src/merge_execution.rs new file mode 100644 index 0000000..0b25265 --- /dev/null +++ b/crates/lance-context-server/src/merge_execution.rs @@ -0,0 +1,362 @@ +//! HTTP handlers only admit/cancel work. The owned executor, not the connection, +//! holds the durable execution fence until the scoped storage future is gone. +use crate::{ + error::AppError, + routes::{generic, rollouts}, + state::AppState, +}; +use axum::{extract::State, http::StatusCode, Json}; +use futures::FutureExt; +use lance_context_merge::{execute_scoped, Coordinator, Execution, Phase}; +use std::{collections::HashMap, sync::Arc, time::Duration}; +use tokio::sync::{watch, Mutex}; + +pub struct Executions { + coordinator: Option, + instance: String, + timeout_secs: u64, + running: Mutex>>, +} + +impl Executions { + pub fn new(coordinator: Option, timeout_secs: u64) -> Self { + Self { + coordinator, + instance: uuid::Uuid::new_v4().to_string(), + timeout_secs, + running: Mutex::new(HashMap::new()), + } + } + pub(crate) fn enabled(&self) -> bool { + self.coordinator.is_some() + } + + /// Invoked after HTTP admission has stopped. Detached executors are not + /// drained by axum's connection shutdown. + pub(crate) async fn shutdown(&self) { + for cancel in self.running.lock().await.values() { + let _ = cancel.send(true); + } + while !self.running.lock().await.is_empty() { + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + fn coordinator(&self) -> Result { + self.coordinator.clone().ok_or_else(|| { + AppError::Overloaded("owned merge execution requires ETCD_ENDPOINTS".into()) + }) + } +} + +pub async fn capabilities( + State(state): State>, +) -> Result, AppError> { + state.merge_executions.coordinator()?; + Ok(Json( + serde_json::json!({"protocol": 1, "instance": state.merge_executions.instance, + "timeout_secs": state.merge_executions.timeout_secs}), + )) +} + +pub async fn start( + State(state): State>, + Json(execution): Json, +) -> Result { + let coordinator = state.merge_executions.coordinator()?; + if execution.instance != state.merge_executions.instance + || execution.phase != Phase::Reserved + || execution.timeout_secs == 0 + || execution.timeout_secs > state.merge_executions.timeout_secs + { + return Err(AppError::InvalidRequest( + "merge executor incarnation or deadline mismatch".into(), + )); + } + let name = execution + .target + .strip_prefix("generic:") + .unwrap_or(&execution.target); + lance_context_core::validate_store_name(name).map_err(AppError::InvalidRequest)?; + let mut running = state.merge_executions.running.lock().await; + if running.contains_key(&execution.id) { + return Ok(StatusCode::ACCEPTED); + } + let (cancel, cancelled) = watch::channel(false); + running.insert(execution.id.clone(), cancel); + // Spawn before any fallible admission I/O. Losing the HTTP connection must + // never leave a Running fence without an executor (or cancel a live write). + let owned = state.clone(); + tokio::spawn(async move { + let result = run(owned.clone(), coordinator, execution.clone(), cancelled).await; + if let Err(error) = result { + tracing::error!(id = %execution.id, target = %execution.target, %error, "merge execution ownership unresolved"); + } + owned + .merge_executions + .running + .lock() + .await + .remove(&execution.id); + }); + Ok(StatusCode::ACCEPTED) +} + +async fn run( + state: Arc, + coordinator: Coordinator, + execution: Execution, + cancelled: watch::Receiver, +) -> Result<(), String> { + // CAS errors are ambiguous: never retry storage work. Reconciliation sees + // the durable fence, so an uncertain admission cannot release ownership. + let running = loop { + match coordinator.start(&execution).await { + Ok(Some(running)) => break running, + Ok(None) => match coordinator.get(&execution.target).await { + Ok(Some(current)) + if current.id == execution.id && current.phase == Phase::Running => + { + break current + } + Ok(_) => return Ok(()), + Err(_) => {} + }, + Err(error) => { + tracing::warn!(id = %execution.id, %error, "reconciling uncertain execution admission") + } + } + tokio::time::sleep(Duration::from_secs(1)).await; + }; + let target = execution.target.clone(); + let mut slot = None; + let work = async { + slot = state.acquire_merge_slot().await; + let result = if let Some(name) = target.strip_prefix("generic:") { + generic::merge_generic_wal_owned( + State(state.clone()), + axum::extract::Path(name.to_string()), + ) + .await + .map(|Json(reply)| reply["reclaimed"].as_u64().unwrap_or(0) as usize) + } else { + rollouts::merge_wal_owned(State(state.clone()), axum::extract::Path(target.clone())) + .await + .map(|Json(reply)| reply.reclaimed) + }; + match result { + Ok(n) => Ok(n), + Err(AppError::NotFound(ref message)) + if message == &format!("Rollout store '{}' does not exist", target) + || message + == &format!( + "Generic store '{}' does not exist", + target.strip_prefix("generic:").unwrap_or(&target) + ) => + { + Ok(0) + } + Err(e) => Err(format!("{e:?}")), + } + }; + let write_scope = lance_context_core::merge_write_scope::MergeWriteScope::new(); + let outcome = std::panic::AssertUnwindSafe(execute_scoped( + write_scope.run(work), + Duration::from_secs(execution.timeout_secs), + cancelled, + )) + .catch_unwind() + .await + .unwrap_or_else(|_| Err("merge executor panicked".into())); + // Top-level cancellation cannot acknowledge a write still in flight at + // object storage. Join those leaf commits before publishing Finished. + write_scope.drain().await; + // Do not admit another memory-heavy merge into this slot while a + // cancelled execution is still joining storage commits. + drop(slot); + // Storage work has ended. An etcd outage may delay acknowledgement, but + // cannot erase the fence or cause the write to be executed twice. + loop { + match coordinator.finish(&running, outcome.clone()).await { + Ok(true) => break, + Ok(false) => { + if coordinator.get(&execution.target).await?.as_ref() + == Some(&running.finished(outcome.clone())) + { + break; + } + return Err("merge execution fence changed unexpectedly".into()); + } + Err(error) => { + tracing::warn!(id = %execution.id, %error, "retrying merge terminal acknowledgement") + } + } + tokio::time::sleep(Duration::from_secs(1)).await; + } + Ok(()) +} + +pub async fn cancel( + State(state): State>, + Json(execution): Json, +) -> Result { + let coordinator = state.merge_executions.coordinator()?; + if coordinator + .cancel_reserved(&execution) + .await + .map_err(AppError::Internal)? + { + return Ok(StatusCode::OK); + } + if let Some(cancel) = state + .merge_executions + .running + .lock() + .await + .get(&execution.id) + { + let _ = cancel.send(true); + return Ok(StatusCode::ACCEPTED); + } + // Absence from the local map is NOT proof that the old process stopped. + // The caller must reconcile the durable result, not interpret HTTP 404 as + // permission for a competing write. + Err(AppError::NotFound( + "execution is not owned by this worker incarnation".into(), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use lance_context_merge::{ClaimProof, EtcdConfig}; + use tokio::io::AsyncWriteExt; + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn disconnected_http_call_is_cancelled_before_ownership_handoff() { + let endpoints = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let prefix = format!("/server-merge-test/{}", uuid::Uuid::new_v4()); + let coordinator = EtcdConfig { + etcd_endpoints: endpoints.split(',').map(str::to_string).collect(), + etcd_prefix: prefix.clone(), + etcd_username: None, + etcd_password: None, + etcd_ca_cert: None, + etcd_client_cert: None, + etcd_client_key: None, + } + .connect() + .await + .unwrap(); + // Test-only claim creation through the isolated etcd connection. + let mut client = + etcd_client::Client::connect(endpoints.split(',').collect::>(), None) + .await + .unwrap(); + let proof = ClaimProof { + key: format!("{prefix}/claim"), + token: "owner".into(), + lease_id: 0, + }; + client + .put(proof.key.clone(), proof.token.clone(), None) + .await + .unwrap(); + client + .put( + lance_context_merge::target_lock_key(&prefix, "blocked"), + proof.token.clone(), + None, + ) + .await + .unwrap(); + let dir = tempfile::tempdir().unwrap(); + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + state.merge_executions = Executions::new(Some(coordinator.clone()), 600); + let state = Arc::new(state); + let name = "blocked"; + // Create via the real API so registry and resident handle agree. + let _ = rollouts::create_rollout_store( + State(state.clone()), + Json(lance_context_api::CreateRolloutStoreRequest { + name: name.into(), + storage_options: None, + }), + ) + .await + .unwrap(); + let store = state.get_or_open_rollout_store(name).await.unwrap(); + let held = store.write().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = crate::routes::router().with_state(state.clone()); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let execution = Execution::new( + name, + &format!("http://{address}"), + &state.merge_executions.instance, + 600, + ); + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + let body = serde_json::to_string(&execution).unwrap(); + let mut connection = tokio::net::TcpStream::connect(address).await.unwrap(); + connection.write_all(format!("POST /api/v1/internal/merge-executor/start HTTP/1.1\r\nHost: {address}\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if coordinator.get(name).await.unwrap().unwrap().phase == Phase::Running { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + drop(connection); + assert_eq!( + coordinator.get(name).await.unwrap().unwrap().phase, + Phase::Running + ); + assert!(!coordinator + .reserve(&proof, &Execution::new(name, "other", "other", 1)) + .await + .unwrap()); + cancel(State(state.clone()), Json(execution.clone())) + .await + .unwrap(); + let finished = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let current = coordinator.get(name).await.unwrap().unwrap(); + if current.phase == Phase::Finished { + break current; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(finished.error.as_ref().unwrap().contains("cancelled")); + assert!(coordinator.release(&proof, &finished).await.unwrap()); + drop(held); + let next = Execution::new(name, &execution.endpoint, &execution.instance, 30); + assert!(coordinator.reserve(&proof, &next).await.unwrap()); + start(State(state.clone()), Json(next)).await.unwrap(); + let done = tokio::time::timeout(Duration::from_secs(30), async { + loop { + let current = coordinator.get(name).await.unwrap().unwrap(); + if current.phase == Phase::Finished { + break current; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(done.error.is_none(), "retry must run: {done:?}"); + assert!(coordinator.release(&proof, &done).await.unwrap()); + server.abort(); + client.delete(proof.key, None).await.unwrap(); + } +} diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index c5fe858..d0f9d87 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -273,7 +273,19 @@ pub async fn merge_generic_wal( State(state): State>, Path(name): Path, ) -> Result, AppError> { + if state.merge_executions.enabled() { + return Err(AppError::Overloaded( + "use the owned merge executor protocol".into(), + )); + } let _slot = state.acquire_merge_slot().await; + merge_generic_wal_owned(State(state), Path(name)).await +} + +pub(crate) async fn merge_generic_wal_owned( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { let store = state.get_or_open_generic_store(&name).await?; // Same prepare/commit split as the sweeper: the object-storage read of the // generations runs under the shared lock so the store keeps serving. diff --git a/crates/lance-context-server/src/routes/mod.rs b/crates/lance-context-server/src/routes/mod.rs index 7f5b2ee..e23c2ab 100644 --- a/crates/lance-context-server/src/routes/mod.rs +++ b/crates/lance-context-server/src/routes/mod.rs @@ -18,6 +18,18 @@ use crate::state::AppState; pub fn router() -> Router> { Router::new() + .route( + "/api/v1/internal/merge-executor", + get(crate::merge_execution::capabilities), + ) + .route( + "/api/v1/internal/merge-executor/start", + post(crate::merge_execution::start), + ) + .route( + "/api/v1/internal/merge-executor/cancel", + post(crate::merge_execution::cancel), + ) .route("/api/v1/health", get(health::health_check)) .route("/api/v1/contexts", post(contexts::create_context)) .route("/api/v1/contexts", get(contexts::list_contexts)) diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index 5f401af..e1ade66 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -728,7 +728,19 @@ pub async fn merge_wal( State(state): State>, Path(name): Path, ) -> Result, AppError> { + if state.merge_executions.enabled() { + return Err(AppError::Overloaded( + "use the owned merge executor protocol".into(), + )); + } let _slot = state.acquire_merge_slot().await; + merge_wal_owned(State(state), Path(name)).await +} + +pub(crate) async fn merge_wal_owned( + State(state): State>, + Path(name): Path, +) -> Result, AppError> { let store_lock = state.get_or_open_rollout_store(&name).await?; // Split by lock scope: seal + read every flushed generation under the // *read* lock so ingest on this store keeps running, then take the write diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index bc3ce77..61632ae 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -76,6 +76,7 @@ impl StoreHandles { } pub struct AppState { + pub merge_executions: crate::merge_execution::Executions, pub stores: RwLock>>>, /// Bounded LRU of resident rollout-store handles. /// @@ -321,7 +322,35 @@ impl AppState { let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) .await .map_err(AppError::from_lance)?; + if !config.merge_etcd.etcd_endpoints.is_empty() + && (config.rollout_merge_after_generations != 0 + || config.rollout_cleanup_interval_secs != 0) + { + return Err(AppError::InvalidRequest( + "owned merge execution requires self-merge sweepers disabled".into(), + )); + } + let merge_coordinator = if config.merge_etcd.etcd_endpoints.is_empty() { + None + } else { + Some( + config + .merge_etcd + .connect() + .await + .map_err(AppError::Internal)?, + ) + }; + if config.merge_execution_timeout_secs == 0 { + return Err(AppError::InvalidRequest( + "MERGE_EXECUTION_TIMEOUT_SECS must be positive".into(), + )); + } Ok(Self { + merge_executions: crate::merge_execution::Executions::new( + merge_coordinator, + config.merge_execution_timeout_secs, + ), stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new(capacity)), rollout_handles: StoreHandles::default(), @@ -391,6 +420,7 @@ impl AppState { .await .expect("open test datagen registry"); Self { + merge_executions: crate::merge_execution::Executions::new(None, 600), stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new( NonZeroUsize::new(DEFAULT_ROLLOUT_CACHE_CAPACITY).unwrap(), diff --git a/docs/merge-recovery.md b/docs/merge-recovery.md new file mode 100644 index 0000000..1a30039 --- /dev/null +++ b/docs/merge-recovery.md @@ -0,0 +1,104 @@ +# Owned merge recovery + +A serial merge fan-out previously had no request deadline and could report +success after one worker failed if another reclaimed generations. A client-side +HTTP timeout alone is unsafe: the worker may still be committing after the +master moves to the next shard. + +This change admits a durable execution before sending the worker request. The +worker owns the merge independently of that request, cancels preparation at its +deadline, and joins already-started base/shard manifest commits before publishing +its terminal result. The master advances only after this acknowledgement. Before +retrying a cancelled commit, the store refreshes its base manifest. WAL drain +remains a relative edit preserving generations flushed during recovery. + +## Ownership + +- `MergeWal`, `Compact`, `IndexId`, and `Repair` share the target lock. +- Admission atomically replaces the lease-backed target lock with a persistent + execution marker. Existing helpers that respect target locks remain excluded. +- A separate leased merge claim ensures only one scheduler reconciles an + execution, including non-deduplicated tasks in dependency chains. +- Worker execution and the persistent target lock survive task lease loss. + Only a terminal execution can restore the target lock to a live task lease. +- Cancelling an unstarted request uses CAS, preventing a delayed POST from + starting after another shard has been admitted. +- Graceful worker shutdown cancels and drains owned executions before closing + resident stores. SIGKILL/OOM is **not** a graceful acknowledgement. + +## Repeated failures + +Failures persist under `merge-failures//` in the master +etcd namespace. The record includes total consecutive failed attempts, diagnostic +class, last error (capped at 4096 characters), next retry time, and attention flag. +Recreating a task or master does not reset this budget. Executed failures are +recorded atomically with releasing terminal execution ownership; a crash in the +master between these two steps cannot lose the failure. + +The first three transient failures use two-second targeted retries, with healthy +shards visited before retrying failed shards. After that, each due task makes +one probe, backing off from 30 seconds to at most 15 minutes. At 15 consecutive +failures, probes are hourly and `needs_attention` is set. Missing data, corruption, +configuration/authentication errors, and unresolved execution ownership set the +attention flag immediately and use hourly probes. Classification uses diagnostic +text conservatively; it never authorizes data removal. + +Successful probes clear the shard's failure record. Failed/deferred shards keep +the overall task failed even if healthy shards reclaim WAL. Per-table merge +cooldown is replaced by per-endpoint backoff so healthy shards can process new +WAL. A 15-second scheduler poll walks at most 256 failure records per page and +enqueues due targets through existing dedupe. It reads metadata only and does +not wait for the 600-second stats sweep. For a large failure ledger, polling +latency grows with page count. Removed endpoints are not automatically probed. + +`GET /api/v1/scheduler/merge-failures?after=` returns at most 256 records +and a `next` cursor. `needs_attention`, `last_error`, and `next_retry_ms` distinguish +failed shards from recovered ones. Existing task error/counter metrics also show +failure. This is an observable circuit state, not an external alert integration. + +The scheduler does **not** automatically enqueue destructive fragment repair for +a merge error. Missing/corrupt files require diagnosis and restoration or an +explicitly reviewed repair; retaining WAL is mandatory. Other maintenance task +repair behavior is unchanged. + +## Limits and required rollout validation + +This is not yet a complete automatic crash-recovery protocol. If a process dies +or a storage commit cannot be joined, there may be no terminal acknowledgement. +After the merge deadline plus a 60-second reconciliation allowance (or 60 seconds +for inherited work), the master records unresolved ownership, fails that task, +and releases its scheduler slot **without deleting the execution or target +lock**. Other tables continue. The affected table remains fenced until the old +executor acknowledges termination. A changed worker UUID, a Ready replacement, +an expired lease, and elapsed time are not sufficient proof. + +Automatic recovery of that orphan requires a storage fencing/barrier protocol +and authoritative termination evidence; neither is implemented here. Do not +manually delete the fence or deploy this as a claimed complete fix for OOM/node +loss. The worker deadline bounds its scoped merge future; draining a backend +that never completes its commit is not bounded by that deadline. + +The execution fence protects scheduler-coordinated maintenance and helpers that +honor target locks. It does not retrofit storage fencing onto arbitrary manual +writers or solve lease loss during the master's own local compaction/index job. + +This wire protocol requires coordinated master/worker rollout. There is no +fallback to unbounded legacy HTTP merges: + +- Workers need `ETCD_ENDPOINTS` and the **same** `ETCD_PREFIX` and credentials as + masters. `MERGE_EXECUTION_TIMEOUT_SECS` defaults to 600; the master caps the + requested deadline at 600. Zero is rejected. +- With the coordinator enabled, legacy merge routes reject admission. Disable + worker self-merge thresholds and cleanup timers; startup checks enforce this. +- `INDEX_BEFORE_MERGE` remains accepted for compatibility, but index preparation + is done by the worker inside the owned execution. +- Drain/reconcile legacy in-flight merges before enabling the new protocol. + Masters must continue their other jobs. Recovery helpers must preserve target + locks; do not clear a keeper's pause while its delegate is alive. +- Keep ordinary workers at six merge slots and 20 GiB. This patch increases + neither memory budgets nor same-table commit concurrency. + +Before production rollout, staging must cover realistic blobs, sustained writes, +WAL pending/reclamation rates, memory/OOM counts, commit conflicts, and p95/p99 +latency. Local cancellation and HTTP-disconnect tests are necessary but cannot +substitute for this soak or the missing orphan recovery protocol. From c8838b856384766d6f86350589491bbd6be25f99 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Thu, 1 Oct 2026 17:42:52 +0000 Subject: [PATCH 2/9] fix: retain ownership for ambiguous manifest commits --- .../src/merge_write_scope.rs | 70 +++++++++++++++++-- crates/lance-context-master/Cargo.toml | 2 +- .../src/merge_execution.rs | 11 ++- crates/lance-context-merge/src/lib.rs | 38 ++++++++++ crates/lance-context-server/Cargo.toml | 2 +- .../src/merge_execution.rs | 38 +++++++--- docs/merge-recovery.md | 7 +- 7 files changed, 147 insertions(+), 21 deletions(-) diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index 061ce0e..9af67d2 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -22,6 +22,7 @@ tokio::task_local! { static CURRENT: Arc; } struct Progress { closed: bool, active: usize, + uncertain: bool, } #[derive(Debug, Default)] @@ -39,6 +40,14 @@ impl MergeWriteScope { CURRENT.scope(self.clone(), future).await } + pub fn has_uncertain_commit(&self) -> bool { + self.progress.lock().unwrap().uncertain + } + + fn mark_uncertain(&self) { + self.progress.lock().unwrap().uncertain = true; + } + /// Call only after dropping/joining the merge future. Prevent new commits /// and join all manifest writes that started before cancellation. pub async fn drain(&self) { @@ -55,11 +64,17 @@ impl MergeWriteScope { } } -struct Active(Arc); +struct Active { + scope: Arc, + completed: bool, +} impl Drop for Active { fn drop(&mut self) { - self.0.progress.lock().unwrap().active -= 1; - self.0.changed.notify_waiters(); + let mut progress = self.scope.progress.lock().unwrap(); + progress.active -= 1; + progress.uncertain |= !self.completed; + drop(progress); + self.scope.changed.notify_waiters(); } } @@ -79,11 +94,19 @@ where } progress.active += 1; } - let active = Active(scope); + let active = Active { + scope, + completed: false, + }; let (send, receive) = oneshot::channel(); tokio::spawn(async move { - let _active = active; + let mut active = active; let result = write.await; + if result.is_err() { + active.scope.mark_uncertain(); + } + active.completed = true; + drop(active); let _ = send.send(result); }); receive @@ -111,6 +134,7 @@ impl CommitHandler for GuardedCommit { let mut owned_manifest = manifest.clone(); let path = base_path.clone(); let store = object_store.clone(); + let scope = CURRENT.try_with(Arc::clone).ok(); let (location, committed) = shield(async move { let location = inner .commit( @@ -123,6 +147,14 @@ impl CommitHandler for GuardedCommit { transaction, ) .await; + // Only a definite conflict proves this conditional commit did + // not happen. A transport/storage error can arrive after the + // remote service accepted it, even though the Rust future ended. + if matches!(&location, Err(CommitError::OtherError(_))) { + if let Some(scope) = scope { + scope.mark_uncertain(); + } + } Ok::<_, CommitError>((location, owned_manifest)) }) .await?; @@ -196,6 +228,33 @@ mod tests { use super::*; use std::sync::atomic::{AtomicBool, Ordering}; + #[tokio::test] + async fn failed_or_panicked_leaf_commit_needs_reconciliation() { + let scope = MergeWriteScope::new(); + let result = scope + .run(shield(async { + Err::<(), _>(Error::io("lost storage response")) + })) + .await; + assert!(result.is_err()); + scope.drain().await; + assert!( + scope.has_uncertain_commit(), + "an error is not proof that remote storage did not commit" + ); + let panicked = MergeWriteScope::new(); + let result = panicked + .run(shield(async { + panic!("commit panic"); + #[allow(unreachable_code)] + Ok::<(), Error>(()) + })) + .await; + assert!(result.is_err()); + panicked.drain().await; + assert!(panicked.has_uncertain_commit()); + } + #[tokio::test] async fn dropping_merge_waits_for_the_surviving_manifest_write() { let scope = MergeWriteScope::new(); @@ -227,6 +286,7 @@ mod tests { release.send(()).unwrap(); drain.await.unwrap(); assert!(committed.load(Ordering::SeqCst)); + assert!(!scope.has_uncertain_commit()); assert!( scope .run(shield(async { Ok::<_, Error>(()) })) diff --git a/crates/lance-context-master/Cargo.toml b/crates/lance-context-master/Cargo.toml index 3ac0027..dc1a0ca 100644 --- a/crates/lance-context-master/Cargo.toml +++ b/crates/lance-context-master/Cargo.toml @@ -13,7 +13,7 @@ name = "lance-context-master" path = "src/main.rs" [dependencies] -lance-context-merge = { path = "../lance-context-merge" } +lance-context-merge = { version = "0.1.0", path = "../lance-context-merge" } lance-context-core = { version = "0.6.3", path = "../lance-context-core" } lance-context-api = { version = "0.6.3", path = "../lance-context-api" } lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" } diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index 1b06226..32713ff 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/crates/lance-context-master/src/merge_execution.rs @@ -27,7 +27,7 @@ pub(crate) async fn run_merge_wal( // Reconcile first, before any index/base-table mutation or new fan-out. if let Some(old) = coordinator.get(&claim.task.target).await? { let endpoint = old.endpoint.clone(); - if old.phase == Phase::Running { + if matches!(old.phase, Phase::Running | Phase::Uncertain) { if let Some(failure) = coordinator.failure(&claim.task.target, &endpoint).await? { if failure.class == lance_context_merge::failure::FailureClass::OwnershipUnresolved && failure.next_retry_ms > lance_context_merge::failure::now_ms() @@ -261,6 +261,15 @@ async fn reconcile_with_grace( continue; } }; + if current.phase == Phase::Uncertain { + let error = current + .error + .unwrap_or_else(|| "merge ownership unresolved: ambiguous storage commit".into()); + coordinator + .record_failure(proof, ¤t.target, ¤t.endpoint, &error) + .await?; + return Err(error); + } if current.phase == Phase::Finished { if !coordinator.release(proof, ¤t).await? { return Err("task claim lost while releasing terminal merge execution".into()); diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index 855a938..e87ebd4 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -23,6 +23,7 @@ pub struct ClaimProof { pub enum Phase { Reserved, Running, + Uncertain, Finished, } @@ -161,6 +162,17 @@ impl Coordinator { self.replace(running, &running.finished(outcome)).await } + /// The Rust write future ended with an ambiguous storage result. This is + /// durable diagnostic evidence, not permission to hand off storage writes. + pub async fn report_uncertain(&self, running: &Execution, error: String) -> Result { + if running.phase != Phase::Running { + return Err("execution is not running".into()); + } + let mut next = running.finished(Err(error)); + next.phase = Phase::Uncertain; + self.replace(running, &next).await + } + /// Cancelling an unstarted request is a CAS. A late POST cannot start it. pub async fn cancel_reserved(&self, execution: &Execution) -> Result { if execution.phase != Phase::Reserved { @@ -445,6 +457,32 @@ mod tests { ) } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn ambiguous_storage_response_cannot_release_the_execution() { + let (coordinator, mut client, proof, lease) = fixture().await; + let operation = Execution::new("table", "worker", "boot", 1); + assert!(coordinator.reserve(&proof, &operation).await.unwrap()); + let running = coordinator.start(&operation).await.unwrap().unwrap(); + assert!(coordinator + .report_uncertain(&running, "storage response lost".into()) + .await + .unwrap()); + let uncertain = coordinator.get("table").await.unwrap().unwrap(); + assert_eq!(uncertain.phase, Phase::Uncertain); + assert!(coordinator.release(&proof, &uncertain).await.is_err()); + assert!(!coordinator.finish(&running, Ok(1)).await.unwrap()); + client.lease_revoke(lease).await.unwrap(); + assert_eq!(coordinator.get("table").await.unwrap(), Some(uncertain)); + client + .delete( + coordinator.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn failure_budget_survives_reconnect_and_lost_claim_cannot_clear_it() { diff --git a/crates/lance-context-server/Cargo.toml b/crates/lance-context-server/Cargo.toml index 743051b..6d838f4 100644 --- a/crates/lance-context-server/Cargo.toml +++ b/crates/lance-context-server/Cargo.toml @@ -13,7 +13,7 @@ name = "lance-context-server" path = "src/main.rs" [dependencies] -lance-context-merge = { path = "../lance-context-merge" } +lance-context-merge = { version = "0.1.0", path = "../lance-context-merge" } lance-context-core = { version = "0.6.7", path = "../lance-context-core" } lance-context-api = { version = "0.6.7", path = "../lance-context-api" } lance-context-metrics = { version = "0.1.0", path = "../lance-context-metrics" } diff --git a/crates/lance-context-server/src/merge_execution.rs b/crates/lance-context-server/src/merge_execution.rs index 0b25265..40b038a 100644 --- a/crates/lance-context-server/src/merge_execution.rs +++ b/crates/lance-context-server/src/merge_execution.rs @@ -174,21 +174,37 @@ async fn run( // Do not admit another memory-heavy merge into this slot while a // cancelled execution is still joining storage commits. drop(slot); - // Storage work has ended. An etcd outage may delay acknowledgement, but - // cannot erase the fence or cause the write to be executed twice. + // A completed Rust future can still have an ambiguous remote result. + // Preserve that distinction durably instead of authorizing another writer. + let uncertain = write_scope.has_uncertain_commit(); + let mut expected = running.finished(outcome.clone()); + if uncertain { + expected.phase = Phase::Uncertain; + expected.reclaimed = 0; + expected.error = Some(format!( + "merge ownership unresolved: manifest commit result unknown; {:?}", + outcome + )); + } loop { - match coordinator.finish(&running, outcome.clone()).await { + let acknowledged = if uncertain { + coordinator + .report_uncertain(&running, expected.error.clone().unwrap()) + .await + } else { + coordinator.finish(&running, outcome.clone()).await + }; + match acknowledged { Ok(true) => break, - Ok(false) => { - if coordinator.get(&execution.target).await?.as_ref() - == Some(&running.finished(outcome.clone())) - { - break; + Ok(false) => match coordinator.get(&execution.target).await { + Ok(Some(current)) if current == expected => break, + Ok(_) => return Err("merge execution fence changed unexpectedly".into()), + Err(error) => { + tracing::warn!(id = %execution.id, %error, "reconciling merge outcome acknowledgement") } - return Err("merge execution fence changed unexpectedly".into()); - } + }, Err(error) => { - tracing::warn!(id = %execution.id, %error, "retrying merge terminal acknowledgement") + tracing::warn!(id = %execution.id, %error, "retrying merge outcome acknowledgement") } } tokio::time::sleep(Duration::from_secs(1)).await; diff --git a/docs/merge-recovery.md b/docs/merge-recovery.md index 1a30039..d9f3928 100644 --- a/docs/merge-recovery.md +++ b/docs/merge-recovery.md @@ -21,6 +21,9 @@ remains a relative edit preserving generations flushed during recovery. execution, including non-deduplicated tasks in dependency chains. - Worker execution and the persistent target lock survive task lease loss. Only a terminal execution can restore the target lock to a live task lease. + Ambiguous manifest errors/panics publish `uncertain`, which cannot release + ownership: an ended Rust future is not evidence that remote storage rejected + its write. Definite conditional-commit conflicts are treated separately. - Cancelling an unstarted request uses CAS, preventing a delayed POST from starting after another shard has been admitted. - Graceful worker shutdown cancels and drains owned executions before closing @@ -63,8 +66,8 @@ repair behavior is unchanged. ## Limits and required rollout validation -This is not yet a complete automatic crash-recovery protocol. If a process dies -or a storage commit cannot be joined, there may be no terminal acknowledgement. +This is not yet a complete automatic crash-recovery protocol. If a process dies, a storage commit cannot be joined, or a commit returns an +ambiguous result, there may be no terminal acknowledgement. After the merge deadline plus a 60-second reconciliation allowance (or 60 seconds for inherited work), the master records unresolved ownership, fails that task, and releases its scheduler slot **without deleting the execution or target From 6f09699a5d0f4f051789186a6cc265dc2b0bcc12 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Thu, 1 Oct 2026 18:22:16 +0000 Subject: [PATCH 3/9] fix: fence admitted manifest versions before recovering stalled merges --- .../src/merge_write_scope.rs | 228 +++++++++++++++++- .../lance-context-core/src/rollout_store.rs | 107 +++++++- crates/lance-context-core/src/store_base.rs | 19 +- .../src/merge_execution.rs | 220 ++++++++++++++--- crates/lance-context-master/src/scheduler.rs | 2 +- crates/lance-context-merge/src/failure.rs | 4 +- crates/lance-context-merge/src/fencing.rs | 177 ++++++++++++++ crates/lance-context-merge/src/lib.rs | 95 +++++++- .../src/merge_execution.rs | 82 ++++++- docs/merge-recovery.md | 216 +++++++++-------- 10 files changed, 991 insertions(+), 159 deletions(-) create mode 100644 crates/lance-context-merge/src/fencing.rs diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index 9af67d2..44075a1 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -12,6 +12,7 @@ use lance_table::{ use object_store::path::Path; use std::{ future::Future, + pin::Pin, sync::{Arc, Mutex}, }; use tokio::sync::{oneshot, Notify}; @@ -25,10 +26,20 @@ struct Progress { uncertain: bool, } +pub trait CommitAuthorizer: std::fmt::Debug + Send + Sync { + fn authorize<'a>( + &'a self, + resource: &'a str, + version: u64, + ) -> Pin> + Send + 'a>>; +} + #[derive(Debug, Default)] pub struct MergeWriteScope { progress: Mutex, changed: Notify, + authorizer: Option>, + leaves: Mutex>, } impl MergeWriteScope { @@ -36,6 +47,23 @@ impl MergeWriteScope { Arc::new(Self::default()) } + pub fn with_authorizer(authorizer: Arc) -> Arc { + Arc::new(Self { + authorizer: Some(authorizer), + ..Self::default() + }) + } + + /// Only after a durable storage barrier fenced every admitted version. + /// Aborting before that evidence would abandon an uncertain storage write. + pub async fn abort_fenced_leaves(&self) { + self.progress.lock().unwrap().closed = true; + for leaf in std::mem::take(&mut *self.leaves.lock().unwrap()) { + leaf.abort(); + } + self.drain().await; + } + pub async fn run(self: &Arc, future: F) -> F::Output { CURRENT.scope(self.clone(), future).await } @@ -64,6 +92,15 @@ impl MergeWriteScope { } } +pub(crate) async fn authorize(resource: &str, version: u64) -> Result<()> { + if let Ok(scope) = CURRENT.try_with(Arc::clone) { + if let Some(authorizer) = &scope.authorizer { + authorizer.authorize(resource, version).await?; + } + } + Ok(()) +} + struct Active { scope: Arc, completed: bool, @@ -94,12 +131,13 @@ where } progress.active += 1; } + let handles = scope.clone(); let active = Active { scope, completed: false, }; let (send, receive) = oneshot::channel(); - tokio::spawn(async move { + let leaf = tokio::spawn(async move { let mut active = active; let result = write.await; if result.is_err() { @@ -109,11 +147,177 @@ where drop(active); let _ = send.send(result); }); + handles.leaves.lock().unwrap().push(leaf.abort_handle()); receive .await .map_err(|_| E::from(Error::io("manifest commit executor panicked")))? } +/// A relative shard drain with admission before *each* immutable version write. +/// Putting one guard around commit_update's retry loop would let a revoked +/// execution obtain new versions after recovery had already fenced it. +pub(crate) async fn drain_generations( + store: lance::dataset::mem_wal::ShardManifestStore, + epoch: u64, + generations: std::collections::HashSet, +) -> Result { + let store = Arc::new(store); + let resource = format!("shard:{}", store.shard_id()); + for _ in 0..10 { + let mut next = store + .read_latest() + .await? + .ok_or_else(|| Error::io("Shard manifest not found"))?; + if next.writer_epoch != epoch { + return Err(Error::io("merge writer epoch changed before drain")); + } + next.version = next + .version + .checked_add(1) + .ok_or_else(|| Error::io("manifest version overflow"))?; + next.flushed_generations + .retain(|g| !generations.contains(&g.generation)); + authorize(&resource, next.version).await?; + let writer = store.clone(); + let scope = CURRENT.try_with(Arc::clone).ok(); + let (next, written) = shield(async move { + let written = writer.write(&next).await; + if written + .as_ref() + .is_err_and(|e| !e.to_string().contains("already exists")) + { + if let Some(scope) = scope { + scope.mark_uncertain(); + } + } + Ok::<_, Error>((next, written)) + }) + .await?; + match written { + Ok(_) => return Ok(next), + Err(error) if error.to_string().contains("already exists") => continue, + Err(error) => return Err(error), + } + } + Err(Error::io("shard drain exceeded conditional-write retries")) +} + +/// The owned protocol is supported only with immutable conditional manifest +/// creation, not Lance's fallback UnsafeCommitHandler or external catalogues. +pub fn supports_version_fencing(uri: &str) -> bool { + let Some((scheme, _)) = uri.split_once(':') else { + return true; + }; + let is_scheme = scheme + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphabetic) + && scheme + .bytes() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, b'+' | b'-' | b'.')); + if !is_scheme || (cfg!(windows) && scheme.len() == 1) { + return true; + } + matches!( + scheme.to_ascii_lowercase().as_str(), + "file" + | "file-object-store" + | "s3" + | "gs" + | "az" + | "abfss" + | "memory" + | "oss" + | "cos" + | "tos" + | "shared-memory" + ) +} + +/// After durable commit admission has closed, occupy versions beyond its +/// watermarks using metadata-only conditional writes. A late old PUT can then +/// only conflict with an occupied version; it cannot change visible table/WAL +/// state. This neither opens a shard writer nor changes its epoch or WAL list. +pub async fn fence_manifest_versions( + uri: &str, + storage_options: Option>, + watermarks: &std::collections::BTreeMap, + recovery_id: &str, +) -> Result<()> { + if !supports_version_fencing(uri) { + return Err(Error::invalid_input( + "storage does not support merge version fencing", + )); + } + if watermarks.is_empty() { + return Ok(()); + } + let mut builder = lance::dataset::builder::DatasetBuilder::from_uri(uri); + if let Some(options) = storage_options { + builder = builder.with_storage_options(options); + } + let mut dataset = builder.load().await?; + if let Some(high) = watermarks.get("base") { + for _ in 0..64 { + dataset.checkout_latest().await?; + if dataset.version().version > *high { + break; + } + let marker = format!("{recovery_id}:{}", dataset.version().version); + dataset + .update_metadata([("lance-context.merge-recovery", marker.as_str())]) + .await?; + } + dataset.checkout_latest().await?; + if dataset.version().version <= *high { + return Err(Error::io( + "base manifest barrier did not cross admitted versions", + )); + } + } + for (resource, high) in watermarks { + if resource == "base" { + continue; + } + let id = resource + .strip_prefix("shard:") + .and_then(|s| uuid::Uuid::parse_str(s).ok()) + .ok_or_else(|| Error::invalid_input("invalid shard manifest watermark"))?; + let store = lance::dataset::mem_wal::ShardManifestStore::new( + dataset.object_store(None).await?, + &dataset.branch_location().path, + id, + 16, + ); + let mut fenced = false; + for _ in 0..64 { + let mut next = store + .read_latest() + .await? + .ok_or_else(|| Error::io("Shard manifest missing during recovery"))?; + if next.version > *high { + fenced = true; + break; + } + next.version = next + .version + .checked_add(1) + .ok_or_else(|| Error::io("manifest version overflow"))?; + match store.write(&next).await { + Ok(_) => {} + Err(error) if error.to_string().contains("already exists") => {} + Err(error) => return Err(error), + } + } + if !fenced { + return Err(Error::io( + "shard manifest barrier did not cross admitted versions", + )); + } + } + Ok(()) +} + #[derive(Debug)] pub(crate) struct GuardedCommit(pub Arc); @@ -130,6 +334,7 @@ impl CommitHandler for GuardedCommit { naming_scheme: ManifestNamingScheme, transaction: Option, ) -> std::result::Result { + authorize("base", manifest.version).await?; let inner = self.0.clone(); let mut owned_manifest = manifest.clone(); let path = base_path.clone(); @@ -228,6 +433,27 @@ mod tests { use super::*; use std::sync::atomic::{AtomicBool, Ordering}; + #[test] + fn recovery_rejects_unsafe_and_external_commit_handlers() { + for uri in [ + "custom:table", + "custom://table", + "https://host/table", + "s3+ddb://bucket/table", + ] { + assert!(!supports_version_fencing(uri), "{uri}"); + } + for uri in [ + "/tmp/table:with-colon", + "relative/table", + "file:/tmp/table", + "S3://bucket/table", + "az://container/table", + ] { + assert!(supports_version_fencing(uri), "{uri}"); + } + } + #[tokio::test] async fn failed_or_panicked_leaf_commit_needs_reconciliation() { let scope = MergeWriteScope::new(); diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 28bc2bf..967ed55 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -2412,12 +2412,43 @@ mod tests { #[tokio::test] async fn cancelled_merge_joins_inflight_commit_then_retries_without_losing_blob_or_new_wal() { + cancellation_recovery_case(false).await; + } + + #[tokio::test] + async fn recovery_barrier_allows_retry_before_late_old_append_returns() { + cancellation_recovery_case(true).await; + } + + async fn cancellation_recovery_case(fence_old: bool) { use crate::merge_write_scope::{GuardedCommit, MergeWriteScope}; use lance_table::format::{IndexMetadata, Manifest, Transaction}; use lance_table::io::commit::{ CommitError, CommitHandler, ManifestLocation, ManifestNamingScheme, ManifestWriter, }; - use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + #[derive(Debug, Default)] + struct Gate { + closed: AtomicBool, + high: AtomicU64, + } + impl crate::merge_write_scope::CommitAuthorizer for Gate { + fn authorize<'a>( + &'a self, + resource: &'a str, + version: u64, + ) -> std::pin::Pin> + Send + 'a>> + { + Box::pin(async move { + if self.closed.load(Ordering::SeqCst) { + return Err(lance::Error::io("execution fenced")); + } + assert_eq!(resource, "base"); + self.high.fetch_max(version, Ordering::SeqCst); + Ok(()) + }) + } + } #[derive(Debug)] struct PausedCommit { @@ -2479,7 +2510,8 @@ mod tests { .await .unwrap(); let store = Arc::new(tokio::sync::Mutex::new(store)); - let scope = MergeWriteScope::new(); + let gate = Arc::new(Gate::default()); + let scope = MergeWriteScope::with_authorizer(gate.clone()); let worker_scope = scope.clone(); let writer = store.clone(); let merge = tokio::spawn(async move { @@ -2505,18 +2537,37 @@ mod tests { .unwrap(); store.flush().await.unwrap(); } + if fence_old { + gate.closed.store(true, Ordering::SeqCst); + let high = gate.high.load(Ordering::SeqCst); + assert!(high > 0); + let plan = std::collections::BTreeMap::from([("base".to_string(), high)]); + crate::merge_write_scope::fence_manifest_versions(uri, None, &plan, "test-recovery") + .await + .unwrap(); + // Recovery must progress while the old storage call is still alive. + assert!(!drain.is_finished()); + let mut next = store.lock().await; + next.base.refresh_latest().await.unwrap(); + assert_eq!(next.base.dataset.count_rows(None).await.unwrap(), 0); + next.cleanup_own_shard().await.unwrap(); + assert_eq!(next.base.dataset.count_rows(None).await.unwrap(), 2); + assert_eq!(flushed_generation_count(&next).await, 0); + } release.notify_one(); tokio::time::timeout(std::time::Duration::from_secs(30), drain) .await .unwrap() .unwrap(); let mut store = store.lock().await; - assert!(flushed_generation_count(&store).await > 0); + if !fence_old { + assert!(flushed_generation_count(&store).await > 0); + } store.base.refresh_latest().await.unwrap(); assert_eq!( store.base.dataset.count_rows(None).await.unwrap(), - 1, - "cancelled caller's leaf append must have committed before retry" + if fence_old { 2 } else { 1 }, + "late old append must be rejected after a storage barrier" ); store.cleanup_own_shard().await.unwrap(); assert_eq!(flushed_generation_count(&store).await, 0); @@ -2530,6 +2581,52 @@ mod tests { store.close().await.unwrap(); } + #[tokio::test] + async fn shard_manifest_barrier_rejects_a_late_drain_and_preserves_new_wal() { + use lance::dataset::mem_wal::ShardManifestStore; + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let mut store = RolloutStore::open(uri).await.unwrap(); + store.add(&[assistant_record("before")]).await.unwrap(); + store.flush().await.unwrap(); + store.add(&[assistant_record("during")]).await.unwrap(); + store.flush().await.unwrap(); + let manifests = ShardManifestStore::new( + store.base.dataset.object_store(None).await.unwrap(), + &store.base.dataset.branch_location().path, + store.base.write_shard, + 16, + ); + let before = manifests.read_latest().await.unwrap().unwrap(); + let mut late_drain = before.clone(); + late_drain.version += 1; + late_drain.flushed_generations.clear(); + let plan = std::collections::BTreeMap::from([( + format!("shard:{}", store.base.write_shard), + late_drain.version, + )]); + crate::merge_write_scope::fence_manifest_versions(uri, None, &plan, "test-recovery") + .await + .unwrap(); + assert!(manifests + .write(&late_drain) + .await + .unwrap_err() + .to_string() + .contains("already exists")); + let after = manifests.read_latest().await.unwrap().unwrap(); + assert!(after.version > late_drain.version); + assert_eq!( + after.writer_epoch, before.writer_epoch, + "recovery must not fence live ingest" + ); + assert_eq!(after.flushed_generations, before.flushed_generations); + store.cleanup_own_shard().await.unwrap(); + assert_eq!(flushed_generation_count(&store).await, 0); + assert_eq!(store.base.dataset.count_rows(None).await.unwrap(), 2); + store.close().await.unwrap(); + } + fn assistant_record(id: &str) -> RolloutRecord { RolloutRecord { id: id.to_string(), diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 6461015..dea7306 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -1047,23 +1047,8 @@ impl StorageBase { ); observe_phase!( "drain", - crate::merge_write_scope::shield(async move { - drain_store - .commit_update(epoch, |current| ShardManifest { - version: current.version + 1, - // Relative edit: retain everything we did not merge. Must - // never become an absolute assignment — see the doc comment. - flushed_generations: current - .flushed_generations - .iter() - .filter(|fg| !merged_generations.contains(&fg.generation)) - .cloned() - .collect(), - ..current.clone() - }) - .await - }) - .await + crate::merge_write_scope::drain_generations(drain_store, epoch, merged_generations) + .await )?; self.delete_merged_generation_dirs(&merged_paths).await?; diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index 32713ff..e5eab2e 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/crates/lance-context-master/src/merge_execution.rs @@ -24,42 +24,89 @@ pub(crate) async fn run_merge_wal( } let coordinator = state.task_store.merge_coordinator(); let proof = state.task_store.merge_claim(claim); - // Reconcile first, before any index/base-table mutation or new fan-out. - if let Some(old) = coordinator.get(&claim.task.target).await? { - let endpoint = old.endpoint.clone(); - if matches!(old.phase, Phase::Running | Phase::Uncertain) { - if let Some(failure) = coordinator.failure(&claim.task.target, &endpoint).await? { - if failure.class == lance_context_merge::failure::FailureClass::OwnershipUnresolved - && failure.next_retry_ms > lance_context_merge::failure::now_ms() - { - return Err(format!( - "merge ownership unresolved; recovery probe at {}; {}", - failure.next_retry_ms, failure.last_error - )); + // Reconcile/fence before any other table mutation. One retry of the fan-out + // after a completed barrier lets healthy shards progress in this task. + for recovery_round in 0..2 { + if let Some(old) = coordinator.get(&claim.task.target).await? { + if old.phase == Phase::Recovering { + if let Some(failure) = coordinator.failure(&old.target, &old.endpoint).await? { + if failure.last_error.contains("recovery barrier failed") + && failure.next_retry_ms > lance_context_merge::failure::now_ms() + { + return Err(format!( + "{}; recovery probe at {}", + failure.last_error, failure.next_retry_ms + )); + } } } - } - match reconcile(&state.http, &coordinator, &proof, old, true).await { - Ok(reclaimed) => { - metrics::counter!("master_merge_wal_generations_reclaimed_total") - .increment(reclaimed as u64); - } - Err(error) => { - if coordinator.get(&claim.task.target).await?.is_some() { - return Err(error); + if recovery_round > 0 || matches!(old.phase, Phase::Recovering | Phase::Uncertain) { + recover_execution(state, &coordinator, &proof, old).await?; + } else if let Err(error) = reconcile(&state.http, &coordinator, &proof, old, true).await + { + if let Some(current) = coordinator.get(&claim.task.target).await? { + recover_execution(state, &coordinator, &proof, current).await?; + } else { + tracing::info!(target = %claim.task.target, %error, "previous merge ended; resuming shards"); } - tracing::info!(target = %claim.task.target, %error, "previous merge terminated; resuming shards"); } } + let outcome = run_workers( + &state.http, + &coordinator, + &proof, + &claim.task.target, + &state.config.worker_endpoints, + ) + .await; + if outcome.is_ok() + || recovery_round == 1 + || coordinator.get(&claim.task.target).await?.is_none() + { + return outcome; + } } - run_workers( - &state.http, - &coordinator, - &proof, - &claim.task.target, - &state.config.worker_endpoints, - ) - .await + unreachable!("bounded merge recovery loop returns on last iteration") +} + +async fn recover_execution( + state: &Arc, + coordinator: &Coordinator, + proof: &ClaimProof, + old: Execution, +) -> Result<(), String> { + let result = async { + let frozen = coordinator.freeze(proof, &old).await? + .ok_or_else(|| "merge execution changed while freezing commit admission".to_string())?; + let watermarks = coordinator.watermarks(&frozen).await?; + let uri = match frozen.target.strip_prefix("generic:") { + Some(name) => state.generic_uri(name), None => state.rollout_uri(&frozen.target), + }; + tokio::time::timeout(Duration::from_secs(120), + lance_context_core::merge_write_scope::fence_manifest_versions(&uri, None, &watermarks.versions, &frozen.id) + ).await.map_err(|_| "storage recovery barrier deadline exceeded".to_string())? + .map_err(|e| e.to_string())?; + if !coordinator.finish_recovery(proof, &frozen).await? { + return Err("claim lost before publishing storage recovery barrier".into()); + } + let recovered = coordinator.get(&frozen.target).await?.ok_or_else(|| "recovered execution disappeared".to_string())?; + if recovered.id != frozen.id || !coordinator.release(proof, &recovered).await? { + return Err("claim lost before releasing recovered execution".into()); + } + metrics::counter!("master_merge_storage_recoveries_total", "result" => "ok").increment(1); + tracing::warn!(target = %frozen.target, endpoint = %frozen.endpoint, id = %frozen.id, "old merge storage commits fenced; resuming healthy shards"); + Ok(()) + }.await; + if let Err(error) = &result { + let message = format!("merge ownership unresolved: recovery barrier failed: {error}"); + coordinator + .record_failure(proof, &old.target, &old.endpoint, &message) + .await?; + metrics::counter!("master_merge_storage_recoveries_total", "result" => "failed") + .increment(1); + return Err(message); + } + result } async fn run_workers( @@ -175,7 +222,7 @@ async fn one( .json() .await .map_err(|e| e.to_string())?; - if capabilities.protocol != 1 || capabilities.timeout_secs == 0 { + if capabilities.protocol != 2 || capabilities.timeout_secs == 0 { return Err("worker lacks bounded owned-merge protocol".into()); } let execution = Execution::new( @@ -261,7 +308,7 @@ async fn reconcile_with_grace( continue; } }; - if current.phase == Phase::Uncertain { + if matches!(current.phase, Phase::Uncertain | Phase::Recovering) { let error = current .error .unwrap_or_else(|| "merge ownership unresolved: ambiguous storage commit".into()); @@ -270,7 +317,7 @@ async fn reconcile_with_grace( .await?; return Err(error); } - if current.phase == Phase::Finished { + if matches!(current.phase, Phase::Finished | Phase::Recovered) { if !coordinator.release(proof, ¤t).await? { return Err("task claim lost while releasing terminal merge execution".into()); } @@ -328,7 +375,7 @@ mod tests { .route( "/api/v1/internal/merge-executor", get(|| async { - Json(serde_json::json!({"protocol":1,"instance":"test","timeout_secs":1})) + Json(serde_json::json!({"protocol":2,"instance":"test","timeout_secs":1})) }), ) .route( @@ -407,6 +454,111 @@ mod tests { (Coordinator::new(client, prefix), proof) } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn orphaned_merge_is_storage_fenced_and_healthy_shard_progresses() { + use clap::Parser; + use lance_context_api::TaskKind; + let dir = tempfile::tempdir().unwrap(); + let endpoints = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let prefix = format!("/merge-orphan-tests/{}", Execution::new("", "", "", 1).id); + let cfg = crate::config::MasterConfig::parse_from([ + "test", + "--data-dir", + dir.path().to_str().unwrap(), + "--etcd-endpoints", + &endpoints, + "--etcd-prefix", + &prefix, + ]); + let mut state = MasterState::new(cfg).await.unwrap(); + let uri = state.rollout_uri("table"); + lance_context_core::RolloutStore::open(&uri) + .await + .unwrap() + .close() + .await + .unwrap(); + let initial_version = lance::Dataset::open(&uri).await.unwrap().version().version; + let coordinator = state.task_store.merge_coordinator(); + let good = Worker { + coordinator: coordinator.clone(), + calls: Arc::new(AtomicUsize::new(0)), + fail: false, + stall_first: false, + name: "healthy", + events: Arc::new(Mutex::new(Vec::new())), + }; + let (healthy, server) = worker(good.clone()).await; + let lost = "http://127.0.0.1:1".to_string(); + Arc::get_mut(&mut state).unwrap().config.worker_endpoints = vec![lost.clone(), healthy]; + state + .task_store + .enqueue(TaskKind::MergeWal, "table", Vec::new()) + .await + .unwrap(); + let claim = state + .task_store + .claim_next_of_kinds(crate::task_store::TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + let proof = state.task_store.merge_claim(&claim); + let execution = Execution::new("table", &lost, "dead-worker", 1); + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + let running = coordinator.start(&execution).await.unwrap().unwrap(); + coordinator + .authorize_commit(&running, &uri, "base", initial_version + 1) + .await + .unwrap(); + assert!(coordinator + .report_uncertain(&running, "lost storage response".into()) + .await + .unwrap()); + let outcome = tokio::time::timeout(Duration::from_secs(30), run_merge_wal(&state, &claim)) + .await + .unwrap(); + assert!(outcome.unwrap_err().contains("merged 7 generations")); + assert_eq!(good.calls.load(Ordering::SeqCst), 1); + assert!(coordinator.get("table").await.unwrap().is_none()); + assert!(lance::Dataset::open(&uri).await.unwrap().version().version > initial_version + 1); + assert!(coordinator + .authorize_commit(&running, &uri, "base", initial_version + 4) + .await + .is_err()); + + // A mismatched worker namespace must remain fenced, never recover by + // committing a barrier to an unrelated same-named table. + let wrong = Execution::new("table", &lost, "misconfigured-worker", 1); + assert!(coordinator.reserve(&proof, &wrong).await.unwrap()); + let wrong = coordinator.start(&wrong).await.unwrap().unwrap(); + coordinator + .authorize_commit(&wrong, "file:///different-table", "base", 7) + .await + .unwrap(); + let error = recover_execution(&state, &coordinator, &proof, wrong) + .await + .unwrap_err(); + assert!(error.contains("dataset URI mismatch")); + assert_eq!( + coordinator.get("table").await.unwrap().unwrap().phase, + Phase::Recovering + ); + state.task_store.finish(claim, Err(error)).await.unwrap(); + server.abort(); + let mut client = + etcd_client::Client::connect(endpoints.split(',').collect::>(), None) + .await + .unwrap(); + client + .delete( + prefix, + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn missing_executor_returns_attention_without_unlocking_surviving_write() { diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 27db7f6..a818035 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -1207,7 +1207,7 @@ mod tests { .route( "/api/v1/internal/merge-executor", get(|| async { - Json(serde_json::json!({"protocol":1,"instance":"stub","timeout_secs":600})) + Json(serde_json::json!({"protocol":2,"instance":"stub","timeout_secs":600})) }), ) .route( diff --git a/crates/lance-context-merge/src/failure.rs b/crates/lance-context-merge/src/failure.rs index 3fdea0b..84f77e6 100644 --- a/crates/lance-context-merge/src/failure.rs +++ b/crates/lance-context-merge/src/failure.rs @@ -74,7 +74,9 @@ impl ShardFailure { ); // First task gets three attempts. Later tasks get one half-open probe, // never a fresh budget of three. Cap probing at once per hour. - let delay_secs = if needs_attention { + let delay_secs = if class == FailureClass::OwnershipUnresolved { + (30u64.saturating_mul(1 << attempts.saturating_sub(1).min(4))).min(300) + } else if needs_attention { 3600 } else if attempts < 3 { 2 diff --git a/crates/lance-context-merge/src/fencing.rs b/crates/lance-context-merge/src/fencing.rs new file mode 100644 index 0000000..4bdff20 --- /dev/null +++ b/crates/lance-context-merge/src/fencing.rs @@ -0,0 +1,177 @@ +//! Close commit admission before fencing the finite set of already admitted +//! manifest versions. Lease expiry and process identity are not storage fences. +use crate::{encode, execution_key, ClaimProof, Coordinator, Execution, Phase, Result}; +use etcd_client::{Compare, CompareOp, TxnOp}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct CommitWatermarks { + #[serde(default)] + pub dataset_uri: Option, + pub versions: BTreeMap, +} + +impl Coordinator { + fn permits_key(&self, execution: &Execution) -> String { + format!( + "{}/merge-commit-permits/{}", + self.prefix.trim_end_matches('/'), + execution.id + ) + } + + /// Every manifest write calls this *after selecting its immutable version* + /// and before issuing storage I/O. A concurrent freeze either observes this + /// permit or prevents it; there is no check-then-write ownership gap. + pub async fn authorize_commit( + &self, + running: &Execution, + dataset_uri: &str, + resource: &str, + version: u64, + ) -> Result<()> { + if running.protocol != 2 || running.phase != Phase::Running { + return Err("merge commit requires fencing protocol 2".into()); + } + if resource != "base" + && resource + .strip_prefix("shard:") + .is_none_or(|s| uuid::Uuid::parse_str(s).is_err()) + { + return Err("invalid merge manifest resource".into()); + } + let key = self.permits_key(running); + for _ in 0..16 { + let response = self + .client + .clone() + .get(key.clone(), None) + .await + .map_err(|e| e.to_string())?; + let previous = response.kvs().first(); + let mut watermarks: CommitWatermarks = previous + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose()? + .unwrap_or_default(); + if watermarks + .dataset_uri + .as_deref() + .is_some_and(|uri| uri != dataset_uri) + { + return Err("merge execution cannot write multiple datasets".into()); + } + watermarks.dataset_uri = Some(dataset_uri.into()); + let watermark = watermarks.versions.entry(resource.into()).or_default(); + *watermark = (*watermark).max(version); + if self + .transact( + vec![ + Compare::value( + execution_key(&self.prefix, &running.target), + CompareOp::Equal, + encode(running), + ), + Compare::mod_revision( + key.clone(), + CompareOp::Equal, + previous.map_or(0, |kv| kv.mod_revision()), + ), + ], + vec![TxnOp::put( + key.clone(), + serde_json::to_vec(&watermarks).unwrap(), + None, + )], + ) + .await? + { + return Ok(()); + } + if self.get(&running.target).await?.as_ref() != Some(running) { + return Err("merge commit ownership revoked for recovery".into()); + } + } + Err("merge commit permission contention exceeded limit".into()) + } + + pub async fn freeze(&self, proof: &ClaimProof, old: &Execution) -> Result> { + if old.protocol != 2 + || !matches!( + old.phase, + Phase::Running | Phase::Uncertain | Phase::Recovering + ) + { + return Err("execution lacks recoverable storage fencing protocol".into()); + } + let mut frozen = old.clone(); + frozen.phase = Phase::Recovering; + Ok(self + .transact( + vec![ + Compare::value(proof.key.as_str(), CompareOp::Equal, proof.token.as_bytes()), + Compare::value( + execution_key(&self.prefix, &old.target), + CompareOp::Equal, + encode(old), + ), + ], + vec![TxnOp::put( + execution_key(&self.prefix, &old.target), + encode(&frozen), + None, + )], + ) + .await? + .then_some(frozen)) + } + + pub async fn watermarks(&self, frozen: &Execution) -> Result { + if frozen.phase != Phase::Recovering { + return Err("commit admission must be frozen first".into()); + } + let response = self + .client + .clone() + .get(self.permits_key(frozen), None) + .await + .map_err(|e| e.to_string())?; + response + .kvs() + .first() + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose() + .map(|v| v.unwrap_or_default()) + } + + /// Only call after storage confirms that every permitted version is fenced. + pub async fn finish_recovery(&self, proof: &ClaimProof, frozen: &Execution) -> Result { + if frozen.phase != Phase::Recovering { + return Err("execution is not recovering".into()); + } + let mut recovered = frozen.finished(Err( + "merge cancelled; admitted manifest versions fenced for recovery".into(), + )); + recovered.phase = Phase::Recovered; + self.transact( + vec![ + Compare::value(proof.key.as_str(), CompareOp::Equal, proof.token.as_bytes()), + Compare::value( + execution_key(&self.prefix, &frozen.target), + CompareOp::Equal, + encode(frozen), + ), + ], + vec![TxnOp::put( + execution_key(&self.prefix, &frozen.target), + encode(&recovered), + None, + )], + ) + .await + } + + pub(crate) fn remove_permits(&self, execution: &Execution) -> TxnOp { + TxnOp::delete(self.permits_key(execution), None) + } +} diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index e87ebd4..2171493 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -5,6 +5,7 @@ //! and reconciles the old execution; it never clears a running fence on timeout. pub mod failure; +pub mod fencing; use etcd_client::{Client, Compare, CompareOp, Txn, TxnOp}; use serde::{Deserialize, Serialize}; @@ -24,6 +25,8 @@ pub enum Phase { Reserved, Running, Uncertain, + Recovering, + Recovered, Finished, } @@ -37,6 +40,12 @@ pub struct Execution { pub phase: Phase, pub reclaimed: usize, pub error: Option, + #[serde(default, skip_serializing_if = "protocol_absent")] + pub protocol: u32, +} + +fn protocol_absent(protocol: &u32) -> bool { + *protocol == 0 } impl Execution { @@ -50,6 +59,7 @@ impl Execution { phase: Phase::Reserved, reclaimed: 0, error: None, + protocol: 2, } } @@ -120,7 +130,10 @@ impl Coordinator { /// Atomic claim check and admission close the delayed-request/lease-loss race. pub async fn reserve(&self, claim: &ClaimProof, execution: &Execution) -> Result { - if execution.phase != Phase::Reserved || execution.timeout_secs == 0 { + if execution.phase != Phase::Reserved + || execution.timeout_secs == 0 + || execution.protocol != 2 + { return Err("invalid merge execution reservation".into()); } let key = execution_key(&self.prefix, &execution.target); @@ -188,7 +201,7 @@ impl Coordinator { /// Never delete on elapsed time, an HTTP error, or claim expiry. Both the /// terminal result and the current task claim must still match atomically. pub async fn release(&self, claim: &ClaimProof, execution: &Execution) -> Result { - if execution.phase != Phase::Finished { + if !matches!(execution.phase, Phase::Finished | Phase::Recovered) { return Err("cannot release a live execution".into()); } let key = execution_key(&self.prefix, &execution.target); @@ -203,6 +216,7 @@ impl Coordinator { ), ]); operations.extend([ + self.remove_permits(execution), TxnOp::delete(key, None), TxnOp::put( target_lock_key(&self.prefix, &execution.target), @@ -457,6 +471,83 @@ mod tests { ) } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn recovery_closes_commit_admission_and_retains_every_allowed_version() { + let (coordinator, mut client, proof, lease) = fixture().await; + let operation = Execution::new("table", "worker", "boot", 1); + assert!(coordinator.reserve(&proof, &operation).await.unwrap()); + let running = coordinator.start(&operation).await.unwrap().unwrap(); + coordinator + .authorize_commit(&running, "file:///table", "base", 7) + .await + .unwrap(); + assert!(coordinator + .authorize_commit(&running, "file:///wrong-table", "base", 8) + .await + .is_err()); + client.lease_revoke(lease).await.unwrap(); + // The old executor remains protected independently of its scheduler. + coordinator + .authorize_commit(&running, "file:///table", "base", 8) + .await + .unwrap(); + assert!(coordinator + .freeze(&proof, &running) + .await + .unwrap() + .is_none()); + let next = ClaimProof { + token: "recovery-owner".into(), + lease_id: 0, + ..proof.clone() + }; + client + .put(next.key.clone(), next.token.clone(), None) + .await + .unwrap(); + let shard = format!("shard:{}", uuid::Uuid::new_v4()); + coordinator + .authorize_commit(&running, "file:///table", &shard, 91) + .await + .unwrap(); + let (admitted, frozen) = tokio::join!( + coordinator.authorize_commit(&running, "file:///table", "base", 9), + coordinator.freeze(&next, &running), + ); + let frozen = frozen.unwrap().unwrap(); + let plan = coordinator.watermarks(&frozen).await.unwrap(); + assert_eq!(plan.dataset_uri.as_deref(), Some("file:///table")); + assert_eq!(plan.versions["base"], if admitted.is_ok() { 9 } else { 8 }); + assert_eq!(plan.versions[&shard], 91); + assert!(coordinator + .authorize_commit(&running, "file:///table", "base", 10) + .await + .is_err()); + assert!(coordinator + .authorize_commit(&running, "file:///table", &shard, 92) + .await + .is_err()); + assert!(coordinator.release(&next, &frozen).await.is_err()); + assert!(!coordinator.finish_recovery(&proof, &frozen).await.unwrap()); + // Storage fencing itself is tested against real manifests in core. + assert!(coordinator.finish_recovery(&next, &frozen).await.unwrap()); + let recovered = coordinator.get("table").await.unwrap().unwrap(); + assert_eq!(recovered.phase, Phase::Recovered); + assert!(coordinator.release(&next, &recovered).await.unwrap()); + assert!(coordinator + .authorize_commit(&running, "file:///table", "base", 11) + .await + .is_err()); + client + .delete( + coordinator.prefix.clone(), + Some(etcd_client::DeleteOptions::new().with_prefix()), + ) + .await + .unwrap(); + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn ambiguous_storage_response_cannot_release_the_execution() { diff --git a/crates/lance-context-server/src/merge_execution.rs b/crates/lance-context-server/src/merge_execution.rs index 40b038a..92501fa 100644 --- a/crates/lance-context-server/src/merge_execution.rs +++ b/crates/lance-context-server/src/merge_execution.rs @@ -11,6 +11,39 @@ use lance_context_merge::{execute_scoped, Coordinator, Execution, Phase}; use std::{collections::HashMap, sync::Arc, time::Duration}; use tokio::sync::{watch, Mutex}; +struct OwnedCommitGuard { + coordinator: Coordinator, + execution: Execution, + dataset_uri: String, +} +impl std::fmt::Debug for OwnedCommitGuard { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("OwnedCommitGuard") + .field("execution", &self.execution.id) + .finish() + } +} +impl lance_context_core::merge_write_scope::CommitAuthorizer for OwnedCommitGuard { + fn authorize<'a>( + &'a self, + resource: &'a str, + version: u64, + ) -> std::pin::Pin< + Box< + dyn std::future::Future> + + Send + + 'a, + >, + > { + Box::pin(async move { + self.coordinator + .authorize_commit(&self.execution, &self.dataset_uri, resource, version) + .await + .map_err(lance_context_core::LanceError::io) + }) + } +} + pub struct Executions { coordinator: Option, instance: String, @@ -54,7 +87,7 @@ pub async fn capabilities( ) -> Result, AppError> { state.merge_executions.coordinator()?; Ok(Json( - serde_json::json!({"protocol": 1, "instance": state.merge_executions.instance, + serde_json::json!({"protocol": 2, "instance": state.merge_executions.instance, "timeout_secs": state.merge_executions.timeout_secs}), )) } @@ -64,7 +97,8 @@ pub async fn start( Json(execution): Json, ) -> Result { let coordinator = state.merge_executions.coordinator()?; - if execution.instance != state.merge_executions.instance + if execution.protocol != 2 + || execution.instance != state.merge_executions.instance || execution.phase != Phase::Reserved || execution.timeout_secs == 0 || execution.timeout_secs > state.merge_executions.timeout_secs @@ -130,7 +164,14 @@ async fn run( }; let target = execution.target.clone(); let mut slot = None; + let uri = match target.strip_prefix("generic:") { + Some(name) => state.generic_uri(name), + None => state.rollout_uri(&target), + }; let work = async { + if !lance_context_core::merge_write_scope::supports_version_fencing(&uri) { + return Err("invalid storage backend for owned merge fencing".into()); + } slot = state.acquire_merge_slot().await; let result = if let Some(name) = target.strip_prefix("generic:") { generic::merge_generic_wal_owned( @@ -159,7 +200,13 @@ async fn run( Err(e) => Err(format!("{e:?}")), } }; - let write_scope = lance_context_core::merge_write_scope::MergeWriteScope::new(); + let write_scope = lance_context_core::merge_write_scope::MergeWriteScope::with_authorizer( + Arc::new(OwnedCommitGuard { + coordinator: coordinator.clone(), + execution: running.clone(), + dataset_uri: uri.clone(), + }), + ); let outcome = std::panic::AssertUnwindSafe(execute_scoped( write_scope.run(work), Duration::from_secs(execution.timeout_secs), @@ -170,7 +217,27 @@ async fn run( .unwrap_or_else(|_| Err("merge executor panicked".into())); // Top-level cancellation cannot acknowledge a write still in flight at // object storage. Join those leaf commits before publishing Finished. - write_scope.drain().await; + loop { + tokio::select! { + biased; + _ = write_scope.drain() => break, + _ = tokio::time::sleep(Duration::from_secs(1)) => { + match coordinator.get(&execution.target).await { + Ok(None) => { + // Only a terminal outcome or a completed storage + // barrier permits removal of this execution record. + write_scope.abort_fenced_leaves().await; + return Ok(()); + } + Ok(Some(current)) if current.id != execution.id || current.phase == Phase::Recovered => { + write_scope.abort_fenced_leaves().await; + return Ok(()); + } + _ => {}, + } + } + } + } // Do not admit another memory-heavy merge into this slot while a // cancelled execution is still joining storage commits. drop(slot); @@ -198,6 +265,13 @@ async fn run( Ok(true) => break, Ok(false) => match coordinator.get(&execution.target).await { Ok(Some(current)) if current == expected => break, + Ok(None) => return Ok(()), + Ok(Some(current)) + if current.id != execution.id + || matches!(current.phase, Phase::Recovering | Phase::Recovered) => + { + return Ok(()) + } Ok(_) => return Err("merge execution fence changed unexpectedly".into()), Err(error) => { tracing::warn!(id = %execution.id, %error, "reconciling merge outcome acknowledgement") diff --git a/docs/merge-recovery.md b/docs/merge-recovery.md index d9f3928..29155c1 100644 --- a/docs/merge-recovery.md +++ b/docs/merge-recovery.md @@ -1,107 +1,135 @@ # Owned merge recovery A serial merge fan-out previously had no request deadline and could report -success after one worker failed if another reclaimed generations. A client-side -HTTP timeout alone is unsafe: the worker may still be committing after the -master moves to the next shard. - -This change admits a durable execution before sending the worker request. The -worker owns the merge independently of that request, cancels preparation at its -deadline, and joins already-started base/shard manifest commits before publishing -its terminal result. The master advances only after this acknowledgement. Before -retrying a cancelled commit, the store refreshes its base manifest. WAL drain -remains a relative edit preserving generations flushed during recovery. - -## Ownership - -- `MergeWal`, `Compact`, `IndexId`, and `Repair` share the target lock. -- Admission atomically replaces the lease-backed target lock with a persistent - execution marker. Existing helpers that respect target locks remain excluded. -- A separate leased merge claim ensures only one scheduler reconciles an - execution, including non-deduplicated tasks in dependency chains. -- Worker execution and the persistent target lock survive task lease loss. - Only a terminal execution can restore the target lock to a live task lease. - Ambiguous manifest errors/panics publish `uncertain`, which cannot release - ownership: an ended Rust future is not evidence that remote storage rejected - its write. Definite conditional-commit conflicts are treated separately. -- Cancelling an unstarted request uses CAS, preventing a delayed POST from - starting after another shard has been admitted. -- Graceful worker shutdown cancels and drains owned executions before closing - resident stores. SIGKILL/OOM is **not** a graceful acknowledgement. +success after one worker failed if another reclaimed generations. An HTTP +timeout alone cannot hand off a table: the worker or object store may still be +committing after the master advances. + +Protocol 2 owns the worker execution independently of its HTTP connection, +records every admitted manifest version, and can fence those versions in storage +before handing off a stuck or crashed executor. Recovery does not depend on a +replacement pod being Ready or on guessing that an old process has died. + +## Admission, cancellation, and recovery + +`MergeWal`, `Compact`, `IndexId`, and `Repair` share the existing target lock. +Merge admission replaces that lease-backed lock with a persistent execution +marker. A separate leased merge claim permits only one scheduler to reconcile +it, including tasks that bypass dedupe through dependency chains. Helpers that +respect the existing target lock stay excluded throughout recovery. + +Before **each** base-manifest commit and **each** shard-drain write attempt, the +worker atomically checks that its exact execution is still Running and records +the resource's maximum admitted immutable version. The permit also binds the +worker's dataset URI. A shard drain guards each conditional write, not a whole +retry loop that could acquire fresh versions after revocation. + +Normal cancellation drops preparation and joins already-started leaf manifest +writes before publishing Finished. Worker merge slots remain held during that +join. A definite conditional-write conflict is distinguished from an ambiguous +storage error or panic. An ambiguous outcome publishes Uncertain and retains +ownership, even though the Rust future has ended. + +For a missing/stuck executor or an uncertain commit result, the master: + +1. Atomically changes the execution to Recovering, closing further commit + admission. A racing permit either precedes this transaction and is included + in the durable maximum, or fails its ownership comparison. +2. Reads the closed set of admitted versions and checks the worker/master URI. +3. Uses metadata-only conditional writes to advance the base and affected shard + manifests beyond those versions. The base barrier updates reserved table + metadata `lance-context.merge-recovery`. Shard barriers copy the current + manifest without changing the WAL list or writer epoch. They do not open a + shard writer, scan payloads, drop fragments, or fence ongoing ingestion. +4. Publishes Recovered only after every barrier is confirmed. Only then can a + live task release the target lock and continue merging. A delayed old PUT + targets an occupied immutable version, and any new old-executor commit is + denied by the closed admission gate. + +A surviving executor may abort/join its pending leaf futures after observing a +completed barrier or a subsequent execution. Before that evidence it retains +those futures. This allows a healthy shard to progress without waiting forever +for a stalled old storage request. The master permits one recovery/restart of +fan-out within a task; further faults are durably queued for another task so a +single target cannot monopolize a scheduler slot indefinitely. + +Reserved work is cancelled by CAS, so a delayed HTTP POST cannot resurrect it. +The base handle refreshes before retrying, and WAL drain remains a relative edit +that preserves generations flushed during recovery. Graceful worker shutdown +cancels/drains owned actors before closing resident stores. ## Repeated failures -Failures persist under `merge-failures//` in the master -etcd namespace. The record includes total consecutive failed attempts, diagnostic -class, last error (capped at 4096 characters), next retry time, and attention flag. -Recreating a task or master does not reset this budget. Executed failures are -recorded atomically with releasing terminal execution ownership; a crash in the -master between these two steps cannot lose the failure. +Per-target/endpoint records survive task recreation and master restart. They +contain consecutive failed attempt counts, diagnostic class, last error (capped +at 4096 characters), next retry time, and `needs_attention`. Executed failures +are recorded atomically with releasing terminal execution ownership; a crash +between release and bookkeeping cannot erase the retry budget. The first three transient failures use two-second targeted retries, with healthy -shards visited before retrying failed shards. After that, each due task makes +shards visited before failed shards are retried. After that, each due task makes one probe, backing off from 30 seconds to at most 15 minutes. At 15 consecutive -failures, probes are hourly and `needs_attention` is set. Missing data, corruption, -configuration/authentication errors, and unresolved execution ownership set the -attention flag immediately and use hourly probes. Classification uses diagnostic -text conservatively; it never authorizes data removal. - -Successful probes clear the shard's failure record. Failed/deferred shards keep +failures, probes are hourly and `needs_attention` is set. Missing data, +corruption, and configuration/authentication errors require attention immediately +and use hourly probes. Diagnostic text classification is conservative and never +authorizes data removal. + +A failed **storage barrier** retains the execution/target lock and releases the +scheduler slot. Metadata recovery probes back off from 30 seconds to at most +five minutes, with attention flagged immediately. They keep trying to establish +a safe handoff when storage recovers. They cannot authorize a new writer while +the barrier is incomplete. Other tables continue running. + +Successful shard merges clear their failure records. Failed/deferred shards keep the overall task failed even if healthy shards reclaim WAL. Per-table merge cooldown is replaced by per-endpoint backoff so healthy shards can process new -WAL. A 15-second scheduler poll walks at most 256 failure records per page and -enqueues due targets through existing dedupe. It reads metadata only and does -not wait for the 600-second stats sweep. For a large failure ledger, polling -latency grows with page count. Removed endpoints are not automatically probed. - -`GET /api/v1/scheduler/merge-failures?after=` returns at most 256 records -and a `next` cursor. `needs_attention`, `last_error`, and `next_retry_ms` distinguish -failed shards from recovered ones. Existing task error/counter metrics also show -failure. This is an observable circuit state, not an external alert integration. - -The scheduler does **not** automatically enqueue destructive fragment repair for -a merge error. Missing/corrupt files require diagnosis and restoration or an -explicitly reviewed repair; retaining WAL is mandatory. Other maintenance task -repair behavior is unchanged. - -## Limits and required rollout validation - -This is not yet a complete automatic crash-recovery protocol. If a process dies, a storage commit cannot be joined, or a commit returns an -ambiguous result, there may be no terminal acknowledgement. -After the merge deadline plus a 60-second reconciliation allowance (or 60 seconds -for inherited work), the master records unresolved ownership, fails that task, -and releases its scheduler slot **without deleting the execution or target -lock**. Other tables continue. The affected table remains fenced until the old -executor acknowledges termination. A changed worker UUID, a Ready replacement, -an expired lease, and elapsed time are not sufficient proof. - -Automatic recovery of that orphan requires a storage fencing/barrier protocol -and authoritative termination evidence; neither is implemented here. Do not -manually delete the fence or deploy this as a claimed complete fix for OOM/node -loss. The worker deadline bounds its scoped merge future; draining a backend -that never completes its commit is not bounded by that deadline. - -The execution fence protects scheduler-coordinated maintenance and helpers that -honor target locks. It does not retrofit storage fencing onto arbitrary manual -writers or solve lease loss during the master's own local compaction/index job. - -This wire protocol requires coordinated master/worker rollout. There is no -fallback to unbounded legacy HTTP merges: - -- Workers need `ETCD_ENDPOINTS` and the **same** `ETCD_PREFIX` and credentials as - masters. `MERGE_EXECUTION_TIMEOUT_SECS` defaults to 600; the master caps the - requested deadline at 600. Zero is rejected. +WAL. A 15-second scheduler poll reads at most 256 failure records per page and +enqueues due targets through existing dedupe. It does not wait for the +600-second stats sweep. Large ledgers add page traversal latency; removed worker +endpoints are not automatically probed. + +`GET /api/v1/scheduler/merge-failures?after=` returns up to 256 records +and a `next` cursor. `needs_attention`, `last_error`, and `next_retry_ms` expose +the circuit state. `master_merge_storage_recoveries_total{result}` records +barrier outcomes alongside existing task and reclaimed-generation metrics. This +is an observable circuit state, not an external notification integration. + +Merge errors do **not** enqueue destructive fragment repair. Missing/corrupt +files require diagnosis and restoration or an explicitly reviewed repair. +Other maintenance tasks retain their existing repair policy. + +## Storage and rollout requirements + +- Masters/workers must use the same physical storage namespace and etcd prefix, + with consistent storage configuration/credentials. URI mismatch fails closed. +- Version fencing supports Lance's immutable conditional manifest handlers: + local files, S3, GCS, Azure, memory, OSS, COS, TOS, and shared memory. Fallback + unsafe handlers and external catalogues such as `s3+ddb` are rejected for + owned merging/recovery. Existing manifest immutability/retention guarantees + must be preserved; never manually remove fencing/version files to unblock a + live writer. +- Workers need `ETCD_ENDPOINTS` and the same `ETCD_PREFIX` as masters. + `MERGE_EXECUTION_TIMEOUT_SECS` defaults to 600 and the master caps the requested + deadline at 600. Zero is rejected. Cancellation reconciliation has a 60-second + allowance; a metadata recovery attempt has a 120-second deadline. +- Enable protocol 2 on both sides. There is no fallback to unbounded legacy + merge HTTP calls. Old executions without version admission cannot be + automatically fenced using an empty watermark set; they remain protected. - With the coordinator enabled, legacy merge routes reject admission. Disable worker self-merge thresholds and cleanup timers; startup checks enforce this. -- `INDEX_BEFORE_MERGE` remains accepted for compatibility, but index preparation - is done by the worker inside the owned execution. -- Drain/reconcile legacy in-flight merges before enabling the new protocol. - Masters must continue their other jobs. Recovery helpers must preserve target - locks; do not clear a keeper's pause while its delegate is alive. -- Keep ordinary workers at six merge slots and 20 GiB. This patch increases - neither memory budgets nor same-table commit concurrency. - -Before production rollout, staging must cover realistic blobs, sustained writes, -WAL pending/reclamation rates, memory/OOM counts, commit conflicts, and p95/p99 -latency. Local cancellation and HTTP-disconnect tests are necessary but cannot -substitute for this soak or the missing orphan recovery protocol. +- `INDEX_BEFORE_MERGE` remains accepted for compatibility; workers prepare their + key indexes inside owned execution. Keep ordinary workers at six merge slots + and 20 GiB; this change increases neither limit. +- Drain/reconcile legacy in-flight merges before enabling protocol 2. Masters + must continue other jobs; recovery helpers must retain target locks. Do not + clear a keeper's pause while its delegate is alive. + +The fence covers scheduler-coordinated maintenance and helpers that honor target +locks. It does not retrofit lease-loss protection onto the master's own local +compaction/index jobs or arbitrary manual writers. It cannot repair missing data +or guarantee a WAL pending ceiling while storage remains unavailable. + +Before production rollout, staging must exercise realistic blobs, sustained +writes, forced worker death and delayed storage responses. Measure pending and +reclaimed generations, memory/OOM counts, commit conflicts, etcd overhead, and +p95/p99 task latency. Local fault tests do not replace that soak. From f019dc182984d8ab23e519b2dcc19196e734aea8 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Thu, 1 Oct 2026 18:32:04 +0000 Subject: [PATCH 4/9] fix: preserve recovery failure causes and enforce dataset identity --- .../src/merge_write_scope.rs | 43 ++++++++++++++++++- .../src/merge_execution.rs | 5 ++- crates/lance-context-merge/src/fencing.rs | 14 ++++-- crates/lance-context-merge/src/lib.rs | 31 +++++++++++++ 4 files changed, 87 insertions(+), 6 deletions(-) diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index 44075a1..5d31748 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -54,8 +54,9 @@ impl MergeWriteScope { }) } - /// Only after a durable storage barrier fenced every admitted version. - /// Aborting before that evidence would abandon an uncertain storage write. + /// Only after dropping the merge future and after a durable storage + /// barrier fenced every admitted version. Aborting before that evidence + /// would abandon an uncertain storage write. pub async fn abort_fenced_leaves(&self) { self.progress.lock().unwrap().closed = true; for leaf in std::mem::take(&mut *self.leaves.lock().unwrap()) { @@ -454,6 +455,44 @@ mod tests { } } + #[tokio::test] + async fn confirmed_fence_can_abort_a_permanently_stalled_leaf() { + struct Writing(Arc); + impl Drop for Writing { + fn drop(&mut self) { + self.0.store(false, Ordering::SeqCst); + } + } + let scope = MergeWriteScope::new(); + let writing = Arc::new(AtomicBool::new(false)); + let flag = writing.clone(); + let (entered, began) = oneshot::channel(); + let worker_scope = scope.clone(); + let caller = tokio::spawn(async move { + worker_scope + .run(shield(async move { + flag.store(true, Ordering::SeqCst); + let _writing = Writing(flag); + entered.send(()).unwrap(); + std::future::pending::>().await + })) + .await + }); + began.await.unwrap(); + caller.abort(); + assert!(caller.await.unwrap_err().is_cancelled()); + assert!(writing.load(Ordering::SeqCst)); + // Real base/shard barriers are exercised by the rollout fault tests. + // Once that proof exists, a stalled local I/O future must release too. + tokio::time::timeout( + std::time::Duration::from_secs(1), + scope.abort_fenced_leaves(), + ) + .await + .unwrap(); + assert!(!writing.load(Ordering::SeqCst)); + } + #[tokio::test] async fn failed_or_panicked_leaf_commit_needs_reconciliation() { let scope = MergeWriteScope::new(); diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index e5eab2e..b2299f8 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/crates/lance-context-master/src/merge_execution.rs @@ -82,6 +82,9 @@ async fn recover_execution( let uri = match frozen.target.strip_prefix("generic:") { Some(name) => state.generic_uri(name), None => state.rollout_uri(&frozen.target), }; + if watermarks.dataset_uri.as_deref().is_some_and(|worker_uri| worker_uri != uri) { + return Err("dataset URI mismatch between worker and recovery master".into()); + } tokio::time::timeout(Duration::from_secs(120), lance_context_core::merge_write_scope::fence_manifest_versions(&uri, None, &watermarks.versions, &frozen.id) ).await.map_err(|_| "storage recovery barrier deadline exceeded".to_string())? @@ -293,7 +296,7 @@ async fn reconcile_with_grace( let mut next_cancel = tokio::time::Instant::now(); loop { if tokio::time::Instant::now() >= handoff_deadline { - let error = "merge ownership unresolved: executor did not acknowledge termination; fence retained, recovery requires old writer termination evidence"; + let error = "merge ownership unresolved: executor did not acknowledge termination; fence retained pending storage version recovery"; coordinator .record_failure(proof, &initial.target, &initial.endpoint, error) .await?; diff --git a/crates/lance-context-merge/src/fencing.rs b/crates/lance-context-merge/src/fencing.rs index 4bdff20..12c6987 100644 --- a/crates/lance-context-merge/src/fencing.rs +++ b/crates/lance-context-merge/src/fencing.rs @@ -149,9 +149,17 @@ impl Coordinator { if frozen.phase != Phase::Recovering { return Err("execution is not recovering".into()); } - let mut recovered = frozen.finished(Err( - "merge cancelled; admitted manifest versions fenced for recovery".into(), - )); + let cause = frozen + .error + .as_deref() + .unwrap_or("merge execution deadline or lost worker"); + let cause = cause + .strip_prefix("merge ownership unresolved: manifest commit result unknown; ") + .or_else(|| cause.strip_prefix("merge ownership unresolved: ")) + .unwrap_or(cause); + let mut recovered = frozen.finished(Err(format!( + "merge storage fenced for recovery; original failure: {cause}" + ))); recovered.phase = Phase::Recovered; self.transact( vec![ diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index 2171493..3d9f097 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -539,6 +539,37 @@ mod tests { .authorize_commit(&running, "file:///table", "base", 11) .await .is_err()); + // A successful barrier must not disguise a permanent data failure as + // another ownership failure with a short metadata-probe backoff. + let operation = Execution::new("table", "worker", "boot", 1); + assert!(coordinator.reserve(&next, &operation).await.unwrap()); + let running = coordinator.start(&operation).await.unwrap().unwrap(); + assert!(coordinator + .report_uncertain( + &running, + "merge ownership unresolved: manifest commit result unknown; schema mismatch" + .into(), + ) + .await + .unwrap()); + let uncertain = coordinator.get("table").await.unwrap().unwrap(); + let frozen = coordinator + .freeze(&next, &uncertain) + .await + .unwrap() + .unwrap(); + assert!(coordinator.finish_recovery(&next, &frozen).await.unwrap()); + let recovered = coordinator.get("table").await.unwrap().unwrap(); + assert!(coordinator.release(&next, &recovered).await.unwrap()); + let failure = coordinator + .failure("table", "worker") + .await + .unwrap() + .unwrap(); + assert_eq!(failure.class, failure::FailureClass::DataOrConfiguration); + assert!(failure.needs_attention); + assert!(failure.last_error.contains("schema mismatch")); + assert_eq!(failure.next_retry_ms - failure.last_failure_ms, 3_600_000); client .delete( coordinator.prefix.clone(), From 8feef55efa8c557081a19eb46dcec652080dc69d Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Thu, 1 Oct 2026 18:34:12 +0000 Subject: [PATCH 5/9] test: enforce persistent merge circuit and recovery probe budgets --- crates/lance-context-merge/src/failure.rs | 59 +++++++++++++++++++++-- 1 file changed, 55 insertions(+), 4 deletions(-) diff --git a/crates/lance-context-merge/src/failure.rs b/crates/lance-context-merge/src/failure.rs index 84f77e6..8d2fa82 100644 --- a/crates/lance-context-merge/src/failure.rs +++ b/crates/lance-context-merge/src/failure.rs @@ -43,6 +43,12 @@ pub fn classify(error: &str) -> FailureClass { "invalid", "403", "401", + "permission denied", + "access denied", + "accessdenied", + "unauthorized", + "unauthenticated", + "forbidden", "protocol", ] .iter() @@ -265,10 +271,20 @@ mod tests { ShardFailure::advance("table", "worker", old, "temporary storage failure", 1000); assert_eq!(next.consecutive_attempts, attempt); assert_eq!(next.needs_attention, attempt >= 15); - assert!(next.next_retry_ms <= 3_601_000); - if attempt == 3 { - assert_eq!(next.next_retry_ms, 31_000); - } + let expected_secs = match attempt { + 1 | 2 => 2, + 3 => 30, + 4 => 60, + 5 => 120, + 6 => 240, + 7 => 480, + 8..=14 => 900, + _ => 3600, + }; + assert_eq!( + next.next_retry_ms - next.last_failure_ms, + expected_secs * 1000 + ); old = Some(next); } let broken = @@ -276,4 +292,39 @@ mod tests { assert!(broken.needs_attention); assert_eq!(broken.next_retry_ms, 3_601_000); } + + #[test] + fn permanent_failures_and_unresolved_barriers_have_distinct_probe_budgets() { + for error in [ + "NotFound: data/fragment.lance", + "corrupt manifest", + "schema mismatch", + "Permission denied", + "AccessDenied", + "unauthenticated", + "HTTP 403 Forbidden", + ] { + let failure = ShardFailure::advance("table", "worker", None, error, 1000); + assert_eq!(failure.class, FailureClass::DataOrConfiguration); + assert!(failure.needs_attention); + assert_eq!(failure.next_retry_ms, 3_601_000); + } + let mut old = None; + for expected_secs in [30, 60, 120, 240, 300, 300] { + let next = ShardFailure::advance( + "table", + "worker", + old, + "merge ownership unresolved: recovery barrier failed: timeout", + 1000, + ); + assert_eq!(next.class, FailureClass::OwnershipUnresolved); + assert!(next.needs_attention); + assert_eq!( + next.next_retry_ms - next.last_failure_ms, + expected_secs * 1000 + ); + old = Some(next); + } + } } From c4946eab7df4383a3ab28c714ab8b951c25d2381 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Fri, 2 Oct 2026 00:57:08 +0000 Subject: [PATCH 6/9] fix: stage owned merges per table and recover only after bounded progress waits --- .../lance-context-core/src/generic_store.rs | 5 + .../src/merge_write_scope.rs | 18 +- crates/lance-context-core/src/metrics.rs | 3 + .../lance-context-core/src/rollout_store.rs | 5 + crates/lance-context-core/src/store_base.rs | 18 ++ crates/lance-context-master/src/config.rs | 2 + .../src/merge_execution.rs | 225 +++++++++++++- crates/lance-context-master/src/routes.rs | 22 +- crates/lance-context-master/src/scheduler.rs | 101 ++++++- crates/lance-context-master/src/state.rs | 2 + crates/lance-context-master/src/task_store.rs | 1 + crates/lance-context-merge/src/lib.rs | 81 ++++- crates/lance-context-merge/src/progress.rs | 53 ++++ crates/lance-context-merge/src/rollout.rs | 117 ++++++++ crates/lance-context-server/src/config.rs | 17 +- .../src/merge_execution.rs | 283 ++++++++++++++++-- .../src/routes/generic.rs | 8 +- .../src/routes/rollouts.rs | 6 +- crates/lance-context-server/src/state.rs | 60 ++-- crates/lance-context-server/src/sweeper.rs | 147 ++++++++- docs/merge-recovery.md | 160 +++++++--- 21 files changed, 1184 insertions(+), 150 deletions(-) create mode 100644 crates/lance-context-merge/src/progress.rs create mode 100644 crates/lance-context-merge/src/rollout.rs diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs index da36431..2282ffd 100644 --- a/crates/lance-context-core/src/generic_store.rs +++ b/crates/lance-context-core/src/generic_store.rs @@ -406,6 +406,11 @@ impl GenericStore { self.base.prepare_count_merge().await } + /// Check the configured count trigger using only this shard's manifest. + pub async fn count_merge_due(&self) -> LanceResult { + self.base.count_merge_due().await + } + /// The shared-lock half of [`Self::cleanup_wal`]: seal, then read a /// budgeted prefix of flushed generations into memory. Callers holding a /// read lock run this while appends continue, then take the write lock diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index 5d31748..5b569d7 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -13,7 +13,10 @@ use object_store::path::Path; use std::{ future::Future, pin::Pin, - sync::{Arc, Mutex}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, + }, }; use tokio::sync::{oneshot, Notify}; @@ -36,6 +39,7 @@ pub trait CommitAuthorizer: std::fmt::Debug + Send + Sync { #[derive(Debug, Default)] pub struct MergeWriteScope { + completed_steps: AtomicU64, progress: Mutex, changed: Notify, authorizer: Option>, @@ -43,6 +47,9 @@ pub struct MergeWriteScope { } impl MergeWriteScope { + pub fn completed_steps(&self) -> u64 { + self.completed_steps.load(Ordering::Relaxed) + } pub fn new() -> Arc { Arc::new(Self::default()) } @@ -93,6 +100,12 @@ impl MergeWriteScope { } } +/// Count completed work, never timer ticks or admission attempts. This is a +/// stall diagnostic, not evidence that WAL has been durably reclaimed. +pub(crate) fn checkpoint() { + let _ = CURRENT.try_with(|scope| scope.completed_steps.fetch_add(1, Ordering::Relaxed)); +} + pub(crate) async fn authorize(resource: &str, version: u64) -> Result<()> { if let Ok(scope) = CURRENT.try_with(Arc::clone) { if let Some(authorizer) = &scope.authorizer { @@ -141,6 +154,9 @@ where let leaf = tokio::spawn(async move { let mut active = active; let result = write.await; + if result.is_ok() { + active.scope.completed_steps.fetch_add(1, Ordering::Relaxed); + } if result.is_err() { active.scope.mark_uncertain(); } diff --git a/crates/lance-context-core/src/metrics.rs b/crates/lance-context-core/src/metrics.rs index 66bb67b..0dffcf9 100644 --- a/crates/lance-context-core/src/metrics.rs +++ b/crates/lance-context-core/src/metrics.rs @@ -162,6 +162,9 @@ macro_rules! observe_phase { ($phase:expr, $body:expr) => {{ let __start = $crate::metrics::timer_start!(); let __result = $body; + if __result.is_ok() { + $crate::merge_write_scope::checkpoint(); + } match &__result { Ok(_) => $crate::metrics::observe_duration!( $crate::metrics::ROLLOUT_WAL_MERGE_DURATION, diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 967ed55..e771aff 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -655,6 +655,11 @@ impl RolloutStore { self.base.prepare_count_merge().await } + /// Check the configured count trigger using only this shard's manifest. + pub async fn count_merge_due(&self) -> LanceResult { + self.base.count_merge_due().await + } + /// [`Self::prepare_merge_if_ready`], but seals the active memtable first — /// the time-triggered behavior of [`Self::cleanup_own_shard`]. pub async fn prepare_cleanup_merge( diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index dea7306..b4fa0b5 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -864,6 +864,23 @@ impl StorageBase { .await } + /// Metadata-only count trigger for a coordinated sweeper. Never prepares + /// batches or takes a merge-memory reservation just to enqueue work. + pub async fn count_merge_due(&self) -> LanceResult { + if self.merge_after_generations == 0 || self.is_version_pinned() || self.deleted { + return Ok(false); + } + let manifest_store = ShardManifestStore::new( + self.dataset.object_store(None).await?, + &self.dataset.branch_location().path, + self.write_shard, + DEFAULT_MANIFEST_SCAN_BATCH_SIZE, + ); + Ok(manifest_store.read_latest().await?.is_some_and(|manifest| { + manifest.flushed_generations.len() >= self.merge_after_generations + })) + } + /// [`Self::prepare_merge_if_ready`], but seals the active memtable *before* /// consulting the manifest — the time-triggered (`threshold = 1`) behavior /// of [`Self::cleanup_own_shard`]. See that method for why the ordering is @@ -1192,6 +1209,7 @@ impl StorageBase { let mut current_batches = Vec::new(); let mut stream = gen_dataset.scan().try_into_stream().await?; while let Some(batch) = stream.try_next().await? { + crate::merge_write_scope::checkpoint(); if batch.num_rows() > 0 { let batch = align_batch_to_schema(batch, merge_schema.clone())?; buffered_bytes = buffered_bytes.saturating_add(batch.get_array_memory_size()); diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index e9b0d5a..fe7ed96 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -6,6 +6,8 @@ use clap::Parser; #[derive(Debug, Clone, Parser)] #[command(name = "lance-context-master", version)] pub struct MasterConfig { + #[command(flatten)] + pub merge_rollout: lance_context_merge::rollout::MergeRollout, /// Data directory / object-store prefix shared with the data-plane server. #[arg(long, env = "DATA_DIR", default_value = "./data")] pub data_dir: String, diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index b2299f8..0d922ce 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/crates/lance-context-master/src/merge_execution.rs @@ -13,6 +13,10 @@ struct Capabilities { protocol: u32, instance: String, timeout_secs: u64, + queue_timeout_secs: u64, + idle_timeout_secs: u64, + progress_protocol: u32, + owned_targets: Vec, } pub(crate) async fn run_merge_wal( @@ -24,6 +28,28 @@ pub(crate) async fn run_merge_wal( } let coordinator = state.task_store.merge_coordinator(); let proof = state.task_store.merge_claim(claim); + let target = &claim.task.target; + if state.config.merge_rollout.draining(target) { + // A drain stops new work, but must still resolve an already-owned + // execution. Otherwise rollback would strand its persistent lock. + if let Some(old) = coordinator.get(target).await? { + if reconcile(&state.http, &coordinator, &proof, old, false) + .await + .is_err() + { + if let Some(current) = coordinator.get(target).await? { + recover_execution(state, &coordinator, &proof, current).await?; + } + } + } + return Err("merge target draining for protocol transition".into()); + } + if !state.config.merge_rollout.owned(target) { + if coordinator.get(target).await?.is_some() { + return Err("owned execution still present; refusing legacy downgrade".into()); + } + return run_legacy(state, target).await; + } // Reconcile/fence before any other table mutation. One retry of the fan-out // after a completed barrier lets healthy shards progress in this task. for recovery_round in 0..2 { @@ -69,6 +95,89 @@ pub(crate) async fn run_merge_wal( unreachable!("bounded merge recovery loop returns on last iteration") } +// Compatibility path is chosen only by explicit configuration, never after a +// failed owned RPC. It retains legacy limitations until that table is drained. +async fn run_legacy(state: &Arc, target: &str) -> Result { + let mut reclaimed = 0u64; + for endpoint in &state.config.worker_endpoints { + let endpoint = endpoint.trim_end_matches('/'); + let url = match target.strip_prefix("generic:") { + Some(name) => format!("{endpoint}/api/v1/generic/{name}/merge-wal"), + None => format!("{endpoint}/api/v1/internal/merge-wal/{target}"), + }; + let started = std::time::Instant::now(); + let result = async { + let response = state + .http + .post(url) + .send() + .await + .map_err(|e| e.to_string())?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(0); + } + let reply: serde_json::Value = response + .error_for_status() + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + reply["reclaimed"] + .as_u64() + .ok_or_else(|| "invalid legacy merge response".to_string()) + } + .await; + metrics::histogram!("master_merge_wal_worker_duration_seconds") + .record(started.elapsed().as_secs_f64()); + metrics::counter!("master_merge_wal_workers_total", "result" => if result.is_ok() { "ok" } else { "failed" }).increment(1); + reclaimed += result?; + } + metrics::counter!("master_merge_wal_generations_reclaimed_total").increment(reclaimed); + Ok(format!( + "merged {reclaimed} generations across {}/{} workers (legacy)", + state.config.worker_endpoints.len(), + state.config.worker_endpoints.len() + )) +} + +struct Deadlines { + queue: tokio::time::Instant, + execution_timeout: Duration, + idle_timeout: Duration, + running: Option, + last_progress: tokio::time::Instant, + sequence: Option, +} + +impl Deadlines { + fn new(execution: &Execution, now: tokio::time::Instant) -> Self { + Self { + queue: now + Duration::from_secs(execution.queue_timeout_secs) + RPC_TIMEOUT, + execution_timeout: Duration::from_secs(execution.timeout_secs), + idle_timeout: Duration::from_secs(execution.idle_timeout_secs) + RPC_TIMEOUT, + running: None, + last_progress: now, + sequence: None, + } + } + fn observe(&mut self, sequence: u64, now: tokio::time::Instant) { + self.running.get_or_insert(now); + if self.sequence.is_none_or(|previous| sequence > previous) { + self.sequence = Some(sequence); + self.last_progress = now; + } + } + fn expired(&self, now: tokio::time::Instant) -> bool { + match self.running { + None => now >= self.queue, + Some(started) => { + now.duration_since(started) >= self.execution_timeout + RPC_TIMEOUT + || now.duration_since(self.last_progress) >= self.idle_timeout + } + } + } +} + async fn recover_execution( state: &Arc, coordinator: &Coordinator, @@ -225,15 +334,23 @@ async fn one( .json() .await .map_err(|e| e.to_string())?; - if capabilities.protocol != 2 || capabilities.timeout_secs == 0 { + if capabilities.protocol != 2 + || capabilities.progress_protocol != 1 + || !capabilities.owned_targets.iter().any(|t| t == target) + || capabilities.timeout_secs == 0 + || capabilities.queue_timeout_secs == 0 + || capabilities.idle_timeout_secs == 0 + { return Err("worker lacks bounded owned-merge protocol".into()); } - let execution = Execution::new( + let mut execution = Execution::new( target, endpoint, &capabilities.instance, - capabilities.timeout_secs.min(600), + capabilities.timeout_secs, ); + execution.queue_timeout_secs = capabilities.queue_timeout_secs; + execution.idle_timeout_secs = capabilities.idle_timeout_secs; // Even a lost reserve response is ambiguous. Recover the fence before // issuing any other storage operation; never infer absence from an error. let admission = coordinator.reserve(proof, &execution).await; @@ -286,16 +403,16 @@ async fn reconcile_with_grace( cancel_immediately: bool, grace: Duration, ) -> Result { - let deadline = - tokio::time::Instant::now() + Duration::from_secs(initial.timeout_secs) + RPC_TIMEOUT; - let handoff_deadline = if cancel_immediately { - tokio::time::Instant::now() - } else { - deadline - } + grace; + let mut deadlines = Deadlines::new(&initial, tokio::time::Instant::now()); + let mut cancel_started = cancel_immediately.then(tokio::time::Instant::now); let mut next_cancel = tokio::time::Instant::now(); loop { - if tokio::time::Instant::now() >= handoff_deadline { + // Losing etcd responses must not turn this into an unbounded scheduler + // wait. Ownership remains durable when the reconciliation slot exits. + if cancel_started.is_none() && deadlines.expired(tokio::time::Instant::now()) { + cancel_started = Some(tokio::time::Instant::now()); + } + if cancel_started.is_some_and(|at| at.elapsed() >= grace) { let error = "merge ownership unresolved: executor did not acknowledge termination; fence retained pending storage version recovery"; coordinator .record_failure(proof, &initial.target, &initial.endpoint, error) @@ -326,9 +443,15 @@ async fn reconcile_with_grace( } return current.error.map_or(Ok(current.reclaimed), Err); } - if (cancel_immediately || tokio::time::Instant::now() >= deadline) - && tokio::time::Instant::now() >= next_cancel - { + if cancel_started.is_none() { + if let Ok(Some(progress)) = coordinator.progress(¤t).await { + deadlines.observe(progress.sequence, tokio::time::Instant::now()); + } + if deadlines.expired(tokio::time::Instant::now()) { + cancel_started = Some(tokio::time::Instant::now()); + } + } + if cancel_started.is_some() && tokio::time::Instant::now() >= next_cancel { next_cancel = tokio::time::Instant::now() + RETRY_DELAY; // CAS reserved work to terminal, or ask its owner to cancel the // actual storage future. HTTP success alone is not acknowledgement. @@ -351,6 +474,41 @@ async fn reconcile_with_grace( #[cfg(test)] mod tests { + #[test] + fn queue_execution_and_real_progress_have_independent_deadlines() { + use super::*; + let mut execution = Execution::new("hot", "worker", "boot", 1800); + execution.queue_timeout_secs = 600; + execution.idle_timeout_secs = 120; + let now = tokio::time::Instant::now(); + let mut deadlines = Deadlines::new(&execution, now); + assert!(!deadlines.expired(now + Duration::from_secs(599))); + // A long queue wait consumes none of the running budget. + let started = now + Duration::from_secs(599); + deadlines.observe(0, started); + for step in 1..=10 { + let current = started + Duration::from_secs(step * 100); + deadlines.observe(step, current); + assert!(!deadlines.expired(current)); + } + // Receiving the same progress repeatedly is only a heartbeat. + let stalled = started + Duration::from_secs(1131); + deadlines.observe(10, stalled); + assert!(deadlines.expired(stalled)); + // Even actual progress cannot extend the configured total ceiling. + deadlines.observe(11, started + Duration::from_secs(1811)); + assert!(deadlines.expired(started + Duration::from_secs(1811))); + } + + #[test] + fn a_worker_that_never_acquires_a_slot_has_a_bounded_queue_wait() { + use super::*; + let mut execution = Execution::new("hot", "worker", "boot", 3600); + execution.queue_timeout_secs = 10; + let now = tokio::time::Instant::now(); + let deadlines = Deadlines::new(&execution, now); + assert!(deadlines.expired(now + Duration::from_secs(21))); + } use super::*; use axum::{ extract::State, @@ -378,7 +536,7 @@ mod tests { .route( "/api/v1/internal/merge-executor", get(|| async { - Json(serde_json::json!({"protocol":2,"instance":"test","timeout_secs":1})) + Json(serde_json::json!({"protocol":2,"instance":"test","timeout_secs":1,"queue_timeout_secs":1,"idle_timeout_secs":1,"progress_protocol":1,"owned_targets":["table"]})) }), ) .route( @@ -473,6 +631,8 @@ mod tests { &endpoints, "--etcd-prefix", &prefix, + "--merge-owned-targets", + "table", ]); let mut state = MasterState::new(cfg).await.unwrap(); let uri = state.rollout_uri("table"); @@ -614,6 +774,41 @@ mod tests { server.abort(); } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn missing_owned_capability_never_falls_back_to_legacy_post() { + let (coordinator, proof) = fixture().await; + let calls = Arc::new(AtomicUsize::new(0)); + let counted = calls.clone(); + let app = Router::new().route( + "/api/v1/internal/merge-wal/table", + post(move || { + let counted = counted.clone(); + async move { + counted.fetch_add(1, Ordering::SeqCst); + Json(serde_json::json!({"reclaimed": 9})) + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + assert!(one( + &reqwest::Client::new(), + &coordinator, + &proof, + "table", + &endpoint + ) + .await + .is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(coordinator.get("table").await.unwrap().is_none()); + server.abort(); + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn stalled_shard_is_terminated_healthy_shards_advance_then_only_failure_retries() { diff --git a/crates/lance-context-master/src/routes.rs b/crates/lance-context-master/src/routes.rs index a575367..c9d7cd8 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -573,26 +573,6 @@ pub async fn enqueue_task( /// lapses, and the last error. A store in this list is broken in a way that /// retrying will not fix (a manifest naming a missing fragment, for example); /// a manual `POST /tasks` still enqueues it. -#[derive(Debug, Deserialize)] -pub struct MergeFailureParams { - pub after: Option, -} - -pub async fn list_merge_failures( - State(state): State>, - Query(params): Query, -) -> Result, MasterError> { - let (failures, next) = state - .task_store - .merge_coordinator() - .failure_page(params.after.as_deref(), 256) - .await - .map_err(MasterError::Internal)?; - Ok(Json( - serde_json::json!({"failures": failures, "next": next}), - )) -} - pub async fn list_cooldowns( State(state): State>, ) -> Result>, MasterError> { @@ -672,7 +652,6 @@ pub fn api_router() -> Router> { .route("/experiments/{name}/rescan", post(rescan_experiment)) .route("/tasks", post(enqueue_task).get(list_tasks)) .route("/scheduler/cooldowns", get(list_cooldowns)) - .route("/scheduler/merge-failures", get(list_merge_failures)) .route("/scheduler/repairs", get(list_repairs)) .route("/tasks/{id}", get(get_task)) .route("/rescan", post(rescan)) @@ -736,6 +715,7 @@ mod tests { fn test_config(dir: &TempDir) -> MasterConfig { MasterConfig { + merge_rollout: Default::default(), data_dir: dir.path().to_string_lossy().to_string(), host: "127.0.0.1".to_string(), port: 0, diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index a818035..5ade72a 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -147,12 +147,17 @@ async fn run_task(state: &Arc, claim: TaskClaim, timing: TaskClaimT .record(timing.permit_wait.as_secs_f64()); let started = std::time::Instant::now(); - let outcome = match task.kind { - TaskKind::Compact => run_compaction(state, &task).await, - TaskKind::MergeWal => crate::merge_execution::run_merge_wal(state, &claim).await, - TaskKind::IndexId => run_index_id(state, &task.target).await, - TaskKind::Repair => run_repair(state, &task).await, - }; + let outcome = + if state.config.merge_rollout.draining(&task.target) && task.kind != TaskKind::MergeWal { + Err("target draining legacy writers for merge protocol transition".into()) + } else { + match task.kind { + TaskKind::Compact => run_compaction(state, &task).await, + TaskKind::MergeWal => crate::merge_execution::run_merge_wal(state, &claim).await, + TaskKind::IndexId => run_index_id(state, &task.target).await, + TaskKind::Repair => run_repair(state, &task).await, + } + }; let work_elapsed = started.elapsed(); let result = if outcome.is_ok() { "success" } else { "failed" }; @@ -583,15 +588,35 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { tokio::spawn(async move { let coordinator = retry_state.task_store.merge_coordinator(); let mut cursor = None; + let mut request_cursor = None; let mut ticker = tokio::time::interval(Duration::from_secs(15)); loop { ticker.tick().await; + match coordinator.request_page(request_cursor.as_deref()).await { + Ok((targets, next)) => { + request_cursor = next; + for target in targets { + if retry_state.config.merge_rollout.owned(&target) { + match enqueue(&retry_state, TaskKind::MergeWal, &target).await { + Ok(_) => { + let _ = coordinator.acknowledge_request(&target).await; + } + Err(error) => { + tracing::warn!(%target, %error, "worker merge demand enqueue failed") + } + } + } + } + } + Err(error) => tracing::warn!(%error, "worker merge demand scan failed"), + } match coordinator.failure_page(cursor.as_deref(), 256).await { Ok((rows, next)) => { cursor = next; let mut targets = std::collections::HashSet::new(); for failure in rows { - if failure.next_retry_ms <= lance_context_merge::failure::now_ms() + if retry_state.config.merge_rollout.owned(&failure.target) + && failure.next_retry_ms <= lance_context_merge::failure::now_ms() && retry_state .config .worker_endpoints @@ -761,6 +786,14 @@ mod tests { fn config(dir: &TempDir) -> MasterConfig { MasterConfig { + merge_rollout: lance_context_merge::rollout::MergeRollout { + owned_targets: ["exp", "generic:gs", "broken"] + .into_iter() + .map(str::to_string) + .chain((0..20).map(|i| format!("exp-{i}"))) + .collect(), + drain_targets: vec![], + }, data_dir: dir.path().to_string_lossy().to_string(), host: "127.0.0.1".to_string(), port: 0, @@ -820,6 +853,54 @@ mod tests { /// Manual enqueue -> dispatcher compacts -> task reaches Done and the stats /// table records a compaction. + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn draining_resolves_existing_owned_merge_and_blocks_new_target_mutations() { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.merge_rollout.owned_targets.retain(|t| t != "exp"); + cfg.merge_rollout.drain_targets.push("exp".into()); + cfg.worker_endpoints = vec!["http://127.0.0.1:1".into()]; + let state = MasterState::new(cfg).await.unwrap(); + enqueue(&state, TaskKind::MergeWal, "exp").await.unwrap(); + let claim = state + .task_store + .claim_next_of_kinds(crate::task_store::TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + let coordinator = state.task_store.merge_coordinator(); + let proof = state.task_store.merge_claim(&claim); + let execution = + lance_context_merge::Execution::new("exp", "http://127.0.0.1:1", "boot", 30); + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + let running = coordinator.start(&execution).await.unwrap().unwrap(); + assert!(coordinator.finish(&running, Ok(0)).await.unwrap()); + let outcome = crate::merge_execution::run_merge_wal(&state, &claim).await; + assert!(outcome.as_ref().unwrap_err().contains("draining")); + assert!(coordinator.get("exp").await.unwrap().is_none()); + state.task_store.finish(claim, outcome).await.unwrap(); + + let compact = enqueue(&state, TaskKind::Compact, "exp").await.unwrap(); + let claim = state + .task_store + .claim_next_of_kinds(crate::task_store::TaskKinds::GENERAL) + .await + .unwrap() + .unwrap(); + run_task( + &state, + claim, + TaskClaimTiming { + claim: Duration::ZERO, + permit_wait: Duration::ZERO, + }, + ) + .await; + let result = state.task_store.get(&compact.id).await.unwrap().unwrap(); + assert!(result.error.as_deref().unwrap().contains("draining")); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn manual_compaction_runs_and_updates_stats() { @@ -1203,11 +1284,13 @@ mod tests { .await .unwrap(); let coordinator = Coordinator::new(client, cfg.etcd_prefix.clone()); + let owned_targets = cfg.merge_rollout.owned_targets.clone(); axum::Router::new() .route( "/api/v1/internal/merge-executor", - get(|| async { - Json(serde_json::json!({"protocol":2,"instance":"stub","timeout_secs":600})) + get(move || { + let owned_targets = owned_targets.clone(); + async move { Json(serde_json::json!({"protocol":2,"instance":"stub","timeout_secs":600,"queue_timeout_secs":600,"idle_timeout_secs":600,"progress_protocol":1,"owned_targets":owned_targets})) } }), ) .route( diff --git a/crates/lance-context-master/src/state.rs b/crates/lance-context-master/src/state.rs index 425d85d..703a40e 100644 --- a/crates/lance-context-master/src/state.rs +++ b/crates/lance-context-master/src/state.rs @@ -120,6 +120,7 @@ pub struct MasterState { impl MasterState { /// Open the registry, stats dataset, and configured durable task store. pub async fn new(config: MasterConfig) -> lance::Result> { + config.merge_rollout.validate().map_err(lance::Error::io)?; 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. @@ -239,6 +240,7 @@ mod tests { fn test_config(dir: &TempDir) -> MasterConfig { MasterConfig { + merge_rollout: Default::default(), data_dir: dir.path().to_string_lossy().to_string(), host: "127.0.0.1".to_string(), port: 0, diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index 66d7a64..31cb9cf 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -1410,6 +1410,7 @@ mod tests { fn config(dir: &TempDir) -> MasterConfig { MasterConfig { + merge_rollout: Default::default(), data_dir: dir.path().to_string_lossy().to_string(), host: "127.0.0.1".to_string(), port: 0, diff --git a/crates/lance-context-merge/src/lib.rs b/crates/lance-context-merge/src/lib.rs index 3d9f097..87b193d 100644 --- a/crates/lance-context-merge/src/lib.rs +++ b/crates/lance-context-merge/src/lib.rs @@ -6,6 +6,8 @@ pub mod failure; pub mod fencing; +pub mod progress; +pub mod rollout; use etcd_client::{Client, Compare, CompareOp, Txn, TxnOp}; use serde::{Deserialize, Serialize}; @@ -37,6 +39,10 @@ pub struct Execution { pub endpoint: String, pub instance: String, pub timeout_secs: u64, + #[serde(default = "default_queue_timeout")] + pub queue_timeout_secs: u64, + #[serde(default = "default_queue_timeout")] + pub idle_timeout_secs: u64, pub phase: Phase, pub reclaimed: usize, pub error: Option, @@ -48,6 +54,10 @@ fn protocol_absent(protocol: &u32) -> bool { *protocol == 0 } +fn default_queue_timeout() -> u64 { + 600 +} + impl Execution { pub fn new(target: &str, endpoint: &str, instance: &str, timeout_secs: u64) -> Self { Self { @@ -56,6 +66,8 @@ impl Execution { endpoint: endpoint.into(), instance: instance.into(), timeout_secs, + queue_timeout_secs: default_queue_timeout(), + idle_timeout_secs: default_queue_timeout(), phase: Phase::Reserved, reclaimed: 0, error: None, @@ -217,6 +229,7 @@ impl Coordinator { ]); operations.extend([ self.remove_permits(execution), + self.remove_progress(execution), TxnOp::delete(key, None), TxnOp::put( target_lock_key(&self.prefix, &execution.target), @@ -249,13 +262,13 @@ impl Coordinator { /// The executor owns this future independently of the request handler. On /// cancellation, drop its scoped storage future *before* publishing Finished. /// Callers must not pass a detached JoinHandle: dropping one does not stop it. -pub async fn execute_scoped( +pub async fn execute_scoped( work: F, timeout: std::time::Duration, mut cancel: tokio::sync::watch::Receiver, -) -> Result +) -> Result where - F: std::future::Future>, + F: std::future::Future>, { // Keep the work in this inner scope: select! only drops branch borrows when // the future was pinned outside it, which would publish completion too soon. @@ -426,7 +439,7 @@ mod tests { async fn pre_cancelled_request_never_polls_storage() { let (cancel, rx) = watch::channel(false); cancel.send(true).unwrap(); - let result = execute_scoped( + let result: Result = execute_scoped( async { panic!("must not execute storage") }, Duration::from_secs(1), rx, @@ -471,6 +484,66 @@ mod tests { ) } + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn worker_demand_coalesces_without_resetting_failure_backoff() { + let (coordinator, _, proof, _) = fixture().await; + let mut last = None; + for _ in 0..4 { + last = Some( + coordinator + .record_failure(&proof, "table", "worker", "storage timeout") + .await + .unwrap(), + ); + } + for _ in 0..10 { + coordinator.request_merge("table").await.unwrap(); + } + let (rows, next) = coordinator.request_page(None).await.unwrap(); + assert_eq!(rows, ["table"]); + assert!(next.is_none()); + let after = coordinator + .failure("table", "worker") + .await + .unwrap() + .unwrap(); + let before = last.unwrap(); + assert_eq!(after.consecutive_attempts, 4); + assert_eq!(after.next_retry_ms, before.next_retry_ms); + coordinator.acknowledge_request("table").await.unwrap(); + assert!(coordinator.request_page(None).await.unwrap().0.is_empty()); + coordinator.request_merge("table").await.unwrap(); + assert_eq!(coordinator.request_page(None).await.unwrap().0, ["table"]); + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn progress_cannot_resurrect_a_frozen_or_released_execution() { + let (coordinator, _, proof, _) = fixture().await; + let execution = Execution::new("table", "worker", "boot", 3600); + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + let running = coordinator.start(&execution).await.unwrap().unwrap(); + assert!(coordinator.publish_progress(&running, 4).await.unwrap()); + assert_eq!( + coordinator + .progress(&running) + .await + .unwrap() + .unwrap() + .sequence, + 4 + ); + let frozen = coordinator.freeze(&proof, &running).await.unwrap().unwrap(); + assert!(!coordinator.publish_progress(&running, 5).await.unwrap()); + // Empty watermarks: this test did not authorize any storage writes. + assert!(coordinator.finish_recovery(&proof, &frozen).await.unwrap()); + let recovered = coordinator.get("table").await.unwrap().unwrap(); + assert!(coordinator.release(&proof, &recovered).await.unwrap()); + assert!(coordinator.progress(&running).await.unwrap().is_none()); + assert!(!coordinator.publish_progress(&running, 6).await.unwrap()); + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn recovery_closes_commit_admission_and_retains_every_allowed_version() { diff --git a/crates/lance-context-merge/src/progress.rs b/crates/lance-context-merge/src/progress.rs new file mode 100644 index 0000000..54024f5 --- /dev/null +++ b/crates/lance-context-merge/src/progress.rs @@ -0,0 +1,53 @@ +//! Execution progress is separate from commit ownership. A heartbeat with an +//! unchanged sequence must never extend a no-progress deadline. +use crate::{encode, execution_key, Coordinator, Execution, Result}; +use etcd_client::{Compare, CompareOp, TxnOp}; + +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct ExecutionProgress { + pub sequence: u64, +} + +impl Coordinator { + fn progress_key(&self, execution: &Execution) -> String { + format!( + "{}/merge-progress/{}", + self.prefix.trim_end_matches('/'), + execution.id + ) + } + + pub async fn publish_progress(&self, execution: &Execution, sequence: u64) -> Result { + self.transact( + vec![Compare::value( + execution_key(&self.prefix, &execution.target), + CompareOp::Equal, + encode(execution), + )], + vec![TxnOp::put( + self.progress_key(execution), + serde_json::to_vec(&ExecutionProgress { sequence }).unwrap(), + None, + )], + ) + .await + } + + pub async fn progress(&self, execution: &Execution) -> Result> { + let response = self + .client + .clone() + .get(self.progress_key(execution), None) + .await + .map_err(|e| e.to_string())?; + response + .kvs() + .first() + .map(|kv| serde_json::from_slice(kv.value()).map_err(|e| e.to_string())) + .transpose() + } + + pub(crate) fn remove_progress(&self, execution: &Execution) -> TxnOp { + TxnOp::delete(self.progress_key(execution), None) + } +} diff --git a/crates/lance-context-merge/src/rollout.rs b/crates/lance-context-merge/src/rollout.rs new file mode 100644 index 0000000..9e7fac2 --- /dev/null +++ b/crates/lance-context-merge/src/rollout.rs @@ -0,0 +1,117 @@ +//! Explicit per-table rollout. Capability deployment alone changes no merge path. +use crate::{execution_key, Coordinator, Result}; +use etcd_client::{Compare, CompareOp, GetOptions, TxnOp}; + +#[derive(Clone, Debug, Default, clap::Args)] +pub struct MergeRollout { + /// Exact scheduler targets using owned merges (generic stores: generic:name). + /// Empty by default. Enable only after the target's legacy writers drain. + #[arg(long, env = "MERGE_OWNED_TARGETS", value_delimiter = ',')] + pub owned_targets: Vec, + /// Drain maintenance for these targets during migration; ingestion and + /// maintenance of other tables continue. Coordinate all replicas and helpers. + #[arg(long, env = "MERGE_DRAIN_TARGETS", value_delimiter = ',')] + pub drain_targets: Vec, +} + +impl MergeRollout { + pub fn owned(&self, target: &str) -> bool { + self.owned_targets.iter().any(|t| t == target) + } + + pub fn draining(&self, target: &str) -> bool { + self.drain_targets.iter().any(|t| t == target) + } + + pub fn validate(&self) -> Result<()> { + if self.owned_targets.iter().any(|t| self.draining(t)) { + return Err("a merge target cannot be both owned and draining".into()); + } + if self + .owned_targets + .iter() + .chain(&self.drain_targets) + .any(|t| t.is_empty() || t == "*" || t.starts_with("datagen:")) + { + return Err("merge rollout requires explicit rollout/generic targets; wildcards and datagen targets are unsupported".into()); + } + Ok(()) + } +} + +impl Coordinator { + fn request_key(&self, target: &str) -> String { + execution_key(&self.prefix, target).replace("/merge-executions/", "/merge-requests/") + } + + /// Coalesce worker count/timer triggers without resetting failure budgets. + pub async fn request_merge(&self, target: &str) -> Result<()> { + let key = self.request_key(target); + self.transact( + vec![Compare::version(key.as_str(), CompareOp::Equal, 0)], + vec![TxnOp::put(key, target, None)], + ) + .await?; + Ok(()) + } + + pub async fn request_page(&self, after: Option<&str>) -> Result<(Vec, Option)> { + let prefix = format!("{}/merge-requests/", self.prefix.trim_end_matches('/')); + let (start, options) = match after { + None => (prefix.clone(), GetOptions::new().with_prefix()), + Some(key) if key.starts_with(&prefix) => { + let mut end = prefix.as_bytes().to_vec(); + *end.last_mut().unwrap() += 1; + (format!("{key}\0"), GetOptions::new().with_range(end)) + } + Some(_) => return Err("invalid merge request cursor".into()), + }; + let response = self + .client + .clone() + .get(start, Some(options.with_limit(256))) + .await + .map_err(|e| e.to_string())?; + let rows = response + .kvs() + .iter() + .map(|kv| String::from_utf8_lossy(kv.value()).into_owned()) + .collect(); + let next = response + .more() + .then(|| { + response + .kvs() + .last() + .map(|kv| String::from_utf8_lossy(kv.key()).into_owned()) + }) + .flatten(); + Ok((rows, next)) + } + + /// Called only after a durable enqueue. A subsequent flush tick reasserts + /// demand if the active task already passed this worker's shard. + pub async fn acknowledge_request(&self, target: &str) -> Result<()> { + self.client + .clone() + .delete(self.request_key(target), None) + .await + .map_err(|e| e.to_string())?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn rollout_is_explicit_and_disjoint() { + let mut config = MergeRollout::default(); + assert!(!config.owned("hot")); + config.owned_targets.push("hot".into()); + assert!(config.owned("hot")); + assert!(!config.owned("hot2")); + config.drain_targets.push("hot".into()); + assert!(config.validate().is_err()); + } +} diff --git a/crates/lance-context-server/src/config.rs b/crates/lance-context-server/src/config.rs index 6ff12b7..aa387ea 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -4,15 +4,26 @@ use clap::Parser; #[command(name = "lance-context-server")] #[command(about = "REST API server for lance-context")] pub struct ServerConfig { - /// Configure etcd to enable the owned merge execution protocol. + /// Connection used lazily by explicitly enabled owned merge targets. #[command(flatten)] pub merge_etcd: lance_context_merge::EtcdConfig, - /// Deadline for scoped merge work, including slot acquisition. Already + #[command(flatten)] + pub merge_rollout: lance_context_merge::rollout::MergeRollout, + + /// Maximum execution time AFTER slot acquisition. Already /// started manifest writes must drain before ownership can be released. - #[arg(long, env = "MERGE_EXECUTION_TIMEOUT_SECS", default_value_t = 600)] + #[arg(long, env = "MERGE_EXECUTION_TIMEOUT_SECS", default_value_t = 3600)] pub merge_execution_timeout_secs: u64, + /// Maximum wait for a merge slot, separate from execution time. + #[arg(long, env = "MERGE_QUEUE_TIMEOUT_SECS", default_value_t = 600)] + pub merge_queue_timeout_secs: u64, + + /// Maximum interval without an observed completed merge step. + #[arg(long, env = "MERGE_IDLE_TIMEOUT_SECS", default_value_t = 600)] + pub merge_idle_timeout_secs: u64, + #[arg(long, default_value = "0.0.0.0")] pub host: String, diff --git a/crates/lance-context-server/src/merge_execution.rs b/crates/lance-context-server/src/merge_execution.rs index 92501fa..e5c2577 100644 --- a/crates/lance-context-server/src/merge_execution.rs +++ b/crates/lance-context-server/src/merge_execution.rs @@ -9,7 +9,7 @@ use axum::{extract::State, http::StatusCode, Json}; use futures::FutureExt; use lance_context_merge::{execute_scoped, Coordinator, Execution, Phase}; use std::{collections::HashMap, sync::Arc, time::Duration}; -use tokio::sync::{watch, Mutex}; +use tokio::sync::{watch, Mutex, OnceCell}; struct OwnedCommitGuard { coordinator: Coordinator, @@ -45,23 +45,68 @@ impl lance_context_core::merge_write_scope::CommitAuthorizer for OwnedCommitGuar } pub struct Executions { - coordinator: Option, + coordinator: OnceCell, + etcd: Option, + pub(crate) rollout: lance_context_merge::rollout::MergeRollout, + queue_timeout_secs: u64, + idle_timeout_secs: u64, instance: String, timeout_secs: u64, running: Mutex>>, } impl Executions { + #[cfg(test)] pub fn new(coordinator: Option, timeout_secs: u64) -> Self { Self { - coordinator, + coordinator: OnceCell::new_with(coordinator), + etcd: None, + rollout: Default::default(), + queue_timeout_secs: 600, + idle_timeout_secs: 600, instance: uuid::Uuid::new_v4().to_string(), timeout_secs, running: Mutex::new(HashMap::new()), } } - pub(crate) fn enabled(&self) -> bool { - self.coordinator.is_some() + pub fn configured( + etcd: lance_context_merge::EtcdConfig, + rollout: lance_context_merge::rollout::MergeRollout, + timeout_secs: u64, + queue_timeout_secs: u64, + idle_timeout_secs: u64, + ) -> Self { + Self { + coordinator: OnceCell::new(), + etcd: Some(etcd), + rollout, + timeout_secs, + queue_timeout_secs, + idle_timeout_secs, + instance: uuid::Uuid::new_v4().to_string(), + running: Mutex::new(HashMap::new()), + } + } + + pub(crate) fn owned(&self, target: &str) -> bool { + self.rollout.owned(target) + } + + pub(crate) fn legacy_allowed(&self, target: &str) -> Result<(), AppError> { + if self.owned(target) || self.rollout.draining(target) { + return Err(AppError::Overloaded( + "table is draining or requires owned merge protocol".into(), + )); + } + Ok(()) + } + + pub(crate) async fn request_merge(&self, target: &str) -> Result<(), String> { + self.coordinator() + .await + .map_err(|e| format!("{e:?}"))? + .request_merge(target) + .await } /// Invoked after HTTP admission has stopped. Detached executors are not @@ -75,20 +120,30 @@ impl Executions { } } - fn coordinator(&self) -> Result { - self.coordinator.clone().ok_or_else(|| { - AppError::Overloaded("owned merge execution requires ETCD_ENDPOINTS".into()) - }) + async fn coordinator(&self) -> Result { + self.coordinator + .get_or_try_init(|| async { + let config = self.etcd.as_ref().ok_or_else(|| { + AppError::Overloaded("owned merge execution requires ETCD_ENDPOINTS".into()) + })?; + config.connect().await.map_err(AppError::Internal) + }) + .await + .cloned() } } pub async fn capabilities( State(state): State>, ) -> Result, AppError> { - state.merge_executions.coordinator()?; Ok(Json( serde_json::json!({"protocol": 2, "instance": state.merge_executions.instance, - "timeout_secs": state.merge_executions.timeout_secs}), + "timeout_secs": state.merge_executions.timeout_secs, + "queue_timeout_secs": state.merge_executions.queue_timeout_secs, + "idle_timeout_secs": state.merge_executions.idle_timeout_secs, + "progress_protocol": 1, + "owned_targets": state.merge_executions.rollout.owned_targets, + "drain_targets": state.merge_executions.rollout.drain_targets}), )) } @@ -96,12 +151,17 @@ pub async fn start( State(state): State>, Json(execution): Json, ) -> Result { - let coordinator = state.merge_executions.coordinator()?; - if execution.protocol != 2 + let coordinator = state.merge_executions.coordinator().await?; + if !state.merge_executions.owned(&execution.target) + || execution.protocol != 2 || execution.instance != state.merge_executions.instance || execution.phase != Phase::Reserved || execution.timeout_secs == 0 || execution.timeout_secs > state.merge_executions.timeout_secs + || execution.queue_timeout_secs == 0 + || execution.queue_timeout_secs > state.merge_executions.queue_timeout_secs + || execution.idle_timeout_secs == 0 + || execution.idle_timeout_secs > state.merge_executions.idle_timeout_secs { return Err(AppError::InvalidRequest( "merge executor incarnation or deadline mismatch".into(), @@ -172,7 +232,6 @@ async fn run( if !lance_context_core::merge_write_scope::supports_version_fencing(&uri) { return Err("invalid storage backend for owned merge fencing".into()); } - slot = state.acquire_merge_slot().await; let result = if let Some(name) = target.strip_prefix("generic:") { generic::merge_generic_wal_owned( State(state.clone()), @@ -207,11 +266,30 @@ async fn run( dataset_uri: uri.clone(), }), ); - let outcome = std::panic::AssertUnwindSafe(execute_scoped( - write_scope.run(work), - Duration::from_secs(execution.timeout_secs), - cancelled, - )) + let outcome = std::panic::AssertUnwindSafe(async { + // No payload work starts before admission to the process-wide slot. + slot = execute_scoped( + async { Ok(state.acquire_merge_slot().await) }, + Duration::from_secs(execution.queue_timeout_secs), + cancelled.clone(), + ) + .await + .map_err(|e| format!("merge queue wait: {e}"))?; + if !coordinator.publish_progress(&running, 0).await? { + return Err("merge ownership revoked before execution".into()); + } + execute_scoped( + async { + tokio::select! { + result = write_scope.run(work) => result, + error = watch_progress(&coordinator, &running, &write_scope) => Err(error), + } + }, + Duration::from_secs(execution.timeout_secs), + cancelled, + ) + .await + }) .catch_unwind() .await .unwrap_or_else(|_| Err("merge executor panicked".into())); @@ -286,11 +364,41 @@ async fn run( Ok(()) } +async fn watch_progress( + coordinator: &Coordinator, + execution: &Execution, + scope: &lance_context_core::merge_write_scope::MergeWriteScope, +) -> String { + let mut sequence = 0; + let mut changed = tokio::time::Instant::now(); + let mut published = 0; + loop { + tokio::time::sleep(Duration::from_secs(1)).await; + let current = scope.completed_steps(); + if current != sequence { + sequence = current; + changed = tokio::time::Instant::now(); + } + if changed.elapsed() >= Duration::from_secs(execution.idle_timeout_secs) { + return "merge no-progress deadline exceeded".into(); + } + if current != published { + match coordinator.publish_progress(execution, current).await { + Ok(true) => published = current, + Ok(false) => return "merge ownership revoked during execution".into(), + Err(error) => { + tracing::warn!(id = %execution.id, %error, "merge progress publication failed") + } + } + } + } +} + pub async fn cancel( State(state): State>, Json(execution): Json, ) -> Result { - let coordinator = state.merge_executions.coordinator()?; + let coordinator = state.merge_executions.coordinator().await?; if coordinator .cancel_reserved(&execution) .await @@ -322,6 +430,136 @@ mod tests { use lance_context_merge::{ClaimProof, EtcdConfig}; use tokio::io::AsyncWriteExt; + #[tokio::test] + async fn unavailable_etcd_does_not_block_startup_or_disable_legacy_self_merge() { + use clap::Parser; + let dir = tempfile::tempdir().unwrap(); + let config = crate::config::ServerConfig::parse_from([ + "test", + "--data-dir", + dir.path().to_str().unwrap(), + "--etcd-endpoints", + "http://127.0.0.1:1", + "--merge-owned-targets", + "hot", + "--rollout-merge-after-generations", + "3", + "--rollout-cleanup-interval-secs", + "30", + ]); + let state = Arc::new(AppState::new(config).await.unwrap()); + assert!(state.merge_executions.coordinator.get().is_none()); + assert_eq!(state.rollout_merge_after_generations, 3); + assert_eq!(state.rollout_cleanup_interval_secs, 30); + assert!(state.merge_executions.legacy_allowed("other").is_ok()); + assert!(state.merge_executions.legacy_allowed("hot").is_err()); + let Json(reply) = capabilities(State(state.clone())).await.unwrap(); + assert_eq!(reply["owned_targets"][0], "hot"); + assert!(state.merge_executions.coordinator.get().is_none()); + } + + async fn deadline_fixture() -> (Arc, Coordinator, ClaimProof, tempfile::TempDir) { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let mut client = etcd_client::Client::connect([endpoint], None) + .await + .unwrap(); + let prefix = format!("/server-deadlines/{}", uuid::Uuid::new_v4()); + let proof = ClaimProof { + key: format!("{prefix}/claim"), + token: "owner".into(), + lease_id: 0, + }; + client + .put(proof.key.as_str(), proof.token.as_str(), None) + .await + .unwrap(); + client + .put( + lance_context_merge::target_lock_key(&prefix, "hot"), + proof.token.as_str(), + None, + ) + .await + .unwrap(); + let coordinator = Coordinator::new(client, prefix); + let dir = tempfile::tempdir().unwrap(); + let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; + state.merge_slots = Some(Arc::new(tokio::sync::Semaphore::new(1))); + state.merge_executions = Executions::new(Some(coordinator.clone()), 3600); + state + .merge_executions + .rollout + .owned_targets + .push("hot".into()); + let state = Arc::new(state); + let _ = rollouts::create_rollout_store( + State(state.clone()), + Json(lance_context_api::CreateRolloutStoreRequest { + name: "hot".into(), + storage_options: None, + }), + ) + .await + .unwrap(); + (state, coordinator, proof, dir) + } + + async fn terminal(coordinator: &Coordinator) -> Execution { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let current = coordinator.get("hot").await.unwrap().unwrap(); + if current.phase == Phase::Finished { + return current; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .unwrap() + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn queue_wait_longer_than_execution_budget_still_runs_after_slot_release() { + let (state, coordinator, proof, _dir) = deadline_fixture().await; + let held = state.acquire_merge_slot().await; + let mut execution = Execution::new("hot", "worker", &state.merge_executions.instance, 1); + execution.queue_timeout_secs = 10; + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + start(State(state.clone()), Json(execution.clone())) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(1300)).await; + assert_eq!( + coordinator.get("hot").await.unwrap().unwrap().phase, + Phase::Running + ); + assert!(coordinator.progress(&execution).await.unwrap().is_none()); + drop(held); + let finished = terminal(&coordinator).await; + assert!(finished.error.is_none(), "{finished:?}"); + assert!(coordinator.release(&proof, &finished).await.unwrap()); + state.merge_executions.shutdown().await; + } + + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn no_progress_cancels_a_blocked_merge_before_the_total_deadline() { + let (state, coordinator, proof, _dir) = deadline_fixture().await; + let store = state.get_or_open_rollout_store("hot").await.unwrap(); + let held = store.write().await; + let mut execution = Execution::new("hot", "worker", &state.merge_executions.instance, 3600); + execution.idle_timeout_secs = 1; + assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + start(State(state.clone()), Json(execution)).await.unwrap(); + let finished = terminal(&coordinator).await; + assert!(finished.error.as_ref().unwrap().contains("no-progress")); + assert!(coordinator.release(&proof, &finished).await.unwrap()); + assert_eq!(state.merge_slots.as_ref().unwrap().available_permits(), 1); + drop(held); + state.merge_executions.shutdown().await; + } + #[tokio::test] #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] async fn disconnected_http_call_is_cancelled_before_ownership_handoff() { @@ -364,6 +602,11 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut state = AppState::new_for_test(dir.path().to_path_buf()).await; state.merge_executions = Executions::new(Some(coordinator.clone()), 600); + state + .merge_executions + .rollout + .owned_targets + .push("blocked".into()); let state = Arc::new(state); let name = "blocked"; // Create via the real API so registry and resident handle agree. diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs index d0f9d87..087f0fd 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -273,11 +273,9 @@ pub async fn merge_generic_wal( State(state): State>, Path(name): Path, ) -> Result, AppError> { - if state.merge_executions.enabled() { - return Err(AppError::Overloaded( - "use the owned merge executor protocol".into(), - )); - } + state + .merge_executions + .legacy_allowed(&format!("generic:{name}"))?; let _slot = state.acquire_merge_slot().await; merge_generic_wal_owned(State(state), Path(name)).await } diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs index e1ade66..cae4e05 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -728,11 +728,7 @@ pub async fn merge_wal( State(state): State>, Path(name): Path, ) -> Result, AppError> { - if state.merge_executions.enabled() { - return Err(AppError::Overloaded( - "use the owned merge executor protocol".into(), - )); - } + state.merge_executions.legacy_allowed(&name)?; let _slot = state.acquire_merge_slot().await; merge_wal_owned(State(state), Path(name)).await } diff --git a/crates/lance-context-server/src/state.rs b/crates/lance-context-server/src/state.rs index 61632ae..ba8e326 100644 --- a/crates/lance-context-server/src/state.rs +++ b/crates/lance-context-server/src/state.rs @@ -322,34 +322,32 @@ impl AppState { let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) .await .map_err(AppError::from_lance)?; - if !config.merge_etcd.etcd_endpoints.is_empty() - && (config.rollout_merge_after_generations != 0 - || config.rollout_cleanup_interval_secs != 0) + config + .merge_rollout + .validate() + .map_err(AppError::InvalidRequest)?; + if !config.merge_rollout.owned_targets.is_empty() + && config.merge_etcd.etcd_endpoints.is_empty() { return Err(AppError::InvalidRequest( - "owned merge execution requires self-merge sweepers disabled".into(), + "owned targets require ETCD_ENDPOINTS".into(), )); } - let merge_coordinator = if config.merge_etcd.etcd_endpoints.is_empty() { - None - } else { - Some( - config - .merge_etcd - .connect() - .await - .map_err(AppError::Internal)?, - ) - }; - if config.merge_execution_timeout_secs == 0 { + if config.merge_execution_timeout_secs == 0 + || config.merge_queue_timeout_secs == 0 + || config.merge_idle_timeout_secs == 0 + { return Err(AppError::InvalidRequest( - "MERGE_EXECUTION_TIMEOUT_SECS must be positive".into(), + "merge execution, queue and idle timeouts must be positive".into(), )); } Ok(Self { - merge_executions: crate::merge_execution::Executions::new( - merge_coordinator, + merge_executions: crate::merge_execution::Executions::configured( + config.merge_etcd.clone(), + config.merge_rollout.clone(), config.merge_execution_timeout_secs, + config.merge_queue_timeout_secs, + config.merge_idle_timeout_secs, ), stores: RwLock::new(std::collections::HashMap::new()), rollout_stores: Mutex::new(LruCache::new(capacity)), @@ -983,17 +981,20 @@ impl AppState { // unswept for hours while their generations piled into the // tens of thousands and every read re-opened all of them. tokio::join!( - sweeper::merge_pass( + sweeper::merge_pass_coordinated( sweeper::resident(&state.rollout_stores).await, - pass_timeout + pass_timeout, + Some(state.clone()), ), - sweeper::merge_pass( + sweeper::merge_pass_coordinated( sweeper::resident(&state.datagen_stores).await, - pass_timeout + pass_timeout, + Some(state.clone()), ), - sweeper::merge_pass( + sweeper::merge_pass_coordinated( sweeper::resident(&state.generic_stores).await, - pass_timeout + pass_timeout, + Some(state.clone()), ), ); } @@ -1044,20 +1045,23 @@ impl AppState { // a no-op in steady state; generic stores default to a deferred // seal and genuinely depend on this. tokio::join!( - sweeper::flush_pass( + sweeper::flush_pass_coordinated( sweeper::resident(&state.rollout_stores).await, pass_timeout, state.merge_slots.clone(), + Some(state.clone()), ), - sweeper::flush_pass( + sweeper::flush_pass_coordinated( sweeper::resident(&state.datagen_stores).await, pass_timeout, state.merge_slots.clone(), + Some(state.clone()), ), - sweeper::flush_pass( + sweeper::flush_pass_coordinated( sweeper::resident(&state.generic_stores).await, pass_timeout, state.merge_slots.clone(), + Some(state.clone()), ), ); } diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs index 3f04f44..2427475 100644 --- a/crates/lance-context-server/src/sweeper.rs +++ b/crates/lance-context-server/src/sweeper.rs @@ -45,6 +45,11 @@ pub(crate) trait Sweepable: Send + Sync + 'static { async { Ok(0) } } + /// Cheap metadata-only predicate used when the master owns execution. + fn count_merge_due(&self) -> impl std::future::Future> + Send { + async { Ok(false) } + } + /// Fold **every** pending flushed generation into the base table; returns /// how many were reclaimed. fn merge_wal(&self) -> impl std::future::Future> + Send; @@ -61,6 +66,14 @@ impl Sweepable for Arc> { guard.flush().await.map_err(|e| e.to_string()) } + async fn count_merge_due(&self) -> Result { + self.read() + .await + .count_merge_due() + .await + .map_err(|e| e.to_string()) + } + async fn merge_if_due(&self) -> Result { let prepared = { let guard = self.read().await; @@ -128,6 +141,14 @@ impl Sweepable for Arc> { guard.flush().await.map_err(|e| e.to_string()) } + async fn count_merge_due(&self) -> Result { + self.read() + .await + .count_merge_due() + .await + .map_err(|e| e.to_string()) + } + async fn merge_if_due(&self) -> Result { // Generic stores default to a deferred seal and take the same // one-row-per-append traffic as rollout, so the count trigger rides @@ -189,10 +210,11 @@ pub(crate) async fn resident(cache: &Mutex>) -> Ve /// and alerts keep working; the new `kind` label is what distinguishes the /// store types. Renaming them would be a silent breakage for anyone graphing /// these today. -pub(crate) async fn flush_pass( +pub(crate) async fn flush_pass_coordinated( stores: Vec<(String, S)>, pass_timeout: Duration, merge_slots: Option>, + state: Option>, ) { let kind = S::kind(); for (name, store) in stores { @@ -211,6 +233,21 @@ pub(crate) async fn flush_pass( // hold up the flush of every store behind this one. When the // slots are full the merge is skipped; the next pass (30 s) // tries again, and the master's sweep covers the store anyway. + if let Some(state) = &state { + match tokio::time::timeout( + pass_timeout, + route_merge(state, &name, &store, true), + ) + .await + { + Ok(Ok(false)) => {} + Ok(Ok(true)) => continue, + outcome => { + tracing::warn!(store = %name, ?outcome, "coordinated count merge request failed"); + continue; + } + } + } let slot = match &merge_slots { Some(slots) => match slots.clone().try_acquire_owned() { Ok(permit) => Some(permit), @@ -250,9 +287,32 @@ pub(crate) async fn flush_pass( } /// Merge every resident store's pending generations, same timeout discipline. -pub(crate) async fn merge_pass(stores: Vec<(String, S)>, pass_timeout: Duration) { +pub(crate) async fn merge_pass_coordinated( + stores: Vec<(String, S)>, + pass_timeout: Duration, + state: Option>, +) { let kind = S::kind(); for (name, store) in stores { + if let Some(state) = &state { + match tokio::time::timeout(pass_timeout, route_merge(state, &name, &store, false)).await + { + Ok(Ok(false)) => {} + Ok(Ok(true)) => continue, + outcome => { + tracing::warn!(store = %name, ?outcome, "coordinated cleanup request failed"); + continue; + } + } + } + let slot = match state.as_ref().and_then(|s| s.merge_slots.clone()) { + Some(slots) => match slots.try_acquire_owned() { + Ok(slot) => Some(slot), + Err(_) => continue, + }, + None => None, + }; + let _slot = slot; report_merge( &name, kind, @@ -261,6 +321,46 @@ pub(crate) async fn merge_pass(stores: Vec<(String, S)>, pass_time } } +/// Return true when routing handled the target (including drain/no-op). +async fn route_merge( + state: &Arc, + name: &str, + store: &S, + count_trigger: bool, +) -> Result { + let target = match S::kind() { + "rollout" => name.to_string(), + "generic" => format!("generic:{name}"), + _ => return Ok(false), + }; + if state.merge_executions.rollout.draining(&target) { + return Ok(true); + } + if !state.merge_executions.owned(&target) { + return Ok(false); + } + if !count_trigger || store.count_merge_due().await? { + state.merge_executions.request_merge(&target).await?; + metrics::counter!("rollout_wal_coordinated_merge_requests_total", "kind" => S::kind()) + .increment(1); + } + Ok(true) +} + +#[cfg(test)] +pub(crate) async fn flush_pass( + stores: Vec<(String, S)>, + timeout: Duration, + slots: Option>, +) { + flush_pass_coordinated(stores, timeout, slots, None).await; +} + +#[cfg(test)] +pub(crate) async fn merge_pass(stores: Vec<(String, S)>, timeout: Duration) { + merge_pass_coordinated(stores, timeout, None).await; +} + /// Record one merge attempt's outcome on the cleanup counters and log. fn report_merge( name: &str, @@ -301,6 +401,49 @@ mod tests { use serde_json::json; use tempfile::TempDir; + #[tokio::test] + #[ignore = "requires isolated local ETCD_TEST_ENDPOINTS"] + async fn owned_count_trigger_enqueues_without_bypassing_slots_or_preparing_payloads() { + let endpoint = std::env::var("ETCD_TEST_ENDPOINTS").unwrap(); + let client = etcd_client::Client::connect([endpoint], None) + .await + .unwrap(); + let coordinator = lance_context_merge::Coordinator::new( + client, + format!("/owned-sweeper/{}", uuid::Uuid::new_v4()), + ); + let state_dir = TempDir::new().unwrap(); + let mut state = crate::state::AppState::new_for_test(state_dir.path().to_path_buf()).await; + state.merge_executions = + crate::merge_execution::Executions::new(Some(coordinator.clone()), 3600); + state + .merge_executions + .rollout + .owned_targets + .push("generic:s".into()); + let state = Arc::new(state); + let dir = TempDir::new().unwrap(); + let store = generic_with_pending(&dir, 2, 2).await; + let slots = Arc::new(Semaphore::new(1)); + let _held = slots.clone().acquire_owned().await.unwrap(); + flush_pass_coordinated( + vec![("s".into(), store.clone())], + Duration::from_secs(5), + Some(slots), + Some(state), + ) + .await; + assert_eq!( + coordinator.request_page(None).await.unwrap().0, + ["generic:s"] + ); + assert_eq!( + store.read().await.pending_wal_generations().await.unwrap(), + 2 + ); + assert_eq!(store.read().await.count_base_rows().await.unwrap(), 0); + } + fn spec() -> SchemaSpec { SchemaSpec::new(vec![( ID_COLUMN.to_string(), diff --git a/docs/merge-recovery.md b/docs/merge-recovery.md index 29155c1..d3c3867 100644 --- a/docs/merge-recovery.md +++ b/docs/merge-recovery.md @@ -58,6 +58,49 @@ The base handle refreshes before retrying, and WAL drain remains a relative edit that preserves generations flushed during recovery. Graceful worker shutdown cancels/drains owned actors before closing resident stores. +## Queue, execution, and no-progress limits + +Admission and slot acquisition have a separate `MERGE_QUEUE_TIMEOUT_SECS` +(default 600). The execution clock starts only after a merge slot is acquired. +`MERGE_EXECUTION_TIMEOUT_SECS` is a configurable total ceiling (default 3600), +not a hard-coded 600-second cap. `MERGE_IDLE_TIMEOUT_SECS` (default 600) limits +the interval without observed completed work. All three must be positive. + +The worker counts completed batch reads, merge phases, and shielded storage +operations. It publishes a changed sequence at most once a second, guarded by +execution ownership. The master tracks the same sequence with bounded RPC +allowance. Repeated heartbeats do not extend the idle limit, and progressing +work cannot extend the total ceiling. A missing progress record means the +executor has not reported acquiring its slot. Readiness and these intermediate +checkpoints are NOT proof of WAL reclamation; only committed shard drains are. + +Progress observation is coarse inside opaque Lance index/build/storage calls. +A single such operation exceeding the idle threshold can still be cancelled. +Measure its p99 on realistic blobs and tune the idle/total limits before enabling +a large target. These defaults are not production throughput measurements. +Cancellation retains the slot until leaf commits drain or a barrier proves them +fenced. It does not restart the pod. An etcd outage cannot indefinitely occupy a +master reconciliation slot; unresolved ownership remains durable for recovery. + +## Count- and time-triggered merge compatibility + +`ROLLOUT_MERGE_AFTER_GENERATIONS` and `ROLLOUT_CLEANUP_INTERVAL_SECS` remain +valid with owned merging. For enabled rollout/generic targets, the flush +sweeper checks its own shard's count using manifest metadata only. At the +threshold it writes a coalesced merge request; the cleanup timer uses the same +request path. The master consumes requests every 15 seconds in pages of at most +256 through normal task dedupe, and executes the table's workers serially under +owned commit protection and their existing merge slots/byte budgets. + +This preserves #289's pending trigger, but changes execution for opted-in tables: +it is scheduled by the master, rather than performed inline by the flush +sweeper. Queue latency therefore remains relevant; 15 seconds is a polling +interval, not a completion SLA. A later flush tick reasserts demand if the active +task has already passed that shard. Requests never clear failure backoff. A busy +merge slot does not block flushing or require preparing any payload to request +work. Unselected targets retain #289's direct, slot-bounded count merge. Datagen +is not a scheduler-owned target and retains its existing sweeper path. + ## Repeated failures Per-target/endpoint records survive task recreation and master restart. They @@ -88,48 +131,91 @@ enqueues due targets through existing dedupe. It does not wait for the 600-second stats sweep. Large ledgers add page traversal latency; removed worker endpoints are not automatically probed. -`GET /api/v1/scheduler/merge-failures?after=` returns up to 256 records -and a `next` cursor. `needs_attention`, `last_error`, and `next_retry_ms` expose -the circuit state. `master_merge_storage_recoveries_total{result}` records -barrier outcomes alongside existing task and reclaimed-generation metrics. This -is an observable circuit state, not an external notification integration. +`needs_attention`, `last_error`, and `next_retry_ms` are persisted in the failure +ledger. `master_merge_storage_recoveries_total{result}` records barrier outcomes +alongside existing task and reclaimed-generation metrics. The read-only failure +inspection HTTP API is a separate follow-up; no external notification integration +is included here. Merge errors do **not** enqueue destructive fragment repair. Missing/corrupt files require diagnosis and restoration or an explicitly reviewed repair. Other maintenance tasks retain their existing repair policy. -## Storage and rollout requirements - -- Masters/workers must use the same physical storage namespace and etcd prefix, - with consistent storage configuration/credentials. URI mismatch fails closed. -- Version fencing supports Lance's immutable conditional manifest handlers: - local files, S3, GCS, Azure, memory, OSS, COS, TOS, and shared memory. Fallback - unsafe handlers and external catalogues such as `s3+ddb` are rejected for - owned merging/recovery. Existing manifest immutability/retention guarantees - must be preserved; never manually remove fencing/version files to unblock a - live writer. -- Workers need `ETCD_ENDPOINTS` and the same `ETCD_PREFIX` as masters. - `MERGE_EXECUTION_TIMEOUT_SECS` defaults to 600 and the master caps the requested - deadline at 600. Zero is rejected. Cancellation reconciliation has a 60-second - allowance; a metadata recovery attempt has a 120-second deadline. -- Enable protocol 2 on both sides. There is no fallback to unbounded legacy - merge HTTP calls. Old executions without version admission cannot be - automatically fenced using an empty watermark set; they remain protected. -- With the coordinator enabled, legacy merge routes reject admission. Disable - worker self-merge thresholds and cleanup timers; startup checks enforce this. -- `INDEX_BEFORE_MERGE` remains accepted for compatibility; workers prepare their - key indexes inside owned execution. Keep ordinary workers at six merge slots - and 20 GiB; this change increases neither limit. -- Drain/reconcile legacy in-flight merges before enabling protocol 2. Masters - must continue other jobs; recovery helpers must retain target locks. Do not - clear a keeper's pause while its delegate is alive. +## Storage requirements + +- Masters/workers use the same physical storage namespace, etcd prefix, and + storage configuration/credentials. URI mismatch fails closed. +- Version fencing requires Lance's immutable conditional manifest handlers: + local files, S3, GCS, Azure, memory, OSS, COS, TOS, and shared memory. Unsafe + fallback handlers and external catalogues such as `s3+ddb` are rejected. + Preserve manifest immutability and retention; never delete fencing files to + unblock a live writer. +- Workers with owned targets need `ETCD_ENDPOINTS`. Connection is lazy and + retried on demand: etcd unavailability does not fail worker startup, flush, or + ingestion. New owned admission and commit authorization fail closed when etcd + is unavailable. Already authorized storage writes may still finish. +- Cancellation reconciliation has a 60-second allowance; each metadata barrier + attempt has a 120-second deadline. These never authorize an unsafe handoff. +- `INDEX_BEFORE_MERGE` remains accepted; owned workers prepare their key indexes + inside execution. Keep ordinary workers at six merge slots / 20 GiB. Neither + limit is increased by this patch. + +## Default-off, per-table rollout + +`MERGE_OWNED_TARGETS` and `MERGE_DRAIN_TARGETS` are comma-separated **exact** +scheduler targets (`name` for rollout, `generic:name` for generic). Both default +to empty; wildcards and overlapping lists are rejected. Merely configuring +etcd or deploying the binaries does not enable the protocol. Unselected tables +use the legacy serial path, including its inability to safely recover a stuck +untracked write. An owned RPC failure never falls back to that path. + +This is an explicit operator-controlled migration, not automatic discovery or +a proof that old writes have drained. All replicas and helpers must follow the +sequence; mixed admission policies for the same table are unsupported. + +1. Deploy capability-aware binaries with both lists empty. Keep self-merge + thresholds/timers and unrelated master jobs running. No worker etcd + connection is opened merely to advertise capabilities. +2. Select one table. Put it in `MERGE_DRAIN_TARGETS` on all masters and workers; + prevent legacy helper admission for that table as well. Workers reject new + legacy merge requests and skip its self-merge; masters reject new maintenance + mutations for that table while still reconciling any already-owned merge. Other tables, ingestion, and master job pools continue. +3. Verify that previously admitted legacy worker/helper writes have actually + completed, including ambiguous remote storage operations. Ready pods, lease + expiry, elapsed time, or a completed HTTP request alone are insufficient. + An untracked legacy write has no admitted-version watermark and cannot be + automatically fenced by this protocol. If it is unresolved, keep **that + table** draining until storage completion is established; do not clear its + ownership or pretend an empty watermark permits recovery. +4. Move the target from draining to owned on every worker. Inspect + `GET /api/v1/internal/merge-executor`: check the exact target lists, + incarnation, progress protocol, and all three timeout values. Keep masters + draining while workers transition. Worker triggers can queue durable demand + but cannot start an unowned merge. Once all writers are ready, move the + target to owned on all masters and update helpers to use owned admission. +5. Soak that table under sustained writes, worker death, delayed storage, and + etcd disruption. Measure committed generations, memory/OOM, commit conflicts, + etcd overhead, queue time, longest opaque phase, and p95/p99 task latency. + Expand one table at a time only after the results justify it. + +Rollback also goes through draining: stop new admission, finish or storage-fence +all owned executions with the capable binaries, then return the selected table +to legacy configuration. Do not simply remove an owned target or roll back a +worker while its execution is unresolved. The master explicitly refuses a +legacy path while a durable execution record exists. Do not clear a helper's +pause while its delegate is alive. There is no fleet-wide stop requirement. + +## Review and deployment boundary + +Timeout, cancellation, commit authorization, storage fencing, and ownership +handoff are one correctness unit. Shipping an HTTP timeout alone would leave +old writes alive and permit conflicting retries. Persistent retry budgets also +remain in the runtime so newly triggered tasks cannot repeatedly reset attempts; +the read-only failure inspection API can be reviewed separately. The fence covers scheduler-coordinated maintenance and helpers that honor target -locks. It does not retrofit lease-loss protection onto the master's own local +locks. It does not retrofit lease-loss protection onto master's local compaction/index jobs or arbitrary manual writers. It cannot repair missing data -or guarantee a WAL pending ceiling while storage remains unavailable. - -Before production rollout, staging must exercise realistic blobs, sustained -writes, forced worker death and delayed storage responses. Measure pending and -reclaimed generations, memory/OOM counts, commit conflicts, etcd overhead, and -p95/p99 task latency. Local fault tests do not replace that soak. +or guarantee a pending ceiling while storage is unavailable. Local fault tests +and CI do not replace the realistic staging soak. This change is not a production +deployment, and its default-off rollout is deliberate. From 5b7618f0683cd4ae62a4037fe5e82fbe43c72874 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Fri, 2 Oct 2026 00:59:47 +0000 Subject: [PATCH 7/9] fix: expose explicitly prefixed merge rollout CLI flags --- crates/lance-context-merge/src/rollout.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/crates/lance-context-merge/src/rollout.rs b/crates/lance-context-merge/src/rollout.rs index 9e7fac2..9f0fcbe 100644 --- a/crates/lance-context-merge/src/rollout.rs +++ b/crates/lance-context-merge/src/rollout.rs @@ -6,11 +6,19 @@ use etcd_client::{Compare, CompareOp, GetOptions, TxnOp}; pub struct MergeRollout { /// Exact scheduler targets using owned merges (generic stores: generic:name). /// Empty by default. Enable only after the target's legacy writers drain. - #[arg(long, env = "MERGE_OWNED_TARGETS", value_delimiter = ',')] + #[arg( + long = "merge-owned-targets", + env = "MERGE_OWNED_TARGETS", + value_delimiter = ',' + )] pub owned_targets: Vec, /// Drain maintenance for these targets during migration; ingestion and /// maintenance of other tables continue. Coordinate all replicas and helpers. - #[arg(long, env = "MERGE_DRAIN_TARGETS", value_delimiter = ',')] + #[arg( + long = "merge-drain-targets", + env = "MERGE_DRAIN_TARGETS", + value_delimiter = ',' + )] pub drain_targets: Vec, } From cbcf8c55aba1f1db808b6e3348933e00ce306b19 Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Fri, 2 Oct 2026 01:12:22 +0000 Subject: [PATCH 8/9] fix: keep draining ownership eligible for bounded recovery probes --- .../src/merge_execution.rs | 30 +++++++++------- crates/lance-context-master/src/scheduler.rs | 36 +++++++++++++++++-- 2 files changed, 52 insertions(+), 14 deletions(-) diff --git a/crates/lance-context-master/src/merge_execution.rs b/crates/lance-context-master/src/merge_execution.rs index 0d922ce..33c5807 100644 --- a/crates/lance-context-master/src/merge_execution.rs +++ b/crates/lance-context-master/src/merge_execution.rs @@ -33,6 +33,7 @@ pub(crate) async fn run_merge_wal( // A drain stops new work, but must still resolve an already-owned // execution. Otherwise rollback would strand its persistent lock. if let Some(old) = coordinator.get(target).await? { + ensure_recovery_due(&coordinator, &old).await?; if reconcile(&state.http, &coordinator, &proof, old, false) .await .is_err() @@ -54,18 +55,7 @@ pub(crate) async fn run_merge_wal( // after a completed barrier lets healthy shards progress in this task. for recovery_round in 0..2 { if let Some(old) = coordinator.get(&claim.task.target).await? { - if old.phase == Phase::Recovering { - if let Some(failure) = coordinator.failure(&old.target, &old.endpoint).await? { - if failure.last_error.contains("recovery barrier failed") - && failure.next_retry_ms > lance_context_merge::failure::now_ms() - { - return Err(format!( - "{}; recovery probe at {}", - failure.last_error, failure.next_retry_ms - )); - } - } - } + ensure_recovery_due(&coordinator, &old).await?; if recovery_round > 0 || matches!(old.phase, Phase::Recovering | Phase::Uncertain) { recover_execution(state, &coordinator, &proof, old).await?; } else if let Err(error) = reconcile(&state.http, &coordinator, &proof, old, true).await @@ -178,6 +168,22 @@ impl Deadlines { } } +async fn ensure_recovery_due(coordinator: &Coordinator, old: &Execution) -> Result<(), String> { + if old.phase == Phase::Recovering { + if let Some(failure) = coordinator.failure(&old.target, &old.endpoint).await? { + if failure.last_error.contains("recovery barrier failed") + && failure.next_retry_ms > lance_context_merge::failure::now_ms() + { + return Err(format!( + "{}; recovery probe at {}", + failure.last_error, failure.next_retry_ms + )); + } + } + } + Ok(()) +} + async fn recover_execution( state: &Arc, coordinator: &Coordinator, diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 5ade72a..7d4f0aa 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -573,6 +573,31 @@ async fn sweep_merge_wal_inner(state: &Arc) -> lance::Result Ok(queued) } +// A draining target still needs metadata recovery while ownership is unresolved. +// Once released, its old failure ledger must not enqueue fresh merge tasks. +async fn should_probe_failure( + state: &Arc, + failure: &lance_context_merge::failure::ShardFailure, +) -> bool { + if state.config.merge_rollout.owned(&failure.target) { + return true; + } + if state.config.merge_rollout.draining(&failure.target) { + match state + .task_store + .merge_coordinator() + .get(&failure.target) + .await + { + Ok(execution) => return execution.is_some(), + Err(error) => { + tracing::warn!(target = %failure.target, %error, "cannot inspect draining merge ownership") + } + } + } + false +} + /// Spawn the scheduler pollers plus the optional periodic auto-sweep. /// /// Returns the handle of the *general* poller only. When a separate WAL-merge @@ -615,13 +640,13 @@ pub fn spawn_scheduler(state: &Arc) -> JoinHandle<()> { cursor = next; let mut targets = std::collections::HashSet::new(); for failure in rows { - if retry_state.config.merge_rollout.owned(&failure.target) - && failure.next_retry_ms <= lance_context_merge::failure::now_ms() + if failure.next_retry_ms <= lance_context_merge::failure::now_ms() && retry_state .config .worker_endpoints .contains(&failure.endpoint) && targets.insert(failure.target.clone()) + && should_probe_failure(&retry_state, &failure).await { if let Err(error) = enqueue(&retry_state, TaskKind::MergeWal, &failure.target).await @@ -871,14 +896,21 @@ mod tests { .unwrap(); let coordinator = state.task_store.merge_coordinator(); let proof = state.task_store.merge_claim(&claim); + let failure = coordinator + .record_failure(&proof, "exp", "http://127.0.0.1:1", "storage timeout") + .await + .unwrap(); + assert!(!should_probe_failure(&state, &failure).await); let execution = lance_context_merge::Execution::new("exp", "http://127.0.0.1:1", "boot", 30); assert!(coordinator.reserve(&proof, &execution).await.unwrap()); + assert!(should_probe_failure(&state, &failure).await); let running = coordinator.start(&execution).await.unwrap().unwrap(); assert!(coordinator.finish(&running, Ok(0)).await.unwrap()); let outcome = crate::merge_execution::run_merge_wal(&state, &claim).await; assert!(outcome.as_ref().unwrap_err().contains("draining")); assert!(coordinator.get("exp").await.unwrap().is_none()); + assert!(!should_probe_failure(&state, &failure).await); state.task_store.finish(claim, outcome).await.unwrap(); let compact = enqueue(&state, TaskKind::Compact, "exp").await.unwrap(); From c0c0b5acc8d2df9c49aab5cbe3b1c4f30f3d583b Mon Sep 17 00:00:00 2001 From: Beinan Wang <> Date: Fri, 2 Oct 2026 01:12:22 +0000 Subject: [PATCH 9/9] test: retry typed index conflicts after concurrent compaction completes --- .../lance-context-core/src/rollout_store.rs | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index e771aff..96feaee 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -4066,9 +4066,10 @@ mod tests { #[test] fn compact_composes_with_concurrent_wal_merge() { - // A base-table compaction (Rewrite) and a WAL merge (Append) are - // non-conflicting in Lance's commit matrix: running them concurrently - // must not fail, and no rows are lost. Instance A compacts while + // Append and Rewrite compose, but preparing the WAL merge can also + // CreateIndex. Lance may ask that transaction to retry after Rewrite. + // One explicit retry once compaction completes must retain every row. + // Instance A compacts while // instance B (a different shard) merges its own generations into the // same base table. use tokio::sync::RwLock; @@ -4136,12 +4137,27 @@ mod tests { }, ); ca.expect("compaction should not fail against a concurrent append"); + let mb = match mb { + Err(LanceError::RetryableCommitConflict { .. }) => { + // The concurrent compactor has completed above. Retry only + // this typed conflict, never arbitrary storage/data errors. + b.write().await.cleanup_own_shard().await + } + result => result, + }; assert_eq!(mb.expect("wal merge should not fail"), 3); // A fresh reader sees all 8 rows exactly once. let reader = RolloutStore::open(&uri).await.unwrap(); let listed = reader.list(None, None).await.unwrap(); assert_eq!(listed.len(), 8); + assert_eq!(reader.base.dataset.count_rows(None).await.unwrap(), 8); + let ids: std::collections::HashSet<_> = listed.iter().map(|r| r.id.clone()).collect(); + let expected = (0..5) + .map(|i| format!("a-{i}")) + .chain((0..3).map(|i| format!("b-{i}"))) + .collect(); + assert_eq!(ids, expected); }); }