feat: add a separate retry loop on object contents download

docs/read_and_write.rst

@@ -700,6 +700,10 @@ These options apply to all object stores.
      - Description
    * - ``allow_http``
      - Allow non-TLS, i.e. non-HTTPS connections. Default, ``False``.
+   * - ``download_retry_count``
+     - Number of times to retry a download. Default, ``3``.  This limit is applied when
+       the HTTP request succeeds but the response is not fully downloaded, typically due
+       to a violation of ``request_timeout``.
    * - ``allow_invalid_certificates``
      - Skip certificate validation on https connections. Default, ``False``.
        Warning: This is insecure and should only be used for testing.

rust/lance-io/src/object_reader.rs

@@ -9,7 +9,7 @@ use bytes::Bytes;
 use deepsize::DeepSizeOf;
 use futures::future::BoxFuture;
 use lance_core::Result;
-use object_store::{path::Path, ObjectStore};
+use object_store::{path::Path, GetOptions, ObjectStore};
 use tokio::sync::OnceCell;
 use tracing::instrument;
 
@@ -28,6 +28,7 @@ pub struct CloudObjectReader {
     size: OnceCell<usize>,
 
     block_size: usize,
+    download_retry_count: usize,
 }
 
 impl DeepSizeOf for CloudObjectReader {
@@ -44,12 +45,14 @@ impl CloudObjectReader {
         path: Path,
         block_size: usize,
         known_size: Option<usize>,
+        download_retry_count: usize,
     ) -> Result<Self> {
         Ok(Self {
             object_store,
             path,
             size: OnceCell::new_with(known_size),
             block_size,
+            download_retry_count,
         })
     }
 
@@ -104,7 +107,40 @@ impl Reader for CloudObjectReader {
 
     #[instrument(level = "debug", skip(self))]
     async fn get_range(&self, range: Range<usize>) -> object_store::Result<Bytes> {
-        self.do_with_retry(|| self.object_store.get_range(&self.path, range.clone()))
-            .await
+        // We have a separate retry loop here.  This is because object_store does not
+        // attempt retries on downloads that fail during streaming of the response body.
+        //
+        // However, this failure is pretty common (e.g. timeout) and we want to retry in these
+        // situations.  In addition, we provide additional logging information in these
+        // failures cases.
+        let mut retries = self.download_retry_count;
+        loop {
+            let get_result = self
+                .do_with_retry(|| {
+                    let options = GetOptions {
+                        range: Some(range.clone().into()),
+                        ..Default::default()
+                    };
+                    self.object_store.get_opts(&self.path, options)
+                })
+                .await?;
+            match get_result.bytes().await {
+                Ok(bytes) => return Ok(bytes),
+                Err(err) => {
+                    if retries == 0 {
+                        log::warn!("Failed to download range {:?} from {} after {} attempts.  This may indicate that cloud storage is overloaded or your timeout settings are too restrictive.  Error details: {:?}", range, self.path, self.download_retry_count, err);
+                        return Err(err);
+                    }
+                    log::debug!(
+                        "Retrying range {:?} from {} (remaining retries: {}).  Error details: {:?}",
+                        range,
+                        self.path,
+                        retries,
+                        err
+                    );
+                    retries -= 1;
+                }
+            }
+        }
     }
 }

rust/lance-io/src/object_store.rs

@@ -46,6 +46,8 @@
 // Cloud disks often need many many threads to saturate the network
 pub const DEFAULT_CLOUD_IO_PARALLELISM: usize = 64;
 
+pub const DEFAULT_DOWNLOAD_RETRY_COUNT: usize = 3;
+
 #[async_trait]
 pub trait ObjectStoreExt {
     /// Returns true if the file exists.
@@ -100,6 +102,8 @@
     /// is true for object stores, but not for local filesystems.
     pub list_is_lexically_ordered: bool,
     io_parallelism: usize,
+    /// Number of times to retry a failed download
+    download_retry_count: usize,
 }
 
 impl DeepSizeOf for ObjectStore {
@@ -440,6 +444,7 @@
                 use_constant_size_upload_parts: false,
                 list_is_lexically_ordered: false,
                 io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM,
+                download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT,
             },
             Path::from_absolute_path(expanded_path.as_path())?,
         ))
