diff --git a/crates/lance-context-core/src/lib.rs b/crates/lance-context-core/src/lib.rs index 836d9b6..efd71da 100644 --- a/crates/lance-context-core/src/lib.rs +++ b/crates/lance-context-core/src/lib.rs @@ -16,7 +16,7 @@ pub mod merge_budget; pub mod merge_write_scope; pub mod metrics; mod namespace; -mod preparation_io; +pub mod preparation_io; mod record; mod registry; mod registry_etcd; diff --git a/crates/lance-context-core/src/merge_write_scope.rs b/crates/lance-context-core/src/merge_write_scope.rs index 7c5e3cb..7b87bd1 100644 --- a/crates/lance-context-core/src/merge_write_scope.rs +++ b/crates/lance-context-core/src/merge_write_scope.rs @@ -105,6 +105,9 @@ pub struct MergeWriteScope { authorizer: Option>, pin_opened_handles: bool, leaves: Mutex>, + /// Test-only storage shim applied **beneath** the progress wrapper, so a + /// delayed or frozen store is still observed as real IO by the watchdog. + io_shim: Option>, } impl MergeWriteScope { @@ -158,6 +161,22 @@ impl MergeWriteScope { }) } + /// Preparation scope whose object stores are wrapped by `shim` under the + /// progress wrapper. For tests that delay or freeze storage beneath a + /// real build to prove the watchdog judges actual IO. + pub fn with_preparation_authorizer_and_io_shim( + authorizer: Arc, + shim: Arc, + ) -> Arc { + Arc::new(Self { + authorizer: Some(authorizer), + pin_opened_handles: true, + observe_file_io: true, + io_shim: Some(shim), + ..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. @@ -273,15 +292,41 @@ pub(crate) fn preparation_store_wrapper( CURRENT .try_with(|scope| { scope.observe_file_io.then(|| { - Arc::new(crate::preparation_io::ProgressWrapper( - scope.completed_steps.clone(), - )) as Arc + let progress: Arc = Arc::new( + crate::preparation_io::ProgressWrapper(scope.completed_steps.clone()), + ); + match &scope.io_shim { + // Progress outermost: it counts the shimmed store's + // completed operations, exactly as it would real storage. + Some(shim) => Arc::new(Layered { + outer: progress, + inner: shim.clone(), + }) + as Arc, + None => progress, + } }) }) .ok() .flatten() } +#[derive(Debug)] +struct Layered { + outer: Arc, + inner: Arc, +} + +impl lance_io::object_store::WrappingObjectStore for Layered { + fn wrap( + &self, + name: &str, + original: Arc, + ) -> Arc { + self.outer.wrap(name, self.inner.wrap(name, original)) + } +} + pub(crate) fn observes_preparation_io() -> bool { CURRENT .try_with(|scope| scope.observe_file_io) diff --git a/crates/lance-context-core/src/preparation_io.rs b/crates/lance-context-core/src/preparation_io.rs index 1ac92ec..93c208a 100644 --- a/crates/lance-context-core/src/preparation_io.rs +++ b/crates/lance-context-core/src/preparation_io.rs @@ -210,3 +210,130 @@ mod tests { assert_eq!(steps.load(Ordering::Relaxed), 4); } } + +/// Test support: an object-store shim that delays or freezes data-file IO +/// beneath the progress wrapper. Lets a test drive a *real* index or +/// compaction build whose storage is slow or stuck and observe what the +/// watchdog does. Not for production use. +#[doc(hidden)] +pub mod test_support { + use super::*; + use std::sync::atomic::AtomicBool; + use std::time::Duration; + + #[derive(Debug, Clone)] + pub struct Throttle { + /// Added to every data-file put/get/get_ranges. + pub delay: Duration, + /// While set, data-file IO parks forever (until cleared). + pub frozen: Arc, + /// Data-file operations started. + pub ops: Arc, + } + + impl Throttle { + pub fn new(delay: Duration) -> Self { + Self { + delay, + frozen: Arc::new(AtomicBool::new(false)), + ops: Arc::new(AtomicU64::new(0)), + } + } + async fn gate(&self, path: &Path) { + if !data_file(path) { + return; + } + self.ops.fetch_add(1, Ordering::Relaxed); + while self.frozen.load(Ordering::Relaxed) { + tokio::time::sleep(Duration::from_millis(50)).await; + } + if !self.delay.is_zero() { + tokio::time::sleep(self.delay).await; + } + } + } + + #[derive(Debug)] + pub struct ThrottleWrapper(pub Throttle); + + impl WrappingObjectStore for ThrottleWrapper { + fn wrap(&self, _: &str, original: Arc) -> Arc { + Arc::new(ThrottledStore { + inner: original, + throttle: self.0.clone(), + }) + } + } + + #[derive(Debug)] + struct ThrottledStore { + inner: Arc, + throttle: Throttle, + } + + impl std::fmt::Display for ThrottledStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "Throttled({})", self.inner) + } + } + + #[async_trait] + impl ObjectStore for ThrottledStore { + async fn put_opts( + &self, + path: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + self.throttle.gate(path).await; + self.inner.put_opts(path, payload, opts).await + } + async fn put_multipart_opts( + &self, + path: &Path, + opts: PutMultipartOptions, + ) -> Result> { + self.throttle.gate(path).await; + self.inner.put_multipart_opts(path, opts).await + } + async fn get_opts(&self, path: &Path, opts: GetOptions) -> Result { + if !opts.head { + self.throttle.gate(path).await; + } + self.inner.get_opts(path, opts).await + } + async fn get_ranges( + &self, + path: &Path, + ranges: &[std::ops::Range], + ) -> Result> { + self.throttle.gate(path).await; + self.inner.get_ranges(path, ranges).await + } + 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 + } + } +} diff --git a/crates/lance-context-master/src/scheduler.rs b/crates/lance-context-master/src/scheduler.rs index b5e405e..070da29 100644 --- a/crates/lance-context-master/src/scheduler.rs +++ b/crates/lance-context-master/src/scheduler.rs @@ -2271,6 +2271,178 @@ mod tests { .is_some_and(|d| d.contains("prepared outside write lock"))); } + /// Review finding 6 on #340: prove the watchdog on a *real* index build + /// with real storage behaviour, not hand-placed checkpoints. + /// + /// (a) Throttled IO: every data/index operation sleeps longer than the + /// idle timeout would allow between manual checkpoints, yet because + /// each completed operation is progress, the build runs past the idle + /// timeout and finishes. + /// (b) Frozen IO: the same build with storage parked part-way is cancelled + /// by the watchdog with the stall error, within [IDLE, IDLE + 4). + /// (c) Isolation: a second, healthy preparation on another table running + /// concurrently with (b) is not cancelled and does not keep (b) alive. + #[tokio::test] + #[ignore = "requires ETCD_TEST_ENDPOINTS"] + async fn real_index_build_survives_slow_io_and_is_cancelled_when_frozen() { + use lance_context_core::merge_write_scope::MergeWriteScope; + use lance_context_core::preparation_io::test_support::{Throttle, ThrottleWrapper}; + const IDLE: u64 = 3; + let dir = TempDir::new().unwrap(); + let mut cfg = config(&dir); + // Plain local paths take Lance's direct file IO and bypass object-store + // wrappers; progress there comes from the io_tracker. The + // `file-object-store` scheme drives the same directory through the + // ObjectStore API, so the shim and the progress wrapper see every + // data/index operation exactly as they would on a cloud store, which + // is the path this test is about. + cfg.data_dir = format!("file-object-store://{}", dir.path().display()); + cfg.maintenance.index_prepare_targets = vec!["*".into()]; + cfg.maintenance.maintenance_idle_timeout_secs = IDLE; + cfg.merge_rollout.owned_targets = vec!["slow".into(), "frozen".into(), "healthy".into()]; + let state = MasterState::new(cfg).await.unwrap(); + // Enough rows to force several data-file operations during the build. + for target in ["slow", "frozen", "healthy"] { + let mut writer = RolloutStore::open(&state.rollout_uri(target)) + .await + .unwrap(); + let records: Vec<_> = (0..64) + .map(|i| rollout_record(&format!("{target}-{i}"))) + .collect(); + writer.add(&records).await.unwrap(); + writer.cleanup_own_shard().await.unwrap(); + } + let claim_for = |state: Arc, target: &'static str| async move { + enqueue(&state, TaskKind::IndexId, target).await.unwrap(); + let claim = state + .task_store + .claim_next_of_kinds(TaskKinds::GENERAL.without_compact()) + .await + .unwrap() + .unwrap(); + assert_eq!(claim.task.target, target); + assert!(claim.preparing_maintenance()); + claim + }; + let prepare = + |state: Arc, scope: Arc, target: &'static str| { + let uri = state.rollout_uri(target); + async move { + 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(|_| "prepared") + .map_err(|e| e.to_string()) + }) + .await + } + }; + + // (a) Throttled: 1.2 s per data-file op, idle timeout 3 s. The build + // must take longer than IDLE in total and still complete. + let claim = claim_for(state.clone(), "slow").await; + let throttle = Throttle::new(Duration::from_millis(1200)); + let scope = MergeWriteScope::with_preparation_authorizer_and_io_shim( + Arc::new(PreparationOnly), + Arc::new(ThrottleWrapper(throttle.clone())), + ); + let began = std::time::Instant::now(); + let outcome = tokio::select! { + done = prepare(state.clone(), scope.clone(), "slow") => done, + error = watch_maintenance_preparation(&state, &claim, &scope) => Err(error), + }; + let elapsed = began.elapsed(); + let ops = throttle.ops.load(std::sync::atomic::Ordering::Relaxed); + assert!( + ops >= 3, + "shim did not intercept data IO (ops={ops}); test is vacuous" + ); + assert_eq!( + outcome, + Ok("prepared"), + "slow-but-alive real build was cancelled" + ); + assert!( + elapsed > Duration::from_secs(IDLE), + "build finished in {elapsed:?}; did not outlive the idle timeout, test is vacuous" + ); + state + .task_store + .finish(claim, Ok("a".into())) + .await + .unwrap(); + + // (b) + (c): freeze "frozen" after its first data op; run "healthy" + // concurrently with no throttle. Only "frozen" is cancelled. + let frozen_claim = claim_for(state.clone(), "frozen").await; + let healthy_claim = claim_for(state.clone(), "healthy").await; + let throttle = Throttle::new(Duration::ZERO); + throttle + .frozen + .store(true, std::sync::atomic::Ordering::Relaxed); + let frozen_scope = MergeWriteScope::with_preparation_authorizer_and_io_shim( + Arc::new(PreparationOnly), + Arc::new(ThrottleWrapper(throttle.clone())), + ); + let healthy_scope = MergeWriteScope::with_preparation_authorizer(Arc::new(PreparationOnly)); + let began = std::time::Instant::now(); + let frozen_run = async { + tokio::select! { + done = prepare(state.clone(), frozen_scope.clone(), "frozen") => done, + error = watch_maintenance_preparation(&state, &frozen_claim, &frozen_scope) => Err(error), + } + }; + let healthy_run = async { + // Give the healthy build a slow start so it overlaps the frozen + // one's stall window rather than finishing instantly. + tokio::time::sleep(Duration::from_millis(500)).await; + tokio::select! { + done = prepare(state.clone(), healthy_scope.clone(), "healthy") => done, + error = watch_maintenance_preparation(&state, &healthy_claim, &healthy_scope) => Err(error), + } + }; + let (frozen_outcome, healthy_outcome) = tokio::join!(frozen_run, healthy_run); + let elapsed = began.elapsed(); + assert_eq!( + frozen_outcome, + Err("maintenance preparation made no progress".to_string()) + ); + assert!( + elapsed >= Duration::from_secs(IDLE) && elapsed < Duration::from_secs(IDLE + 4), + "frozen build cancelled after {elapsed:?}" + ); + assert_eq!( + healthy_outcome, + Ok("prepared"), + "healthy build was collateral damage" + ); + assert!( + throttle.ops.load(std::sync::atomic::Ordering::Relaxed) >= 1, + "frozen shim never saw IO; test is vacuous" + ); + throttle + .frozen + .store(false, std::sync::atomic::Ordering::Relaxed); + state + .task_store + .finish(frozen_claim, Err("stalled".into())) + .await + .unwrap(); + state + .task_store + .finish(healthy_claim, Ok("h".into())) + .await + .unwrap(); + } + #[tokio::test] #[ignore = "requires ETCD_TEST_ENDPOINTS"] async fn prepared_index_task_runs_with_fenced_publication() {