acknowledge and use aws-retry-after when received (#6527)

* acknowledge and use aws-retry-after when received

* bound retry delay and handle aws sdk errors we wrapped

quickwit/quickwit-aws/src/error.rs

@@ -29,6 +29,20 @@ use aws_smithy_types::retry::ErrorKind;
 
 use crate::retry::AwsRetryable;
 
+/// Extracts the `x-amz-retry-after` header from an SDK error response, if present.
+/// The header value is in milliseconds.
+pub fn retry_after_from_sdk_error<E>(error: &SdkError<E>) -> Option<std::time::Duration> {
+    let headers = match error {
+        SdkError::ServiceError(e) => e.raw().headers(),
+        SdkError::ResponseError(e) => e.raw().headers(),
+        _ => return None,
+    };
+    headers
+        .get("x-amz-retry-after")
+        .and_then(|v| v.parse::<u64>().ok())
+        .map(std::time::Duration::from_millis)
+}
+
 impl<E> AwsRetryable for SdkError<E>
 where E: AwsRetryable
 {
@@ -53,6 +67,10 @@ where E: AwsRetryable
             _ => false,
         }
     }
+
+    fn retry_after(&self) -> Option<std::time::Duration> {
+        retry_after_from_sdk_error(self)
+    }
 }
 
 fn is_retryable(meta: &aws_sdk_s3::error::ErrorMetadata) -> bool {

quickwit/quickwit-aws/src/retry.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::fmt::Debug;
+use std::time::Duration;
 
 use futures::{Future, TryFutureExt};
 use quickwit_common::retry::{
@@ -23,15 +24,27 @@ pub trait AwsRetryable {
     fn is_retryable(&self) -> bool {
         false
     }
+
+    /// Returns a server-suggested delay before the next retry, if provided by the response.
+    fn retry_after(&self) -> Option<Duration> {
+        None
+    }
 }
 
-impl<E> AwsRetryable for Retry<E> {
+impl<E: AwsRetryable> AwsRetryable for Retry<E> {
     fn is_retryable(&self) -> bool {
         match self {
             Retry::Transient(_) => true,
             Retry::Permanent(_) => false,
         }
     }
+
+    fn retry_after(&self) -> Option<Duration> {
+        match self {
+            Retry::Transient(e) => e.retry_after(),
+            Retry::Permanent(_) => None,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -43,6 +56,10 @@ where E: AwsRetryable
     fn is_retryable(&self) -> bool {
         self.0.is_retryable()
     }
+
+    fn retry_after(&self) -> Option<Duration> {
+        self.0.retry_after()
+    }
 }
 
 pub async fn aws_retry<U, E, Fut>(retry_params: &RetryParams, f: impl Fn() -> Fut) -> Result<U, E>

quickwit/quickwit-common/src/retry.rs

@@ -24,6 +24,12 @@ pub trait Retryable {
     fn is_retryable(&self) -> bool {
         false
     }
+
+    /// Returns a server-suggested delay before the next retry, if provided.
+    /// When present, this overrides the computed exponential backoff delay.
+    fn retry_after(&self) -> Option<Duration> {
+        None
+    }
 }
 
 #[derive(Debug, Eq, PartialEq)]
@@ -166,7 +172,10 @@ where
             );
             return Err(error);
         }
-        let delay = retry_params.compute_delay(num_attempts);
+        let delay = match error.retry_after() {
+            Some(hint) => hint.min(retry_params.max_delay),
+            None => retry_params.compute_delay(num_attempts),
+        };
         debug!(
             num_attempts=%num_attempts,
             delay_ms=%delay.as_millis(),

quickwit/quickwit-index-management/src/garbage_collection.rs

@@ -550,6 +550,7 @@ pub async fn delete_splits_from_storage_and_metastore(
 }
 
 #[cfg(test)]
+#[allow(clippy::result_large_err)] // BulkDeleteError is large; acceptable in mock closures
 mod tests {
     use std::time::Duration;
 

quickwit/quickwit-index-management/src/parquet_garbage_collection.rs

@@ -407,6 +407,7 @@ async fn delete_parquet_splits_from_storage_and_metastore(
 }
 
 #[cfg(test)]
+#[allow(clippy::result_large_err)] // BulkDeleteError is large; acceptable in mock closures
 mod tests {
     use std::collections::HashMap;
     use std::path::PathBuf;

quickwit/quickwit-janitor/src/actors/garbage_collector.rs

@@ -342,6 +342,7 @@ impl Handler<Loop> for GarbageCollector {
 }
 
 #[cfg(test)]
+#[allow(clippy::result_large_err)] // BulkDeleteError is large; acceptable in mock closures
 mod tests {
     use std::collections::HashSet;
     use std::ops::Bound;

quickwit/quickwit-storage/src/error.rs

@@ -15,6 +15,7 @@
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Duration;
 use std::{fmt, io};
 
 use serde::{Deserialize, Serialize};
@@ -71,6 +72,7 @@ impl StorageErrorKind {
         StorageError {
             kind: self,
             source: Arc::new(source.into()),
+            retry_after: None,
         }
     }
 }
@@ -94,6 +96,8 @@ pub struct StorageError {
     pub kind: StorageErrorKind,
     #[source]
     source: Arc<anyhow::Error>,
+    /// Server-suggested delay before the next retry, if provided (e.g. from `x-amz-retry-after`).
+    pub(crate) retry_after: Option<Duration>,
 }
 
 /// Generic Result type for storage operations.
@@ -106,6 +110,7 @@ impl StorageError {
         StorageError {
             kind: self.kind,
             source: Arc::new(anyhow::anyhow!("{ctx}").context(self.source)),
+            retry_after: self.retry_after,
         }
     }
 
@@ -113,6 +118,14 @@ impl StorageError {
     pub fn kind(&self) -> StorageErrorKind {
         self.kind
     }
+
+    /// Attaches a server-suggested retry delay to this error.
+    pub fn with_retry_after(self, retry_after: Option<Duration>) -> Self {
+        StorageError {
+            retry_after,
+            ..self
+        }
+    }
 }
 
 impl From<io::Error> for StorageError {

quickwit/quickwit-storage/src/object_storage/error.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::time::Duration;
+
 use aws_sdk_s3::error::{DisplayErrorContext, ProvideErrorMetadata, SdkError};
 use aws_sdk_s3::operation::abort_multipart_upload::AbortMultipartUploadError;
 use aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError;
@@ -22,6 +24,8 @@ use aws_sdk_s3::operation::get_object::GetObjectError;
 use aws_sdk_s3::operation::head_object::HeadObjectError;
 use aws_sdk_s3::operation::put_object::PutObjectError;
 use aws_sdk_s3::operation::upload_part::UploadPartError;
+use quickwit_aws::error::retry_after_from_sdk_error;
+use quickwit_aws::retry::AwsRetryable;
 use quickwit_metrics::counter;
 
 use crate::metrics::OBJECT_STORAGE_GET_ERRORS_TOTAL;
@@ -53,8 +57,16 @@ where E: std::error::Error + ToStorageErrorKind + Send + Sync + 'static
             SdkError::TimeoutError(_) => StorageErrorKind::Timeout,
             _ => StorageErrorKind::Internal,
         };
+        // Extract the server-suggested retry delay before consuming the error.
+        let retry_after = retry_after_from_sdk_error(&error);
         let source = anyhow::anyhow!("{}", DisplayErrorContext(error));
-        error_kind.with_error(source)
+        error_kind.with_error(source).with_retry_after(retry_after)
+    }
+}
+
+impl AwsRetryable for StorageError {
+    fn retry_after(&self) -> Option<Duration> {
+        self.retry_after
     }
 }
 

quickwit/quickwit-storage/src/storage.rs

@@ -41,7 +41,11 @@ impl<W: AsyncWrite + Send + Unpin> SendableAsync for W {}
 /// object storage treat them. This means when directory separators a present
 /// in the storage operation path, the storage implementation should create and remove transparently
 /// these intermediate directories.
-#[cfg_attr(any(test, feature = "testsuite"), mockall::automock)]
+#[cfg_attr(
+    any(test, feature = "testsuite"),
+    allow(clippy::result_large_err), // BulkDeleteError is large but only in mock/test code
+    mockall::automock
+)]
 #[async_trait]
 pub trait Storage: fmt::Debug + Send + Sync + 'static {
     /// Check storage connection if applicable