@@ -466,6 +471,7 @@
             use_constant_size_upload_parts: false,
             list_is_lexically_ordered: false,
             io_parallelism: DEFAULT_LOCAL_IO_PARALLELISM,
+            download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT,
         }
     }
 
@@ -478,6 +484,7 @@
             use_constant_size_upload_parts: false,
             list_is_lexically_ordered: true,
             io_parallelism: get_num_compute_intensive_cpus(),
+            download_retry_count: DEFAULT_DOWNLOAD_RETRY_COUNT,
         }
     }
 
@@ -516,6 +523,7 @@
                 path.clone(),
                 self.block_size,
                 None,
+                self.download_retry_count,
             )?)),
         }
     }
@@ -533,6 +541,7 @@
                 path.clone(),
                 self.block_size,
                 Some(known_size),
+                self.download_retry_count,
             )?)),
         }
     }
@@ -620,7 +629,7 @@
     pub fn remove_stream<'a>(
         &'a self,
         locations: BoxStream<'a, Result<Path>>,
     ) -> BoxStream<Result<Path>> {
         self.inner
             .delete_stream(locations.err_into::<ObjectStoreError>().boxed())
             .err_into::<Error>()
@@ -641,6 +650,28 @@
         Ok(self.inner.head(path).await?.size)
     }
 }
+
+/// Options that can be set for multiple object stores
+#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy)]
+pub enum LanceConfigKey {
+    /// Number of times to retry a download that fails
+    DownloadRetryCount,
+}
+
+impl FromStr for LanceConfigKey {
+    type Err = Error;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s.to_ascii_lowercase().as_str() {
+            "download_retry_count" => Ok(Self::DownloadRetryCount),
+            _ => Err(Error::InvalidInput {
+                source: format!("Invalid LanceConfigKey: {}", s).into(),
+                location: location!(),
+            }),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct StorageOptions(pub HashMap<String, String>);
 
@@ -709,6 +740,15 @@
         })
     }
 
+    /// Number of times to retry a download that fails
+    pub fn download_retry_count(&self) -> usize {
+        self.0
+            .iter()
+            .find(|(key, _)| key.to_ascii_lowercase() == "download_retry_count")
+            .map(|(_, value)| value.parse::<usize>().unwrap_or(3))
+            .unwrap_or(3)
+    }
+
     /// Subset of options relevant for azure storage
     pub fn as_azure_options(&self) -> HashMap<AzureConfigKey, String> {
         self.0
@@ -755,6 +795,7 @@
     options: ObjectStoreParams,
 ) -> Result<ObjectStore> {
     let mut storage_options = StorageOptions(options.storage_options.clone().unwrap_or_default());
+    let download_retry_count = storage_options.download_retry_count();
     let mut url = ensure_table_uri(url)?;
     // Block size: On local file systems, we use 4KB block size. On cloud
     // object stores, we use 64KB block size. This is generally the largest
@@ -813,6 +854,7 @@
                 use_constant_size_upload_parts,
                 list_is_lexically_ordered: true,
                 io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
+                download_retry_count,
             })
         }
         "gs" => {
@@ -831,6 +873,7 @@
                 use_constant_size_upload_parts: false,
                 list_is_lexically_ordered: true,
                 io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
+                download_retry_count,
             })
         }
         "az" => {
@@ -845,6 +888,7 @@
                 use_constant_size_upload_parts: false,
                 list_is_lexically_ordered: true,
                 io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
+                download_retry_count,
             })
         }
         // we have a bypass logic to use `tokio::fs` directly to lower overhead
