diff --git a/Cargo.lock b/Cargo.lock index e44a4c5b..1cd26f28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -412,18 +412,18 @@ checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] name = "async-trait" -version = "0.1.89" +version = "0.1.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 3.0.6", ] [[package]] @@ -780,7 +780,7 @@ checksum = "8d7396fd9500589e62e460e987ecb671bad374934e55ec3b5f498cc7a8a8a7b7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -955,7 +955,7 @@ dependencies = [ "regex", "rustc-hash", "shlex 1.3.0", - "syn", + "syn 2.0.117", ] [[package]] @@ -1069,7 +1069,7 @@ checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1229,7 +1229,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1576,7 +1576,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn", + "syn 2.0.117", ] [[package]] @@ -1587,7 +1587,7 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2037,7 +2037,7 @@ checksum = "1a3614234dd93578c92428cb4f408e020874f0d2b7e6c90c928d9d28b5df2ceb" dependencies = [ "datafusion-doc", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2285,7 +2285,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2589,7 +2589,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2838,7 +2838,7 @@ checksum = "53010ccb100b96a67bc32c0175f0ed1426b31b655d562898e57325f81c023ac0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3578,7 +3578,7 @@ checksum = "782d32378dddf207193ac91cefb848ad41abb58195c95168e1291227a0832b47" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3623,7 +3623,7 @@ dependencies = [ "quote", "rustc_version", "simd_cesu8", - "syn", + "syn 2.0.117", ] [[package]] @@ -3642,7 +3642,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" dependencies = [ "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -3885,6 +3885,7 @@ dependencies = [ "arrow-select", "async-trait", "base64", + "bytes", "chrono", "clap", "datafusion", @@ -4108,7 +4109,7 @@ checksum = "bf46656b359786f0b73f13936193583972b7af6e53ccb542da87830be18e94b2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5048,7 +5049,7 @@ dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5445,7 +5446,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5660,7 +5661,7 @@ checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5767,7 +5768,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.117", ] [[package]] @@ -5836,7 +5837,7 @@ dependencies = [ "pulldown-cmark", "pulldown-cmark-to-cmark", "regex", - "syn", + "syn 2.0.117", "tempfile", ] @@ -5850,7 +5851,7 @@ dependencies = [ "itertools 0.14.0", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5928,7 +5929,7 @@ dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -5941,7 +5942,7 @@ dependencies = [ "proc-macro2", "pyo3-build-config", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -6252,7 +6253,7 @@ checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -6882,7 +6883,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -6917,7 +6918,7 @@ checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -6961,7 +6962,7 @@ dependencies = [ "darling", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7140,7 +7141,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7199,7 +7200,7 @@ checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7270,7 +7271,7 @@ dependencies = [ "proc-macro2", "quote", "rustversion", - "syn", + "syn 2.0.117", ] [[package]] @@ -7296,6 +7297,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "syn" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8593e8e72159ed2257d083c7a454a85cbf854f37a0966d8d483aff8c8a3ebcee" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -7313,7 +7325,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7408,7 +7420,7 @@ checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7419,7 +7431,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7530,7 +7542,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7648,7 +7660,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -7673,7 +7685,7 @@ dependencies = [ "prost-build", "prost-types", "quote", - "syn", + "syn 2.0.117", "tempfile", "tonic-build", ] @@ -7771,7 +7783,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -8074,7 +8086,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wasm-bindgen-shared", ] @@ -8282,7 +8294,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -8293,7 +8305,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -8551,7 +8563,7 @@ dependencies = [ "heck", "indexmap 2.14.0", "prettyplease", - "syn", + "syn 2.0.117", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -8567,7 +8579,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -8827,7 +8839,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -8848,7 +8860,7 @@ checksum = "0b631b19d36a892ab55420c92dbc83ccd79274f25be714855d3074aa71cab639" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -8868,7 +8880,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -8910,7 +8922,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] diff --git a/crates/lance-context-core/Cargo.toml b/crates/lance-context-core/Cargo.toml index c24653ad..ac86edcd 100644 --- a/crates/lance-context-core/Cargo.toml +++ b/crates/lance-context-core/Cargo.toml @@ -18,11 +18,12 @@ default = ["metrics"] metrics = ["dep:metrics"] [dependencies] -async-trait = "0.1" +async-trait = "0.1.92" lance-table = "9.0.0" lance-io = "9.0.0" object_store = "0.13.2" base64 = "0.22" +bytes = "1" arrow-array = "58" arrow-ipc = "58" arrow-json = "58" diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 4435f3ec..485f71e8 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -16,6 +16,7 @@ pub mod merge_budget; pub mod merge_write_scope; pub mod metrics; mod namespace; +mod preparation_io; mod record; mod registry; mod registry_etcd; @@ -103,4 +104,4 @@ pub use lance::Error as LanceError; // capacity-bounded cache session across all resident rollout stores. pub use lance::session::Session; -pub use store_base::PreparedCompaction; +pub use store_base::{PreparedCompaction, PreparedKeyIndex}; diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index e1f1d550..248c4625 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -45,7 +45,9 @@ pub trait CommitAuthorizer: std::fmt::Debug + Send + Sync { #[derive(Debug, Default)] pub struct MergeWriteScope { - completed_steps: AtomicU64, + completed_steps: Arc, + observe_file_io: bool, + local_io: Mutex>, progress: Mutex, changed: Notify, authorizer: Option>, @@ -55,7 +57,18 @@ pub struct MergeWriteScope { impl MergeWriteScope { pub fn completed_steps(&self) -> u64 { - self.completed_steps.load(Ordering::Relaxed) + self.local_io.lock().unwrap().iter().fold( + self.completed_steps.load(Ordering::Relaxed), + |steps, (tracker, baseline)| { + let stats = tracker.stats(); + steps.saturating_add( + stats + .read_bytes + .saturating_add(stats.written_bytes) + .saturating_sub(*baseline), + ) + }, + ) } pub fn new() -> Arc { Arc::new(Self::default()) @@ -79,6 +92,17 @@ impl MergeWriteScope { }) } + /// Preparation reads/writes immutable files. Observe successful file IO so + /// long scans do not appear idle when Lance omits training callbacks. + pub fn with_preparation_authorizer(authorizer: Arc) -> Arc { + Arc::new(Self { + authorizer: Some(authorizer), + pin_opened_handles: true, + observe_file_io: true, + ..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. @@ -140,6 +164,47 @@ pub(crate) fn write_progress() -> lance::dataset::write::WriteProgressFn { }) } +/// Capture only the preparation scope. Cached handles from ordinary writers +/// must never acquire another task's diagnostic counter. +pub(crate) fn preparation_store_wrapper( +) -> Option> { + CURRENT + .try_with(|scope| { + scope.observe_file_io.then(|| { + Arc::new(crate::preparation_io::ProgressWrapper( + scope.completed_steps.clone(), + )) as Arc + }) + }) + .ok() + .flatten() +} + +pub(crate) fn observes_preparation_io() -> bool { + CURRENT + .try_with(|scope| scope.observe_file_io) + .unwrap_or(false) +} + +/// Lance's direct local/uring readers and writers bypass ObjectStore wrappers. +/// Their tracker records completed byte IO. The unique preparation wrapper in +/// the store params isolates this handle's tracker from other cached stores. +pub(crate) fn register_local_preparation_io(store: &ObjectStore) { + if !store.is_local() { + return; + } + let _ = CURRENT.try_with(|scope| { + if scope.observe_file_io { + let tracker = store.io_tracker().clone(); + let stats = tracker.stats(); + scope.local_io.lock().unwrap().push(( + tracker, + stats.read_bytes.saturating_add(stats.written_bytes), + )); + } + }); +} + 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 { diff --git a/crates/lance-context-core/src/preparation_io.rs b/crates/lance-context-core/src/preparation_io.rs new file mode 100644 index 00000000..1ac92ec1 --- /dev/null +++ b/crates/lance-context-core/src/preparation_io.rs @@ -0,0 +1,212 @@ +//! Diagnostics for immutable preparation IO, never write authorization. +use async_trait::async_trait; +use bytes::Bytes; +use futures::{stream::BoxStream, StreamExt}; +use lance_io::object_store::WrappingObjectStore; +use object_store::{path::Path, *}; +use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, +}; + +#[derive(Debug)] +pub(crate) struct ProgressWrapper(pub Arc); + +impl WrappingObjectStore for ProgressWrapper { + fn wrap(&self, _: &str, original: Arc) -> Arc { + Arc::new(ProgressStore { + inner: original, + steps: self.0.clone(), + }) + } +} + +#[derive(Debug)] +struct ProgressStore { + inner: Arc, + steps: Arc, +} + +impl std::fmt::Display for ProgressStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "PreparationProgress({})", self.inner) + } +} + +// Metadata/lease polling must not mask stalled data work. Match path segments +// so both bucket-relative and dataset-relative stores work. +fn data_file(path: &Path) -> bool { + path.as_ref() + .split('/') + .any(|s| matches!(s, "data" | "_indices")) +} + +#[async_trait] +impl ObjectStore for ProgressStore { + async fn put_opts( + &self, + path: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + let count = data_file(path) && payload.content_length() > 0; + let result = self.inner.put_opts(path, payload, opts).await?; + if count { + self.steps.fetch_add(1, Ordering::Relaxed); + } + Ok(result) + } + + async fn put_multipart_opts( + &self, + path: &Path, + opts: PutMultipartOptions, + ) -> Result> { + let inner = self.inner.put_multipart_opts(path, opts).await?; + if !data_file(path) { + return Ok(inner); + } + Ok(Box::new(ProgressUpload { + inner, + steps: self.steps.clone(), + })) + } + + async fn get_opts(&self, path: &Path, opts: GetOptions) -> Result { + let count = data_file(path) && !opts.head; + let result = self.inner.get_opts(path, opts).await?; + if !count { + return Ok(result); + } + let meta = result.meta.clone(); + let range = result.range.clone(); + let attributes = result.attributes.clone(); + let steps = self.steps.clone(); + // Count body chunks after they arrive, not successful response headers. + // into_stream also observes actual local file reads without buffering. + let payload = GetResultPayload::Stream( + result + .into_stream() + .inspect(move |chunk| { + if chunk.as_ref().is_ok_and(|b| !b.is_empty()) { + steps.fetch_add(1, Ordering::Relaxed); + } + }) + .boxed(), + ); + Ok(GetResult { + payload, + meta, + range, + attributes, + }) + } + + async fn get_ranges(&self, path: &Path, ranges: &[std::ops::Range]) -> Result> { + let result = self.inner.get_ranges(path, ranges).await?; + if data_file(path) && result.iter().any(|b| !b.is_empty()) { + self.steps.fetch_add(1, Ordering::Relaxed); + } + Ok(result) + } + + fn delete_stream( + &self, + paths: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + self.inner.delete_stream(paths) + } + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + self.inner.list(prefix) + } + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + self.inner.list_with_offset(prefix, offset) + } + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + self.inner.list_with_delimiter(prefix).await + } + async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> Result<()> { + self.inner.copy_opts(from, to, opts).await + } + async fn rename_opts(&self, from: &Path, to: &Path, opts: RenameOptions) -> Result<()> { + self.inner.rename_opts(from, to, opts).await + } +} + +#[derive(Debug)] +struct ProgressUpload { + inner: Box, + steps: Arc, +} +#[async_trait] +impl MultipartUpload for ProgressUpload { + fn put_part(&mut self, payload: PutPayload) -> UploadPart { + let nonempty = payload.content_length() > 0; + let part = self.inner.put_part(payload); + let steps = self.steps.clone(); + Box::pin(async move { + part.await?; + if nonempty { + steps.fetch_add(1, Ordering::Relaxed); + } + Ok(()) + }) + } + async fn complete(&mut self) -> Result { + self.inner.complete().await + } + async fn abort(&mut self) -> Result<()> { + self.inner.abort().await + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[tokio::test] + async fn counts_data_bodies_and_upload_parts_but_not_metadata_or_failures() { + let steps = Arc::new(AtomicU64::new(0)); + let store = ProgressStore { + inner: Arc::new(memory::InMemory::new()), + steps: steps.clone(), + }; + let data = Path::from("table/data/part.lance"); + let metadata = Path::from("table/_versions/1.manifest"); + store.put(&metadata, "metadata".into()).await.unwrap(); + store.get(&metadata).await.unwrap().bytes().await.unwrap(); + assert_eq!(steps.load(Ordering::Relaxed), 0); + store.put(&data, "payload".into()).await.unwrap(); + let result = store.get(&data).await.unwrap(); + store.head(&data).await.unwrap(); + assert_eq!( + steps.load(Ordering::Relaxed), + 1, + "headers and HEAD are not data progress" + ); + assert_eq!(result.bytes().await.unwrap(), "payload"); + assert_eq!(steps.load(Ordering::Relaxed), 2); + assert!(store.get(&Path::from("table/data/missing")).await.is_err()); + assert!(store + .put_opts(&data, "conflict".into(), PutOptions::from(PutMode::Create)) + .await + .is_err()); + assert_eq!(steps.load(Ordering::Relaxed), 2); + let mut upload = store + .put_multipart(&Path::from("table/_indices/new/index.idx")) + .await + .unwrap(); + upload.put_part("index".into()).await.unwrap(); + upload.complete().await.unwrap(); + assert_eq!(steps.load(Ordering::Relaxed), 3); + let ranges = store.get_ranges(&data, &[0..3, 3..7]).await.unwrap(); + assert_eq!( + ranges, + vec![Bytes::from_static(b"pay"), Bytes::from_static(b"load")] + ); + assert_eq!(steps.load(Ordering::Relaxed), 4); + } +} diff --git a/crates/lance-context-core/src/rollout_append.rs b/crates/lance-context-core/src/rollout_append.rs index 13b02ca3..a8ac9825 100644 --- a/crates/lance-context-core/src/rollout_append.rs +++ b/crates/lance-context-core/src/rollout_append.rs @@ -983,6 +983,224 @@ mod tests { .await } + async fn prepare_index_without_commit( + uri: &str, + kind: crate::KeyIndexType, + ) -> crate::PreparedKeyIndex { + let scope = crate::merge_write_scope::MergeWriteScope::with_preparation_authorizer( + Arc::new(NoPreparationCommit), + ); + let prepared = scope + .run(async { + RolloutStore::open_existing_with_options( + uri, + RolloutStoreOptions { + key_index_type: kind, + ..Default::default() + }, + ) + .await + .unwrap() + .prepare_id_key_index() + .await + .unwrap() + }) + .await; + assert!( + scope.completed_steps() > 0, + "index file IO must advance the idle watchdog" + ); + prepared + } + + #[tokio::test] + async fn prepared_index_preserves_append_watermarks_and_uncovered_reads() { + use lance::index::DatasetIndexExt; + for kind in [crate::KeyIndexType::Btree, crate::KeyIndexType::Zonemap] { + for merge_first in [true, false] { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (a, mut coordinator) = prepared_compaction_fixture(uri).await; + let before = Dataset::open(uri).await.unwrap().version().version; + let prepared = prepare_index_without_commit(uri, kind).await; + assert_eq!(Dataset::open(uri).await.unwrap().version().version, before); + put(&a, "during-index", 8192).await; + let (plans, _) = coordinator + .plan(&["a".into()], 64, 1024 * 1024) + .await + .unwrap(); + let staged = stage(uri, plans[0].clone(), budget(), None).await.unwrap(); + let high = staged.plan.generations.last().unwrap().number; + let shard = staged.plan.shard; + if merge_first { + coordinator.commit(vec![staged.clone()]).await.unwrap(); + } + let mut publisher = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + publisher + .commit_prepared_id_key_index(prepared) + .await + .unwrap(); + if !merge_first { + coordinator.commit(vec![staged.clone()]).await.unwrap(); + } + assert_eq!(coordinator.commit(vec![staged]).await.unwrap(), 0); + let current = Dataset::open(uri).await.unwrap(); + assert_eq!(watermarks(¤t).await.unwrap().get(&shard), Some(&high)); + assert_eq!(rows(uri).await, 5); + // A fresh indexed reader must find both covered and uncovered + // keys, including the append that raced with preparation. + let reader = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + assert_eq!( + reader.get_blob("during-index").await.unwrap().unwrap(), + vec![42; 8192] + ); + assert!(reader.get_blob("seed-0").await.unwrap().is_some()); + let indices = current.load_indices().await.unwrap(); + let key = indices + .iter() + .find(|i| i.name == crate::store_base::ID_INDEX_NAME) + .unwrap(); + assert_eq!(key.fragment_bitmap.as_ref().unwrap().len(), 4); + assert_eq!(current.manifest().fragments.len(), 5); + } + } + } + + #[tokio::test] + async fn prepared_index_progress_isolated_from_other_cached_readers() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (_writer, _) = prepared_compaction_fixture(uri).await; + let options = RolloutStoreOptions { + session: Some(RolloutStore::build_session(1024 * 1024, 1024 * 1024)), + ..Default::default() + }; + let scope = crate::merge_write_scope::MergeWriteScope::with_preparation_authorizer( + Arc::new(NoPreparationCommit), + ); + let builder = scope + .run(RolloutStore::open_existing_with_options( + uri, + options.clone(), + )) + .await + .unwrap(); + let before = scope.completed_steps(); + let other = RolloutStore::open_existing_with_options(uri, options) + .await + .unwrap(); + assert!(other.get_blob("seed-0").await.unwrap().is_some()); + assert_eq!( + scope.completed_steps(), + before, + "another task's reads cannot keep a stalled builder alive" + ); + scope.run(builder.prepare_id_key_index()).await.unwrap(); + assert!( + scope.completed_steps() > before, + "local index IO must advance its own scope" + ); + } + + #[tokio::test] + async fn prepared_index_rejects_changed_sources_schema_or_replacement() { + for mutation in ["delete", "compact", "index", "schema"] { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (_a, _) = prepared_compaction_fixture(uri).await; + let prepared = prepare_index_without_commit(uri, crate::KeyIndexType::Btree).await; + let mut publisher = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + match mutation { + "delete" => { + Dataset::open(uri) + .await + .unwrap() + .delete("id = 'seed-0'") + .await + .unwrap(); + } + "compact" => { + publisher + .compact(Some(crate::CompactionConfig { + target_rows_per_fragment: 1024, + ..Default::default() + })) + .await + .unwrap(); + } + "index" => publisher.create_id_key_index().await.unwrap(), + "schema" => { + Dataset::open(uri) + .await + .unwrap() + .add_columns( + lance::dataset::NewColumnTransform::AllNulls(Arc::new( + arrow_schema::Schema::new(vec![arrow_schema::Field::new( + "new_column", + arrow_schema::DataType::Int64, + true, + )]), + )), + None, + None, + ) + .await + .unwrap(); + } + _ => unreachable!(), + } + let before = Dataset::open(uri).await.unwrap().version().version; + let error = publisher + .commit_prepared_id_key_index(prepared) + .await + .unwrap_err(); + assert!( + error.to_string().contains("reprepare"), + "{mutation}: {error}" + ); + assert_eq!(Dataset::open(uri).await.unwrap().version().version, before); + } + } + + #[tokio::test] + async fn prepared_index_cannot_publish_with_revoked_authorization() { + let dir = tempfile::tempdir().unwrap(); + let uri = dir.path().to_str().unwrap(); + let (_a, _) = prepared_compaction_fixture(uri).await; + let prepared = prepare_index_without_commit(uri, crate::KeyIndexType::Btree).await; + let before = Dataset::open(uri).await.unwrap().version().version; + let error = crate::merge_write_scope::MergeWriteScope::with_pinned_authorizer(Arc::new( + NoPreparationCommit, + )) + .run(async { + let mut publisher = + RolloutStore::open_existing_with_options(uri, RolloutStoreOptions::default()) + .await + .unwrap(); + publisher + .commit_prepared_id_key_index(prepared) + .await + .unwrap_err() + }) + .await; + assert!( + error + .to_string() + .contains("unexpected preparation manifest commit"), + "{error}" + ); + assert_eq!(Dataset::open(uri).await.unwrap().version().version, before); + } + #[tokio::test] async fn prepared_compaction_preserves_append_and_watermarks_in_both_orders() { for merge_first in [true, false] { diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index 669948d0..a1aedb90 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -737,6 +737,19 @@ impl RolloutStore { self.base.create_configured_key_index().await } + /// Build immutable key-index files without publishing a table version. + pub async fn prepare_id_key_index(&self) -> LanceResult { + self.base.prepare_key_index().await + } + + /// Publish prepared key-index files under fresh table write ownership. + pub async fn commit_prepared_id_key_index( + &mut self, + prepared: crate::PreparedKeyIndex, + ) -> LanceResult<()> { + self.base.commit_prepared_key_index(prepared).await + } + /// Explicitly create a BTree, independent of the configured merge policy. pub async fn create_id_btree_index(&mut self) -> LanceResult<()> { self.base.create_key_btree_index().await diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs index 0bb82edd..12af8f04 100644 --- a/crates/lance-context-core/src/store_base.rs +++ b/crates/lance-context-core/src/store_base.rs @@ -131,6 +131,18 @@ impl PreparedCompaction { } } +/// Immutable index files built against one snapshot. The publishing handle +/// must be opened separately under commit ownership; no authorization or +/// mutable dataset handle escapes preparation. +#[derive(Debug)] +pub struct PreparedKeyIndex { + uri: String, + schema: Schema, + fragments: Vec, + index: lance_table::format::IndexMetadata, + replaced: Vec, +} + /// Execute only the first `max_source_fragments` from a Lance compaction plan. /// /// Lance's built-in `max_source_fragments` stops before a whole planned task @@ -1437,6 +1449,101 @@ impl StorageBase { self.reload().await } + /// Build files only. Callers must open this handle in a preparation-only + /// write scope so even an unexpected metadata write is rejected. + pub async fn prepare_key_index(&self) -> LanceResult { + self.ensure_writable()?; + let mut snapshot = self.dataset.clone(); + let replaced = snapshot + .load_indices() + .await? + .iter() + .filter(|index| index.name == ID_INDEX_NAME) + .cloned() + .collect(); + let kind = match self.key_index_type { + crate::KeyIndexType::Btree => IndexType::BTree, + crate::KeyIndexType::Zonemap => IndexType::ZoneMap, + }; + let index = snapshot + .create_index_builder( + &[self.key_column.as_str()], + kind, + &ScalarIndexParams::default(), + ) + .name(ID_INDEX_NAME.to_string()) + .replace(true) + .execute_uncommitted() + .await?; + Ok(PreparedKeyIndex { + uri: snapshot.uri().to_owned(), + schema: Schema::from(snapshot.schema()), + fragments: snapshot.manifest().fragments.as_ref().clone(), + index, + replaced, + }) + } + + /// Publish on a fresh fenced handle. Appends may have added uncovered + /// fragments; changes to the indexed snapshot or same-name index require + /// rebuilding. Keep the original read version for Lance conflict checks. + pub async fn commit_prepared_key_index( + &mut self, + prepared: PreparedKeyIndex, + ) -> LanceResult<()> { + use lance::dataset::transaction::{Operation, Transaction}; + use lance::dataset::write::CommitBuilder; + + self.ensure_writable()?; + self.dataset.checkout_latest().await?; + if prepared.uri != self.dataset.uri() + || prepared.schema != Schema::from(self.dataset.schema()) + { + return Err(LanceError::invalid_input( + "index dataset or schema changed; reprepare", + )); + } + let sources: HashMap<_, _> = self + .dataset + .manifest() + .fragments + .iter() + .map(|f| (f.id, f)) + .collect(); + if prepared + .fragments + .iter() + .any(|old| sources.get(&old.id).copied() != Some(old)) + { + return Err(LanceError::invalid_input( + "index source fragment changed; reprepare", + )); + } + let current = self.dataset.load_indices().await?; + let replaced: Vec<_> = current + .iter() + .filter(|index| index.name == prepared.index.name) + .cloned() + .collect(); + if replaced != prepared.replaced { + return Err(LanceError::invalid_input( + "index metadata changed; reprepare", + )); + } + let transaction = Transaction::new( + prepared.index.dataset_version, + Operation::CreateIndex { + new_indices: vec![prepared.index], + removed_indices: prepared.replaced, + }, + None, + ); + self.dataset = CommitBuilder::new(Arc::new(self.dataset.clone())) + .execute(transaction) + .await?; + self.reload().await + } + /// Evolve an older base table to the store's latest additive schema. /// /// Missing nullable columns are added as all-null arrays. Existing unknown @@ -2188,23 +2295,38 @@ impl StorageBase { storage_options: Option>, session: Option>, ) -> LanceResult { - let store_params = storage_options.clone().map(|options| ObjectStoreParams { + let mut store_params = storage_options.clone().map(|options| ObjectStoreParams { storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( options, ))), ..Default::default() }); + if let Some(wrapper) = crate::merge_write_scope::preparation_store_wrapper() { + store_params + .get_or_insert_with(ObjectStoreParams::default) + .object_store_wrapper = Some(wrapper); + } 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::new(handler), - )); + let mut builder = DatasetBuilder::from_uri(uri) + .with_read_params(lance::dataset::ReadParams { + store_options: store_params, + ..Default::default() + }) + .with_commit_handler(Arc::new(crate::merge_write_scope::GuardedCommit::new( + handler, + ))); if let Some(options) = storage_options { builder = builder.with_storage_options(options); } if let Some(session) = session { builder = builder.with_session(session); } - builder.load().await + let dataset = builder.load().await?; + if crate::merge_write_scope::observes_preparation_io() { + let store = dataset.object_store(None).await?; + crate::merge_write_scope::register_local_preparation_io(&store); + } + Ok(dataset) } /// Create an empty Lance dataset with `schema` at `uri`. diff --git a/crates/lance-context-master/src/config.rs b/crates/lance-context-master/src/config.rs index fe639a54..890971be 100644 --- a/crates/lance-context-master/src/config.rs +++ b/crates/lance-context-master/src/config.rs @@ -327,6 +327,10 @@ pub struct MaintenanceConfig { /// source fragments invalidate preparation instead of being overwritten. #[arg(long, env = "COMPACTION_PREPARE_TARGETS", value_delimiter = ',')] pub compaction_prepare_targets: Vec, + /// Build rollout key-index files outside the table write lock. Exact + /// targets or '*'; empty retains legacy indexing. Commit remains fenced. + #[arg(long, env = "INDEX_PREPARE_TARGETS", value_delimiter = ',')] + pub index_prepare_targets: Vec, /// Exact owned rollout targets whose dedicated publishers all use the /// common fenced task/target/merge-claim protocol. Allows Compact/IndexId /// beside persistent catch-up identity; empty retains legacy exclusion. @@ -337,6 +341,10 @@ pub struct MaintenanceConfig { /// This limits local capacity retention; it never cancels the current writer. #[arg(long, env = "COMPACTION_COMMIT_WAIT_SECS", default_value_t = 120)] pub compaction_commit_wait_secs: u64, + /// Maximum wait for index publication ownership after files are ready. + /// Expiry releases preparation capacity, never interrupts the other writer. + #[arg(long, env = "INDEX_COMMIT_WAIT_SECS", default_value_t = 120)] + pub index_commit_wait_secs: u64, /// Legacy wire field; owned work now uses the real-progress idle timeout. #[arg(long, env = "MAINTENANCE_TIMEOUT_SECS", default_value_t = 3600)] pub maintenance_timeout_secs: u64, @@ -351,8 +359,10 @@ impl Default for MaintenanceConfig { fn default() -> Self { Self { compaction_prepare_targets: Vec::new(), + index_prepare_targets: Vec::new(), maintenance_catchup_targets: Vec::new(), compaction_commit_wait_secs: 120, + index_commit_wait_secs: 120, maintenance_timeout_secs: 3600, maintenance_idle_timeout_secs: 600, maintenance_drain_timeout_secs: 30, diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index 5c8106d5..81a0e83c 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -15,10 +15,9 @@ //! The master cannot do this itself without fencing the live shard writer, so //! it fans out to every configured worker endpoint and each worker merges its //! own shard (`POST /api/v1/internal/merge-wal/{name}`). -//! - **IndexId** — builds a ZoneMap scalar index on the base table's `id` column -//! (runs locally on the master). It commits a `CreateIndex`, which can conflict -//! with a concurrent `Compact` `Rewrite` on the same dataset, so it shares the -//! per-name task-store lock with `Compact`. +//! - **IndexId** — builds the configured scalar index on `id` locally. Opted-in +//! targets prepare immutable files without table write ownership, then acquire +//! the shared writer fence to validate and publish the index metadata. //! //! Each task runs in its own `tokio::spawn`, so one task's failure never affects //! another. A global [`Semaphore`] bounds how many run at once. @@ -148,7 +147,9 @@ async fn run_task(state: &Arc, mut claim: TaskClaim, timing: TaskCl let started = std::time::Instant::now(); let mut refresh_compaction_stats = false; - let outcome = if claim.preparing_compaction() { + let outcome = if claim.preparing_maintenance() && task.kind == TaskKind::IndexId { + run_prepared_index(state, &mut claim).await + } else if claim.preparing_maintenance() { run_prepared_compaction(state, &mut claim) .await .map(|(detail, changed)| { @@ -359,11 +360,11 @@ impl lance_context_core::merge_write_scope::CommitAuthorizer for PreparationOnly _: &'a str, _: u64, ) -> std::pin::Pin> + Send + 'a>> { - Box::pin(async { Err(lance::Error::io("compaction preparation cannot commit")) }) + Box::pin(async { Err(lance::Error::io("maintenance preparation cannot commit")) }) } } -async fn watch_compaction_preparation( +async fn watch_maintenance_preparation( state: &Arc, claim: &TaskClaim, scope: &lance_context_core::merge_write_scope::MergeWriteScope, @@ -374,7 +375,7 @@ async fn watch_compaction_preparation( tokio::time::sleep(Duration::from_secs(2)).await; match state.task_store.preparation_owned(claim).await { Ok(true) => {} - Ok(false) => return "compaction preparation claim lost".into(), + Ok(false) => return "maintenance preparation claim lost".into(), Err(error) => return error.to_string(), } let current = scope.completed_steps(); @@ -385,23 +386,27 @@ async fn watch_compaction_preparation( if changed.elapsed() >= Duration::from_secs(state.config.maintenance.maintenance_idle_timeout_secs) { - return "compaction preparation made no progress".into(); + return "maintenance preparation made no progress".into(); } } } -async fn wait_for_compaction_commit( +async fn wait_for_maintenance_commit( state: &Arc, claim: &mut TaskClaim, ) -> Result<(), String> { let waiting = std::time::Instant::now(); + let wait_budget = match claim.task.kind { + TaskKind::IndexId => state.config.maintenance.index_commit_wait_secs, + _ => state.config.maintenance.compaction_commit_wait_secs, + }; if !state .task_store - .request_compaction_commit(claim) + .request_maintenance_commit(claim) .await .map_err(|e| e.to_string())? { - return Err("compaction preparation claim lost before requesting commit".into()); + return Err("maintenance preparation claim lost before requesting commit".into()); } loop { if !state @@ -410,11 +415,11 @@ async fn wait_for_compaction_commit( .await .map_err(|e| e.to_string())? { - return Err("compaction preparation claim lost before commit".into()); + return Err("maintenance preparation claim lost before commit".into()); } if state .task_store - .promote_compaction(claim) + .promote_maintenance(claim) .await .map_err(|e| e.to_string())? { @@ -422,10 +427,8 @@ async fn wait_for_compaction_commit( } // Check only after promotion has definitively declined. Never cancel // an in-flight acquisition and mistake an accepted claim for a timeout. - if waiting.elapsed() - >= Duration::from_secs(state.config.maintenance.compaction_commit_wait_secs) - { - return Err("compaction commit ownership wait budget exhausted; reprepare".into()); + if waiting.elapsed() >= Duration::from_secs(wait_budget) { + return Err("maintenance commit ownership wait budget exhausted; reprepare".into()); } tokio::time::sleep(Duration::from_millis(500)).await; } @@ -438,7 +441,7 @@ async fn run_prepared_compaction( let target = claim.task.target.clone(); let (_, name) = parse_target(&target); let uri = state.rollout_uri(name); - let scope = lance_context_core::merge_write_scope::MergeWriteScope::with_pinned_authorizer( + let scope = lance_context_core::merge_write_scope::MergeWriteScope::with_preparation_authorizer( Arc::new(PreparationOnly), ); let began = std::time::Instant::now(); @@ -456,7 +459,7 @@ async fn run_prepared_compaction( }); tokio::select! { result = work => result?, - error = watch_compaction_preparation(state, claim, &scope) => return Err(error), + error = watch_maintenance_preparation(state, claim, &scope) => return Err(error), } }; metrics::histogram!("master_compaction_phase_duration_seconds", "phase" => "prepare") @@ -476,7 +479,7 @@ async fn run_prepared_compaction( } let waiting = std::time::Instant::now(); tracing::info!(task = %claim.task.id, target = %claim.task.target, seconds = began.elapsed().as_secs_f64(), "compaction files ready; waiting for commit ownership"); - wait_for_compaction_commit(state, claim).await?; + wait_for_maintenance_commit(state, claim).await?; metrics::histogram!("master_compaction_phase_duration_seconds", "phase" => "commit_wait") .record(waiting.elapsed().as_secs_f64()); // No preparation handle is reused: only this scope may authorize metadata @@ -545,7 +548,7 @@ async fn finish_compaction( )) } -/// Build a ZoneMap scalar index on one experiment's `id` column. Shares the +/// Build the configured scalar index on one experiment's `id` column. Shares the /// per-name base-table write gate with [`run_compaction`] so an `IndexId` and a /// `Compact` for the same experiment never commit concurrently (`CreateIndex` /// vs `Rewrite` can conflict). Distinct experiments index concurrently. @@ -572,6 +575,63 @@ async fn index_id_inner(state: &Arc, name: &str) -> Result, + claim: &mut TaskClaim, +) -> Result { + let uri = state.rollout_uri(&claim.task.target); + let scope = lance_context_core::merge_write_scope::MergeWriteScope::with_preparation_authorizer( + Arc::new(PreparationOnly), + ); + let began = std::time::Instant::now(); + tracing::info!(task = %claim.task.id, target = %claim.task.target, + "index file preparation started without table ownership"); + let prepared = { + let work = scope.run(async { + let store = + RolloutStore::open_existing_with_options(&uri, state.rollout_store_options()) + .await + .map_err(|e| e.to_string())?; + store + .prepare_id_key_index() + .await + .map_err(|e| e.to_string()) + }); + tokio::select! { + result = work => result?, + error = watch_maintenance_preparation(state, claim, &scope) => return Err(error), + } + }; + metrics::histogram!("master_index_phase_duration_seconds", "phase" => "prepare") + .record(began.elapsed().as_secs_f64()); + tracing::info!(task = %claim.task.id, target = %claim.task.target, + seconds = began.elapsed().as_secs_f64(), "index files ready; waiting for commit ownership"); + let waiting = std::time::Instant::now(); + wait_for_maintenance_commit(state, claim).await?; + metrics::histogram!("master_index_phase_duration_seconds", "phase" => "commit_wait") + .record(waiting.elapsed().as_secs_f64()); + crate::maintenance_execution::reconcile_previous(state, claim).await?; + let started = std::time::Instant::now(); + let result = crate::maintenance_execution::run(state, claim, async { + let mut store = + RolloutStore::open_existing_with_options(&uri, state.rollout_store_options()) + .await + .map_err(|e| e.to_string())?; + store + .commit_prepared_id_key_index(prepared) + .await + .map_err(|e| e.to_string())?; + Ok(format!( + "built {} index on id (prepared outside write lock)", + state.config.key_index_type.as_str() + )) + }) + .await; + metrics::histogram!("master_index_phase_duration_seconds", "phase" => "commit") + .record(started.elapsed().as_secs_f64()); + result +} + async fn compact_inner( state: &Arc, name: &str, @@ -1215,7 +1275,7 @@ mod tests { .unwrap() .unwrap(); let merge_id = merge.task.id.clone(); - let error = wait_for_compaction_commit(&state, &mut prep) + let error = wait_for_maintenance_commit(&state, &mut prep) .await .unwrap_err(); assert!(error.contains("wait budget exhausted")); @@ -1253,7 +1313,9 @@ mod tests { .await .unwrap() .unwrap(); - wait_for_compaction_commit(&state, &mut next).await.unwrap(); + wait_for_maintenance_commit(&state, &mut next) + .await + .unwrap(); state .task_store .finish(next, Ok("retry acquired free writer".into())) @@ -1437,7 +1499,7 @@ mod tests { .await .unwrap() .unwrap(); - assert!(claim.preparing_compaction()); + assert!(claim.preparing_maintenance()); run_task( &state, claim, @@ -2025,6 +2087,51 @@ mod tests { worker.abort(); } + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn prepared_index_task_runs_with_fenced_publication() { + for kind in [ + lance_context_core::KeyIndexType::Btree, + lance_context_core::KeyIndexType::Zonemap, + ] { + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + cfg.maintenance.index_prepare_targets = vec!["exp".into()]; + cfg.merge_rollout.owned_targets = vec!["exp".into()]; + cfg.key_index_type = kind; + let state = MasterState::new(cfg).await.unwrap(); + let uri = state.rollout_uri("exp"); + let mut writer = RolloutStore::open(&uri).await.unwrap(); + writer.add(&[rollout_record("seed")]).await.unwrap(); + writer.cleanup_own_shard().await.unwrap(); + let task = enqueue(&state, TaskKind::IndexId, "exp").await.unwrap(); + let claim = state + .task_store + .claim_next_of_kinds(TaskKinds::GENERAL.without_compact()) + .await + .unwrap() + .unwrap(); + assert!(claim.preparing_maintenance()); + run_task(&state, claim, TaskClaimTiming::default()).await; + let status = state.task_store.get(&task.id).await.unwrap().unwrap(); + assert_eq!(status.state, TaskState::Done, "{status:?}"); + assert_eq!( + status.detail, + Some(format!( + "built {} index on id (prepared outside write lock)", + kind.as_str() + )) + ); + assert!(state + .task_store + .merge_coordinator() + .get("exp") + .await + .unwrap() + .is_none()); + } + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn resident_merge_waiting_for_memory_keeps_general_tasks_running() { diff --git a/crates/lance-context-master/src/task_store.rs b/crates/lance-context-master/src/task_store.rs index e28bba89..dfca7ba7 100644 --- a/crates/lance-context-master/src/task_store.rs +++ b/crates/lance-context-master/src/task_store.rs @@ -119,6 +119,7 @@ struct EtcdTaskStore { prefix: String, lease_ttl: i64, prepare_targets: Vec, + index_prepare_targets: Vec, maintenance_catchup_targets: Vec, rollout: lance_context_merge::rollout::MergeRollout, } @@ -131,7 +132,7 @@ pub struct TaskClaim { } impl TaskClaim { - pub(crate) fn preparing_compaction(&self) -> bool { + pub(crate) fn preparing_maintenance(&self) -> bool { self.backend.preparation_key.is_some() && self.backend.target_key.is_none() } } @@ -186,7 +187,10 @@ impl TaskStore { /// Ask subsequent merge admissions to yield after files are ready. This /// marker shares the preparation lease: crash/finish removes the request. /// It grants no write ownership and never cancels an admitted merge. - pub(crate) async fn request_compaction_commit(&self, claim: &TaskClaim) -> lance::Result { + pub(crate) async fn request_maintenance_commit( + &self, + claim: &TaskClaim, + ) -> lance::Result { let Some(preparation) = &claim.backend.preparation_key else { return Ok(false); }; @@ -208,13 +212,13 @@ impl TaskStore { ) .await .map(|r| r.succeeded()) - .map_err(etcd_error("request compaction commit turn")) + .map_err(etcd_error("request maintenance commit turn")) } /// Acquire the existing table writer protocol only after immutable files /// are ready. An ambiguous response can adopt only this claim's own token. - pub(crate) async fn promote_compaction(&self, claim: &mut TaskClaim) -> lance::Result { - assert!(claim.preparing_compaction()); + pub(crate) async fn promote_maintenance(&self, claim: &mut TaskClaim) -> lance::Result { + assert!(claim.preparing_maintenance()); let b = &claim.backend; let target = &claim.task.target; let execution_key = lance_context_merge::execution_key(&self.inner.prefix, target); @@ -296,7 +300,7 @@ impl TaskStore { { true } else { - return Err(etcd_error("promote compaction")(error)); + return Err(etcd_error("promote maintenance")(error)); } } }; @@ -686,6 +690,7 @@ impl EtcdTaskStore { prefix: config.etcd.prefix().to_string(), lease_ttl: config.etcd_lease_ttl_secs, prepare_targets: config.maintenance.compaction_prepare_targets.clone(), + index_prepare_targets: config.maintenance.index_prepare_targets.clone(), maintenance_catchup_targets: config.maintenance.maintenance_catchup_targets.clone(), rollout: config.merge_rollout.clone(), }) @@ -840,14 +845,21 @@ impl EtcdTaskStore { } } - let preparing = task.kind == TaskKind::Compact + let preparation_targets: &[String] = match task.kind { + TaskKind::Compact => &self.prepare_targets, + TaskKind::IndexId => &self.index_prepare_targets, + _ => &[], + }; + let preparing = matches!(task.kind, TaskKind::Compact | TaskKind::IndexId) && !task.target.starts_with("generic:") && !self.rollout.draining(&task.target) - && self - .prepare_targets + && preparation_targets .iter() .any(|t| t == "*" || t == &task.target); let write_claim = requires_target_lock(task.kind) && !preparing; + // Keep the existing wire keys so older masters also yield to + // prepared index commits. Compact/IndexId share one preparation + // slot per table; append merges need neither preparation slot. let preparation_key = preparing.then(|| { format!( "{}/compact-preparations/{}", @@ -2485,6 +2497,110 @@ mod tests { (dir, cfg, store) } + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn index_preparation_allows_merge_and_shares_compact_commit_turn() { + let (_dir, mut cfg, _) = preparation_test_store().await; + cfg.maintenance.index_prepare_targets = vec!["hot".into()]; + let store = TaskStore::open(&cfg).await.unwrap(); + store + .enqueue(TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + let merge = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + store + .enqueue(TaskKind::IndexId, "hot", vec![]) + .await + .unwrap(); + let mut prep = store + .claim_next_of_kinds(TaskKinds::GENERAL.without_compact()) + .await + .unwrap() + .unwrap(); + assert!(prep.preparing_maintenance()); + // Readers/file builders can start even when merge already owns the table. + assert!(!store.promote_maintenance(&mut prep).await.unwrap()); + store + .enqueue(TaskKind::Compact, "hot", vec![]) + .await + .unwrap(); + assert!(store + .claim_next_of_kinds(TaskKinds::COMPACT) + .await + .unwrap() + .is_none()); + assert!(store.request_maintenance_commit(&prep).await.unwrap()); + store + .finish(merge, Ok("merge advanced during index build".into())) + .await + .unwrap(); + store + .enqueue(TaskKind::MergeWal, "hot", vec![]) + .await + .unwrap(); + // This uses the old compact wire keys: old masters also yield the turn. + assert!(store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .is_none()); + assert!(store.promote_maintenance(&mut prep).await.unwrap()); + assert!(store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .is_none()); + store + .finish(prep, Ok("index published".into())) + .await + .unwrap(); + let merge = store + .claim_next_of_kinds(TaskKinds::MERGE_WAL) + .await + .unwrap() + .unwrap(); + store.finish(merge, Ok("continued".into())).await.unwrap(); + let compact = store + .claim_next_of_kinds(TaskKinds::COMPACT) + .await + .unwrap() + .unwrap(); + store.finish(compact, Ok("continued".into())).await.unwrap(); + } + + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn expired_index_preparation_cannot_publish_or_release_replacement() { + let (_dir, mut cfg, _) = preparation_test_store().await; + cfg.maintenance.index_prepare_targets = vec!["*".into()]; + let store = TaskStore::open(&cfg).await.unwrap(); + let task = store + .enqueue(TaskKind::IndexId, "hot", vec![]) + .await + .unwrap(); + let kinds = TaskKinds::GENERAL.without_compact(); + let mut old = store.claim_next_of_kinds(kinds).await.unwrap().unwrap(); + assert!(store.request_maintenance_commit(&old).await.unwrap()); + store + .inner + .revoke_lease(old.backend.lease_id) + .await + .unwrap(); + store.recover_orphaned().await.unwrap(); + let mut replacement = store.claim_next_of_kinds(kinds).await.unwrap().unwrap(); + assert_eq!(replacement.task.id, task.id); + assert!(!store.request_maintenance_commit(&old).await.unwrap()); + assert!(!store.promote_maintenance(&mut old).await.unwrap()); + assert!(store.finish(old, Ok("stale".into())).await.is_err()); + assert!(store.preparation_owned(&replacement).await.unwrap()); + assert!(store.promote_maintenance(&mut replacement).await.unwrap()); + store.finish(replacement, Ok("valid".into())).await.unwrap(); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn compaction_preparation_allows_merge_but_commit_is_exclusive() { @@ -2498,7 +2614,7 @@ mod tests { .await .unwrap() .unwrap(); - assert!(prep.preparing_compaction()); + assert!(prep.preparing_maintenance()); assert!(store .inner .get_text(&store.inner.target_lock_key("hot")) @@ -2514,13 +2630,13 @@ mod tests { .await .unwrap() .unwrap(); - assert!(!store.promote_compaction(&mut prep).await.unwrap()); + assert!(!store.promote_maintenance(&mut prep).await.unwrap()); store .finish(merge, Ok("append while rewriting".into())) .await .unwrap(); - assert!(store.promote_compaction(&mut prep).await.unwrap()); - assert!(!prep.preparing_compaction()); + assert!(store.promote_maintenance(&mut prep).await.unwrap()); + assert!(!prep.preparing_maintenance()); store .enqueue(TaskKind::MergeWal, "hot", Vec::new()) .await @@ -2576,12 +2692,12 @@ mod tests { .unwrap() .unwrap(); assert!( - prep.preparing_compaction(), + prep.preparing_maintenance(), "preparation must not block native merge admission" ); - assert!(store.request_compaction_commit(&prep).await.unwrap()); + assert!(store.request_maintenance_commit(&prep).await.unwrap()); assert!( - !store.promote_compaction(&mut prep).await.unwrap(), + !store.promote_maintenance(&mut prep).await.unwrap(), "running merge must finish cooperatively" ); store @@ -2600,7 +2716,7 @@ mod tests { .is_none(), "new pass must yield to ready compact" ); - assert!(store.promote_compaction(&mut prep).await.unwrap()); + assert!(store.promote_maintenance(&mut prep).await.unwrap()); store .finish(prep, Ok("compact committed".into())) .await @@ -2635,13 +2751,13 @@ mod tests { .await .unwrap() .unwrap(); - assert!(store.request_compaction_commit(&prep).await.unwrap()); + assert!(store.request_maintenance_commit(&prep).await.unwrap()); store .inner .revoke_lease(prep.backend.lease_id) .await .unwrap(); - assert!(!store.request_compaction_commit(&prep).await.unwrap()); + assert!(!store.request_maintenance_commit(&prep).await.unwrap()); assert!(store .inner .get_text(&store.inner.compaction_commit_key("hot")) @@ -2692,7 +2808,7 @@ mod tests { .reserve(&store.merge_claim(&merge), &execution) .await .unwrap()); - assert!(store.request_compaction_commit(&prep).await.unwrap()); + assert!(store.request_maintenance_commit(&prep).await.unwrap()); store.abandon_claim_for_test(merge).await.unwrap(); let recovery = store .claim_next_of_kinds(TaskKinds::MERGE_WAL) @@ -2751,13 +2867,13 @@ mod tests { .await .unwrap() .unwrap(); - assert!(prep.preparing_compaction()); - assert!(!store.promote_compaction(&mut prep).await.unwrap()); + assert!(prep.preparing_maintenance()); + assert!(!store.promote_maintenance(&mut prep).await.unwrap()); store .finish(native, Ok("native pass joined".into())) .await .unwrap(); - assert!(store.promote_compaction(&mut prep).await.unwrap()); + assert!(store.promote_maintenance(&mut prep).await.unwrap()); store .enqueue(TaskKind::MergeWal, "hot", Vec::new()) .await @@ -2882,7 +2998,7 @@ mod tests { .put(active.clone(), "replacement", None) .await .unwrap(); - assert!(!store.promote_compaction(&mut prep).await.unwrap()); + assert!(!store.promote_maintenance(&mut prep).await.unwrap()); store .finish(prep, Err("native identity changed".into())) .await @@ -2933,7 +3049,7 @@ mod tests { .put(active.clone(), "new-dedicated", None) .await .unwrap(); - assert!(!store.promote_compaction(&mut prep).await.unwrap()); + assert!(!store.promote_maintenance(&mut prep).await.unwrap()); store.finish(prep, Err("superseded".into())).await.unwrap(); assert_eq!( store.inner.get_text(&active).await.unwrap().as_deref(), @@ -3018,10 +3134,10 @@ mod tests { .unwrap(); assert_eq!(replacement.task.id, task.id); assert!(!store.preparation_owned(&old).await.unwrap()); - assert!(!store.promote_compaction(&mut old).await.unwrap()); + assert!(!store.promote_maintenance(&mut old).await.unwrap()); assert!(store.finish(old, Ok("stale".into())).await.is_err()); assert!(store.preparation_owned(&replacement).await.unwrap()); - assert!(store.promote_compaction(&mut replacement).await.unwrap()); + assert!(store.promote_maintenance(&mut replacement).await.unwrap()); store.finish(replacement, Ok("valid".into())).await.unwrap(); } diff --git a/docs/concurrent-indexing.md b/docs/concurrent-indexing.md new file mode 100644 index 00000000..a4da3942 --- /dev/null +++ b/docs/concurrent-indexing.md @@ -0,0 +1,56 @@ +# Index preparation alongside WAL merge + +`INDEX_PREPARE_TARGETS` selects rollout tables (comma-separated exact names or +`*`) whose IndexId tasks build files outside the table write lock. An empty +list keeps legacy indexing; generic and draining targets keep their existing +behavior. Both configured index types, BTree and ZoneMap, use this path. + +The task takes a lease-backed preparation claim and calls Lance's +`execute_uncommitted` against a fixed snapshot. Handles opened in preparation +have a pinned authorizer that rejects manifest commits. Append merge can keep +publishing while the index scans keys and writes immutable files. Existing +general task concurrency and Lance memory limits still bound indexing; this +change does not add another indexing pool or increase parallelism. + +Once files are ready, the task requests the next writer turn, then promotes its +exact task/preparation token into the existing table writer protocol. It never +cancels the active writer. `INDEX_COMMIT_WAIT_SECS` (default 120) bounds how long +prepared files retain a local task slot waiting for ownership. Expiry fails that +attempt and leaves the other writer untouched. Lease loss invalidates promotion; +a recovered task builds fresh files. Abandoned files remain unpublished. + +Index and compact share the existing `compact-preparations` and +`compact-commit-ready` wire keys deliberately. There is one preparation slot per +table, independent of merge. Older masters already honor this commit request, +so a rolling upgrade preserves scheduling fairness without a queue migration. +Old full-lock maintenance may run during new preparation; validation still +rejects stale output. Persistent dedicated-owner guards retain their existing +opt-in requirements (`MAINTENANCE_CATCHUP_TARGETS`). + +Publication opens a fresh handle inside the existing maintenance execution +scope, retaining owned-target fencing and ambiguous-commit recovery. Before +committing it checks the dataset URI/schema, every original source fragment, +and the same-name index metadata. A delete, rewrite, schema change, or concurrent +index replacement invalidates preparation. The CreateIndex transaction retains +the original snapshot read version, allowing Lance to check intervening commits +as well. Newly appended fragments remain unindexed and query planning scans them; +index publication must preserve merge watermarks and newly appended rows. + +The progress watchdog counts completed data/index reads and upload parts, +including IO from child tasks. Response headers, metadata polling and failed IO +do not count. Lance 9 does not report fine-grained progress from its basic BTree +or ZoneMap trainer, so training callbacks alone are insufficient. Streaming +reads remain streaming and use no new payload buffer. Local/uring readers bypass +ObjectStore wrappers, so preparation also samples their handle-specific completed +byte counters. The preparation-specific store parameters isolate those counters +from other tasks using the same session. Storage commit progress +continues to use the existing execution scope. The +`master_index_phase_duration_seconds` histogram separates `prepare`, `commit_wait`, +and `commit`. Logs include task/target identities and phase transitions without +adding table labels to metrics. + +Validation covers both index types and append commit orders, indexed/uncovered +reads, WAL watermark retention, stale source/schema/index rejection, forbidden +publication, etcd admission during an active merge, exclusive publication, and +replacement after preparation lease loss. Production rollout should compare +these phase durations and pending generations while ordinary jobs continue.