Use the paged memory pool in Mountpoint

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

Author: passaro

mountpoint-s3-fs/examples/mount_from_config.rs

@@ -19,6 +19,7 @@ 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, TargetThroughputSetting},
@@ -234,14 +235,15 @@ fn mount_filesystem(
 
     // Create the client and runtime
     let client_config = config.build_client_config()?;
+    let pool = PagedPool::new([client_config.part_config.read_size_bytes]);
     let client = client_config
-        .create_client(None)
+        .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.
@@ -273,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.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,6 +19,7 @@ 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::{ClientConfig, S3Path, TargetThroughputSetting};
 
@@ -29,6 +30,7 @@ fn main() -> anyhow::Result<()> {
 
 pub fn create_mock_client(
     client_config: ClientConfig,
+    _pool: PagedPool,
     s3_path: &S3Path,
     personality: Option<S3Personality>,
 ) -> anyhow::Result<(Arc<ThroughputMockClient>, Runtime, S3Personality)> {

mountpoint-s3/src/run.rs

@@ -10,6 +10,7 @@ use mountpoint_s3_client::{ObjectClient, S3CrtClient};
 use mountpoint_s3_fs::data_cache::{DataCacheConfig, ManagedCacheDir};
 use mountpoint_s3_fs::fuse::session::FuseSession;
 use mountpoint_s3_fs::logging::init_logging;
+use mountpoint_s3_fs::memory::PagedPool;
 use mountpoint_s3_fs::s3::S3Personality;
 use mountpoint_s3_fs::s3::config::{ClientConfig, S3Path};
 use mountpoint_s3_fs::{MountpointConfig, Runtime, metrics};
@@ -174,8 +175,17 @@ fn mount(args: CliArgs, client_builder: impl ClientBuilder) -> anyhow::Result<Fu
 
     let client_config = args.client_config(build_info::FULL_VERSION);
 
+    // Set up a paged memory pool
+    let pool = PagedPool::new([
+        1024 * 1024,
+        client_config.part_config.read_size_bytes,
+        client_config.part_config.write_size_bytes,
+    ]);
+    pool.schedule_trim(Duration::from_secs(60));
+
     let s3_path = args.s3_path()?;
-    let (client, runtime, s3_personality) = client_builder.build(client_config, &s3_path, args.personality())?;
+    let (client, runtime, s3_personality) =
+        client_builder.build(client_config, pool.clone(), &s3_path, args.personality())?;
 
     let bucket_description = args.bucket_description()?;
     tracing::debug!("using S3 personality {s3_personality:?} for {bucket_description}");
@@ -189,7 +199,7 @@ fn mount(args: CliArgs, client_builder: impl ClientBuilder) -> anyhow::Result<Fu
     let mount_point_path = format!("{}", fuse_session_config.mount_point());
 
     let mut fuse_session = MountpointConfig::new(fuse_session_config, filesystem_config, data_cache_config)
-        .create_fuse_session(s3_path, client, runtime)?;
+        .create_fuse_session(s3_path, client, runtime, Some(pool))?;
     tracing::info!("successfully mounted {} at {}", bucket_description, mount_point_path);
 
     if let Some(managed_cache_dir) = managed_cache_dir {
@@ -208,36 +218,39 @@ pub trait ClientBuilder {
     fn build(
         self,
         client_config: ClientConfig,
+        pool: PagedPool,
         s3_path: &S3Path,
         personality: Option<S3Personality>,
     ) -> anyhow::Result<(Self::Client, Runtime, S3Personality)>;
 }
 
 impl<F, C> ClientBuilder for F
 where
-    F: FnOnce(ClientConfig, &S3Path, Option<S3Personality>) -> anyhow::Result<(C, Runtime, S3Personality)>,
+    F: FnOnce(ClientConfig, PagedPool, &S3Path, Option<S3Personality>) -> anyhow::Result<(C, Runtime, S3Personality)>,
     C: ObjectClient + Clone + Send + Sync + 'static,
 {
     type Client = C;
 
     fn build(
         self,
         client_config: ClientConfig,
+        pool: PagedPool,
         s3_path: &S3Path,
         personality: Option<S3Personality>,
     ) -> anyhow::Result<(Self::Client, Runtime, S3Personality)> {
-        self(client_config, s3_path, personality)
+        self(client_config, pool, s3_path, personality)
     }
 }
 
 // Create a real S3 client
 pub fn create_s3_client(
     client_config: ClientConfig,
+    pool: PagedPool,
     s3_path: &S3Path,
     personality: Option<S3Personality>,
 ) -> anyhow::Result<(S3CrtClient, Runtime, S3Personality)> {
     let client = client_config
-        .create_client(Some(s3_path))
+        .create_client(pool, Some(s3_path))
         .context("Failed to create S3 client")?;
 
     let runtime = Runtime::new(client.event_loop_group());