@@ -862,6 +906,7 @@
             use_constant_size_upload_parts: false,
             list_is_lexically_ordered: true,
             io_parallelism: get_num_compute_intensive_cpus(),
+            download_retry_count,
         }),
         unknown_scheme => {
             if let Some(provider) = registry.providers.get(unknown_scheme) {
@@ -878,6 +923,7 @@
 }
 
 impl ObjectStore {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         store: Arc<DynObjectStore>,
         location: Url,
@@ -886,6 +932,7 @@
         use_constant_size_upload_parts: bool,
         list_is_lexically_ordered: bool,
         io_parallelism: usize,
+        download_retry_count: usize,
     ) -> Self {
         let scheme = location.scheme();
         let block_size = block_size.unwrap_or_else(|| infer_block_size(scheme));
@@ -902,6 +949,7 @@
             use_constant_size_upload_parts,
             list_is_lexically_ordered,
             io_parallelism,
+            download_retry_count,
         }
     }
 }

rust/lance-io/src/scheduler.rs

@@ -660,7 +660,7 @@ mod tests {
     use tokio::{runtime::Handle, time::timeout};
     use url::Url;
 
-    use crate::testing::MockObjectStore;
+    use crate::{object_store::DEFAULT_DOWNLOAD_RETRY_COUNT, testing::MockObjectStore};
 
     use super::*;
 
@@ -743,6 +743,7 @@ mod tests {
             false,
             false,
             1,
+            DEFAULT_DOWNLOAD_RETRY_COUNT,
         ));
 
         let config = SchedulerConfig {
@@ -831,6 +832,7 @@ mod tests {
             false,
             false,
             1,
+            DEFAULT_DOWNLOAD_RETRY_COUNT,
         ));
 
         let config = SchedulerConfig {

rust/lance-io/src/utils.rs

@@ -183,7 +183,7 @@ mod tests {
 
     use crate::{
         object_reader::CloudObjectReader,
-        object_store::ObjectStore,
+        object_store::{ObjectStore, DEFAULT_DOWNLOAD_RETRY_COUNT},
         object_writer::ObjectWriter,
         traits::{ProtoStruct, WriteExt, Writer},
         utils::read_struct,
@@ -226,7 +226,9 @@ mod tests {
         assert_eq!(pos, 0);
         object_writer.shutdown().await.unwrap();
 
-        let object_reader = CloudObjectReader::new(store.inner, path, 1024, None).unwrap();
+        let object_reader =
+            CloudObjectReader::new(store.inner, path, 1024, None, DEFAULT_DOWNLOAD_RETRY_COUNT)
+                .unwrap();
         let actual: BytesWrapper = read_struct(&object_reader, pos).await.unwrap();
         assert_eq!(some_message, actual);
     }

rust/lance/src/dataset/builder.rs

@@ -4,7 +4,8 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
 
 use lance_file::datatypes::populate_schema_dictionary;
 use lance_io::object_store::{
-    ObjectStore, ObjectStoreParams, ObjectStoreRegistry, DEFAULT_CLOUD_IO_PARALLELISM,
+    ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptions,
+    DEFAULT_CLOUD_IO_PARALLELISM,
 };
 use lance_table::{
     format::Manifest,
@@ -220,6 +221,14 @@ impl DatasetBuilder {
             None => commit_handler_from_url(&self.table_uri, &Some(self.options.clone())).await,
         }?;
 
+        let storage_options = self
+            .options
+            .storage_options
+            .clone()
+            .map(StorageOptions::new)
+            .unwrap_or_default();
+        let download_retry_count = storage_options.download_retry_count();
+
         match &self.options.object_store {
             Some(store) => Ok((
                 ObjectStore::new(
@@ -232,6 +241,7 @@ impl DatasetBuilder {
                     // If user supplied an object store then we just assume it's probably
                     // cloud-like
                     DEFAULT_CLOUD_IO_PARALLELISM,
+                    download_retry_count,
                 ),
                 Path::from(store.1.path()),
                 commit_handler,