Use pool in cache and uploader and fix memory limiter

Signed-off-by: Alessandro Passaro <alexpax@amazon.co.uk>

Author: passaro

mountpoint-s3-fs/benches/cache_serialization.rs

@@ -8,6 +8,7 @@ use criterion::{BenchmarkGroup, Criterion, criterion_group, criterion_main};
 
 use mountpoint_s3_client::types::ETag;
 use mountpoint_s3_fs::data_cache::{ChecksummedBytes, DataCache, DiskDataCache, DiskDataCacheConfig};
+use mountpoint_s3_fs::memory::PagedPool;
 use mountpoint_s3_fs::object::ObjectId;
 use rand::RngCore;
 use tempfile::TempDir;
@@ -39,7 +40,8 @@ fn cache_read_benchmark(group: &mut BenchmarkGroup<'_, WallTime>, dir_path: &Pat
         block_size: BLOCK_SIZE,
         limit: mountpoint_s3_fs::data_cache::CacheLimit::Unbounded,
     };
-    let cache = DiskDataCache::new(config);
+    let pool = PagedPool::new([BLOCK_SIZE as usize]);
+    let cache = DiskDataCache::new(config, pool);
     let cache_key = ObjectId::new("a".into(), ETag::for_tests());
     let bytes = ChecksummedBytes::new(data.to_owned().into());
 

mountpoint-s3-fs/examples/fs_benchmark.rs

@@ -8,6 +8,7 @@ use fuser::{BackgroundSession, MountOption, Session};
 use mountpoint_s3_client::S3CrtClient;
 use mountpoint_s3_client::config::{EndpointConfig, RustLogAdapter, S3ClientConfig};
 use mountpoint_s3_fs::fuse::S3FuseFilesystem;
+use mountpoint_s3_fs::memory::PagedPool;
 use mountpoint_s3_fs::prefetch::Prefetcher;
 use mountpoint_s3_fs::{Runtime, S3Filesystem, S3FilesystemConfig};
 use tempfile::tempdir;
