Use the paged memory pool + ClientBuilder changes

Signed-off-by: Alessandro Passaro <[email protected]>

Author: passaro

mountpoint-s3-fs/examples/mount_from_config.rs

@@ -19,9 +19,10 @@ use mountpoint_s3_fs::{
     },
     logging::{LoggingConfig, LoggingHandle, error_logger::FileErrorLogger, init_logging},
     manifest::{Manifest, ingest_manifest},
+    memory::PagedPool,
     metrics::{self, MetricsSinkHandle},
     prefix::Prefix,
-    s3::config::{ClientConfig, PartConfig, Region, S3Path},
+    s3::config::{ClientConfig, PartConfig, Region, S3Path, TargetThroughputSetting},
 };
 use nix::sys::signal::{self, Signal};
 use nix::unistd::Pid;
@@ -121,17 +122,17 @@ impl ConfigOptions {
             Some(prefix) => format!("{prefix}/mp-exmpl"),
             None => "mountpoint-s3-example/mp-exmpl".to_string(),
         };
-        let target_throughput = self.determine_throughput()?;
+        let throughput_target = self.determine_throughput()?;
         Ok(ClientConfig {
             region: Region::new_user_specified(self.region.clone()),
             endpoint_url: self.endpoint_url.clone(),
             addressing_style: AddressingStyle::Automatic,
             dual_stack: false,
             transfer_acceleration: false,
-            auth_config: mountpoint_s3_client::config::S3ClientAuthConfig::Default,
+            auth_config: Default::default(),
             requester_pays: false,
             expected_bucket_owner: self.expected_bucket_owner.clone(),
-            throughput_target_gbps: target_throughput,
+            throughput_target,
             bind: None,
             part_config: PartConfig::with_part_size(self.part_size.unwrap_or(8388608)),
             user_agent: UserAgent::new(Some(user_agent_string)),
@@ -157,23 +158,28 @@ impl ConfigOptions {
         DataCacheConfig::default()
     }
 
-    fn determine_throughput(&self) -> Result<f64> {
+    fn determine_throughput(&self) -> Result<TargetThroughputSetting> {
         match &self.throughput_config {
             // TODO(chagem): Remove some code duplication, by moving this logic into fs crate.
-            ThroughputConfig::Explicit { throughput } => Ok(*throughput),
+            ThroughputConfig::Explicit { throughput } => Ok(TargetThroughputSetting::User { gbps: *throughput }),
             ThroughputConfig::IMDSAutoConfigure => {
-                const DEFAULT_THROUGHPUT: f64 = 10.0;
                 let instance_info = InstanceInfo::new();
                 match autoconfigure::network_throughput(&instance_info) {
-                    Ok(throughput) => Ok(throughput),
+                    Ok(throughput) => Ok(TargetThroughputSetting::Instance { gbps: throughput }),
                     Err(e) => {
-                        tracing::warn!("Failed to detect network throughput. Using {DEFAULT_THROUGHPUT} gbps: {e:?}");
-                        Ok(DEFAULT_THROUGHPUT)
+                        tracing::warn!(
+                            "Failed to detect network throughput. Using {} gbps: {:?}",
+                            TargetThroughputSetting::DEFAULT_TARGET_THROUGHPUT,
+                            e
+                        );
+                        Ok(TargetThroughputSetting::Default)
                     }
                 }
             }
             ThroughputConfig::IMDSLookUp { ec2_instance_type } => {
-                autoconfigure::get_maximum_network_throughput(ec2_instance_type).context("Unrecognized instance ID")
+                let target = autoconfigure::get_maximum_network_throughput(ec2_instance_type)
+                    .context("Unrecognized instance ID")?;
+                Ok(TargetThroughputSetting::Instance { gbps: target })
             }
         }
     }
@@ -228,15 +234,16 @@ fn mount_filesystem(
     let s3_path = config.build_s3_path()?;
 
     // Create the client and runtime
-    let client = config
-        .build_client_config()?
-        .create_client(None)
+    let client_config = config.build_client_config()?;
+    let pool = PagedPool::new([client_config.part_config.read_size_bytes]);
+    let client = client_config
+        .create_client(pool.clone(), None)
         .context("Failed to create S3 client")?;
     let runtime = Runtime::new(client.event_loop_group());
 
     // Create and run the FUSE session
     let fuse_session = mp_config
-        .create_fuse_session(s3_path, client, runtime)
+        .create_fuse_session(s3_path, client, runtime, Some(pool))
         .context("Failed to create FUSE session")?;
 
     Ok(fuse_session)

mountpoint-s3-fs/src/config.rs

@@ -6,6 +6,7 @@ use crate::data_cache::{DataCacheConfig, DiskDataCache, ExpressDataCache, Multil
 use crate::fuse::config::FuseSessionConfig;
 use crate::fuse::session::FuseSession;
 use crate::fuse::{ErrorLogger, S3FuseFilesystem};
+use crate::memory::PagedPool;
 use crate::prefetch::{Prefetcher, PrefetcherBuilder};
 use crate::s3::config::S3Path;
 use crate::sync::Arc;
@@ -46,11 +47,12 @@ impl MountpointConfig {
         s3_path: S3Path,
         client: Client,
         runtime: Runtime,
+        memory_pool: Option<PagedPool>,
     ) -> anyhow::Result<FuseSession>
     where
         Client: ObjectClient + Clone + Send + Sync + 'static,
     {
-        let prefetcher_builder = create_prefetcher_builder(self.data_cache_config, &client, &runtime)?;
+        let prefetcher_builder = create_prefetcher_builder(self.data_cache_config, &client, &runtime, memory_pool)?;
         tracing::trace!(filesystem_config=?self.filesystem_config, "creating file system");
         let fs = S3Filesystem::new(
             client,
@@ -72,11 +74,14 @@ fn create_prefetcher_builder<Client>(
     data_cache_config: DataCacheConfig,
     client: &Client,
     runtime: &Runtime,
+    memory_pool: Option<PagedPool>,
 ) -> anyhow::Result<PrefetcherBuilder<Client>>
 where
     Client: ObjectClient + Clone + Send + Sync + 'static,
 {
-    let disk_cache = data_cache_config.disk_cache_config.map(DiskDataCache::new);
+    let disk_cache = data_cache_config
+        .disk_cache_config
+        .map(|config| DiskDataCache::new_with_pool(config, memory_pool));
     let express_cache = match data_cache_config.express_cache_config {
         None => None,
         Some(config) => {

mountpoint-s3-fs/src/data_cache/disk_data_cache.rs

@@ -20,6 +20,7 @@ use tracing::{trace, warn};
 
 use crate::checksums::IntegrityError;
 use crate::data_cache::DataCacheError;
+use crate::memory::PagedPool;
 use crate::object::ObjectId;
 use crate::sync::Mutex;
 
@@ -36,6 +37,7 @@ pub struct DiskDataCache {
     config: DiskDataCacheConfig,
     /// Tracks blocks usage. `None` when no cache limit was set.
     usage: Option<Mutex<UsageInfo<DiskBlockKey>>>,
+    pool: Option<PagedPool>,
 }
 
 /// Configuration for a [DiskDataCache].
@@ -246,20 +248,27 @@ impl DiskBlock {
     }
 
     /// Deserialize an instance from `reader`.
-    fn read(reader: &mut impl Read, block_size: u64) -> Result<Self, DiskBlockReadWriteError> {
+    fn read(
+        reader: &mut impl Read,
+        block_size: u64,
+        pool: Option<&PagedPool>,
+    ) -> Result<Self, DiskBlockReadWriteError> {
         let header: DiskBlockHeader = bincode::decode_from_std_read(reader, BINCODE_CONFIG)?;
 
         if header.block_len > block_size {
             return Err(DiskBlockReadWriteError::InvalidBlockLength(header.block_len));
         }
 
-        let mut buffer = vec![0u8; header.block_len as usize];
-        reader.read_exact(&mut buffer)?;
+        let size = header.block_len as usize;
+        let data = if let Some(pool) = pool {
+            pool.read_exact(reader, size)?
+        } else {
+            let mut buffer = vec![0u8; size];
+            reader.read_exact(&mut buffer)?;
+            buffer.into()
+        };
 
-        Ok(Self {
-            header,
-            data: buffer.into(),
-        })
+        Ok(Self { header, data })
     }
 
     /// Serialize this instance to `writer` and return the number of bytes written on success.
@@ -303,14 +312,28 @@ impl From<DiskBlockReadWriteError> for DataCacheError {
     }
 }
 
+fn usage_info(limit: &CacheLimit) -> Option<Mutex<UsageInfo<DiskBlockKey>>> {
+    match limit {
+        CacheLimit::Unbounded => None,
+        CacheLimit::TotalSize { .. } | CacheLimit::AvailableSpace { .. } => Some(Mutex::new(UsageInfo::new())),
+    }
+}
+
 impl DiskDataCache {
     /// Create a new instance of an [DiskDataCache] with the specified configuration.
     pub fn new(config: DiskDataCacheConfig) -> Self {
-        let usage = match config.limit {
-            CacheLimit::Unbounded => None,
-            CacheLimit::TotalSize { .. } | CacheLimit::AvailableSpace { .. } => Some(Mutex::new(UsageInfo::new())),
-        };
-        DiskDataCache { config, usage }
+        let usage = usage_info(&config.limit);
+        DiskDataCache {
+            config,
+            usage,
+            pool: None,
+        }
+    }
+
+    /// Create a new instance of an [DiskDataCache] with the specified configuration.
+    pub fn new_with_pool(config: DiskDataCacheConfig, pool: Option<PagedPool>) -> Self {
+        let usage = usage_info(&config.limit);
+        DiskDataCache { config, usage, pool }
     }
 
     /// Get the relative path for the given block.
@@ -349,7 +372,7 @@ impl DiskDataCache {
             return Err(DataCacheError::InvalidBlockContent);
         }
 
-        let block = DiskBlock::read(&mut file, self.block_size())
+        let block = DiskBlock::read(&mut file, self.block_size(), self.pool.as_ref())
             .inspect_err(|e| warn!(path = ?path.as_ref(), "block could not be deserialized: {:?}", e))?;
         let bytes = block
             .data(cache_key, block_idx, block_offset)
@@ -1063,7 +1086,7 @@ mod tests {
         // "Corrupt" the serialized value with an invalid length.
         replace_u64_at(&mut buf, offset, u64::MAX);
 
-        let err = DiskBlock::read(&mut Cursor::new(buf), MAX_LENGTH).expect_err("deserialization should fail");
+        let err = DiskBlock::read(&mut Cursor::new(buf), MAX_LENGTH, None).expect_err("deserialization should fail");
         match length_to_corrupt {
             "key" | "etag" => assert!(matches!(
                 err,

mountpoint-s3-fs/src/s3/config.rs

@@ -11,6 +11,7 @@ use mountpoint_s3_client::{ObjectClient, S3CrtClient, S3RequestError};
 use regex::Regex;
 use thiserror::Error;
 
+use crate::memory::PagedPool;
 use crate::prefix::{Prefix, PrefixError};
 
 /// Configuration for the S3 Client to use in Mountpoint.
@@ -41,7 +42,7 @@ pub struct ClientConfig {
     pub expected_bucket_owner: Option<String>,
 
     /// Target throughput in Gbps
-    pub throughput_target_gbps: f64,
+    pub throughput_target: TargetThroughputSetting,
 
     /// One or more network interfaces to use when accessing S3
     pub bind: Option<Vec<String>>,
@@ -182,10 +183,10 @@ fn matches_bucket_regex(bucket_name: &str) -> bool {
 #[derive(Debug)]
 pub struct PartConfig {
     /// Part size for GET in bytes
-    read_size_bytes: usize,
+    pub read_size_bytes: usize,
 
     /// Part size for multi-part PUT in bytes
-    write_size_bytes: usize,
+    pub write_size_bytes: usize,
 }
 
 impl PartConfig {
@@ -238,6 +239,26 @@ impl Display for Region {
     }
 }
 
+/// Target throughput setting.
+#[derive(Debug, Clone, Copy)]
+pub enum TargetThroughputSetting {
+    Default,
+    User { gbps: f64 },
+    Instance { gbps: f64 },
+}
+
+impl TargetThroughputSetting {
+    pub const DEFAULT_TARGET_THROUGHPUT: f64 = 10.0;
+
+    pub fn value(&self) -> f64 {
+        match self {
+            TargetThroughputSetting::Default => Self::DEFAULT_TARGET_THROUGHPUT,
+            TargetThroughputSetting::User { gbps } => *gbps,
+            TargetThroughputSetting::Instance { gbps } => *gbps,
+        }
+    }
+}
+
 // This is a weird looking number! We really want our first request size to be 1MiB,
 // which is a common IO size. But Linux's readahead will try to read an extra 128k on on
 // top of a 1MiB read, which we'd have to wait for a second request to service. Because
@@ -253,15 +274,20 @@ const INITIAL_READ_WINDOW_SIZE: usize = 1024 * 1024 + 128 * 1024;
 
 impl ClientConfig {
     /// Create an [S3CrtClient]
-    pub fn create_client(self, validate_on_s3_path: Option<&S3Path>) -> anyhow::Result<S3CrtClient> {
+    pub fn create_client(
+        self,
+        memory_pool: PagedPool,
+        validate_on_s3_path: Option<&S3Path>,
+    ) -> anyhow::Result<S3CrtClient> {
         let mut client_config = S3ClientConfig::new()
             .auth_config(self.auth_config)
-            .throughput_target_gbps(self.throughput_target_gbps)
+            .throughput_target_gbps(self.throughput_target.value())
             .read_part_size(self.part_config.read_size_bytes)
             .write_part_size(self.part_config.write_size_bytes)
             .read_backpressure(true)
             .initial_read_window(INITIAL_READ_WINDOW_SIZE)
-            .user_agent(self.user_agent);
+            .user_agent(self.user_agent)
+            .memory_pool(memory_pool);
         if let Some(interfaces) = self.bind {
             client_config = client_config.network_interface_names(interfaces);
         }

mountpoint-s3/src/bin/mock-mount-s3.rs

@@ -19,56 +19,55 @@ use mountpoint_s3_client::mock_client::throughput_client::ThroughputMockClient;
 use mountpoint_s3_client::mock_client::{MockClient, MockObject};
 use mountpoint_s3_client::types::ETag;
 use mountpoint_s3_fs::Runtime;
+use mountpoint_s3_fs::memory::PagedPool;
 use mountpoint_s3_fs::s3::S3Personality;
-use mountpoint_s3_fs::s3::config::BucketNameOrS3Uri;
+use mountpoint_s3_fs::s3::config::{ClientConfig, S3Path, TargetThroughputSetting};
 
 fn main() -> anyhow::Result<()> {
     let cli_args = CliArgs::parse();
     mountpoint_s3::run(create_mock_client, cli_args)
 }
 
-fn create_mock_client(args: &CliArgs) -> anyhow::Result<(Arc<ThroughputMockClient>, Runtime, S3Personality)> {
-    let bucket_name: String = match &args.bucket_name {
-        BucketNameOrS3Uri::BucketName(bucket_name) => bucket_name.clone().into(),
-        BucketNameOrS3Uri::S3Uri(_) => panic!("mock-mount-s3 bucket names do not support S3 URIs"),
-    };
-
+pub fn create_mock_client(
+    client_config: ClientConfig,
+    _pool: PagedPool,
+    s3_path: S3Path,
+    personality: Option<S3Personality>,
+) -> anyhow::Result<(Arc<ThroughputMockClient>, Runtime, S3Personality)> {
     // An extra little safety thing to make sure we can distinguish the real mount-s3 binary and
     // this one. Buckets starting with "sthree-" are always invalid against real S3:
     // https://docs.aws.amazon.com/AmazonS3/latest/userguide/bucketnamingrules.html
     anyhow::ensure!(
-        &bucket_name.starts_with("sthree-"),
+        &s3_path.bucket_name.starts_with("sthree-"),
         "mock-mount-s3 bucket names must start with `sthree-`"
     );
 
-    tracing::warn!("using mock client");
-
     // TODO: Actually update the mock client to support different part sizes
     let part_size = {
-        if args.read_part_size.is_some() || args.write_part_size.is_some() {
+        if client_config.part_config.read_size_bytes != client_config.part_config.write_size_bytes {
             tracing::warn!("mock client does not support separate part sizes for reading and writing, ignoring");
         }
-        args.part_size
+        client_config.part_config.read_size_bytes
     };
 
-    let s3_personality = if let Some(bucket_type) = &args.bucket_type {
-        bucket_type.to_personality()
-    } else {
-        S3Personality::Standard
-    };
+    let s3_personality = personality.unwrap_or(S3Personality::Standard);
 
     let config = MockClient::config()
-        .bucket(&bucket_name)
+        .bucket(&s3_path.bucket_name)
         .part_size(part_size as usize)
         .unordered_list_seed(None)
         .enable_backpressure(true)
         .initial_read_window_size(1024 * 1024 + 128 * 1024) // matching real MP
         .enable_rename(s3_personality.supports_rename_object());
 
-    let client = if let Some(max_throughput_gbps) = args.maximum_throughput_gbps {
+    let client = if let TargetThroughputSetting::User {
+        gbps: max_throughput_gbps,
+    } = client_config.throughput_target
+    {
         tracing::info!("mock client limited to {max_throughput_gbps} Gb/s download throughput");
-        ThroughputMockClient::new(config, max_throughput_gbps as f64)
+        ThroughputMockClient::new(config, max_throughput_gbps)
     } else {
+        tracing::info!("mock client with no throughput limit");
         ThroughputMockClient::new_unlimited_throughput(config)
     };