fix(cleanup): add retry logic to manifest reads and per-manifest error tolerance

rust/lance-io/src/object_store.rs

@@ -556,6 +556,11 @@ impl ObjectStore {
             .unwrap_or(self.io_parallelism)
     }
 
+    /// Number of times to retry a download that fails during body streaming.
+    pub fn download_retry_count(&self) -> usize {
+        self.download_retry_count
+    }
+
     /// Get the IO tracker for this object store
     ///
     /// The IO tracker can be used to get statistics about read/write operations

rust/lance-table/src/io/manifest.rs

@@ -26,15 +26,59 @@ use crate::format::{DataStorageFormat, IndexMetadata, MAGIC, Manifest, Transacti
 
 use super::commit::ManifestLocation;
 
+/// Wrapper around `object_store.inner.get_range()` that retries on transient
+/// errors (e.g. mid-stream body download failures from S3). The upstream
+/// `object_store` crate only retries on the initial HTTP request, not on
+/// failures during body streaming. This mirrors the retry logic used by
+/// `CloudObjectReader` in `lance-io`.
+async fn get_range_with_retry(
+    object_store: &ObjectStore,
+    path: &Path,
+    range: Range<u64>,
+    retry_count: usize,
+) -> Result<Bytes> {
+    let mut retries = retry_count;
+    loop {
+        match object_store.inner.get_range(path, range.clone()).await {
+            Ok(bytes) => return Ok(bytes),
+            Err(err) => {
+                if retries == 0 {
+                    log::warn!(
+                        "Failed to read manifest range {:?} from {} after {} attempts. \
+                         Error details: {:?}",
+                        range,
+                        path,
+                        retry_count + 1,
+                        err,
+                    );
+                    return Err(err.into());
+                }
+                log::info!(
+                    "Retrying manifest read range {:?} from {} \
+                     (remaining retries: {}). Error: {:?}",
+                    range,
+                    path,
+                    retries,
+                    err,
+                );
+                retries -= 1;
+            }
+        }
+    }
+}
+
 /// Read Manifest on URI.
 ///
 /// This only reads manifest files. It does not read data files.
+/// Retries transient S3 errors (including mid-stream body failures)
+/// up to `object_store.download_retry_count()` times per range request.
 #[instrument(level = "debug", skip(object_store))]
 pub async fn read_manifest(
     object_store: &ObjectStore,
     path: &Path,
     known_size: Option<u64>,
 ) -> Result<Manifest> {
+    let retry_count = object_store.download_retry_count();
     let file_size = if let Some(known_size) = known_size {
         known_size
     } else {
@@ -46,7 +90,7 @@ pub async fn read_manifest(
         start: initial_start,
         end: file_size,
     };
-    let buf = object_store.inner.get_range(path, range).await?;
+    let buf = get_range_with_retry(object_store, path, range, retry_count).await?;
 
     // In case of corruption, the known_size might be wrong. We can retry without
     // the size to be more robust.
@@ -75,18 +119,15 @@ pub async fn read_manifest(
     } else {
         // The prefetch only captured part of the manifest. We need to make an
         // additional range request to read the remainder.
-        let mut buf2: BytesMut = object_store
-            .inner
-            .get_range(
-                path,
-                Range {
-                    start: manifest_pos as u64,
-                    end: file_size - PREFETCH_SIZE,
-                },
-            )
-            .await?
-            .into_iter()
-            .collect();
+        let remainder_range = Range {
+            start: manifest_pos as u64,
+            end: file_size - PREFETCH_SIZE,
+        };
+        let mut buf2: BytesMut =
+            get_range_with_retry(object_store, path, remainder_range, retry_count)
+                .await?
+                .into_iter()
+                .collect();
         buf2.extend_from_slice(&buf);
         buf2.freeze()
     };

rust/lance/src/dataset/cleanup.rs

@@ -194,19 +194,116 @@ impl<'a> CleanupTask<'a> {
         Ok(final_stats)
     }
 
+    /// Maximum number of retry rounds for manifests that fail during
+    /// `process_manifests`. Each round re-attempts all manifests that failed
+    /// in the previous round.
+    const MANIFEST_RETRY_ROUNDS: usize = 3;
+
     #[instrument(level = "debug", skip_all)]
     async fn process_manifests(
         &'a self,
         tagged_versions: &HashSet<u64>,
     ) -> Result<CleanupInspection> {
         let inspection = Mutex::new(CleanupInspection::default());
-        self.dataset
+        let io_parallelism = self.dataset.object_store.io_parallelism();
+
+        // Collect all manifest locations first so we can retry failures.
+        let locations: Vec<ManifestLocation> = self
+            .dataset
             .commit_handler
             .list_manifest_locations(&self.dataset.base, &self.dataset.object_store, false)
-            .try_for_each_concurrent(self.dataset.object_store.io_parallelism(), |location| {
-                self.process_manifest_file(location, &inspection, tagged_versions)
-            })
+            .try_collect()
             .await?;
+
+        let total = locations.len();
+        info!(
+            "Processing {} manifest files (concurrency={})",
+            total, io_parallelism
+        );
+
+        let failed: Mutex<Vec<(ManifestLocation, String)>> = Mutex::new(Vec::new());
+
+        // First pass: process all manifests, collecting failures instead of aborting.
+        stream::iter(locations)
+            .for_each_concurrent(io_parallelism, |location| {
+                let failed_ref = &failed;
+                let inspection_ref = &inspection;
+                async move {
+                    if let Err(err) = self
+                        .process_manifest_file(location.clone(), inspection_ref, tagged_versions)
+                        .await
+                    {
+                        log::warn!("Failed to process manifest {}: {:?}", location.path, err,);
+                        failed_ref
+                            .lock()
+                            .unwrap()
+                            .push((location, format!("{:?}", err)));
+                    }
+                }
+            })
+            .await;
+
+        // Retry rounds for failed manifests.
+        let mut pending_failures = std::mem::take(&mut *failed.lock().unwrap());
+        for round in 1..=Self::MANIFEST_RETRY_ROUNDS {
+            if pending_failures.is_empty() {
+                break;
+            }
+            let retry_count = pending_failures.len();
+            info!(
+                "Retry round {}/{}: re-processing {} failed manifest(s)",
+                round,
+                Self::MANIFEST_RETRY_ROUNDS,
+                retry_count,
+            );
+            let round_failed: Mutex<Vec<(ManifestLocation, String)>> = Mutex::new(Vec::new());
+
+            stream::iter(pending_failures.into_iter().map(|(loc, _)| loc))
+                .for_each_concurrent(io_parallelism, |location| {
+                    let round_failed_ref = &round_failed;
+                    let inspection_ref = &inspection;
+                    async move {
+                        if let Err(err) = self
+                            .process_manifest_file(
+                                location.clone(),
+                                inspection_ref,
+                                tagged_versions,
+                            )
+                            .await
+                        {
+                            log::warn!(
+                                "Retry round {}: failed to process manifest {}: {:?}",
+                                round,
+                                location.path,
+                                err,
+                            );
+                            round_failed_ref
+                                .lock()
+                                .unwrap()
+                                .push((location, format!("{:?}", err)));
+                        }
+                    }
+                })
+                .await;
+
+            pending_failures = std::mem::take(&mut *round_failed.lock().unwrap());
+        }
+
+        if !pending_failures.is_empty() {
+            let failed_paths: Vec<String> = pending_failures
+                .iter()
+                .map(|(loc, err)| format!("{}: {}", loc.path, err))
+                .collect();
+            return Err(Error::io(format!(
+                "Failed to read {} manifest(s) after {} retry rounds: [{}]",
+                pending_failures.len(),
+                Self::MANIFEST_RETRY_ROUNDS,
+                failed_paths.join(", "),
+            )));
+        }
+
+        let succeeded = total;
+        info!("Successfully processed all {} manifests", succeeded);
         Ok(inspection.into_inner().unwrap())
     }