Recover acknowledged own-shard WAL before cleanup after restart

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

@@ -700,6 +700,15 @@ impl GenericStore {
         self.base.prepare_cleanup_merge().await
     }
 
+    /// Recover this instance's unsealed WAL, then prepare cleanup. The caller
+    /// must exclusively own the ingestion writer ID, as when appending. Never
+    /// use this to merge a live worker's shard from a maintenance master.
+    pub async fn prepare_recovered_cleanup_merge(
+        &self,
+    ) -> LanceResult<Option<(ShardManifestStore, ShardManifest, PreparedMerge)>> {
+        self.base.prepare_recovered_cleanup_merge().await
+    }
+
     /// Commit a merge prepared by [`Self::prepare_cleanup_merge`].
     pub async fn commit_prepared_merge(
         &mut self,
@@ -1383,6 +1392,110 @@ mod tests {
         });
     }
 
+    #[tokio::test]
+    async fn cold_maintenance_cleanup_does_not_claim_a_live_ingestion_epoch() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        let options = GenericStoreOptions {
+            shard_id: Some("live-ingestion-writer".into()),
+            ..Default::default()
+        };
+        let writer = GenericStore::open(uri, spec(), options.clone())
+            .await
+            .unwrap();
+        writer
+            .add(&[row(json!({"id":"shared", "score":1.0}))])
+            .await
+            .unwrap();
+        writer.flush().await.unwrap();
+        let before = writer
+            .own_shard_watermark()
+            .await
+            .unwrap()
+            .unwrap()
+            .writer_epoch;
+        let mut maintenance = GenericStore::open_existing(uri, options).await.unwrap();
+        let (manifests, manifest, prepared) =
+            maintenance.prepare_cleanup_merge().await.unwrap().unwrap();
+        maintenance
+            .commit_prepared_merge(&manifests, &manifest, prepared)
+            .await
+            .unwrap();
+        writer
+            .add(&[row(json!({"id":"shared", "score":2.0}))])
+            .await
+            .unwrap();
+        writer.flush().await.unwrap();
+        assert_eq!(
+            writer
+                .own_shard_watermark()
+                .await
+                .unwrap()
+                .unwrap()
+                .writer_epoch,
+            before
+        );
+        assert_eq!(
+            writer.get("shared", None).await.unwrap().unwrap()["score"],
+            json!(2.0)
+        );
+    }
+
+    #[test]
+    fn cleanup_recovers_own_wal_updates_without_new_ingest_after_runtime_exit() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let options = GenericStoreOptions {
+            shard_id: Some("restarted-generic-writer".into()),
+            ..Default::default()
+        };
+        let old_runtime = tokio::runtime::Runtime::new().unwrap();
+        let old = old_runtime.block_on(async {
+            let store = GenericStore::open(&uri, spec(), options.clone())
+                .await
+                .unwrap();
+            store
+                .add(&[row(json!({"id":"shared", "score":1.0}))])
+                .await
+                .unwrap();
+            store.flush().await.unwrap();
+            store
+                .add(&[row(json!({"id":"shared", "score":2.0}))])
+                .await
+                .unwrap();
+            store
+        });
+        drop(old_runtime);
+        drop(old);
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+        runtime.block_on(async {
+            let mut recovered = GenericStore::open_existing(&uri, options).await.unwrap();
+            assert_eq!(
+                recovered.get("shared", None).await.unwrap().unwrap()["score"],
+                json!(1.0)
+            );
+            let (manifests, manifest, prepared) = recovered
+                .prepare_recovered_cleanup_merge()
+                .await
+                .unwrap()
+                .expect("own WAL must recover without a new append");
+            assert!(
+                recovered
+                    .commit_prepared_merge(&manifests, &manifest, prepared)
+                    .await
+                    .unwrap()
+                    > 0
+            );
+            assert_eq!(
+                recovered.get("shared", None).await.unwrap().unwrap()["score"],
+                json!(2.0)
+            );
+            assert_eq!(recovered.list(None, None).await.unwrap().len(), 1);
+            assert_eq!(recovered.cleanup_wal().await.unwrap(), 0);
+            recovered.close().await.unwrap();
+        });
+    }
+
     #[test]
     fn seal_mode_survives_a_reopen() {
         // The mode is a property of the store, not of whoever opens it. Before

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

@@ -702,6 +702,15 @@ impl RolloutStore {
         self.base.prepare_cleanup_merge().await
     }
 
+    /// Recover this instance's unsealed WAL, then prepare cleanup. The caller
+    /// must exclusively own the ingestion writer ID, as when appending. Never
+    /// use this to merge a live worker's shard from a maintenance master.
+    pub async fn prepare_recovered_cleanup_merge(
+        &self,
+    ) -> LanceResult<Option<(ShardManifestStore, ShardManifest, PreparedMerge)>> {
+        self.base.prepare_recovered_cleanup_merge().await
+    }
+
     /// Commit a merge prepared by [`Self::prepare_merge_if_ready`].
     pub async fn commit_prepared_merge(
         &mut self,
@@ -4686,6 +4695,98 @@ mod tests {
         });
     }
 
