Merge remote-tracking branch 'origin/main' into arpad/scrubber_azure

Cargo.lock

@@ -4861,6 +4861,7 @@ dependencies = [
  "once_cell",
  "pin-project-lite",
  "rand 0.8.5",
+ "reqwest 0.11.19",
  "scopeguard",
  "serde",
  "serde_json",

libs/remote_storage/Cargo.toml

@@ -18,6 +18,7 @@ camino = { workspace = true, features = ["serde1"] }
 humantime-serde.workspace = true
 hyper = { workspace = true, features = ["client"] }
 futures.workspace = true
+reqwest = { version = "0.11", default-features = false, features = ["rustls-tls"] }
 serde.workspace = true
 serde_json.workspace = true
 tokio = { workspace = true, features = ["sync", "fs", "io-util"] }

libs/remote_storage/src/azure_blob.rs

@@ -15,8 +15,11 @@ use std::time::SystemTime;
 use super::REMOTE_STORAGE_PREFIX_SEPARATOR;
 use anyhow::Result;
 use azure_core::request_options::{IfMatchCondition, MaxResults, Metadata, Range};
+use azure_core::HttpClient;
+use azure_core::TransportOptions;
 use azure_core::{Continuable, RetryOptions};
 use azure_identity::DefaultAzureCredential;
+use azure_storage::CloudLocation;
 use azure_storage::StorageCredentials;
 use azure_storage_blobs::blob::CopyStatus;
 use azure_storage_blobs::prelude::ClientBuilder;
@@ -72,8 +75,17 @@ impl AzureBlobStorage {
             StorageCredentials::token_credential(Arc::new(token_credential))
         };
 
-        // we have an outer retry
-        let builder = ClientBuilder::new(account, credentials).retry(RetryOptions::none());
+        let location = match &azure_config.endpoint {
+            None => CloudLocation::Public { account },
+            Some(endpoint) => CloudLocation::Custom {
+                account,
+                uri: endpoint.clone(),
+            },
+        };
+        let builder = ClientBuilder::with_location(location, credentials)
+            // we have an outer retry
+            .retry(RetryOptions::none())
+            .transport(TransportOptions::new(reqwest_client(true)));
 
         let client = builder.container_client(azure_config.container_name.to_owned());
 
@@ -246,6 +258,15 @@ impl AzureBlobStorage {
     }
 }
 
+fn reqwest_client(allow_idle_connections: bool) -> Arc<dyn HttpClient> {
+    let max_idle = if allow_idle_connections { 8 } else { 0 };
+    let client = reqwest::ClientBuilder::new()
+        .pool_max_idle_per_host(max_idle)
+        .build()
+        .expect("failed to build `reqwest` client");
+    Arc::new(client)
+}
+
 fn to_azure_metadata(metadata: StorageMetadata) -> Metadata {
     let mut res = Metadata::new();
     for (k, v) in metadata.0.into_iter() {

libs/remote_storage/src/config.rs

@@ -125,6 +125,8 @@ pub struct AzureConfig {
     pub container_region: String,
     /// A "subfolder" in the container, to use the same container separately by multiple remote storage users at once.
     pub prefix_in_container: Option<String>,
+    /// The endpoint to use. Use the default if None.
+    pub endpoint: Option<String>,
     /// Azure has various limits on its API calls, we need not to exceed those.
     /// See [`DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT`] for more details.
     #[serde(default = "default_remote_storage_azure_concurrency_limit")]
@@ -144,6 +146,7 @@ impl Debug for AzureConfig {
             .field("storage_account", &self.storage_account)
             .field("bucket_region", &self.container_region)
             .field("prefix_in_container", &self.prefix_in_container)
+            .field("endpoint", &self.endpoint)
             .field("concurrency_limit", &self.concurrency_limit)
             .field(
                 "max_keys_per_list_response",
@@ -296,6 +299,7 @@ timeout = '5s'";
                     storage_account: None,
                     container_region: "westeurope".into(),
                     prefix_in_container: None,
+                    endpoint: None,
                     concurrency_limit: default_remote_storage_azure_concurrency_limit(),
                     max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
                 }),

libs/remote_storage/tests/test_real_azure.rs

@@ -216,6 +216,7 @@ async fn create_azure_client(
             storage_account: None,
             container_region: remote_storage_azure_region,
             prefix_in_container: Some(format!("test_{millis}_{random:08x}/")),
+            endpoint: None,
             concurrency_limit: NonZeroUsize::new(100).unwrap(),
             max_keys_per_list_response,
         }),