LCOV - code coverage report
Current view: top level - libs/remote_storage/src - config.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 70.4 % 199 140
Test Date: 2025-03-12 00:01:28 Functions: 28.8 % 80 23

            Line data    Source code
       1              : use std::fmt::Debug;
       2              : use std::num::NonZeroUsize;
       3              : use std::str::FromStr;
       4              : use std::time::Duration;
       5              : 
       6              : use aws_sdk_s3::types::StorageClass;
       7              : use camino::Utf8PathBuf;
       8              : use serde::{Deserialize, Serialize};
       9              : 
      10              : use crate::{
      11              :     DEFAULT_MAX_KEYS_PER_LIST_RESPONSE, DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT,
      12              :     DEFAULT_REMOTE_STORAGE_LOCALFS_CONCURRENCY_LIMIT, DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT,
      13              : };
      14              : 
      15              : /// External backup storage configuration, enough for creating a client for that storage.
      16           37 : #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
      17              : pub struct RemoteStorageConfig {
      18              :     /// The storage connection configuration.
      19              :     #[serde(flatten)]
      20              :     pub storage: RemoteStorageKind,
      21              :     /// A common timeout enforced for all requests after concurrency limiter permit has been
      22              :     /// acquired.
      23              :     #[serde(
      24              :         with = "humantime_serde",
      25              :         default = "default_timeout",
      26              :         skip_serializing_if = "is_default_timeout"
      27              :     )]
      28              :     pub timeout: Duration,
      29              :     /// Alternative timeout used for metadata objects which are expected to be small
      30              :     #[serde(
      31              :         with = "humantime_serde",
      32              :         default = "default_small_timeout",
      33              :         skip_serializing_if = "is_default_small_timeout"
      34              :     )]
      35              :     pub small_timeout: Duration,
      36              : }
      37              : 
      38              : impl RemoteStorageKind {
      39            0 :     pub fn bucket_name(&self) -> Option<&str> {
      40            0 :         match self {
      41            0 :             RemoteStorageKind::LocalFs { .. } => None,
      42            0 :             RemoteStorageKind::AwsS3(config) => Some(&config.bucket_name),
      43            0 :             RemoteStorageKind::AzureContainer(config) => Some(&config.container_name),
      44              :         }
      45            0 :     }
      46              : }
      47              : 
      48              : impl RemoteStorageConfig {
      49              :     /// Helper to fetch the configured concurrency limit.
      50            0 :     pub fn concurrency_limit(&self) -> usize {
      51            0 :         match &self.storage {
      52            0 :             RemoteStorageKind::LocalFs { .. } => DEFAULT_REMOTE_STORAGE_LOCALFS_CONCURRENCY_LIMIT,
      53            0 :             RemoteStorageKind::AwsS3(c) => c.concurrency_limit.into(),
      54            0 :             RemoteStorageKind::AzureContainer(c) => c.concurrency_limit.into(),
      55              :         }
      56            0 :     }
      57              : }
      58              : 
      59            1 : fn default_timeout() -> Duration {
      60            1 :     RemoteStorageConfig::DEFAULT_TIMEOUT
      61            1 : }
      62              : 
      63           10 : fn default_small_timeout() -> Duration {
      64           10 :     RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT
      65           10 : }
      66              : 
      67            0 : fn is_default_timeout(d: &Duration) -> bool {
      68            0 :     *d == RemoteStorageConfig::DEFAULT_TIMEOUT
      69            0 : }
      70              : 
      71            0 : fn is_default_small_timeout(d: &Duration) -> bool {
      72            0 :     *d == RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT
      73            0 : }
      74              : 
      75              : /// A kind of a remote storage to connect to, with its connection configuration.
      76           28 : #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
      77              : #[serde(untagged)]
      78              : pub enum RemoteStorageKind {
      79              :     /// Storage based on local file system.
      80              :     /// Specify a root folder to place all stored files into.
      81              :     LocalFs { local_path: Utf8PathBuf },
      82              :     /// AWS S3 based storage, storing all files in the S3 bucket
      83              :     /// specified by the config
      84              :     AwsS3(S3Config),
      85              :     /// Azure Blob based storage, storing all files in the container
      86              :     /// specified by the config
      87              :     AzureContainer(AzureConfig),
      88              : }
      89              : 
      90              : /// AWS S3 bucket coordinates and access credentials to manage the bucket contents (read and write).
      91           25 : #[derive(Clone, PartialEq, Eq, Deserialize, Serialize)]
      92              : pub struct S3Config {
      93              :     /// Name of the bucket to connect to.
      94              :     pub bucket_name: String,
      95              :     /// The region where the bucket is located at.
      96              :     pub bucket_region: String,
      97              :     /// A "subfolder" in the bucket, to use the same bucket separately by multiple remote storage users at once.
      98              :     pub prefix_in_bucket: Option<String>,
      99              :     /// A base URL to send S3 requests to.
     100              :     /// By default, the endpoint is derived from a region name, assuming it's
     101              :     /// an AWS S3 region name, erroring on wrong region name.
     102              :     /// Endpoint provides a way to support other S3 flavors and their regions.
     103              :     ///
     104              :     /// Example: `http://127.0.0.1:5000`
     105              :     pub endpoint: Option<String>,
     106              :     /// AWS S3 has various limits on its API calls, we need not to exceed those.
     107              :     /// See [`DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT`] for more details.
     108              :     #[serde(default = "default_remote_storage_s3_concurrency_limit")]
     109              :     pub concurrency_limit: NonZeroUsize,
     110              :     #[serde(default = "default_max_keys_per_list_response")]
     111              :     pub max_keys_per_list_response: Option<i32>,
     112              :     #[serde(
     113              :         deserialize_with = "deserialize_storage_class",
     114              :         serialize_with = "serialize_storage_class",
     115              :         default
     116              :     )]
     117              :     pub upload_storage_class: Option<StorageClass>,
     118              : }
     119              : 
     120            7 : fn default_remote_storage_s3_concurrency_limit() -> NonZeroUsize {
     121            7 :     DEFAULT_REMOTE_STORAGE_S3_CONCURRENCY_LIMIT
     122            7 :         .try_into()
     123            7 :         .unwrap()
     124            7 : }
     125              : 
     126            7 : fn default_max_keys_per_list_response() -> Option<i32> {
     127            7 :     DEFAULT_MAX_KEYS_PER_LIST_RESPONSE
     128            7 : }
     129              : 
     130            0 : fn default_azure_conn_pool_size() -> usize {
     131            0 :     // By default, the Azure SDK does no connection pooling, due to historic reports of hard-to-reproduce issues
     132            0 :     // (https://github.com/hyperium/hyper/issues/2312)
     133            0 :     //
     134            0 :     // However, using connection pooling is important to avoid exhausting client ports when
     135            0 :     // doing huge numbers of requests (https://github.com/neondatabase/cloud/issues/20971)
     136            0 :     //
     137            0 :     // We therefore enable a modest pool size by default: this may be configured to zero if
     138            0 :     // issues like the alleged upstream hyper issue appear.
     139            0 :     8
     140            0 : }
     141              : 
     142              : impl Debug for S3Config {
     143            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     144            0 :         f.debug_struct("S3Config")
     145            0 :             .field("bucket_name", &self.bucket_name)
     146            0 :             .field("bucket_region", &self.bucket_region)
     147            0 :             .field("prefix_in_bucket", &self.prefix_in_bucket)
     148            0 :             .field("concurrency_limit", &self.concurrency_limit)
     149            0 :             .field(
     150            0 :                 "max_keys_per_list_response",
     151            0 :                 &self.max_keys_per_list_response,
     152            0 :             )
     153            0 :             .finish()
     154            0 :     }
     155              : }
     156              : 
     157              : /// Azure  bucket coordinates and access credentials to manage the bucket contents (read and write).
     158           12 : #[derive(Clone, PartialEq, Eq, Serialize, Deserialize)]
     159              : pub struct AzureConfig {
     160              :     /// Name of the container to connect to.
     161              :     pub container_name: String,
     162              :     /// Name of the storage account the container is inside of
     163              :     pub storage_account: Option<String>,
     164              :     /// The region where the bucket is located at.
     165              :     pub container_region: String,
     166              :     /// A "subfolder" in the container, to use the same container separately by multiple remote storage users at once.
     167              :     pub prefix_in_container: Option<String>,
     168              :     /// Azure has various limits on its API calls, we need not to exceed those.
     169              :     /// See [`DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT`] for more details.
     170              :     #[serde(default = "default_remote_storage_azure_concurrency_limit")]
     171              :     pub concurrency_limit: NonZeroUsize,
     172              :     #[serde(default = "default_max_keys_per_list_response")]
     173              :     pub max_keys_per_list_response: Option<i32>,
     174              :     #[serde(default = "default_azure_conn_pool_size")]
     175              :     pub conn_pool_size: usize,
     176              : }
     177              : 
     178            6 : fn default_remote_storage_azure_concurrency_limit() -> NonZeroUsize {
     179            6 :     NonZeroUsize::new(DEFAULT_REMOTE_STORAGE_AZURE_CONCURRENCY_LIMIT).unwrap()
     180            6 : }
     181              : 
     182              : impl Debug for AzureConfig {
     183            0 :     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     184            0 :         f.debug_struct("AzureConfig")
     185            0 :             .field("bucket_name", &self.container_name)
     186            0 :             .field("storage_account", &self.storage_account)
     187            0 :             .field("bucket_region", &self.container_region)
     188            0 :             .field("prefix_in_container", &self.prefix_in_container)
     189            0 :             .field("concurrency_limit", &self.concurrency_limit)
     190            0 :             .field(
     191            0 :                 "max_keys_per_list_response",
     192            0 :                 &self.max_keys_per_list_response,
     193            0 :             )
     194            0 :             .finish()
     195            0 :     }
     196              : }
     197              : 
     198           15 : fn deserialize_storage_class<'de, D: serde::Deserializer<'de>>(
     199           15 :     deserializer: D,
     200           15 : ) -> Result<Option<StorageClass>, D::Error> {
     201           15 :     Option::<String>::deserialize(deserializer).and_then(|s| {
     202           15 :         if let Some(s) = s {
     203              :             use serde::de::Error;
     204           12 :             let storage_class = StorageClass::from_str(&s).expect("infallible");
     205              :             #[allow(deprecated)]
     206           12 :             if matches!(storage_class, StorageClass::Unknown(_)) {
     207            0 :                 return Err(D::Error::custom(format!(
     208            0 :                     "Specified storage class unknown to SDK: '{s}'. Allowed values: {:?}",
     209            0 :                     StorageClass::values()
     210            0 :                 )));
     211           12 :             }
     212           12 :             Ok(Some(storage_class))
     213              :         } else {
     214            3 :             Ok(None)
     215              :         }
     216           15 :     })
     217           15 : }
     218              : 
     219            9 : fn serialize_storage_class<S: serde::Serializer>(
     220            9 :     val: &Option<StorageClass>,
     221            9 :     serializer: S,
     222            9 : ) -> Result<S::Ok, S::Error> {
     223            9 :     let val = val.as_ref().map(StorageClass::as_str);
     224            9 :     Option::<&str>::serialize(&val, serializer)
     225            9 : }
     226              : 
     227              : impl RemoteStorageConfig {
     228              :     pub const DEFAULT_TIMEOUT: Duration = std::time::Duration::from_secs(120);
     229              :     pub const DEFAULT_SMALL_TIMEOUT: Duration = std::time::Duration::from_secs(30);
     230              : 
     231           10 :     pub fn from_toml(toml: &toml_edit::Item) -> anyhow::Result<RemoteStorageConfig> {
     232           10 :         Ok(utils::toml_edit_ext::deserialize_item(toml)?)
     233           10 :     }
     234              : 
     235            9 :     pub fn from_toml_str(input: &str) -> anyhow::Result<RemoteStorageConfig> {
     236            9 :         let toml_document = toml_edit::DocumentMut::from_str(input)?;
     237            9 :         if let Some(item) = toml_document.get("remote_storage") {
     238            0 :             return Self::from_toml(item);
     239            9 :         }
     240            9 :         Self::from_toml(toml_document.as_item())
     241            9 :     }
     242              : }
     243              : 
     244              : #[cfg(test)]
     245              : mod tests {
     246              :     use super::*;
     247              : 
     248            9 :     fn parse(input: &str) -> anyhow::Result<RemoteStorageConfig> {
     249            9 :         RemoteStorageConfig::from_toml_str(input)
     250            9 :     }
     251              : 
     252              :     #[test]
     253            3 :     fn parse_localfs_config_with_timeout() {
     254            3 :         let input = "local_path = '.'
     255            3 : timeout = '5s'";
     256            3 : 
     257            3 :         let config = parse(input).unwrap();
     258            3 : 
     259            3 :         assert_eq!(
     260            3 :             config,
     261            3 :             RemoteStorageConfig {
     262            3 :                 storage: RemoteStorageKind::LocalFs {
     263            3 :                     local_path: Utf8PathBuf::from(".")
     264            3 :                 },
     265            3 :                 timeout: Duration::from_secs(5),
     266            3 :                 small_timeout: RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT
     267            3 :             }
     268            3 :         );
     269            3 :     }
     270              : 
     271              :     #[test]
     272            3 :     fn test_s3_parsing() {
     273            3 :         let toml = "\
     274            3 :     bucket_name = 'foo-bar'
     275            3 :     bucket_region = 'eu-central-1'
     276            3 :     upload_storage_class = 'INTELLIGENT_TIERING'
     277            3 :     timeout = '7s'
     278            3 :     ";
     279            3 : 
     280            3 :         let config = parse(toml).unwrap();
     281            3 : 
     282            3 :         assert_eq!(
     283            3 :             config,
     284            3 :             RemoteStorageConfig {
     285            3 :                 storage: RemoteStorageKind::AwsS3(S3Config {
     286            3 :                     bucket_name: "foo-bar".into(),
     287            3 :                     bucket_region: "eu-central-1".into(),
     288            3 :                     prefix_in_bucket: None,
     289            3 :                     endpoint: None,
     290            3 :                     concurrency_limit: default_remote_storage_s3_concurrency_limit(),
     291            3 :                     max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
     292            3 :                     upload_storage_class: Some(StorageClass::IntelligentTiering),
     293            3 :                 }),
     294            3 :                 timeout: Duration::from_secs(7),
     295            3 :                 small_timeout: RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT
     296            3 :             }
     297            3 :         );
     298            3 :     }
     299              : 
     300              :     #[test]
     301            3 :     fn test_storage_class_serde_roundtrip() {
     302            3 :         let classes = [
     303            3 :             None,
     304            3 :             Some(StorageClass::Standard),
     305            3 :             Some(StorageClass::IntelligentTiering),
     306            3 :         ];
     307           12 :         for class in classes {
     308            9 :             #[derive(Serialize, Deserialize)]
     309            9 :             struct Wrapper {
     310              :                 #[serde(
     311              :                     deserialize_with = "deserialize_storage_class",
     312              :                     serialize_with = "serialize_storage_class"
     313              :                 )]
     314              :                 class: Option<StorageClass>,
     315              :             }
     316            9 :             let wrapped = Wrapper {
     317            9 :                 class: class.clone(),
     318            9 :             };
     319            9 :             let serialized = serde_json::to_string(&wrapped).unwrap();
     320            9 :             let deserialized: Wrapper = serde_json::from_str(&serialized).unwrap();
     321            9 :             assert_eq!(class, deserialized.class);
     322              :         }
     323            3 :     }
     324              : 
     325              :     #[test]
     326            3 :     fn test_azure_parsing() {
     327            3 :         let toml = "\
     328            3 :     container_name = 'foo-bar'
     329            3 :     container_region = 'westeurope'
     330            3 :     upload_storage_class = 'INTELLIGENT_TIERING'
     331            3 :     timeout = '7s'
     332            3 :     conn_pool_size = 8
     333            3 :     ";
     334            3 : 
     335            3 :         let config = parse(toml).unwrap();
     336            3 : 
     337            3 :         assert_eq!(
     338            3 :             config,
     339            3 :             RemoteStorageConfig {
     340            3 :                 storage: RemoteStorageKind::AzureContainer(AzureConfig {
     341            3 :                     container_name: "foo-bar".into(),
     342            3 :                     storage_account: None,
     343            3 :                     container_region: "westeurope".into(),
     344            3 :                     prefix_in_container: None,
     345            3 :                     concurrency_limit: default_remote_storage_azure_concurrency_limit(),
     346            3 :                     max_keys_per_list_response: DEFAULT_MAX_KEYS_PER_LIST_RESPONSE,
     347            3 :                     conn_pool_size: 8,
     348            3 :                 }),
     349            3 :                 timeout: Duration::from_secs(7),
     350            3 :                 small_timeout: RemoteStorageConfig::DEFAULT_SMALL_TIMEOUT
     351            3 :             }
     352            3 :         );
     353            3 :     }
     354              : }
        

Generated by: LCOV version 2.1-beta