diff --git a/crates/lance-context-core/src/rollout_store.rs b/crates/lance-context-core/src/rollout_store.rs index fbbfbfd..d0bd4a7 100644 --- a/crates/lance-context-core/src/rollout_store.rs +++ b/crates/lance-context-core/src/rollout_store.rs @@ -425,6 +425,15 @@ pub struct RolloutStoreOptions { /// typed read APIs built on top of the base's scanners. pub struct RolloutStore { base: StorageBase, + #[cfg(test)] + blob_read_pause: std::sync::Mutex>>, +} + +#[cfg(test)] +#[derive(Default)] +struct BlobReadPause { + opened: tokio::sync::Notify, + resume: tokio::sync::Notify, } impl RolloutStore { @@ -522,7 +531,11 @@ impl RolloutStore { // Encoding follows the base schema. Evolve before accepting writes, // otherwise values for newly supported fields are silently discarded. base.ensure_latest_schema().await?; - Ok(Self { base }) + Ok(Self { + base, + #[cfg(test)] + blob_read_pause: Default::default(), + }) } /// URI of the underlying Lance dataset. #[must_use] @@ -1592,15 +1605,42 @@ impl RolloutStore { /// both), so base-first cannot shadow a newer WAL copy and the parallel /// fallback needs no ordering. /// - /// The fallback is tolerant of a generation that a concurrent merge drained - /// and deleted between snapshot and open: such an open fails with a - /// not-found error, which is skipped (the row's data is already covered by - /// the base table we checked first). This removes the transient 500s that - /// the previous fail-fast fallback surfaced under concurrent auto-merge. + /// A concurrent merge can delete a generation after its dataset is open, + /// while its IDs or payload are being read. Missing-file errors trigger at + /// most two retries from a refreshed base snapshot. Persistent failures + /// remain errors, and explicitly pinned reads never advance versions. + /// Generations missing at open are skipped; a final miss is also retried + /// by the HTTP caller after refreshing its store. pub async fn get_blob(&self, id: &str) -> LanceResult>> { + let mut base = self.base.dataset.clone(); + let mut retries = 0; + loop { + match self.get_blob_from_snapshot(&base, id).await { + Err(error) + if is_not_found_error(&error) + && !self.base.is_version_pinned() + && retries < 2 => + { + // Merge can delete a generation after its manifest was + // opened, during the ID scan or payload take. Re-read from + // the committed base instead of turning that race into a + // 500. Persistent missing files still fail after 3 tries. + retries += 1; + base.checkout_latest().await?; + } + result => return result, + } + } + } + + async fn get_blob_from_snapshot( + &self, + base: &Dataset, + id: &str, + ) -> LanceResult>> { // Base-table-first: an already-merged row is found here with no MemWAL // manifest reads and no per-generation opens. - if let Some(payload) = Self::get_blob_from_dataset(&self.base.dataset, id).await? { + if let Some(payload) = Self::get_blob_from_dataset(base, id).await? { return Ok(payload); } @@ -1610,9 +1650,9 @@ impl RolloutStore { // Refresh a clone so shared-lock readers never mutate the resident // dataset, and explicitly pinned historical reads remain historical. if !self.base.is_version_pinned() { - let mut latest = self.base.dataset.clone(); + let mut latest = base.clone(); latest.checkout_latest().await?; - if latest.manifest.version != self.base.dataset.manifest.version { + if latest.manifest.version != base.manifest.version { if let Some(payload) = Self::get_blob_from_dataset(&latest, id).await? { return Ok(payload); } @@ -1635,10 +1675,20 @@ impl RolloutStore { let mut hits = stream::iter(uris) .map(|uri| async move { match self.open_flushed_dataset(&uri).await { - Ok(dataset) => Self::get_blob_from_dataset(&dataset, id).await, + Ok(dataset) => { + #[cfg(test)] + { + let pause = self.blob_read_pause.lock().unwrap().take(); + if let Some(pause) = pause { + pause.opened.notify_one(); + pause.resume.notified().await; + } + } + Self::get_blob_from_dataset(&dataset, id).await + } // A generation drained + deleted by a concurrent merge - // between snapshot and open: its rows are already in the base - // table (checked above), so skip it rather than 500. + // between snapshot and open: its rows are in a committed + // base version. The HTTP caller refreshes on a final miss. Err(err) if is_not_found_error(&err) => Ok(None), Err(err) => Err(err), } @@ -3210,6 +3260,7 @@ mod tests { // Model a writer from the old release without running today's open- // time migration, so the fixtures really contain the legacy schema. let mut store = RolloutStore { + blob_read_pause: Default::default(), base: StorageBase::open( &uri, StorageBaseOptions { @@ -6259,6 +6310,106 @@ mod tests { }); } + #[tokio::test] + async fn blob_lookup_recovers_when_open_wal_is_deleted() { + for requested_id in ["target", "mid-absent"] { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_str().unwrap(); + let open = |shard: &str| { + RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some(shard.to_owned()), + ..Default::default() + }, + ) + }; + let mut writer = open("race-writer").await.unwrap(); + let reader = open("race-reader").await.unwrap(); + let original_version = reader.version(); + let bytes = b"blob-moved-after-generation-open"; + writer + .add(&[ + assistant_record("a"), + artifact_record("target", bytes), + assistant_record("z"), + ]) + .await + .unwrap(); + writer.flush().await.unwrap(); + let pause = Arc::new(BlobReadPause::default()); + *reader.blob_read_pause.lock().unwrap() = Some(pause.clone()); + let reading = tokio::spawn(async move { + let result = reader.get_blob(requested_id).await; + (result, reader.version()) + }); + tokio::time::timeout(std::time::Duration::from_secs(30), pause.opened.notified()) + .await + .unwrap(); + // The generation dataset is already open, but its payload has not + // been read. Commit it to the base and delete its source files. + assert_eq!(writer.cleanup_own_shard().await.unwrap(), 1); + assert_eq!(flushed_generation_count(&writer).await, 0); + pause.resume.notify_one(); + let (result, version) = + tokio::time::timeout(std::time::Duration::from_secs(30), reading) + .await + .unwrap() + .unwrap(); + assert_eq!( + result.unwrap().as_deref(), + (requested_id == "target").then_some(bytes.as_slice()), + ); + assert_eq!(version, original_version); + } + } + + #[tokio::test] + async fn blob_lookup_does_not_hide_persistent_missing_data() { + let dir = TempDir::new().unwrap(); + let uri = dir.path().to_str().unwrap(); + let writer = RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some("missing-writer".to_owned()), + ..Default::default() + }, + ) + .await + .unwrap(); + writer + .add(&[artifact_record("target", b"not-merged")]) + .await + .unwrap(); + writer.flush().await.unwrap(); + let reader = RolloutStore::open_with_options( + uri, + RolloutStoreOptions { + shard_id: Some("missing-reader".to_owned()), + ..Default::default() + }, + ) + .await + .unwrap(); + let snapshots = writer.wal_shard_snapshots().await.unwrap(); + let generation_uri = + writer.flushed_generation_uri(snapshots[0].shard_id, &snapshots[0].sstables[0].path); + // Retain the generation manifest. Unlike a completed merge, no base + // commit can recover this row: bounded retries must preserve the error. + let path = generation_uri + .strip_prefix("file://") + .unwrap_or(&generation_uri); + std::fs::remove_dir_all(std::path::Path::new(path).join("data")).unwrap(); + let error = tokio::time::timeout( + std::time::Duration::from_secs(30), + reader.get_blob("target"), + ) + .await + .unwrap() + .unwrap_err(); + assert!(is_not_found_error(&error), "unexpected error: {error}"); + } + #[tokio::test] async fn blob_lookup_refreshes_base_before_opening_unrelated_wal() { let dir = TempDir::new().unwrap();