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/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/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..5b569d7 --- /dev/null +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -0,0 +1,579 @@ +//! 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, + pin::Pin, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, + }, +}; +use tokio::sync::{oneshot, Notify}; + +tokio::task_local! { static CURRENT: Arc; } + +#[derive(Debug, Default)] +struct Progress { + closed: bool, + active: usize, + 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 { + completed_steps: AtomicU64, + progress: Mutex, + changed: Notify, + authorizer: Option>, + leaves: Mutex>, +} + +impl MergeWriteScope { + pub fn completed_steps(&self) -> u64 { + self.completed_steps.load(Ordering::Relaxed) + } + pub fn new() -> Arc { + Arc::new(Self::default()) + } + + pub fn with_authorizer(authorizer: Arc) -> Arc { + Arc::new(Self { + authorizer: Some(authorizer), + ..Self::default() + }) + } + + /// 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()) { + leaf.abort(); + } + self.drain().await; + } + + pub async fn run(self: &Arc, future: F) -> F::Output { + 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) { + 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; + } + } +} + +/// 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 { + authorizer.authorize(resource, version).await?; + } + } + Ok(()) +} + +struct Active { + scope: Arc, + completed: bool, +} +impl Drop for Active { + fn drop(&mut self) { + let mut progress = self.scope.progress.lock().unwrap(); + progress.active -= 1; + progress.uncertain |= !self.completed; + drop(progress); + self.scope.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 handles = scope.clone(); + let active = Active { + scope, + completed: false, + }; + let (send, receive) = oneshot::channel(); + 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(); + } + active.completed = true; + 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); + +#[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 { + authorize("base", manifest.version).await?; + let inner = self.0.clone(); + 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( + &mut owned_manifest, + indices, + &path, + &store, + manifest_writer, + naming_scheme, + 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?; + *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}; + + #[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 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(); + 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(); + 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.has_uncertain_commit()); + assert!( + scope + .run(shield(async { Ok::<_, Error>(()) })) + .await + .is_err(), + "cancelled scope cannot issue another commit" + ); + } +} 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 b15c0db..96feaee 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( @@ -2410,6 +2415,223 @@ 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, 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 { + 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 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 { + 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(); + } + 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; + 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(), + 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); + 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(); + } + + #[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(), @@ -3844,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; @@ -3914,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); }); } diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index d9a4230..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 @@ -1015,6 +1032,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,21 +1056,15 @@ 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() - }) + crate::merge_write_scope::drain_generations(drain_store, epoch, merged_generations) .await )?; @@ -1194,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()); @@ -2016,7 +2032,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 +2076,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..dc1a0ca 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 = { 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" } @@ -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..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, @@ -106,9 +108,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..33c5807 --- /dev/null +++ b/crates/lance-context-master/src/merge_execution.rs @@ -0,0 +1,928 @@ +//! 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, + queue_timeout_secs: u64, + idle_timeout_secs: u64, + progress_protocol: u32, + owned_targets: Vec, +} + +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); + 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? { + ensure_recovery_due(&coordinator, &old).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 { + if let Some(old) = coordinator.get(&claim.task.target).await? { + 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 + { + 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"); + } + } + } + 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; + } + } + 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 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, + 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), + }; + 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())? + .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( + 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 != 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 mut execution = Execution::new( + target, + endpoint, + &capabilities.instance, + 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; + 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 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 { + // 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) + .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 matches!(current.phase, Phase::Uncertain | Phase::Recovering) { + 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 matches!(current.phase, Phase::Finished | Phase::Recovered) { + 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_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. + 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 { + #[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, + 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":2,"instance":"test","timeout_secs":1,"queue_timeout_secs":1,"idle_timeout_secs":1,"progress_protocol":1,"owned_targets":["table"]})) + }), + ) + .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 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, + "--merge-owned-targets", + "table", + ]); + 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() { + 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 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() { + 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..c9d7cd8 100644 --- a/crates/lance-context-master/src/routes.rs +++ b/crates/lance-context-master/src/routes.rs @@ -715,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 09d5fa0..7d4f0aa 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 => run_merge_wal(state, &task.target).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" }; @@ -168,7 +173,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 +377,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`. /// @@ -774,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 @@ -783,6 +607,60 @@ 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 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() + && 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 + { + 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 { @@ -933,6 +811,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, @@ -992,6 +878,61 @@ 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 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(); + 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() { @@ -1144,20 +1085,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 +1303,70 @@ 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()); + let owned_targets = cfg.merge_rollout.owned_targets.clone(); + axum::Router::new() + .route( + "/api/v1/internal/merge-executor", + 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( + "/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 +1375,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 +1391,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 +1436,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 +1466,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 +1480,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 +1489,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 +1540,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 +1559,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 +1569,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 +1652,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 +1797,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 +1822,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 +1838,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 +1863,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 +1881,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 +1919,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/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 91567da..31cb9cf 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 ) } @@ -1323,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, @@ -1364,6 +1452,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..8d2fa82 --- /dev/null +++ b/crates/lance-context-merge/src/failure.rs @@ -0,0 +1,330 @@ +//! 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", + "permission denied", + "access denied", + "accessdenied", + "unauthorized", + "unauthenticated", + "forbidden", + "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 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 + } 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); + 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 = + ShardFailure::advance("table", "worker", None, "Not found: /data/1.lance", 1000); + 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); + } + } +} diff --git a/crates/lance-context-merge/src/fencing.rs b/crates/lance-context-merge/src/fencing.rs new file mode 100644 index 0000000..12c6987 --- /dev/null +++ b/crates/lance-context-merge/src/fencing.rs @@ -0,0 +1,185 @@ +//! 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 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![ + 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 new file mode 100644 index 0000000..87b193d --- /dev/null +++ b/crates/lance-context-merge/src/lib.rs @@ -0,0 +1,833 @@ +//! 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; +pub mod fencing; +pub mod progress; +pub mod rollout; + +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, + Uncertain, + Recovering, + Recovered, + 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, + #[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, + #[serde(default, skip_serializing_if = "protocol_absent")] + pub protocol: u32, +} + +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 { + id: uuid::Uuid::new_v4().to_string(), + target: target.into(), + 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, + protocol: 2, + } + } + + 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 + || execution.protocol != 2 + { + 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 + } + + /// 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 { + 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 !matches!(execution.phase, Phase::Finished | Phase::Recovered) { + 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([ + self.remove_permits(execution), + self.remove_progress(execution), + 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: 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 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() { + 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()); + // 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(), + 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() { + 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() { + 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-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..9f0fcbe --- /dev/null +++ b/crates/lance-context-merge/src/rollout.rs @@ -0,0 +1,125 @@ +//! 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 = "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 = "merge-drain-targets", + 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/Cargo.toml b/crates/lance-context-server/Cargo.toml index bebeed0..6d838f4 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 = { 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" } @@ -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..aa387ea 100644 --- a/crates/lance-context-server/src/config.rs +++ b/crates/lance-context-server/src/config.rs @@ -4,6 +4,26 @@ use clap::Parser; #[command(name = "lance-context-server")] #[command(about = "REST API server for lance-context")] pub struct ServerConfig { + /// Connection used lazily by explicitly enabled owned merge targets. + #[command(flatten)] + pub merge_etcd: lance_context_merge::EtcdConfig, + + #[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 = 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/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..e5c2577 --- /dev/null +++ b/crates/lance-context-server/src/merge_execution.rs @@ -0,0 +1,695 @@ +//! 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, OnceCell}; + +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: 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: 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 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 + /// 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; + } + } + + 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> { + Ok(Json( + serde_json::json!({"protocol": 2, "instance": state.merge_executions.instance, + "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}), + )) +} + +pub async fn start( + State(state): State>, + Json(execution): Json, +) -> Result { + 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(), + )); + } + 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 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()); + } + 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::with_authorizer( + Arc::new(OwnedCommitGuard { + coordinator: coordinator.clone(), + execution: running.clone(), + dataset_uri: uri.clone(), + }), + ); + 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())); + // Top-level cancellation cannot acknowledge a write still in flight at + // object storage. Join those leaf commits before publishing Finished. + 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); + // 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 { + 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) => 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") + } + }, + Err(error) => { + tracing::warn!(id = %execution.id, %error, "retrying merge outcome acknowledgement") + } + } + tokio::time::sleep(Duration::from_secs(1)).await; + } + 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().await?; + 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] + 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() { + 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); + 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. + 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..087f0fd 100644 --- a/crates/lance-context-server/src/routes/generic.rs +++ b/crates/lance-context-server/src/routes/generic.rs @@ -273,7 +273,17 @@ pub async fn merge_generic_wal( State(state): State>, Path(name): Path, ) -> Result, AppError> { + state + .merge_executions + .legacy_allowed(&format!("generic:{name}"))?; 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..cae4e05 100644 --- a/crates/lance-context-server/src/routes/rollouts.rs +++ b/crates/lance-context-server/src/routes/rollouts.rs @@ -728,7 +728,15 @@ pub async fn merge_wal( State(state): State>, Path(name): Path, ) -> Result, AppError> { + state.merge_executions.legacy_allowed(&name)?; 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..ba8e326 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,33 @@ impl AppState { let datagen_registry = RolloutRegistry::open_or_create(&datagen_registry_uri, None) .await .map_err(AppError::from_lance)?; + 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 targets require ETCD_ENDPOINTS".into(), + )); + } + 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, queue and idle timeouts must be positive".into(), + )); + } Ok(Self { + 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)), rollout_handles: StoreHandles::default(), @@ -391,6 +418,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(), @@ -953,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()), ), ); } @@ -1014,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 new file mode 100644 index 0000000..d3c3867 --- /dev/null +++ b/docs/merge-recovery.md @@ -0,0 +1,221 @@ +# 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. 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. + +## 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 +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 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, 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 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. + +`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 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 master's local +compaction/index jobs or arbitrary manual writers. It cannot repair missing data +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.