test(master): real index builds under throttled, frozen and concurrent IO

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;

crates/lance-context-core/src/merge_write_scope.rs

@@ -105,6 +105,9 @@ pub struct MergeWriteScope {
     authorizer: Option<Arc<dyn CommitAuthorizer>>,
     pin_opened_handles: bool,
     leaves: Mutex<Vec<tokio::task::AbortHandle>>,
+    /// 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<Arc<dyn lance_io::object_store::WrappingObjectStore>>,
 }
 
 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<dyn CommitAuthorizer>,
+        shim: Arc<dyn lance_io::object_store::WrappingObjectStore>,
+    ) -> Arc<Self> {
+        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<dyn lance_io::object_store::WrappingObjectStore>
+                let progress: Arc<dyn lance_io::object_store::WrappingObjectStore> = 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<dyn lance_io::object_store::WrappingObjectStore>,
+                    None => progress,
+                }
             })
         })
         .ok()
         .flatten()
 }
 
+#[derive(Debug)]
+struct Layered {
+    outer: Arc<dyn lance_io::object_store::WrappingObjectStore>,
+    inner: Arc<dyn lance_io::object_store::WrappingObjectStore>,
+}
+
+impl lance_io::object_store::WrappingObjectStore for Layered {
+    fn wrap(
+        &self,
+        name: &str,
+        original: Arc<dyn object_store::ObjectStore>,
+    ) -> Arc<dyn object_store::ObjectStore> {
+        self.outer.wrap(name, self.inner.wrap(name, original))
+    }
+}
+
 pub(crate) fn observes_preparation_io() -> bool {
     CURRENT
         .try_with(|scope| scope.observe_file_io)

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<AtomicBool>,
+        /// Data-file operations started.
+        pub ops: Arc<AtomicU64>,
+    }
+
+    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<dyn ObjectStore>) -> Arc<dyn ObjectStore> {
+            Arc::new(ThrottledStore {
+                inner: original,
+                throttle: self.0.clone(),
+            })
+        }
+    }
+
+    #[derive(Debug)]
+    struct ThrottledStore {
+        inner: Arc<dyn ObjectStore>,
+        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<PutResult> {
+            self.throttle.gate(path).await;
+            self.inner.put_opts(path, payload, opts).await
+        }
+        async fn put_multipart_opts(
+            &self,
+            path: &Path,
+            opts: PutMultipartOptions,
+        ) -> Result<Box<dyn MultipartUpload>> {
+            self.throttle.gate(path).await;
+            self.inner.put_multipart_opts(path, opts).await
+        }
+        async fn get_opts(&self, path: &Path, opts: GetOptions) -> Result<GetResult> {
+            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<u64>],
+        ) -> Result<Vec<Bytes>> {
+            self.throttle.gate(path).await;
+            self.inner.get_ranges(path, ranges).await
+        }
+        fn delete_stream(
+            &self,
+            paths: BoxStream<'static, Result<Path>>,
+        ) -> BoxStream<'static, Result<Path>> {
+            self.inner.delete_stream(paths)
+        }
+        fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
+            self.inner.list(prefix)
+        }
+        fn list_with_offset(
+            &self,
+            prefix: Option<&Path>,
+            offset: &Path,
+        ) -> BoxStream<'static, Result<ObjectMeta>> {
+            self.inner.list_with_offset(prefix, offset)
+        }
+        async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
+            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
+        }
+    }
+}

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<MasterState>, 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<MasterState>, scope: Arc<MergeWriteScope>, 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() {