fix: keep blob misses under shared store locks

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

@@ -434,6 +434,7 @@ pub struct RolloutStore {
 struct BlobReadPause {
     opened: tokio::sync::Notify,
     resume: tokio::sync::Notify,
+    before_wal_snapshot: bool,
 }
 
 impl RolloutStore {
@@ -1600,17 +1601,18 @@ impl RolloutStore {
     /// **concurrently** (bounded by [`DEFAULT_OBSERVE_CONCURRENCY`]) with the
     /// first match winning, rather than one-at-a-time newest-first.
     ///
-    /// Because rollout rows are immutable and an `id` is never re-appended, a
-    /// row lives in exactly one place (the base table *or* one generation, never
-    /// both), so base-first cannot shadow a newer WAL copy and the parallel
+    /// Because rollout rows are immutable and an `id` is never re-appended,
+    /// any WAL copy still visible after a merge commits is identical to the
+    /// base row. Base-first cannot shadow a newer value, and the parallel
     /// fallback needs no ordering.
     ///
     /// 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.
+    /// After a WAL miss, check the latest base once more to cover a generation
+    /// drained between the earlier base check and WAL discovery. All refreshes
+    /// use clones, so even misses can complete under a shared store guard.
     pub async fn get_blob(&self, id: &str) -> LanceResult<Option<Vec<u8>>> {
         let mut base = self.base.dataset.clone();
         let mut retries = 0;
@@ -1649,8 +1651,8 @@ impl RolloutStore {
         // (or refusing the request because those generations exceed the cap).
         // Refresh a clone so shared-lock readers never mutate the resident
         // dataset, and explicitly pinned historical reads remain historical.
+        let mut latest = base.clone();
         if !self.base.is_version_pinned() {
-            let mut latest = base.clone();
             latest.checkout_latest().await?;
             if latest.manifest.version != base.manifest.version {
                 if let Some(payload) = Self::get_blob_from_dataset(&latest, id).await? {
@@ -1659,6 +1661,25 @@ impl RolloutStore {
             }
         }
 
+        #[cfg(test)]
+        {
+            let pause = {
+                let mut guard = self.blob_read_pause.lock().unwrap();
+                if guard
+                    .as_ref()
+                    .is_some_and(|pause| pause.before_wal_snapshot)
+                {
+                    guard.take()
+                } else {
+                    None
+                }
+            };
+            if let Some(pause) = pause {
+                pause.opened.notify_one();
+                pause.resume.notified().await;
+            }
+        }
+
         // Base miss: fall back to the flushed generations. Open them with
         // bounded concurrency and take the first hit; tolerate a generation that
         // was concurrently merged away (NotFound) instead of failing the request.
@@ -1688,7 +1709,7 @@ impl RolloutStore {
                     }
                     // A generation drained + deleted by a concurrent merge
                     // between snapshot and open: its rows are in a committed
-                    // base version. The HTTP caller refreshes on a final miss.
+                    // base version, covered by the final base check below.
                     Err(err) if is_not_found_error(&err) => Ok(None),
                     Err(err) => Err(err),
                 }
@@ -1703,6 +1724,20 @@ impl RolloutStore {
             }
         }
 
+        if !self.base.is_version_pinned() {
+            let checked_version = latest.manifest.version;
+            latest.checkout_latest().await?;
+            if latest.manifest.version != checked_version {
+                // Any generation already drained during the completed WAL
+                // scan was committed before this checkout. A final base check
+                // therefore covers it without an exclusive resident refresh or
+                // waiting for merging to stop. Concurrent newer appends need
+                // not be part of this read's captured view.
+                return Self::get_blob_from_dataset(&latest, id)
+                    .await
+                    .map(Option::flatten);
+            }
+        }
         Ok(None)
     }
 
@@ -6364,6 +6399,60 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn blob_lookup_recovers_when_wal_drains_before_discovery() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        let mut writer = RolloutStore::open_with_options(
+            uri,
+            RolloutStoreOptions {
+                shard_id: Some("snapshot-writer".to_owned()),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+        let reader = RolloutStore::open_with_options(
+            uri,
+            RolloutStoreOptions {
+                shard_id: Some("snapshot-reader".to_owned()),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+        let original_version = reader.version();
+        let bytes = b"present-before-read-starts";
+        writer
+            .add(&[artifact_record("target", bytes)])
+            .await
+            .unwrap();
+        writer.flush().await.unwrap();
+        let pause = Arc::new(BlobReadPause {
+            before_wal_snapshot: true,
+            ..Default::default()
+        });
+        *reader.blob_read_pause.lock().unwrap() = Some(pause.clone());
+        let reading = tokio::spawn(async move {
+            let result = reader.get_blob("target").await;
+            (result, reader.version())
+        });
+        tokio::time::timeout(std::time::Duration::from_secs(30), pause.opened.notified())
+            .await
+            .unwrap();
+        // The row existed before the read. Move it after the base check but
+        // before WAL discovery, so opening the WAL returns no generations.
+        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(), 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();

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

@@ -557,25 +557,14 @@ async fn get_rollout_refreshing_on_miss(
     store.get_by_id(id).await.map_err(AppError::from_lance)
 }
 
-/// Blob lookup can use a cloned latest base under the shared lock. Retain the
-/// exclusive refresh/retry for a final miss, including a merge racing that
-/// cloned snapshot; successful payload hits never upgrade the store lock.
-async fn get_rollout_blob_refreshing_on_miss(
+/// Core blob lookup refreshes cloned base snapshots around WAL fallback and
+/// retries missing-file races. Even a final miss stays under the shared guard,
+/// so it does not queue an exclusive refresh behind concurrent append/flush.
+async fn get_rollout_blob_shared(
     store_lock: &RwLock<RolloutStore>,
     id: &str,
 ) -> Result<Option<Vec<u8>>, AppError> {
-    {
-        let store = store_lock.read().await;
-        let payload = store.get_blob(id).await.map_err(AppError::from_lance)?;
-        if payload.is_some() || store.is_version_pinned() {
-            return Ok(payload);
-        }
-    }
-
-    let mut store = store_lock.write().await;
-    if !store.is_version_pinned() {
-        store.refresh_latest().await.map_err(AppError::from_lance)?;
-    }
+    let store = store_lock.read().await;
     store.get_blob(id).await.map_err(AppError::from_lance)
 }
 
@@ -606,11 +595,8 @@ pub async fn fetch_rollout_blob(
 ) -> Result<Response, AppError> {
     let store_lock = state.get_or_open_rollout_store(&name).await?;
 
-    let (bytes, reservation) = load_blob_with_budget(
-        &state,
-        get_rollout_blob_refreshing_on_miss(&store_lock, &id),
-    )
-    .await?;
+    let (bytes, reservation) =
+        load_blob_with_budget(&state, get_rollout_blob_shared(&store_lock, &id)).await?;
     let bytes =
         bytes.ok_or_else(|| AppError::NotFound(format!("Rollout '{}' has no payload", id)))?;
 
@@ -803,6 +789,31 @@ mod tests {
     use super::*;
     use crate::state::AppState;
 
+    #[tokio::test]
+    async fn blob_miss_does_not_require_exclusive_store_lock() {
+        let dir = TempDir::new().unwrap();
+        let store = RolloutStore::open_with_options(
+            dir.path().to_str().unwrap(),
+            RolloutStoreOptions::default(),
+        )
+        .await
+        .unwrap();
+        let store_lock = RwLock::new(store);
+        // Appends hold a shared guard while adding/flushing. A missing blob
+        // must finish without upgrading to an exclusive guard behind them.
+        let concurrent_request = store_lock.read().await;
+        let version = concurrent_request.version();
+        let payload = tokio::time::timeout(
+            std::time::Duration::from_secs(2),
+            get_rollout_blob_shared(&store_lock, "absent"),
+        )
+        .await
+        .expect("a missing blob must not wait for another shared store guard")
+        .unwrap();
+        assert!(payload.is_none());
+        assert_eq!(concurrent_request.version(), version);
+    }
+
     #[test]
     fn flush_query_param_is_parsed_conservatively() {
         // Opt-in only: the default (and any non-affirmative spelling) must keep