fix: preserve completed read snapshots during WAL merges

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

@@ -2410,6 +2410,86 @@ mod tests {
     use serde_json::json;
     use tempfile::TempDir;
 
+    #[tokio::test]
+    async fn list_retries_wrapped_missing_generation_errors() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let dir = tempfile::tempdir().unwrap();
+        let store = RolloutStore::open(dir.path().to_str().unwrap())
+            .await
+            .unwrap();
+        let attempts = AtomicUsize::new(0);
+        let result = store
+            .base
+            .read_consistent(ListSource::All, |_, _| {
+                let first = attempts.fetch_add(1, Ordering::SeqCst) == 0;
+                async move {
+                    if first {
+                        // Same nested error surface observed in Azure while
+                        // GC removed a generation's data file during a take.
+                        Err(LanceError::io(
+                            "Wrapped error: Not found: _mem_wal/shard/gen_9/data/file.lance",
+                        ))
+                    } else {
+                        Ok(())
+                    }
+                }
+            })
+            .await;
+        assert!(result.is_ok(), "{result:?}");
+        assert_eq!(attempts.load(Ordering::SeqCst), 2);
+    }
+
+    #[tokio::test]
+    async fn completed_list_snapshot_survives_later_merge_commits() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        let dir = tempfile::tempdir().unwrap();
+        let uri = dir.path().to_str().unwrap();
+        let writer = RolloutStore::open_with_options(
+            uri,
+            RolloutStoreOptions {
+                shard_id: Some("snapshot-writer".into()),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap();
+        writer.add(&[assistant_record("initial")]).await.unwrap();
+        writer.flush().await.unwrap();
+        let reader = RolloutStore::open(uri).await.unwrap();
+        let writer = tokio::sync::Mutex::new(writer);
+        let attempts = AtomicUsize::new(0);
+
+        let rows = reader
+            .base
+            .read_consistent(ListSource::All, |dataset, snapshots| {
+                let reader = &reader;
+                let writer = &writer;
+                let attempts = &attempts;
+                async move {
+                    let rows = reader.list_non_blob_snapshot(dataset, snapshots).await?;
+                    // Publish a later write only after this read has collected
+                    // its complete snapshot. It must not invalidate those rows.
+                    let sequence = attempts.fetch_add(1, Ordering::SeqCst);
+                    let mut writer = writer.lock().await;
+                    writer
+                        .add(&[assistant_record(&format!("later-{sequence}"))])
+                        .await
+                        .unwrap();
+                    writer.flush().await.unwrap();
+                    writer.cleanup_own_shard().await.unwrap();
+                    Ok(rows)
+                }
+            })
+            .await
+            .unwrap();
+
+        assert_eq!(attempts.load(Ordering::SeqCst), 1);
+        assert_eq!(rows.len(), 1);
+        assert_eq!(rows[0].id, "initial");
+    }
+
     #[tokio::test]
     async fn list_retries_when_a_generation_is_deleted_during_the_read() {
         use std::sync::atomic::{AtomicUsize, Ordering};

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

@@ -2282,10 +2282,12 @@ impl StorageBase {
 
     /// Read a fresh base/WAL view and retry if a merge invalidates it. A merge
     /// commits base before draining/deleting its generations, so an unchanged
-    /// base version brackets a successful read. Retry missing generations even
-    /// when the base version is unchanged. Never return a partial result after
-    /// a generation disappears. Ordinary WAL appends do not change base and do
-    /// not force retries. Pinned reads use only their selected base version.
+    /// base version must bracket the WAL manifest capture. Validate that pair
+    /// before reading payloads. A later commit does not invalidate a completed
+    /// read of that snapshot. Retry missing generations even when the base
+    /// version is unchanged; never return a partial result after one disappears.
+    /// Ordinary WAL appends do not force retries. Pinned reads use only their
+    /// selected base version.
     pub async fn read_consistent<T, F, Fut>(
         &self,
         source: ListSource,
@@ -2306,6 +2308,16 @@ impl StorageBase {
             } else {
                 self.wal_shard_snapshots().await?
             };
+            if !self.is_version_pinned() {
+                let mut latest = dataset.clone();
+                latest.checkout_latest().await?;
+                if latest.manifest.version != version {
+                    // A merge could have drained WAL after we opened base.
+                    // Retry this metadata pair before scanning any payloads,
+                    // otherwise the view could omit its newly merged rows.
+                    continue;
+                }
+            }
             let result = read(dataset.clone(), snapshots).await;
             if self.is_version_pinned() {
                 return result;
@@ -2319,10 +2331,12 @@ impl StorageBase {
                 // the base version. Refresh both snapshots even in that case.
                 continue;
             }
-            dataset.checkout_latest().await?;
-            if dataset.manifest.version == version {
-                return result;
-            }
+            // The immutable base/WAL view was complete at capture. A commit
+            // after capture may force a missing-file retry, but cannot make
+            // successfully read rows disappear. Requiring the base to stay
+            // unchanged throughout payload decoding starves reads under
+            // continuous merging and repeats all their decoding work.
+            return result;
         }
         Err(LanceError::io(
             "dataset changed during read; retry the request",
@@ -2531,21 +2545,19 @@ pub(crate) fn is_fenced_error(err: &LanceError) -> bool {
 
 /// Whether a Lance error means the target dataset/path does not exist.
 ///
-/// Used by WAL read fallbacks to skip a flushed generation that a concurrent
-/// merge drained from the shard manifest and deleted between the manifest
-/// snapshot and the open — the row's data is already in the base table, so
-/// skipping is safe and avoids a transient failure. Both the structured
-/// not-found variants and the object-store "not found" surface are matched; the
-/// latter is text-based because it arrives as a generic wrapped IO error.
+/// WAL reads can race with a merge draining and deleting a generation. A read
+/// must then refresh its base/WAL view, not return a partial result. Both the
+/// structured variants and the object-store message are matched; the latter
+/// can arrive as a generic wrapped IO error with different capitalization.
 pub(crate) fn is_not_found_error(err: &LanceError) -> bool {
     if matches!(
         err,
         LanceError::DatasetNotFound { .. } | LanceError::NotFound { .. }
     ) {
         return true;
     }
-    let text = err.to_string();
-    text.contains("was not found") || text.contains("not found") || text.contains("NotFound")
+    let text = err.to_string().to_ascii_lowercase();
+    text.contains("not found") || text.contains("notfound")
 }
 
 /// Keep the newest prepared row for each merge key.