+    #[test]
+    fn cleanup_recovers_own_wal_without_new_ingest_after_runtime_exit() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let options = RolloutStoreOptions {
+            shard_id: Some("restarted-writer".into()),
+            ..Default::default()
+        };
+        let payload = vec![37; 64 * 1024];
+        let old_runtime = tokio::runtime::Runtime::new().unwrap();
+        let old = old_runtime.block_on(async {
+            let store = RolloutStore::open_with_options(&uri, options.clone())
+                .await
+                .unwrap();
+            store
+                .add(&[artifact_record("durable-before-exit", &payload)])
+                .await
+                .unwrap();
+            assert_eq!(flushed_generation_count(&store).await, 0);
+            store
+        });
+        // Stop every old background task before dropping the store outside a
+        // runtime. Do not run its graceful close (which would hide this bug).
+        drop(old_runtime);
+        drop(old);
+
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+        runtime.block_on(async {
+            let other = RolloutStore::open_with_options(
+                &uri,
+                RolloutStoreOptions {
+                    shard_id: Some("different-live-writer".into()),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+            assert!(other
+                .prepare_recovered_cleanup_merge()
+                .await
+                .unwrap()
+                .is_none());
+            assert!(other.own_shard_watermark().await.unwrap().is_none());
+            assert!(other
+                .get_by_id("durable-before-exit")
+                .await
+                .unwrap()
+                .is_none());
+
+            let mut recovered = RolloutStore::open_with_options(&uri, options)
+                .await
+                .unwrap();
+            recovered.flush().await.unwrap();
+            assert!(recovered
+                .get_by_id("durable-before-exit")
+                .await
+                .unwrap()
+                .is_none());
+            // Exercise the split-lock route used by the actual HTTP merge.
+            let (manifests, manifest, prepared) = recovered
+                .prepare_recovered_cleanup_merge()
+                .await
+                .unwrap()
+                .expect("cleanup must replay acknowledged but unsealed WAL");
+            assert!(
+                recovered
+                    .commit_prepared_merge(&manifests, &manifest, prepared)
+                    .await
+                    .unwrap()
+                    > 0
+            );
+            assert_eq!(
+                recovered.get_blob("durable-before-exit").await.unwrap(),
+                Some(payload)
+            );
+            assert_eq!(recovered.cleanup_own_shard().await.unwrap(), 0);
+            let epoch = shard_writer_epoch(&recovered).await;
+            recovered
+                .add(&[assistant_record("after-recovery")])
+                .await
+                .unwrap();
+            recovered.cleanup_own_shard().await.unwrap();
+            assert_eq!(
+                shard_writer_epoch(&recovered).await,
+                epoch,
+                "live writer must be reused"
+            );
+            assert_eq!(recovered.list(None, None).await.unwrap().len(), 2);
+            recovered.close().await.unwrap();
+        });
+    }
+
     #[test]
     fn flush_is_noop_without_writes() {
         // flush() with no resident writer / nothing buffered must be a harmless

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

@@ -81,6 +81,10 @@ pub(crate) const DEFAULT_MANIFEST_SCAN_BATCH_SIZE: usize = 16;
 /// concurrently while collecting observability metrics.
 pub(crate) const DEFAULT_OBSERVE_CONCURRENCY: usize = 16;
 
+// Cold WAL replay decodes payloads. Serialize maintenance-triggered recovery
+// across tables rather than opening a writer for every cold store at once.
+static OWN_WAL_RECOVERY_SLOT: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(1);
+
 /// Flushed generations folded into the base table by one merge pass, by default.
 ///
 /// A merge buffers every row of every generation it takes before appending, and
@@ -763,6 +767,51 @@ impl StorageBase {
         result.map(|_| ())
     }
 
+    /// Recover a restarted instance's durable WAL before its cleanup merge.
+    /// Ordinary flush sweeps stay cheap: they only touch resident writers.
+    /// This path is for this handle's OWN write shard, never a foreign shard;
+    /// as with `put`, the caller must have exclusive ownership of that writer ID.
+    /// Opening a writer claims its epoch, so this is not a mechanism for taking
+    /// over an instance that is still alive.
+    /// Returns false when there is no own shard to merge.
+    async fn recover_and_flush_own_wal(&self) -> LanceResult<bool> {
+        if self.is_version_pinned() || self.deleted {
+            return Ok(false);
+        }
+        if self.write_writer.lock().await.is_some() {
+            self.flush().await?;
+            return Ok(true);
+        }
+        let manifest_store = ShardManifestStore::new(
+            self.dataset.object_store(None).await?,
+            &self.dataset.branch_location().path,
+            self.write_shard,
+            DEFAULT_MANIFEST_SCAN_BATCH_SIZE,
+        );
+        // A worker with no own shard must not queue behind another table's
+        // recovery. Serial master fan-out visits many such workers.
+        if manifest_store.refresh_latest().await?.is_none() {
+            return Ok(false);
+        }
+        let _recovery_slot = OWN_WAL_RECOVERY_SLOT
+            .acquire()
+            .await
+            .expect("WAL recovery semaphore is never closed");
+        // Manifest WAL-tip hints can lag acknowledged PUTs, so an existing
+        // own shard needs replay even if its hint says it is empty. The normal
+        // writer open rechecks under its mutex if ingest won the race. Reuse its
+        // normal replay limits (including Lance's streaming seals); changing
+        // those here would also change future ingest's generation sizes because
+        // this writer remains resident. One indivisible entry can exceed the
+        // configured soft memory ceiling.
+        info!(shard = %self.write_shard, "recovering own WAL before cleanup merge");
+        self.resident_writer().await?;
+        // Keep the recovered writer resident so concurrent ingest shares it;
+        // a later cheap flush can also finish sealing if this caller is cancelled.
+        self.flush().await?;
+        Ok(true)
+    }
+
     async fn flush_inner(&self) -> LanceResult<FlushOutcome> {
         let writer = {
             let guard = self.write_writer.lock().await;
@@ -983,6 +1032,19 @@ impl StorageBase {
         self.prepare_merge_if_ready_inner(1, true).await
     }
 
+    /// Recover and seal this instance's durable WAL before preparing cleanup.
+    /// The caller MUST own the ingestion writer ID exclusively. A maintenance
+    /// master merging a live worker's sealed shard must use the ordinary
+    /// preparation methods, which never acquire an ingestion epoch.
+    pub async fn prepare_recovered_cleanup_merge(
+        &self,
+    ) -> LanceResult<Option<(ShardManifestStore, ShardManifest, PreparedMerge)>> {
+        if !self.recover_and_flush_own_wal().await? {
+            return Ok(None);
+        }
+        self.prepare_merge_if_ready_inner(1, false).await
+    }
+
     async fn prepare_merge_if_ready_inner(
         &self,
         threshold: usize,
@@ -993,7 +1055,6 @@ impl StorageBase {
             return Ok(None);
         }
         if seal_first {
-            // Materialize anything buffered so it is eligible for this pass.
             self.flush().await?;
         }
         let object_store = self.dataset.object_store(None).await?;
@@ -2739,3 +2800,28 @@ pub fn derive_shard_id(instance_id: Option<&str>) -> Uuid {
     let input = instance_id.unwrap_or("default");
     Uuid::new_v5(&Uuid::NAMESPACE_OID, input.as_bytes())
 }
+
+#[cfg(test)]
+mod cold_recovery_tests {
+    use super::OWN_WAL_RECOVERY_SLOT;
+    use crate::RolloutStore;
+    use std::time::Duration;
+
+    #[tokio::test]
+    async fn empty_shard_cleanup_does_not_wait_for_other_recovery() {
+        let dir = tempfile::TempDir::new().unwrap();
+        let store = RolloutStore::open(dir.path().to_str().unwrap())
+            .await
+            .unwrap();
+        let _occupied = OWN_WAL_RECOVERY_SLOT.acquire().await.unwrap();
+        let prepared = tokio::time::timeout(
+            Duration::from_secs(2),
+            store.prepare_recovered_cleanup_merge(),
+        )
+        .await
+        .expect("absent own shard must bypass cold recovery admission")
+        .unwrap();
+        assert!(prepared.is_none());
+        assert!(store.own_shard_watermark().await.unwrap().is_none());
+    }
+}

crates/lance-context-server/src/routes/generic.rs

@@ -338,7 +338,7 @@ pub(crate) async fn merge_generic_wal_owned(
     let prepared = {
         let guard = store.read().await;
         guard
-            .prepare_cleanup_merge()
+            .prepare_recovered_cleanup_merge()
             .await
             .map_err(AppError::from_lance)?
     };

crates/lance-context-server/src/routes/rollouts.rs

@@ -756,7 +756,7 @@ pub(crate) async fn merge_wal_owned(
         let store = store_lock.read().await;
         ::metrics::histogram!("rollout_wal_merge_lock_wait_seconds")
             .record(lock_start.elapsed().as_secs_f64());
-        match store.prepare_cleanup_merge().await {
+        match store.prepare_recovered_cleanup_merge().await {
             Ok(prepared) => prepared,
             Err(e) => {
                 ::metrics::counter!("rollout_wal_cleanup_total", "result" => "failed").increment(1);

crates/lance-context-server/src/sweeper.rs

@@ -116,7 +116,7 @@ impl Sweepable for Arc<RwLock<RolloutStore>> {
         let prepared = {
             let guard = self.read().await;
             guard
-                .prepare_cleanup_merge()
+                .prepare_recovered_cleanup_merge()
                 .await
                 .map_err(|e| e.to_string())?
         };
@@ -205,7 +205,7 @@ impl Sweepable for Arc<RwLock<GenericStore>> {
         let prepared = {
             let guard = self.read().await;
             guard
-                .prepare_cleanup_merge()
+                .prepare_recovered_cleanup_merge()
                 .await
                 .map_err(|e| e.to_string())?
         };