diff --git a/crates/lance-context-core/src/generic_store.rs b/crates/lance-context-core/src/generic_store.rs
index 4ddff56..b966499 100644
--- a/crates/lance-context-core/src/generic_store.rs
+++ b/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> {
+ 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
diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs
index 940fd9c..1be83be 100644
--- a/crates/lance-context-core/src/rollout_store.rs
+++ b/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 > {
+ 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
diff --git a/crates/lance-context-core/src/store_base.rs b/crates/lance-context-core/src/store_base.rs
index c4766bd..ed29a90 100644
--- a/crates/lance-context-core/src/store_base.rs
+++ b/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 {
+ 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 {
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> {
+ 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());
+ }
+}
diff --git a/crates/lance-context-server/src/routes/generic.rs b/crates/lance-context-server/src/routes/generic.rs
index 287ba02..eaece28 100644
--- a/crates/lance-context-server/src/routes/generic.rs
+++ b/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)?
};
diff --git a/crates/lance-context-server/src/routes/rollouts.rs b/crates/lance-context-server/src/routes/rollouts.rs
index 1b00fc1..6e01e72 100644
--- a/crates/lance-context-server/src/routes/rollouts.rs
+++ b/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);
diff --git a/crates/lance-context-server/src/sweeper.rs b/crates/lance-context-server/src/sweeper.rs
index 783980e..3ca5690 100644
--- a/crates/lance-context-server/src/sweeper.rs
+++ b/crates/lance-context-server/src/sweeper.rs
@@ -116,7 +116,7 @@ impl Sweepable for Arc> {
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> {
let prepared = {
let guard = self.read().await;
guard
- .prepare_cleanup_merge()
+ .prepare_recovered_cleanup_merge()
.await
.map_err(|e| e.to_string())?
};