feat: add the s3 retry config options for storage option (#3268)

Add `client_max_retries` and `client_retry_timeout` of `RetryConfig` for S3 client. If there are some server error of object store server, the `object store` module of `arrow-rs` will retry `client_max_retries` times and also the total execute time is not over `client_retry_timeout`. Closes #3182 --------- Co-authored-by: Will Jones <[email protected]>
lancedb · Dec 18, 2024 · 6cd6ae8 · 6cd6ae8
1 parent a07717a
commit 6cd6ae8
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 2 deletions.
diff --git a/docs/read_and_write.rst b/docs/read_and_write.rst
@@ -727,7 +727,10 @@ These options apply to all object stores.
        and IP masks. Any subdomain of the provided domain will be bypassed. For 
        example, ``example.com, 192.168.1.0/24`` would bypass ``https://api.example.com``,
        ``https://www.example.com``, and any IP in the range ``192.168.1.0/24``.
-
+   * - ``client_max_retries``
+     - Number of times for a s3 client to retry the request. Default, ``10``.
+   * - ``client_retry_timeout``
+     - Timeout for a s3 client to retry the request in seconds. Default, ``180``.
 
 S3 Configuration
 ~~~~~~~~~~~~~~~~

diff --git a/rust/lance-io/src/object_store.rs b/rust/lance-io/src/object_store.rs
@@ -26,7 +26,9 @@ use object_store::{
     aws::AmazonS3Builder, azure::AzureConfigKey, gcp::GoogleConfigKey, local::LocalFileSystem,
     memory::InMemory, CredentialProvider, Error as ObjectStoreError, Result as ObjectStoreResult,
 };
-use object_store::{parse_url_opts, ClientOptions, DynObjectStore, StaticCredentialProvider};
+use object_store::{
+    parse_url_opts, ClientOptions, DynObjectStore, RetryConfig, StaticCredentialProvider,
+};
 use object_store::{path::Path, ObjectMeta, ObjectStore as OSObjectStore};
 use shellexpand::tilde;
 use snafu::{location, Location};
@@ -787,6 +789,24 @@ impl StorageOptions {
             .unwrap_or(3)
     }
 
+    /// Max retry times to set in RetryConfig for s3 client
+    pub fn client_max_retries(&self) -> usize {
+        self.0
+            .iter()
+            .find(|(key, _)| key.to_ascii_lowercase() == "client_max_retries")
+            .and_then(|(_, value)| value.parse::<usize>().ok())
+            .unwrap_or(10)
+    }
+
+    /// Seconds of timeout to set in RetryConfig for s3 client
+    pub fn client_retry_timeout(&self) -> u64 {
+        self.0
+            .iter()
+            .find(|(key, _)| key.to_ascii_lowercase() == "client_retry_timeout")
+            .and_then(|(_, value)| value.parse::<u64>().ok())
+            .unwrap_or(180)
+    }
+
     /// Subset of options relevant for azure storage
     pub fn as_azure_options(&self) -> HashMap<AzureConfigKey, String> {
         self.0
@@ -850,6 +870,13 @@ async fn configure_store(
             //     });
             // }
 
+            let max_retries = storage_options.client_max_retries();
+            let retry_timeout = storage_options.client_retry_timeout();
+            let retry_config = RetryConfig {
+                backoff: Default::default(),
+                max_retries,
+                retry_timeout: Duration::from_secs(retry_timeout),
+            };
             let storage_options = storage_options.as_s3_options();
             let region = resolve_s3_region(&url, &storage_options).await?;
             let (aws_creds, region) = build_aws_credential(
@@ -882,6 +909,7 @@ async fn configure_store(
             builder = builder
                 .with_url(url.as_ref())
                 .with_credentials(aws_creds)
+                .with_retry(retry_config)
                 .with_region(region);
             let store = builder.build()?;