Recover stalled local maintenance before resuming WAL merges

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

@@ -43,6 +43,7 @@ pub struct MergeWriteScope {
     progress: Mutex<Progress>,
     changed: Notify,
     authorizer: Option<Arc<dyn CommitAuthorizer>>,
+    pin_opened_handles: bool,
     leaves: Mutex<Vec<tokio::task::AbortHandle>>,
 }
 
@@ -61,6 +62,17 @@ impl MergeWriteScope {
         })
     }
 
+    /// Bind newly opened dataset handles to this execution, including commits
+    /// from child tasks. Use only for ephemeral maintenance handles: worker
+    /// caches must keep resolving the current execution at each call instead.
+    pub fn with_pinned_authorizer(authorizer: Arc<dyn CommitAuthorizer>) -> Arc<Self> {
+        Arc::new(Self {
+            authorizer: Some(authorizer),
+            pin_opened_handles: 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.
@@ -336,7 +348,21 @@ pub async fn fence_manifest_versions(
 }
 
 #[derive(Debug)]
-pub(crate) struct GuardedCommit(pub Arc<dyn CommitHandler>);
+pub(crate) struct GuardedCommit {
+    inner: Arc<dyn CommitHandler>,
+    scope: Option<Arc<MergeWriteScope>>,
+}
+impl GuardedCommit {
+    pub(crate) fn new(inner: Arc<dyn CommitHandler>) -> Self {
+        Self {
+            inner,
+            scope: CURRENT
+                .try_with(Arc::clone)
+                .ok()
+                .filter(|scope| scope.pin_opened_handles),
+        }
+    }
+}
 
 #[async_trait::async_trait]
 #[allow(clippy::too_many_arguments)]
@@ -351,45 +377,51 @@ impl CommitHandler for GuardedCommit {
         naming_scheme: ManifestNamingScheme,
         transaction: Option<Transaction>,
     ) -> std::result::Result<ManifestLocation, CommitError> {
-        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();
+        let work = async {
+            authorize("base", manifest.version).await?;
+            let inner = self.inner.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
+                Ok::<_, CommitError>((location, owned_manifest))
+            })
+            .await?;
+            *manifest = committed;
+            location
+        };
+        match &self.scope {
+            Some(scope) => scope.run(work).await,
+            None => work.await,
+        }
     }
 
     async fn resolve_latest_location(
         &self,
         base_path: &Path,
         object_store: &ObjectStore,
     ) -> Result<ManifestLocation> {
-        self.0
+        self.inner
             .resolve_latest_location(base_path, object_store)
             .await
     }
@@ -399,7 +431,7 @@ impl CommitHandler for GuardedCommit {
         version: u64,
         object_store: &dyn object_store::ObjectStore,
     ) -> Result<ManifestLocation> {
-        self.0
+        self.inner
             .resolve_version_location(base_path, version, object_store)
             .await
     }
@@ -410,7 +442,7 @@ impl CommitHandler for GuardedCommit {
         object_store: &dyn object_store::ObjectStore,
         naming: ManifestNamingScheme,
     ) -> Result<bool> {
-        self.0
+        self.inner
             .version_exists(base_path, version, object_store, naming)
             .await
     }
@@ -419,7 +451,7 @@ impl CommitHandler for GuardedCommit {
         base_path: &Path,
         object_store: &'a ObjectStore,
     ) -> BoxStream<'a, Result<ManifestLocation>> {
-        self.0
+        self.inner
             .list_detached_manifest_locations(base_path, object_store)
     }
     fn list_manifest_locations<'a>(
@@ -428,7 +460,7 @@ impl CommitHandler for GuardedCommit {
         object_store: &'a ObjectStore,
         sorted: bool,
     ) -> BoxStream<'a, Result<ManifestLocation>> {
-        self.0
+        self.inner
             .list_manifest_locations(base_path, object_store, sorted)
     }
     fn list_manifest_locations_since<'a>(
@@ -437,11 +469,11 @@ impl CommitHandler for GuardedCommit {
         object_store: &'a ObjectStore,
         since: u64,
     ) -> BoxStream<'a, Result<ManifestLocation>> {
-        self.0
+        self.inner
             .list_manifest_locations_since(base_path, object_store, since)
     }
     async fn delete(&self, base_path: &Path) -> Result<()> {
-        self.0.delete(base_path).await
+        self.inner.delete(base_path).await
     }
 }
 

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

@@ -2372,6 +2372,94 @@ mod tests {
         assert_eq!(attempts.load(Ordering::SeqCst), 1);
     }
 
