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())? };