fix: robust against body read failures in list endpoint (#3644)

Addressing errors like:

```
LanceError(IO): Generic GCS error: Error getting list response body: error decoding response body, /src/lance/rust/lance-table/src/io/commit.rs:216:17
```

Author: wjones127

rust/lance-io/src/object_store.rs

@@ -6,6 +6,7 @@
 use std::collections::HashMap;
 use std::ops::Range;
 use std::path::PathBuf;
+use std::pin::Pin;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::{Duration, SystemTime};
@@ -16,9 +17,11 @@ use aws_credential_types::provider::ProvideCredentials;
 use bytes::Bytes;
 use chrono::{DateTime, Utc};
 use deepsize::DeepSizeOf;
+use futures::Stream;
 use futures::{future, stream::BoxStream, StreamExt, TryStreamExt};
 use lance_core::utils::parse::str_is_truthy;
 use lance_core::utils::tokio::get_num_compute_intensive_cpus;
+use list_retry::ListRetryStream;
 use object_store::aws::{
     AmazonS3ConfigKey, AwsCredential as ObjectStoreAwsCredential, AwsCredentialProvider,
 };
@@ -37,6 +40,7 @@ use tokio::sync::RwLock;
 use url::Url;
 
 use super::local::LocalObjectReader;
+mod list_retry;
 mod tracing;
 use self::tracing::ObjectStoreTracingExt;
 use crate::object_writer::WriteResult;
@@ -620,6 +624,13 @@ impl ObjectStore {
             .collect())
     }
 
+    pub fn list(
+        &self,
+        path: Option<Path>,
+    ) -> Pin<Box<dyn Stream<Item = Result<ObjectMeta>> + Send>> {
+        Box::pin(ListRetryStream::new(self.inner.clone(), path, 5).map(|m| m.map_err(|e| e.into())))
+    }
+
     /// Read all files (start from base directory) recursively
     ///
     /// unmodified_since can be specified to only return files that have not been modified since the given time.

rust/lance-io/src/object_store/list_retry.rs

@@ -0,0 +1,155 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::{pin::Pin, sync::Arc, task::Poll};
+
+use futures::{Stream, StreamExt};
+use object_store::{path::Path, ObjectMeta, ObjectStore};
+use tokio::task::JoinHandle;
+
+/// ObjectStore::list() and ObjectStore::list_with_offset() return a stream
+/// where the lifetime is tied to the object store. This makes it hard to wrap.
+/// So here we put it inside a tokio task and return a channel receiver.
+struct StaticListStream {
+    rx: tokio::sync::mpsc::Receiver<Result<ObjectMeta, object_store::Error>>,
+    handle: JoinHandle<()>,
+}
+
+impl StaticListStream {
+    fn new(object_store: Arc<dyn ObjectStore>, prefix: Option<Path>, offset: Option<Path>) -> Self {
+        let (tx, rx) = tokio::sync::mpsc::channel(100);
+        let handle = tokio::spawn(async move {
+            let mut stream = if let Some(offset) = offset {
+                object_store.list_with_offset(prefix.as_ref(), &offset)
+            } else {
+                object_store.list(prefix.as_ref())
+            };
+            while let Some(item) = stream.next().await {
+                if tx.send(item).await.is_err() {
+                    break;
+                }
+            }
+        });
+        Self { rx, handle }
+    }
+
+    fn abort(&self) {
+        self.handle.abort();
+    }
+}
+
+impl Stream for StaticListStream {
+    type Item = Result<ObjectMeta, object_store::Error>;
+
+    fn poll_next(
+        self: Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        match this.rx.poll_recv(cx) {
+            Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending if this.handle.is_finished() => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+/// A stream that does outer retries on list operations.
+///
+/// This is to handle request responses that ObjectStore doesn't handle, such as
+/// the error `error decoding response body` from queries to GCS.
+pub struct ListRetryStream {
+    object_store: Arc<dyn ObjectStore>,
+    current_stream: StaticListStream,
+    prefix: Option<Path>,
+    last_successful_key: Option<Path>,
+    max_retries: usize,
+    current_retries: usize,
+}
+
+impl ListRetryStream {
+    pub fn new(
+        object_store: Arc<dyn ObjectStore>,
+        prefix: Option<Path>,
+        max_retries: usize,
+    ) -> Self {
+        let current_stream = StaticListStream::new(object_store.clone(), prefix.clone(), None);
+        Self {
+            object_store,
+            current_stream,
+            prefix,
+            last_successful_key: None,
+            max_retries,
+            current_retries: 0,
+        }
+    }
+
+    fn is_retryable(error: &object_store::Error) -> bool {
+        !matches!(
+            error,
+            object_store::Error::NotFound { .. }
+                | object_store::Error::InvalidPath { .. }
+                | object_store::Error::NotSupported { .. }
+                | object_store::Error::NotImplemented
+        )
+    }
+}
+
+impl Stream for ListRetryStream {
+    type Item = Result<ObjectMeta, object_store::Error>;
+
+    fn poll_next(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        loop {
+            match this.current_stream.poll_next_unpin(cx) {
+                Poll::Ready(Some(Ok(meta))) => {
+                    this.last_successful_key = Some(meta.location.clone());
+                    return Poll::Ready(Some(Ok(meta)));
+                }
+                Poll::Ready(None) => {
+                    // If the stream is done, return None
+                    return Poll::Ready(None);
+                }
+                Poll::Ready(Some(Err(error))) if Self::is_retryable(&error) => {
+                    if this.current_retries < this.max_retries {
+                        this.current_retries += 1;
+
+                        this.current_stream.abort();
+                        this.current_stream = StaticListStream::new(
+                            this.object_store.clone(),
+                            this.prefix.clone(),
+                            this.last_successful_key.clone(),
+                        );
+
+                        continue;
+                    } else {
+                        return Poll::Ready(Some(Err(error)));
+                    }
+                }
+                Poll::Ready(Some(Err(error))) => {
+                    return Poll::Ready(Some(Err(error)));
+                }
+                Poll::Pending => {
+                    return Poll::Pending;
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn assert_send<T: Send>() {}
+
+    #[test]
+    fn test_list_retry_stream_send() {
+        // Ensure that ListRetryStream is Send
+        assert_send::<ListRetryStream>();
+    }
+}

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

@@ -236,7 +236,7 @@ async fn current_manifest_path(
         }
     }
 
-    let manifest_files = object_store.inner.list(Some(&base.child(VERSIONS_DIR)));
+    let manifest_files = object_store.list(Some(base.child(VERSIONS_DIR)));
 
     let mut valid_manifests = manifest_files.try_filter_map(|res| {
         if let Some(scheme) = ManifestNamingScheme::detect_scheme(res.location.filename().unwrap())