@@ -139,11 +140,13 @@ fn mount_file_system(
     region: &str,
     throughput_target_gbps: Option<f64>,
 ) -> BackgroundSession {
+    let pool = PagedPool::new([8 * 1024 * 1024]);
     let mut config = S3ClientConfig::new().endpoint_config(EndpointConfig::new(region));
     let initial_read_window_size = 1024 * 1024 + 128 * 1024;
     config = config
         .read_backpressure(true)
-        .initial_read_window(initial_read_window_size);
+        .initial_read_window(initial_read_window_size)
+        .memory_pool(pool.clone());
     if let Some(throughput_target_gbps) = throughput_target_gbps {
         config = config.throughput_target_gbps(throughput_target_gbps);
     }
@@ -164,6 +167,7 @@ fn mount_file_system(
     let fs = S3Filesystem::new(
         client,
         prefetcher_builder,
+        pool,
         runtime,
         bucket_name,
         &Default::default(),

mountpoint-s3-fs/examples/mount_from_config.rs

@@ -243,7 +243,7 @@ fn mount_filesystem(
 
     // 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, pool)
         .context("Failed to create FUSE session")?;
 
     Ok(fuse_session)

mountpoint-s3-fs/examples/prefetch_benchmark.rs

@@ -13,6 +13,7 @@ use mountpoint_s3_client::types::HeadObjectParams;
 use mountpoint_s3_client::{ObjectClient, S3CrtClient};
 use mountpoint_s3_fs::Runtime;
 use mountpoint_s3_fs::mem_limiter::MemoryLimiter;
+use mountpoint_s3_fs::memory::PagedPool;
 use mountpoint_s3_fs::object::ObjectId;
 use mountpoint_s3_fs::prefetch::{PrefetchGetObject, Prefetcher, PrefetcherConfig};
 use serde_json::{json, to_writer};
@@ -164,9 +165,10 @@ fn main() -> anyhow::Result<()> {
     let args = CliArgs::parse();
 
     let bucket = args.bucket.as_str();
-    let client_config = args.s3_client_config();
+    let pool = PagedPool::new([args.part_size.unwrap_or(8 * 1024 * 1024) as usize]);
+    let client_config = args.s3_client_config().memory_pool(pool.clone());
     let client = S3CrtClient::new(client_config).context("failed to create S3 CRT client")?;
-    let mem_limiter = Arc::new(MemoryLimiter::new(client.clone(), args.memory_target()));
+    let mem_limiter = Arc::new(MemoryLimiter::new(pool, args.memory_target()));
     let runtime = Runtime::new(client.event_loop_group());
 
     // Verify if all objects exist and collect metadata

mountpoint-s3-fs/examples/upload_benchmark.rs

@@ -6,6 +6,7 @@ use mountpoint_s3_client::config::{Allocator, EndpointConfig, RustLogAdapter, S3
 use mountpoint_s3_client::types::ChecksumAlgorithm;
 use mountpoint_s3_client::{ObjectClient, S3CrtClient};
 use mountpoint_s3_fs::mem_limiter::MemoryLimiter;
+use mountpoint_s3_fs::memory::PagedPool;
 use mountpoint_s3_fs::upload::{Uploader, UploaderConfig};
 use mountpoint_s3_fs::{Runtime, ServerSideEncryption};
 use sysinfo::{RefreshKind, System};
@@ -100,10 +101,12 @@ fn main() {
         let endpoint_uri = Uri::new_from_str(&Allocator::default(), url).expect("Failed to parse endpoint URL");
         endpoint_config = endpoint_config.endpoint(endpoint_uri);
     }
+    let pool = PagedPool::new([args.write_part_size]);
     let mut config = S3ClientConfig::new()
         .endpoint_config(endpoint_config)
         .throughput_target_gbps(args.throughput_target_gbps as f64)
-        .write_part_size(args.write_part_size);
+        .write_part_size(args.write_part_size)
+        .memory_pool(pool.clone());
     if let Some(crt_mem_limit_gib) = args.crt_memory_limit_gib {
         config = config.memory_limit_in_bytes(crt_mem_limit_gib * 1024 * 1024 * 1024);
     }
@@ -118,7 +121,7 @@ fn main() {
             let sys = System::new_with_specifics(RefreshKind::everything());
             (sys.total_memory() as f64 * 0.95) as u64
         };
-        let mem_limiter = Arc::new(MemoryLimiter::new(client.clone(), max_memory_target));
+        let mem_limiter = Arc::new(MemoryLimiter::new(pool.clone(), max_memory_target));
 
         let buffer_size = args.write_part_size;
         let server_side_encryption = ServerSideEncryption::new(args.sse.clone(), args.sse_kms_key_id.clone());
@@ -134,6 +137,7 @@ fn main() {
         let uploader = Uploader::new(
             client.clone(),
             runtime.clone(),
+            pool.clone(),
             mem_limiter,
             UploaderConfig::new(buffer_size)
                 .server_side_encryption(server_side_encryption)

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,15 +47,18 @@ impl MountpointConfig {
         s3_path: S3Path,
         client: Client,
         runtime: Runtime,
+        memory_pool: 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.clone())?;
         tracing::trace!(filesystem_config=?self.filesystem_config, "creating file system");
         let fs = S3Filesystem::new(
             client,
             prefetcher_builder,
+            memory_pool,
             runtime,
             &s3_path.bucket_name,
             &s3_path.prefix,
@@ -72,11 +76,14 @@ fn create_prefetcher_builder<Client>(
     data_cache_config: DataCacheConfig,
     client: &Client,
     runtime: &Runtime,
+    memory_pool: 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(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::{BufferKind, PagedPool};
 use crate::object::ObjectId;
 use crate::sync::Mutex;
 
@@ -34,6 +35,7 @@ const HASHED_DIR_SPLIT_INDEX: usize = 2;
 /// On-disk implementation of [DataCache].
 pub struct DiskDataCache {
     config: DiskDataCacheConfig,
+    pool: PagedPool,
     /// Tracks blocks usage. `None` when no cache limit was set.
     usage: Option<Mutex<UsageInfo<DiskBlockKey>>>,
 }
@@ -246,20 +248,22 @@ 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: &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 = {
+            let mut buffer = pool.get_buffer_mut(size, BufferKind::DiskCache);
+            buffer.set_len_uninit(size);
+            reader.read_exact(buffer.as_mut())?;
+            buffer.into_bytes()
+        };
 
-        Ok(Self {
-            header,
-            data: buffer.into(),
-        })
+        Ok(Self { header, data })
     }
 
     /// Serialize this instance to `writer` and return the number of bytes written on success.
@@ -305,12 +309,12 @@ impl From<DiskBlockReadWriteError> for DataCacheError {
 
 impl DiskDataCache {
     /// Create a new instance of an [DiskDataCache] with the specified configuration.
-    pub fn new(config: DiskDataCacheConfig) -> Self {
-        let usage = match config.limit {
+    pub fn new(config: DiskDataCacheConfig, pool: PagedPool) -> Self {
+        let usage = match &config.limit {
             CacheLimit::Unbounded => None,
             CacheLimit::TotalSize { .. } | CacheLimit::AvailableSpace { .. } => Some(Mutex::new(UsageInfo::new())),
         };
-        DiskDataCache { config, usage }
+        DiskDataCache { config, pool, usage }
     }
 
     /// Get the relative path for the given block.
@@ -349,7 +353,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)
             .inspect_err(|e| warn!(path = ?path.as_ref(), "block could not be deserialized: {:?}", e))?;
         let bytes = block
             .data(cache_key, block_idx, block_offset)
@@ -659,11 +663,15 @@ mod tests {
     #[test]
     fn get_path_for_block_key() {
         let cache_dir = PathBuf::from("mountpoint-cache/");
-        let data_cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_dir,
-            block_size: 1024,
-            limit: CacheLimit::Unbounded,
-        });
+        let pool = PagedPool::new([1024]);
+        let data_cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_dir,
+                block_size: 1024,
+                limit: CacheLimit::Unbounded,
+            },
+            pool,
+        );
 
         let s3_key = "a".repeat(266);
         let etag = ETag::for_tests();
@@ -687,11 +695,15 @@ mod tests {
     #[test]
     fn get_path_for_block_key_huge_block_index() {
         let cache_dir = PathBuf::from("mountpoint-cache/");
-        let data_cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_dir,
-            block_size: 1024,
-            limit: CacheLimit::Unbounded,
-        });
+        let pool = PagedPool::new([1024]);
+        let data_cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_dir,
+                block_size: 1024,
+                limit: CacheLimit::Unbounded,
+            },
+            pool,
+        );
 
         let s3_key = "a".repeat(266);
         let etag = ETag::for_tests();
@@ -723,11 +735,15 @@ mod tests {
 
         let block_size = 8 * 1024 * 1024;
         let cache_directory = tempfile::tempdir().unwrap();
-        let cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_directory.path().to_path_buf(),
-            block_size,
-            limit: CacheLimit::Unbounded,
-        });
+        let pool = PagedPool::new([block_size as usize]);
+        let cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_directory.path().to_path_buf(),
+                block_size,
+                limit: CacheLimit::Unbounded,
+            },
+            pool,
+        );
         let cache_key_1 = ObjectId::new("a".into(), ETag::for_tests());
         let cache_key_2 = ObjectId::new(
             "long-key_".repeat(100), // at least 900 chars, exceeding easily 255 chars (UNIX filename limit)
@@ -806,11 +822,15 @@ mod tests {
         let slice = data.slice(1..5);
 
         let cache_directory = tempfile::tempdir().unwrap();
-        let cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_directory.path().to_path_buf(),
-            block_size: 8 * 1024 * 1024,
-            limit: CacheLimit::Unbounded,
-        });
+        let pool = PagedPool::new([8 * 1024 * 1024]);
+        let cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_directory.path().to_path_buf(),
+                block_size: 8 * 1024 * 1024,
+                limit: CacheLimit::Unbounded,
+            },
+            pool,
+        );
         let cache_key = ObjectId::new("a".into(), ETag::for_tests());
 
         cache
@@ -884,11 +904,15 @@ mod tests {
         let small_object_key = ObjectId::new("small".into(), ETag::for_tests());
 
         let cache_directory = tempfile::tempdir().unwrap();
-        let cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_directory.path().to_path_buf(),
-            block_size: BLOCK_SIZE as u64,
-            limit: CacheLimit::TotalSize { max_size: CACHE_LIMIT },
-        });
+        let pool = PagedPool::new([BLOCK_SIZE]);
+        let cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_directory.path().to_path_buf(),
+                block_size: BLOCK_SIZE as u64,
+                limit: CacheLimit::TotalSize { max_size: CACHE_LIMIT },
+            },
+            pool,
+        );
 
         // Put all of large_object
         for (block_idx, bytes) in large_object_blocks.iter().enumerate() {
@@ -1063,7 +1087,8 @@ 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 pool = PagedPool::new([MAX_LENGTH as usize]);
+        let err = DiskBlock::read(&mut Cursor::new(buf), MAX_LENGTH, &pool).expect_err("deserialization should fail");
         match length_to_corrupt {
             "key" | "etag" => assert!(matches!(
                 err,
@@ -1078,11 +1103,15 @@ mod tests {
     fn test_concurrent_access() {
         let block_size = 1024 * 1024;
         let cache_directory = tempfile::tempdir().unwrap();
-        let data_cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_directory.path().to_path_buf(),
-            block_size: block_size as u64,
-            limit: CacheLimit::Unbounded,
-        });
+        let pool = PagedPool::new([block_size]);
+        let data_cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_directory.path().to_path_buf(),
+                block_size: block_size as u64,
+                limit: CacheLimit::Unbounded,
+            },
+            pool,
+        );
         let data_cache = Arc::new(data_cache);
 
         let cache_key = ObjectId::new("foo".to_owned(), ETag::for_tests());

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

@@ -116,6 +116,7 @@ mod tests {
     use super::*;
     use crate::checksums::ChecksummedBytes;
     use crate::data_cache::{CacheLimit, DiskDataCache, DiskDataCacheConfig, ExpressDataCache, ExpressDataCacheConfig};
+    use crate::memory::PagedPool;
 
     use futures::executor::ThreadPool;
     use mountpoint_s3_client::mock_client::MockClient;
@@ -128,11 +129,15 @@ mod tests {
 
     fn create_disk_cache() -> (TempDir, Arc<DiskDataCache>) {
         let cache_directory = tempfile::tempdir().unwrap();
-        let cache = DiskDataCache::new(DiskDataCacheConfig {
-            cache_directory: cache_directory.path().to_path_buf(),
-            block_size: BLOCK_SIZE,
-            limit: CacheLimit::Unbounded,
-        });
+        let pool = PagedPool::new([BLOCK_SIZE as usize, PART_SIZE]);
+        let cache = DiskDataCache::new(
+            DiskDataCacheConfig {
+                cache_directory: cache_directory.path().to_path_buf(),
+                block_size: BLOCK_SIZE,
+                limit: CacheLimit::Unbounded,
+            },
+            pool,
+        );
         (cache_directory, Arc::new(cache))
     }