feat: Time out S3 attempts individually and retry throttled requests longer

quickwit/quickwit-aws/src/error.rs

@@ -16,7 +16,7 @@
 
 use aws_config::retry::ErrorKind;
 use aws_runtime::retries::classifiers::{THROTTLING_ERRORS, TRANSIENT_ERRORS};
-use aws_sdk_s3::error::SdkError;
+use aws_sdk_s3::error::{ProvideErrorMetadata, SdkError};
 use aws_sdk_s3::operation::abort_multipart_upload::AbortMultipartUploadError;
 use aws_sdk_s3::operation::complete_multipart_upload::CompleteMultipartUploadError;
 use aws_sdk_s3::operation::create_multipart_upload::CreateMultipartUploadError;
@@ -30,7 +30,7 @@ use aws_sdk_s3::operation::upload_part::UploadPartError;
 use crate::retry::AwsRetryable;
 
 impl<E> AwsRetryable for SdkError<E>
-where E: AwsRetryable
+where E: AwsRetryable + ProvideErrorMetadata
 {
     fn is_retryable(&self) -> bool {
         match self {
@@ -55,6 +55,10 @@ where E: AwsRetryable
             _ => false,
         }
     }
+
+    fn is_throttling(&self) -> bool {
+        matches!(self.code(), Some(code) if THROTTLING_ERRORS.contains(&code))
+    }
 }
 
 fn is_retryable(meta: &aws_sdk_s3::error::ErrorMetadata) -> bool {

quickwit/quickwit-aws/src/retry.rs

@@ -23,14 +23,19 @@ pub trait AwsRetryable {
     fn is_retryable(&self) -> bool {
         false
     }
+
+    fn is_throttling(&self) -> bool {
+        false
+    }
 }
 
 impl<E> AwsRetryable for Retry<E> {
     fn is_retryable(&self) -> bool {
-        match self {
-            Retry::Transient(_) => true,
-            Retry::Permanent(_) => false,
-        }
+        Retryable::is_retryable(self)
+    }
+
+    fn is_throttling(&self) -> bool {
+        Retryable::is_throttling(self)
     }
 }
 
@@ -43,6 +48,10 @@ where E: AwsRetryable
     fn is_retryable(&self) -> bool {
         self.0.is_retryable()
     }
+
+    fn is_throttling(&self) -> bool {
+        self.0.is_throttling()
+    }
 }
 
 pub async fn aws_retry<U, E, Fut>(retry_params: &RetryParams, f: impl Fn() -> Fut) -> Result<U, E>

quickwit/quickwit-cli/src/lib.rs

@@ -35,7 +35,7 @@ use quickwit_metastore::{IndexMetadataResponseExt, MetastoreResolver};
 use quickwit_proto::metastore::{IndexMetadataRequest, MetastoreService, MetastoreServiceClient};
 use quickwit_rest_client::models::Timeout;
 use quickwit_rest_client::rest_client::{DEFAULT_BASE_URL, QuickwitClient, QuickwitClientBuilder};
-use quickwit_storage::{StorageResolver, load_file};
+use quickwit_storage::{S3NodeOptions, StorageResolver, load_file};
 use reqwest::Url;
 use tabled::settings::object::Rows;
 use tabled::settings::panel::Header;
@@ -247,38 +247,44 @@ async fn load_node_config(config_uri: &Uri) -> anyhow::Result<NodeConfig> {
     Ok(config)
 }
 
