diff --git a/config/config.yaml b/config/config.yaml index ba49474c9f..f1a3c53a22 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -37,6 +37,11 @@ storage: # secret_key: "" # # Custom endpoint, for MinIO and similar # endpoint_url: "" + # # Timeouts and retries, in seconds + # timeout_sec: 30 + # connect_timeout_sec: 5 + # max_retries: 10 + # retry_timeout_sec: 180 # Google Cloud Storage. Without credential fields, Application Default # Credentials are used: GOOGLE_* environment, gcloud config or the @@ -51,6 +56,11 @@ storage: # application_credentials_path: "" # # Custom endpoint, for fake-gcs-server and similar # endpoint_url: "" + # # Timeouts and retries, same as in s3_config + # timeout_sec: 30 + # connect_timeout_sec: 5 + # max_retries: 10 + # retry_timeout_sec: 180 # Azure Blob Storage. Without credential fields, the default chain is # used: AZURE_* environment, managed identity or the Azure CLI. @@ -67,6 +77,11 @@ storage: # tenant_id: "" # # Custom endpoint, for Azurite and similar # endpoint_url: "" + # # Timeouts and retries, same as in s3_config + # timeout_sec: 30 + # connect_timeout_sec: 5 + # max_retries: 10 + # retry_timeout_sec: 180 # Where to store temporary files # If null, temporary snapshots are stored in: storage/snapshots_temp/ diff --git a/lib/collection/src/common/snapshots_manager.rs b/lib/collection/src/common/snapshots_manager.rs index b56f79e01c..172894d720 100644 --- a/lib/collection/src/common/snapshots_manager.rs +++ b/lib/collection/src/common/snapshots_manager.rs @@ -1,15 +1,16 @@ +use std::num::NonZeroU64; use std::path::{Path, PathBuf}; +use std::time::Duration; use common::defaults::APP_USER_AGENT; use common::tempfile_ext::MaybeTempPath; use fs_err as fs; use fs_err::tokio as tokio_fs; -use http::HeaderValue; -use object_store::aws::AmazonS3Builder; +use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey}; use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder}; -use object_store::gcp::GoogleCloudStorageBuilder; +use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey}; use object_store::prefix::PrefixStore; -use object_store::{ClientOptions, ObjectStoreExt}; +use object_store::{ClientConfigKey, ObjectStoreExt, RetryConfig}; use serde::Deserialize; use tempfile::TempPath; use tokio::io::AsyncWriteExt; @@ -55,6 +56,14 @@ pub struct S3Config { pub access_key: Option, pub secret_key: Option, pub endpoint_url: Option, + /// Request timeout. Overrides `AWS_TIMEOUT`. + pub timeout_sec: Option, + /// Connect timeout. Overrides `AWS_CONNECT_TIMEOUT`. + pub connect_timeout_sec: Option, + /// Retries for a failed request. 0 disables retries. + pub max_retries: Option, + /// How long to keep retrying after the first attempt. + pub retry_timeout_sec: Option, } /// Google Cloud Storage connection parameters. @@ -75,6 +84,14 @@ pub struct GcsConfig { pub application_credentials_path: Option, /// Custom base URL, for emulators such as fake-gcs-server. pub endpoint_url: Option, + /// See [`S3Config::timeout_sec`]. Overrides `GOOGLE_TIMEOUT`. + pub timeout_sec: Option, + /// See [`S3Config::connect_timeout_sec`]. Overrides `GOOGLE_CONNECT_TIMEOUT`. + pub connect_timeout_sec: Option, + /// See [`S3Config::max_retries`]. + pub max_retries: Option, + /// See [`S3Config::retry_timeout_sec`]. + pub retry_timeout_sec: Option, } /// Azure Blob Storage connection parameters. @@ -99,6 +116,14 @@ pub struct AzureConfig { pub tenant_id: Option, /// Custom endpoint URL, for emulators such as Azurite. pub endpoint_url: Option, + /// See [`S3Config::timeout_sec`]. Overrides `AZURE_TIMEOUT`. + pub timeout_sec: Option, + /// See [`S3Config::connect_timeout_sec`]. Overrides `AZURE_CONNECT_TIMEOUT`. + pub connect_timeout_sec: Option, + /// See [`S3Config::max_retries`]. + pub max_retries: Option, + /// See [`S3Config::retry_timeout_sec`]. + pub retry_timeout_sec: Option, } pub struct SnapshotStorageCloud { @@ -137,23 +162,108 @@ pub enum SnapshotStorageManager { Cloud(SnapshotStorageCloud), } -/// HTTP client options shared by all object store backends. -/// -/// Identifies Qdrant with its user agent so providers can attribute snapshot traffic. -/// Plain HTTP is only allowed for explicitly configured `http://` endpoints, which -/// are typically local emulators. -fn client_options(endpoint_url: Option<&str>) -> ClientOptions { - let allow_http = endpoint_url.is_some_and(|url| url.starts_with("http://")); - let mut options = ClientOptions::new().with_allow_http(allow_http); - if let Ok(user_agent) = HeaderValue::from_str(APP_USER_AGENT.as_str()) { - options = options.with_user_agent(user_agent); +/// HTTP client settings that every object store backend accepts. +#[derive(Clone, Copy, Debug, Default)] +struct ClientSettings<'a> { + endpoint_url: Option<&'a str>, + timeout_sec: Option, + connect_timeout_sec: Option, + max_retries: Option, + retry_timeout_sec: Option, +} + +impl<'a> From<&'a S3Config> for ClientSettings<'a> { + fn from(config: &'a S3Config) -> Self { + Self { + endpoint_url: config.endpoint_url.as_deref(), + timeout_sec: config.timeout_sec, + connect_timeout_sec: config.connect_timeout_sec, + max_retries: config.max_retries, + retry_timeout_sec: config.retry_timeout_sec, + } + } +} + +impl<'a> From<&'a GcsConfig> for ClientSettings<'a> { + fn from(config: &'a GcsConfig) -> Self { + Self { + endpoint_url: config.endpoint_url.as_deref(), + timeout_sec: config.timeout_sec, + connect_timeout_sec: config.connect_timeout_sec, + max_retries: config.max_retries, + retry_timeout_sec: config.retry_timeout_sec, + } + } +} + +impl<'a> From<&'a AzureConfig> for ClientSettings<'a> { + fn from(config: &'a AzureConfig) -> Self { + Self { + endpoint_url: config.endpoint_url.as_deref(), + timeout_sec: config.timeout_sec, + connect_timeout_sec: config.connect_timeout_sec, + max_retries: config.max_retries, + retry_timeout_sec: config.retry_timeout_sec, + } + } +} + +impl ClientSettings<'_> { + /// HTTP client settings as config keys. + /// + /// Identifies Qdrant with its user agent so providers can attribute snapshot traffic. + /// Plain HTTP is allowed for explicitly configured `http://` endpoints, which are + /// typically local emulators. + /// + /// The settings are config keys rather than a `ClientOptions`, because + /// `with_client_options` would replace the client settings that `from_env` read + /// from the environment, such as `AWS_TIMEOUT`. + fn client_config(&self) -> Vec<(ClientConfigKey, String)> { + let mut config = vec![(ClientConfigKey::UserAgent, APP_USER_AGENT.clone())]; + if self + .endpoint_url + .is_some_and(|url| url.starts_with("http://")) + { + config.push((ClientConfigKey::AllowHttp, true.to_string())); + } + if let Some(timeout_sec) = self.timeout_sec { + config.push((ClientConfigKey::Timeout, format!("{timeout_sec}s"))); + } + if let Some(connect_timeout_sec) = self.connect_timeout_sec { + config.push(( + ClientConfigKey::ConnectTimeout, + format!("{connect_timeout_sec}s"), + )); + } + config + } + + /// Retry settings, with the `object_store` defaults for unset fields. + fn retry_config(&self) -> RetryConfig { + let mut retry = RetryConfig::default(); + if let Some(max_retries) = self.max_retries { + retry.max_retries = max_retries; + } + if let Some(retry_timeout_sec) = self.retry_timeout_sec { + retry.retry_timeout = Duration::from_secs(retry_timeout_sec); + } + retry } - options } fn build_s3_client(config: Option<&S3Config>) -> CollectionResult { - let endpoint_url = config.and_then(|config| config.endpoint_url.as_deref()); - let mut builder = AmazonS3Builder::from_env().with_client_options(client_options(endpoint_url)); + s3_builder(AmazonS3Builder::from_env(), config) + .build() + .map_err(|e| CollectionError::service_error(format!("Failed to create S3 client: {e}"))) +} + +/// Apply `config` on top of `builder`, which normally comes from `AmazonS3Builder::from_env`. +fn s3_builder(mut builder: AmazonS3Builder, config: Option<&S3Config>) -> AmazonS3Builder { + let settings = config.map(ClientSettings::from).unwrap_or_default(); + for (key, value) in settings.client_config() { + builder = builder.with_config(AmazonS3ConfigKey::Client(key), value); + } + builder = builder.with_retry(settings.retry_config()); if let Some(config) = config { builder = builder.with_bucket_name(&config.bucket); @@ -172,16 +282,17 @@ fn build_s3_client(config: Option<&S3Config>) -> CollectionResult, ) -> CollectionResult { - let endpoint_url = config.and_then(|config| config.endpoint_url.as_deref()); - let mut builder = - GoogleCloudStorageBuilder::from_env().with_client_options(client_options(endpoint_url)); + let settings = config.map(ClientSettings::from).unwrap_or_default(); + let mut builder = GoogleCloudStorageBuilder::from_env(); + for (key, value) in settings.client_config() { + builder = builder.with_config(GoogleConfigKey::Client(key), value); + } + builder = builder.with_retry(settings.retry_config()); if let Some(config) = config { builder = builder.with_bucket_name(&config.bucket); @@ -207,9 +318,12 @@ fn build_gcs_client( fn build_azure_client( config: Option<&AzureConfig>, ) -> CollectionResult { - let endpoint_url = config.and_then(|config| config.endpoint_url.as_deref()); - let mut builder = - MicrosoftAzureBuilder::from_env().with_client_options(client_options(endpoint_url)); + let settings = config.map(ClientSettings::from).unwrap_or_default(); + let mut builder = MicrosoftAzureBuilder::from_env(); + for (key, value) in settings.client_config() { + builder = builder.with_config(AzureConfigKey::Client(key), value); + } + builder = builder.with_retry(settings.retry_config()); if let Some(config) = config { builder = builder @@ -709,6 +823,10 @@ mod tests { access_key: Some("test-access-key".into()), secret_key: Some("test-secret-key".into()), endpoint_url: Some("http://localhost:9000".into()), + timeout_sec: None, + connect_timeout_sec: None, + max_retries: None, + retry_timeout_sec: None, }), ..Default::default() }) @@ -1152,4 +1270,131 @@ mod tests { fs::canonicalize(my_collection.join("my.snapshot")).unwrap(), ); } + + /// `from_env` maps `AWS_TIMEOUT` to the client timeout key. Building on top of it + /// must keep that value unless the config file sets its own. + #[test] + fn s3_builder_keeps_env_timeouts_unless_configured() { + let env_builder = || { + let key: AmazonS3ConfigKey = "aws_timeout".parse().unwrap(); + assert_eq!(key, AmazonS3ConfigKey::Client(ClientConfigKey::Timeout)); + AmazonS3Builder::new().with_config(key, "77s") + }; + let timeout = |builder: &AmazonS3Builder| { + builder.get_config_value(&AmazonS3ConfigKey::Client(ClientConfigKey::Timeout)) + }; + let mut config = S3Config { + bucket: "test-bucket".into(), + prefix: None, + region: None, + access_key: None, + secret_key: None, + endpoint_url: None, + timeout_sec: None, + connect_timeout_sec: None, + max_retries: None, + retry_timeout_sec: None, + }; + + let builder = s3_builder(env_builder(), Some(&config)); + assert_eq!(timeout(&builder).as_deref(), Some("77s")); + let builder = s3_builder(env_builder(), None); + assert_eq!(timeout(&builder).as_deref(), Some("77s")); + + config.timeout_sec = NonZeroU64::new(120); + let builder = s3_builder(env_builder(), Some(&config)); + assert_eq!(timeout(&builder).as_deref(), Some("120s")); + } + + #[test] + fn retry_config_applies_limits() { + let retry = |max_retries, retry_timeout_sec| { + let settings = ClientSettings { + endpoint_url: None, + timeout_sec: None, + connect_timeout_sec: None, + max_retries, + retry_timeout_sec, + }; + let retry = settings.retry_config(); + (retry.max_retries, retry.retry_timeout) + }; + let default = RetryConfig::default(); + + assert_eq!( + retry(None, None), + (default.max_retries, default.retry_timeout), + ); + assert_eq!(retry(Some(3), Some(60)), (3, Duration::from_secs(60))); + // Each field overrides only its own default. + assert_eq!(retry(Some(0), None), (0, default.retry_timeout)); + assert_eq!( + retry(None, Some(600)), + (default.max_retries, Duration::from_secs(600)), + ); + } + + #[test] + fn gcs_and_azure_configs_deserialize_client_settings() { + let gcs: GcsConfig = serde_json::from_value(serde_json::json!({ + "bucket": "b", + "timeout_sec": 120, + "connect_timeout_sec": 10, + "max_retries": 20, + "retry_timeout_sec": 600, + })) + .unwrap(); + let azure: AzureConfig = serde_json::from_value(serde_json::json!({ + "account": "a", + "container": "c", + "timeout_sec": 120, + "connect_timeout_sec": 10, + "max_retries": 20, + "retry_timeout_sec": 600, + })) + .unwrap(); + + for settings in [ClientSettings::from(&gcs), ClientSettings::from(&azure)] { + assert_eq!(settings.timeout_sec, NonZeroU64::new(120)); + assert_eq!(settings.connect_timeout_sec, NonZeroU64::new(10)); + assert_eq!(settings.max_retries, Some(20)); + assert_eq!(settings.retry_timeout_sec, Some(600)); + } + } + + #[test] + fn s3_builder_applies_client_settings() { + let config = S3Config { + bucket: "test-bucket".into(), + prefix: None, + region: None, + access_key: None, + secret_key: None, + endpoint_url: Some("http://localhost:9000".into()), + timeout_sec: NonZeroU64::new(120), + connect_timeout_sec: NonZeroU64::new(10), + max_retries: Some(3), + retry_timeout_sec: Some(60), + }; + let builder = s3_builder(AmazonS3Builder::new(), Some(&config)); + let client_value = |key| builder.get_config_value(&AmazonS3ConfigKey::Client(key)); + + assert_eq!( + client_value(ClientConfigKey::Timeout).as_deref(), + Some("120s") + ); + assert_eq!( + client_value(ClientConfigKey::ConnectTimeout).as_deref(), + Some("10s"), + ); + assert_eq!( + client_value(ClientConfigKey::AllowHttp).as_deref(), + Some("true") + ); + assert_eq!( + client_value(ClientConfigKey::UserAgent), + Some(APP_USER_AGENT.clone()), + ); + builder.build().unwrap(); + } }