+    #[tokio::test]
+    async fn owned_store_commit_guard_survives_transfer_to_a_child_task() {
+        use crate::merge_write_scope::{CommitAuthorizer, MergeWriteScope};
+        #[derive(Debug)]
+        struct Deny;
+        impl CommitAuthorizer for Deny {
+            fn authorize<'a>(
+                &'a self,
+                _: &'a str,
+                _: u64,
+            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = lance::Result<()>> + Send + 'a>>
+            {
+                Box::pin(async { Err(lance::Error::io("old execution fenced")) })
+            }
+        }
+        let dir = tempfile::tempdir().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        let store = RolloutStore::open(uri).await.unwrap();
+        let version = store.version();
+        let scope = MergeWriteScope::with_pinned_authorizer(Arc::new(Deny));
+        let mut dataset = scope
+            .run(StorageBase::load_with_options(uri, None, None))
+            .await
+            .unwrap();
+        let error =
+            tokio::spawn(async move { dataset.update_metadata([("attempt", "late-child")]).await })
+                .await
+                .unwrap()
+                .unwrap_err();
+        assert!(error.to_string().contains("old execution fenced"));
+        scope.drain().await;
+        assert_eq!(
+            RolloutStore::open_existing_with_options(uri, Default::default())
+                .await
+                .unwrap()
+                .version(),
+            version
+        );
+    }
+
+    #[tokio::test]
+    async fn cached_store_commit_guard_uses_each_new_execution() {
+        use crate::merge_write_scope::{CommitAuthorizer, MergeWriteScope};
+        #[derive(Debug)]
+        struct Deny;
+        impl CommitAuthorizer for Deny {
+            fn authorize<'a>(
+                &'a self,
+                _: &'a str,
+                _: u64,
+            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = lance::Result<()>> + Send + 'a>>
+            {
+                Box::pin(async { Err(lance::Error::io("current execution fenced")) })
+            }
+        }
+        let dir = tempfile::tempdir().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        RolloutStore::open(uri).await.unwrap();
+        let old = MergeWriteScope::with_authorizer(Arc::new(Deny));
+        let mut dataset = old
+            .run(StorageBase::load_with_options(uri, None, None))
+            .await
+            .unwrap();
+        old.drain().await;
+        // A cached handle must not retain the old, now closed execution.
+        let next = MergeWriteScope::new();
+        next.run(async {
+            dataset
+                .update_metadata([("attempt", "next-execution")])
+                .await
+        })
+        .await
+        .unwrap();
+        next.drain().await;
+        // The same handle must still enforce a later execution's authorization.
+        let denied = MergeWriteScope::with_authorizer(Arc::new(Deny));
+        let error = denied
+            .run(async {
+                dataset
+                    .update_metadata([("attempt", "denied-execution")])
+                    .await
+            })
+            .await
+            .unwrap_err();
+        assert!(error.to_string().contains("current execution fenced"));
+        denied.drain().await;
+    }
+
     #[test]
     fn rollout_filters_parse_supported_fields() {
         let filters = RolloutFilters::from_json_value(json!({
@@ -2509,7 +2597,7 @@ mod tests {
         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 {
+        let handler = Arc::new(GuardedCommit::new(Arc::new(PausedCommit {
             inner: lance_table::io::commit::commit_handler_from_url(uri, &None)
                 .await
                 .unwrap(),

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

@@ -2039,8 +2039,9 @@ impl StorageBase {
             ..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)));
+        let mut builder = DatasetBuilder::from_uri(uri).with_commit_handler(Arc::new(
+            crate::merge_write_scope::GuardedCommit::new(handler),
+        ));
         if let Some(options) = storage_options {
             builder = builder.with_storage_options(options);
         }
@@ -2078,7 +2079,9 @@ impl StorageBase {
         params.session = session;
         let handler =
             lance_table::io::commit::commit_handler_from_url(uri, &params.store_params).await?;
-        params.commit_handler = Some(Arc::new(crate::merge_write_scope::GuardedCommit(handler)));
+        params.commit_handler = Some(Arc::new(crate::merge_write_scope::GuardedCommit::new(
+            handler,
+        )));
 
         Dataset::write(batches, uri, Some(params)).await
     }

crates/lance-context-master/src/config.rs

@@ -8,6 +8,8 @@ use clap::Parser;
 pub struct MasterConfig {
     #[command(flatten)]
     pub merge_rollout: lance_context_merge::rollout::MergeRollout,
+    #[command(flatten)]
+    pub maintenance: MaintenanceConfig,
     /// Data directory / object-store prefix shared with the data-plane server.
     #[arg(long, env = "DATA_DIR", default_value = "./data")]
     pub data_dir: String,
@@ -308,3 +310,26 @@ mod tests {
         assert_eq!(config.compaction_max_bytes_per_file, 1024 * 1024 * 1024);
     }
 }
+
+/// Deadlines for local table mutations on explicitly owned targets.
+#[derive(Debug, Clone, clap::Args)]
+pub struct MaintenanceConfig {
+    /// Maximum local index/compact/repair execution time on owned targets.
+    #[arg(long, env = "MAINTENANCE_TIMEOUT_SECS", default_value_t = 3600)]
+    pub maintenance_timeout_secs: u64,
+    /// Maximum time without a completed storage/processing step.
+    #[arg(long, env = "MAINTENANCE_IDLE_TIMEOUT_SECS", default_value_t = 600)]
+    pub maintenance_idle_timeout_secs: u64,
+    /// Grace for in-flight manifest commits before storage fencing.
+    #[arg(long, env = "MAINTENANCE_DRAIN_TIMEOUT_SECS", default_value_t = 30)]
+    pub maintenance_drain_timeout_secs: u64,
+}
+impl Default for MaintenanceConfig {
+    fn default() -> Self {
+        Self {
+            maintenance_timeout_secs: 3600,
+            maintenance_idle_timeout_secs: 600,
+            maintenance_drain_timeout_secs: 30,
+        }
+    }
+}

crates/lance-context-master/src/lib.rs

@@ -3,6 +3,7 @@
 pub mod config;
 pub mod discovery;
 pub mod error;
+mod maintenance_execution;
 mod merge_execution;
 pub mod routes;
 pub mod scanner;