-/// Builds the storage and metastore resolvers.
-///
-/// When `s3_assembly_thread_pool_opt` is set, S3 storages assemble large downloads on that pool
-/// instead of the tokio runtime. Only searcher-only nodes set it (see
-/// [`s3_assembly_thread_pool_for_services`]).
+/// Builds the storage and metastore resolvers. `s3_node_options` are the S3 storage settings that
+/// depend on the node rather than on the storage configs (see [`s3_node_options`]).
 fn get_resolvers(
     storage_configs: &StorageConfigs,
     metastore_configs: &MetastoreConfigs,
-    s3_assembly_thread_pool_opt: Option<ThreadPoolWithPriority>,
+    s3_node_options: S3NodeOptions,
 ) -> (StorageResolver, MetastoreResolver) {
     // The CLI tests rely on the unconfigured singleton resolvers, so it's better to return them if
     // the storage and metastore configs are not set.
-    if storage_configs.is_empty()
-        && metastore_configs.is_empty()
-        && s3_assembly_thread_pool_opt.is_none()
-    {
+    if storage_configs.is_empty() && metastore_configs.is_empty() && s3_node_options.is_empty() {
         return (
             StorageResolver::unconfigured(),
             MetastoreResolver::unconfigured(),
         );
     }
-    let storage_resolver = match s3_assembly_thread_pool_opt {
-        Some(thread_pool) => {
-            StorageResolver::configured_with_s3_assembly_thread_pool(storage_configs, thread_pool)
-        }
-        None => StorageResolver::configured(storage_configs),
-    };
+    let storage_resolver =
+        StorageResolver::configured_with_s3_node_options(storage_configs, s3_node_options);
     let metastore_resolver =
         MetastoreResolver::configured(storage_resolver.clone(), metastore_configs);
     (storage_resolver, metastore_resolver)
 }
 
+/// Returns the S3 storage settings of a node running `node_config`.
+///
+/// The searcher's storage timeout policy is applied by the S3 storages themselves, to each attempt
+/// of a request, so it is handed over to them whenever the node runs the searcher.
+fn s3_node_options(node_config: &NodeConfig) -> S3NodeOptions {
+    let timeout_policy = if node_config.is_service_enabled(QuickwitService::Searcher) {
+        node_config.searcher_config.storage_timeout_policy.clone()
+    } else {
+        None
+    };
+    S3NodeOptions {
+        assembly_thread_pool: s3_assembly_thread_pool_for_services(&node_config.enabled_services),
+        timeout_policy,
+    }
+}
+
 /// Returns the thread pool on which S3 storages should assemble large downloads, if any.
 ///
 /// Only nodes that run the searcher service alone get one: the search thread pool. On such nodes
@@ -553,7 +559,10 @@ mod tests {
         };
         let storage_configs = StorageConfigs::new(vec![s3_storage_config.into()]);
         let metastore_configs = MetastoreConfigs::default();
-        let (_storage_resolver, _metastore_resolver) =
-            get_resolvers(&storage_configs, &metastore_configs, None);
+        let (_storage_resolver, _metastore_resolver) = get_resolvers(
+            &storage_configs,
+            &metastore_configs,
+            S3NodeOptions::default(),
+        );
     }
 }

quickwit/quickwit-cli/src/service.rs

