From 6e897bf3fda21d559bb6c42e35c856fdb42994bf Mon Sep 17 00:00:00 2001
From: Nikhil Sinha <131262146+nikhilsinhaparseable@users.noreply.github.com>
Date: Wed, 23 Oct 2024 15:24:41 +0530
Subject: [PATCH] chore: retry for S3 and blob store calls (#968)

add retry config for S3 and blob-store configs
retry config:
- max retries - 5
- retry timeout - 120 secs
- backoff config (default by object_store crate) -
	initial backoff - 100 ms (initial delay before 1st retry)
	max backoff - 15 secs (max delay between retries)
	base - 2 (exponential backoff factor i.e. 2s, 4s, 8s etc)
---
 server/src/storage/azure_blob.rs | 11 +++++++++--
 server/src/storage/s3.rs         | 10 ++++++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/server/src/storage/azure_blob.rs b/server/src/storage/azure_blob.rs
index 1c3a27883..6475b8fdc 100644
--- a/server/src/storage/azure_blob.rs
+++ b/server/src/storage/azure_blob.rs
@@ -30,7 +30,7 @@ use datafusion::datasource::object_store::{
 };
 use datafusion::execution::runtime_env::RuntimeConfig;
 use object_store::azure::{MicrosoftAzure, MicrosoftAzureBuilder};
-use object_store::{ClientOptions, ObjectStore, PutPayload};
+use object_store::{BackoffConfig, ClientOptions, ObjectStore, PutPayload, RetryConfig};
 use relative_path::{RelativePath, RelativePathBuf};
 use std::path::Path as StdPath;
 use url::Url;
@@ -120,10 +120,17 @@ impl AzureBlobConfig {
             .with_connect_timeout(Duration::from_secs(CONNECT_TIMEOUT_SECS))
             .with_timeout(Duration::from_secs(REQUEST_TIMEOUT_SECS));
 
+        let retry_config = RetryConfig {
+            max_retries: 5,
+            retry_timeout: Duration::from_secs(120),
+            backoff: BackoffConfig::default(),
+        };
+
         let mut builder = MicrosoftAzureBuilder::new()
             .with_endpoint(self.endpoint_url.clone())
             .with_account(&self.account)
-            .with_container_name(&self.container);
+            .with_container_name(&self.container)
+            .with_retry(retry_config);
 
         if let Some(access_key) = self.access_key.clone() {
             builder = builder.with_access_key(access_key)
diff --git a/server/src/storage/s3.rs b/server/src/storage/s3.rs
index 0c9eb982e..0d6513437 100644
--- a/server/src/storage/s3.rs
+++ b/server/src/storage/s3.rs
@@ -28,7 +28,7 @@ use futures::{StreamExt, TryStreamExt};
 use object_store::aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, Checksum};
 use object_store::limit::LimitStore;
 use object_store::path::Path as StorePath;
-use object_store::{ClientOptions, ObjectStore, PutPayload};
+use object_store::{BackoffConfig, ClientOptions, ObjectStore, PutPayload, RetryConfig};
 use relative_path::{RelativePath, RelativePathBuf};
 
 use std::collections::BTreeMap;
@@ -218,13 +218,19 @@ impl S3Config {
         if self.skip_tls {
             client_options = client_options.with_allow_invalid_certificates(true)
         }
+        let retry_config = RetryConfig {
+            max_retries: 5,
+            retry_timeout: Duration::from_secs(30),
+            backoff: BackoffConfig::default(),
+        };
 
         let mut builder = AmazonS3Builder::new()
             .with_region(&self.region)
             .with_endpoint(&self.endpoint_url)
             .with_bucket_name(&self.bucket_name)
             .with_virtual_hosted_style_request(!self.use_path_style)
-            .with_allow_http(true);
+            .with_allow_http(true)
+            .with_retry(retry_config);
 
         if self.set_checksum {
             builder = builder.with_checksum_algorithm(Checksum::SHA256)