@@ -30,8 +30,7 @@ use tracing::{debug, info};
 
 use crate::checklist::{BLUE_COLOR, RED_COLOR};
 use crate::{
-    config_cli_arg, get_resolvers, load_node_config, s3_assembly_thread_pool_for_services,
-    start_actor_runtimes,
+    config_cli_arg, get_resolvers, load_node_config, s3_node_options, start_actor_runtimes,
 };
 
 pub fn build_run_command() -> Command {
@@ -118,7 +117,7 @@ impl RunCliCommand {
         let (storage_resolver, metastore_resolver) = get_resolvers(
             &node_config.storage_configs,
             &node_config.metastore_configs,
-            s3_assembly_thread_pool_for_services(&node_config.enabled_services),
+            s3_node_options(&node_config),
         );
         crate::busy_detector::set_enabled(true);
         // TODO move in serve quickwit?

quickwit/quickwit-cli/src/tool.rs

@@ -55,7 +55,7 @@ use quickwit_search::{SearchResponseRest, single_node_search};
 use quickwit_serve::{
     BodyFormat, SearchRequestQueryString, SortBy, search_request_from_api_request,
 };
-use quickwit_storage::{BundleStorage, Storage};
+use quickwit_storage::{BundleStorage, S3NodeOptions, Storage};
 use tantivy::Index;
 use tantivy::directory::FileSlice;
 use tantivy::schema::FieldType;
@@ -65,7 +65,7 @@ use tracing::{debug, info};
 use crate::checklist::{GREEN_COLOR, RED_COLOR};
 use crate::{
     THROUGHPUT_WINDOW_SIZE, config_cli_arg, get_resolvers, load_node_config, run_index_checklist,
-    start_actor_runtimes,
+    s3_node_options, start_actor_runtimes,
 };
 
 pub fn build_tool_command() -> Command {
@@ -585,8 +585,11 @@ pub async fn local_ingest_docs_cli(args: LocalIngestDocsArgs) -> anyhow::Result<
     println!("❯ Ingesting documents locally...");
 
     let config = load_node_config(&args.config_uri).await?;
-    let (storage_resolver, metastore_resolver) =
-        get_resolvers(&config.storage_configs, &config.metastore_configs, None);
+    let (storage_resolver, metastore_resolver) = get_resolvers(
+        &config.storage_configs,
+        &config.metastore_configs,
+        S3NodeOptions::default(),
+    );
     let mut metastore = metastore_resolver.resolve(&config.metastore_uri).await?;
 
     let source_params = if let Some(uri) = args.input_path_opt.as_ref() {
@@ -715,8 +718,11 @@ pub async fn local_search_cli(args: LocalSearchArgs) -> anyhow::Result<()> {
     debug!(args=?args, "local-search");
     println!("❯ Searching directly on the index storage (without calling REST API)...");
     let config = load_node_config(&args.config_uri).await?;
-    let (storage_resolver, metastore_resolver) =
-        get_resolvers(&config.storage_configs, &config.metastore_configs, None);
+    let (storage_resolver, metastore_resolver) = get_resolvers(
+        &config.storage_configs,
+        &config.metastore_configs,
+        s3_node_options(&config),
+    );
     let metastore: MetastoreServiceClient =
         metastore_resolver.resolve(&config.metastore_uri).await?;
     let aggs = args
@@ -754,8 +760,11 @@ pub async fn merge_cli(args: MergeArgs) -> anyhow::Result<()> {
     debug!(args=?args, "run-merge-operations");
     println!("❯ Merging splits locally...");
     let config = load_node_config(&args.config_uri).await?;
-    let (storage_resolver, metastore_resolver) =
-        get_resolvers(&config.storage_configs, &config.metastore_configs, None);
+    let (storage_resolver, metastore_resolver) = get_resolvers(
+        &config.storage_configs,
+        &config.metastore_configs,
+        S3NodeOptions::default(),
+    );
     let mut metastore = metastore_resolver.resolve(&config.metastore_uri).await?;
     run_index_checklist(&mut metastore, &storage_resolver, &args.index_id, None).await?;
     // The indexing service needs to update its cluster chitchat state so that the control plane is
@@ -842,8 +851,11 @@ pub async fn merge_mature_cli(args: MatureMergeArgs) -> anyhow::Result<()> {
     info!(merge_config=?args.merge_config, "merge-mature configuration");
     println!("❯ Scanning all indexes for mature merge opportunities...");
     let config = load_node_config(&args.config_uri).await?;
-    let (storage_resolver, metastore_resolver) =
-        get_resolvers(&config.storage_configs, &config.metastore_configs, None);
+    let (storage_resolver, metastore_resolver) = get_resolvers(
+        &config.storage_configs,
+        &config.metastore_configs,
+        S3NodeOptions::default(),
+    );
     let metastore = metastore_resolver.resolve(&config.metastore_uri).await?;
 
     let runtimes_config = RuntimesConfig::default();
@@ -874,8 +886,11 @@ pub async fn garbage_collect_index_cli(args: GarbageCollectIndexArgs) -> anyhow:
     println!("❯ Garbage collecting index...");
 
     let config = load_node_config(&args.config_uri).await?;
-    let (storage_resolver, metastore_resolver) =
-        get_resolvers(&config.storage_configs, &config.metastore_configs, None);
+    let (storage_resolver, metastore_resolver) = get_resolvers(
+        &config.storage_configs,
+        &config.metastore_configs,
+        S3NodeOptions::default(),
+    );
     let metastore = metastore_resolver.resolve(&config.metastore_uri).await?;
     let mut index_service = IndexService::new(metastore, storage_resolver);
     let removal_info = index_service
@@ -939,8 +954,11 @@ async fn extract_split_cli(args: ExtractSplitArgs) -> anyhow::Result<()> {
     println!("❯ Extracting split...");
 
     let config = load_node_config(&args.config_uri).await?;
-    let (storage_resolver, metastore_resolver) =
-        get_resolvers(&config.storage_configs, &config.metastore_configs, None);
+    let (storage_resolver, metastore_resolver) = get_resolvers(
+        &config.storage_configs,
+        &config.metastore_configs,
+        S3NodeOptions::default(),
+    );
     let metastore = metastore_resolver.resolve(&config.metastore_uri).await?;
     let index_metadata = metastore
         .index_metadata(IndexMetadataRequest::for_index_id(args.index_id))

quickwit/quickwit-common/src/retry.rs

@@ -24,19 +24,25 @@ pub trait Retryable {
     fn is_retryable(&self) -> bool {
         false
     }
+
+    fn is_throttling(&self) -> bool {
+        false
+    }
 }
 
 #[derive(Debug, Eq, PartialEq)]
 pub enum Retry<E> {
     Permanent(E),
     Transient(E),
+    Throttled(E),
 }
 
 impl<E> Retry<E> {
     pub fn into_inner(self) -> E {
         match self {
             Self::Transient(error) => error,
             Self::Permanent(error) => error,
+            Self::Throttled(error) => error,
         }
     }
 }
@@ -45,11 +51,20 @@ impl<E> Retryable for Retry<E> {
     fn is_retryable(&self) -> bool {
         match self {
             Retry::Permanent(_) => false,
-            Retry::Transient(_) => true,
+            Retry::Transient(_) | Retry::Throttled(_) => true,
         }
     }
+
+    fn is_throttling(&self) -> bool {
+        matches!(self, Retry::Throttled(_))
+    }
 }
 
+/// Number of attempts granted to an operation that keeps being throttled, on top of
+/// [`RetryParams::max_attempts`]. Throttling can last longer than the other transient errors, and
+/// the exponential backoff spaces these last attempts out, so they add little load.
+pub const THROTTLING_EXTRA_ATTEMPTS: usize = 2;
+
 #[derive(Debug, Clone, Copy)]
 pub struct RetryParams {
     pub base_delay: Duration,
@@ -159,7 +174,13 @@ where
         }
         num_attempts += 1;
 
-        if num_attempts >= retry_params.max_attempts {
+        // Throttled operations get extra attempts, unless retries are disabled altogether.
+        let max_attempts = if error.is_throttling() && retry_params.max_attempts > 1 {
+            retry_params.max_attempts + THROTTLING_EXTRA_ATTEMPTS
+        } else {
+            retry_params.max_attempts
+        };
+        if num_attempts >= max_attempts {
             warn!(
                 num_attempts=%num_attempts,
                 "request failed"
@@ -188,11 +209,14 @@ where
 #[cfg(test)]
 mod tests {
     use std::sync::RwLock;
+    use std::sync::atomic::{AtomicUsize, Ordering};
     use std::time::Duration;
 
     use futures::future::ready;
 
-    use super::{MockableSleep, RetryParams, Retryable, retry_with_mockable_sleep};
+    use super::{
+        MockableSleep, RetryParams, Retryable, THROTTLING_EXTRA_ATTEMPTS, retry_with_mockable_sleep,
+    };
 
     #[derive(Debug, Eq, PartialEq)]
     pub enum Retry<E> {
@@ -266,6 +290,44 @@ mod tests {
         );
     }
 
+    /// Returns the number of attempts of an operation that always fails with `error`.
+    async fn count_attempts(retry_params: RetryParams, error: fn() -> super::Retry<()>) -> usize {
+        let num_attempts = AtomicUsize::new(0);
+        retry_with_mockable_sleep(
+            &retry_params,
+            || {
+                num_attempts.fetch_add(1, Ordering::Relaxed);
+                ready(Err::<(), _>(error()))
+            },
+            NoopSleep,
+        )
+        .await
+        .unwrap_err();
+        num_attempts.into_inner()
+    }
+
+    #[tokio::test]
+    async fn test_retry_gives_throttled_operations_extra_attempts() {
+        let retry_params = RetryParams {
+            base_delay: Duration::from_millis(1),
+            max_delay: Duration::from_millis(2),
+            max_attempts: 3,
+        };
+        assert_eq!(
+            count_attempts(retry_params, || super::Retry::Transient(())).await,
+            3
+        );
+        assert_eq!(
+            count_attempts(retry_params, || super::Retry::Throttled(())).await,
+            3 + THROTTLING_EXTRA_ATTEMPTS
+        );
+        // Disabled retries stay disabled for throttled operations.
+        assert_eq!(
+            count_attempts(RetryParams::no_retries(), || super::Retry::Throttled(())).await,
+            1
+        );
+    }
+
     #[tokio::test]
     async fn test_retry_retries_up_to_max_attempts_times() {
         let retry_sequence: Vec<_> = (0..29)