Decouple mem_limiter from client

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

Author: passaro

mountpoint-s3-fs/examples/prefetch_benchmark.rs

@@ -124,15 +124,37 @@ fn parse_duration(arg: &str) -> Result<Duration, String> {
         .map_err(|e| format!("Invalid duration: {e}"))
 }
 
-fn create_memory_limiter(args: &CliArgs, client: &S3CrtClient) -> Arc<MemoryLimiter<S3CrtClient>> {
-    let max_memory_target = if let Some(target) = args.max_memory_target {
-        target * 1024 * 1024
-    } else {
-        // Default to 95% of total system memory
-        let sys = System::new_with_specifics(RefreshKind::everything());
-        (sys.total_memory() as f64 * 0.95) as u64
-    };
-    Arc::new(MemoryLimiter::new(client.clone(), max_memory_target))
+impl CliArgs {
+    fn memory_target(&self) -> u64 {
+        if let Some(target) = self.max_memory_target {
+            target * 1024 * 1024
+        } else {
+            // Default to 95% of total system memory
+            let sys = System::new_with_specifics(RefreshKind::everything());
+            (sys.total_memory() as f64 * 0.95) as u64
+        }
+    }
+
+    fn s3_client_config(&self) -> S3ClientConfig {
+        let initial_read_window_size = 1024 * 1024 + 128 * 1024;
+        let mut client_config = S3ClientConfig::new()
+            .read_backpressure(true)
+            .initial_read_window(initial_read_window_size)
+            .endpoint_config(EndpointConfig::new(self.region.as_str()));
+        if let Some(throughput_target_gbps) = self.maximum_throughput_gbps {
+            client_config = client_config.throughput_target_gbps(throughput_target_gbps as f64);
+        }
+        if let Some(limit_gib) = self.crt_memory_limit_gib {
+            client_config = client_config.memory_limit_in_bytes(limit_gib * 1024 * 1024 * 1024);
+        }
+        if let Some(part_size) = self.part_size {
+            client_config = client_config.part_size(part_size as usize);
+        }
+        if let Some(nics) = &self.bind {
+            client_config = client_config.network_interface_names(nics.to_vec());
+        }
+        client_config
+    }
 }
 
 fn main() -> anyhow::Result<()> {
@@ -142,8 +164,9 @@ fn main() -> anyhow::Result<()> {
     let args = CliArgs::parse();
 
     let bucket = args.bucket.as_str();
-    let client = make_s3_client_from_args(&args).context("failed to create S3 CRT client")?;
-    let mem_limiter = create_memory_limiter(&args, &client);
+    let client_config = args.s3_client_config();
+    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 runtime = Runtime::new(client.event_loop_group());
 
     // Verify if all objects exist and collect metadata
@@ -258,24 +281,3 @@ async fn wait_for_download(
     }
     Ok(total_bytes_read)
 }
-
-fn make_s3_client_from_args(args: &CliArgs) -> anyhow::Result<S3CrtClient> {
-    let initial_read_window_size = 1024 * 1024 + 128 * 1024;
-    let mut client_config = S3ClientConfig::new()
-        .read_backpressure(true)
-        .initial_read_window(initial_read_window_size)
-        .endpoint_config(EndpointConfig::new(args.region.as_str()));
-    if let Some(throughput_target_gbps) = args.maximum_throughput_gbps {
-        client_config = client_config.throughput_target_gbps(throughput_target_gbps as f64);
-    }
-    if let Some(limit_gib) = args.crt_memory_limit_gib {
-        client_config = client_config.memory_limit_in_bytes(limit_gib * 1024 * 1024 * 1024);
-    }
-    if let Some(part_size) = args.part_size {
-        client_config = client_config.part_size(part_size as usize);
-    }
-    if let Some(nics) = &args.bind {
-        client_config = client_config.network_interface_names(nics.to_vec());
-    }
-    Ok(S3CrtClient::new(client_config)?)
-}

mountpoint-s3-fs/src/mem_limiter.rs

@@ -51,8 +51,7 @@ impl BufferArea {
 ///                                     mem reserved by the backpressure controller
 ///                                     (on `BackpressureFeedbackEvent`)
 ///
-#[derive(Debug)]
-pub struct MemoryLimiter<Client: ObjectClient> {
+pub struct MemoryLimiter {
     mem_limit: u64,
     /// Reserved memory for allocations we are tracking, such as buffers we allocate for prefetching.
     /// The memory may not be used yet but has been reserved.
@@ -63,11 +62,24 @@ pub struct MemoryLimiter<Client: ObjectClient> {
     // prefetch takes control over the entire read path but we don't record or control
     // memory usage on the write path today, so we will rely on the client's stats
     // for "other buffers" and adjust the prefetcher read window accordingly.
-    client: Client,
+    get_client_mem: Box<dyn Fn() -> u64 + Send + Sync + 'static>,
+}
+
+impl std::fmt::Debug for MemoryLimiter {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("MemoryLimiter")
+            .field("mem_limit", &self.mem_limit)
+            .field("mem_reserved", &self.mem_reserved)
+            .field("additional_mem_reserved", &self.additional_mem_reserved)
+            .finish_non_exhaustive()
+    }
 }
 
-impl<Client: ObjectClient> MemoryLimiter<Client> {
-    pub fn new(client: Client, mem_limit: u64) -> Self {
+impl MemoryLimiter {
+    pub fn new<Client>(client: Client, mem_limit: u64) -> Self
+    where
+        Client: ObjectClient + Send + Sync + 'static,
+    {
         let min_reserved = 128 * 1024 * 1024;
         let reserved_mem = (mem_limit / 8).max(min_reserved);
         let formatter = make_format(humansize::BINARY);
@@ -77,10 +89,14 @@ impl<Client: ObjectClient> MemoryLimiter<Client> {
             formatter(reserved_mem)
         );
         Self {
-            client,
             mem_limit,
             mem_reserved: AtomicU64::new(0),
             additional_mem_reserved: reserved_mem,
+            get_client_mem: Box::new(move || {
+                client
+                    .mem_usage_stats()
+                    .map_or(0, |stats| stats.primary_allocated.saturating_add(stats.secondary_used))
+            }),
         }
     }
 
@@ -147,8 +163,6 @@ impl<Client: ObjectClient> MemoryLimiter<Client> {
     // where memory is allocated exactly equal to the used memory. So total allocated memory for the CRT client would
     // be `primary_allocated` + `secondary_used`.
     fn client_mem_allocated(&self) -> u64 {
-        self.client
-            .mem_usage_stats()
-            .map_or(0, |stats| stats.primary_allocated.saturating_add(stats.secondary_used))
+        (self.get_client_mem)()
     }
 }

mountpoint-s3-fs/src/prefetch/backpressure_controller.rs

@@ -3,7 +3,6 @@ use std::sync::Arc;
 
 use async_channel::{Receiver, RecvError, Sender, unbounded};
 use humansize::make_format;
-use mountpoint_s3_client::ObjectClient;
 use tracing::trace;
 
 use crate::mem_limiter::{BufferArea, MemoryLimiter};
@@ -35,7 +34,7 @@ pub struct BackpressureConfig {
 /// It is used to send feedback ([Self::send_feedback]) to its corresponding [BackpressureLimiter],
 /// the counterpart which should be leveraged by the stream producer.
 #[derive(Debug)]
-pub struct BackpressureController<Client: ObjectClient> {
+pub struct BackpressureController {
     /// Sender for the [BackpressureLimiter] to receive size increments from the controller.
     read_window_updater: Sender<usize>,
     /// Amount by which the producer should be producing data ahead of [Self::next_read_offset].
@@ -58,7 +57,7 @@ pub struct BackpressureController<Client: ObjectClient> {
     /// Memory limiter is used to guide decisions on how much data to prefetch.
     ///
     /// For example, when memory is low we should scale down [Self::preferred_read_window_size].
-    mem_limiter: Arc<MemoryLimiter<Client>>,
+    mem_limiter: Arc<MemoryLimiter>,
 }
 
 /// The [BackpressureLimiter] is used on producer side of a stream, for example,
@@ -81,10 +80,10 @@ pub struct BackpressureLimiter {
 ///
 /// This pair allows a consumer to send feedback ([BackpressureFeedbackEvent]) when starved or bytes are consumed,
 /// informing a producer (a holder of the [BackpressureLimiter]) when it should provide data more aggressively.
-pub fn new_backpressure_controller<Client: ObjectClient>(
+pub fn new_backpressure_controller(
     config: BackpressureConfig,
-    mem_limiter: Arc<MemoryLimiter<Client>>,
-) -> (BackpressureController<Client>, BackpressureLimiter) {
+    mem_limiter: Arc<MemoryLimiter>,
+) -> (BackpressureController, BackpressureLimiter) {
     // Minimum window size multiplier as the scaling up and down won't work if the multiplier is 1.
     const MIN_WINDOW_SIZE_MULTIPLIER: usize = 2;
     let read_window_end_offset = config.request_range.start + config.initial_read_window_size as u64;
@@ -114,7 +113,7 @@ pub fn new_backpressure_controller<Client: ObjectClient>(
     (controller, limiter)
 }
 
-impl<Client: ObjectClient> BackpressureController<Client> {
+impl BackpressureController {
     pub fn read_window_end_offset(&self) -> u64 {
         self.read_window_end_offset
     }
@@ -234,7 +233,7 @@ impl<Client: ObjectClient> BackpressureController<Client> {
     }
 }
 
-impl<Client: ObjectClient> Drop for BackpressureController<Client> {
+impl Drop for BackpressureController {
     fn drop(&mut self) {
         debug_assert!(
             self.next_read_offset <= self.request_end_offset,
@@ -435,7 +434,7 @@ mod tests {
 
     fn new_backpressure_controller_for_test(
         backpressure_config: BackpressureConfig,
-    ) -> (BackpressureController<MockClient>, BackpressureLimiter) {
+    ) -> (BackpressureController, BackpressureLimiter) {
         let client = MockClient::config()
             .bucket("test-bucket")
             .part_size(8 * 1024 * 1024)

mountpoint-s3-fs/src/prefetch/builder.rs

@@ -38,7 +38,7 @@ where
     pub fn build(
         self,
         runtime: Runtime,
-        mem_limiter: Arc<MemoryLimiter<Client>>,
+        mem_limiter: Arc<MemoryLimiter>,
         prefetcher_config: PrefetcherConfig,
     ) -> Prefetcher<Client> {
         self.inner.build(runtime, mem_limiter, prefetcher_config)
@@ -58,7 +58,7 @@ where
     fn build(
         self: Box<Self>,
         runtime: Runtime,
-        mem_limiter: Arc<MemoryLimiter<Client>>,
+        mem_limiter: Arc<MemoryLimiter>,
         prefetcher_config: PrefetcherConfig,
     ) -> Prefetcher<Client>;
 }
@@ -74,7 +74,7 @@ where
     fn build(
         self: Box<Self>,
         runtime: Runtime,
-        mem_limiter: Arc<MemoryLimiter<Client>>,
+        mem_limiter: Arc<MemoryLimiter>,
         prefetcher_config: PrefetcherConfig,
     ) -> Prefetcher<Client> {
         let part_stream = ClientPartStream::new(runtime, self.client, mem_limiter);
@@ -95,7 +95,7 @@ where
     fn build(
         self: Box<Self>,
         runtime: Runtime,
-        mem_limiter: Arc<MemoryLimiter<Client>>,
+        mem_limiter: Arc<MemoryLimiter>,
         prefetcher_config: PrefetcherConfig,
     ) -> Prefetcher<Client> {
         let part_stream = CachingPartStream::new(runtime, self.client, mem_limiter, self.cache);

mountpoint-s3-fs/src/prefetch/caching_stream.rs

@@ -29,11 +29,11 @@ pub struct CachingPartStream<Cache, Client: ObjectClient + Clone + Send + Sync +
     cache: Arc<Cache>,
     runtime: Runtime,
     client: Client,
-    mem_limiter: Arc<MemoryLimiter<Client>>,
+    mem_limiter: Arc<MemoryLimiter>,
 }
 
 impl<Cache, Client: ObjectClient + Clone + Send + Sync + 'static> CachingPartStream<Cache, Client> {
-    pub fn new(runtime: Runtime, client: Client, mem_limiter: Arc<MemoryLimiter<Client>>, cache: Cache) -> Self {
+    pub fn new(runtime: Runtime, client: Client, mem_limiter: Arc<MemoryLimiter>, cache: Cache) -> Self {
         Self {
             cache: Arc::new(cache),
             runtime,

mountpoint-s3-fs/src/prefetch/part_queue.rs

@@ -23,7 +23,7 @@ pub struct PartQueue<Client: ObjectClient> {
     failed: bool,
     /// The total number of bytes sent to the underlying queue of `self.receiver`
     bytes_received: Arc<AtomicUsize>,
-    mem_limiter: Arc<MemoryLimiter<Client>>,
+    mem_limiter: Arc<MemoryLimiter>,
 }
 
 /// Producer side of the queue of [Part]s.
@@ -36,7 +36,7 @@ pub struct PartQueueProducer<E: std::error::Error> {
 
 /// Creates an unbounded [PartQueue] and its related [PartQueueProducer].
 pub fn unbounded_part_queue<Client: ObjectClient>(
-    mem_limiter: Arc<MemoryLimiter<Client>>,
+    mem_limiter: Arc<MemoryLimiter>,
 ) -> (PartQueue<Client>, PartQueueProducer<Client::ClientError>) {
     let (sender, receiver) = unbounded();
     let bytes_counter = Arc::new(AtomicUsize::new(0));

mountpoint-s3-fs/src/prefetch/part_stream.rs

@@ -195,11 +195,11 @@ impl<Client> Debug for PartStream<Client> {
 pub struct ClientPartStream<Client: ObjectClient + Clone + Send + Sync + 'static> {
     runtime: Runtime,
     client: Client,
-    mem_limiter: Arc<MemoryLimiter<Client>>,
+    mem_limiter: Arc<MemoryLimiter>,
 }
 
 impl<Client: ObjectClient + Clone + Send + Sync + 'static> ClientPartStream<Client> {
-    pub fn new(runtime: Runtime, client: Client, mem_limiter: Arc<MemoryLimiter<Client>>) -> Self {
+    pub fn new(runtime: Runtime, client: Client, mem_limiter: Arc<MemoryLimiter>) -> Self {
         Self {
             runtime,
             client,

mountpoint-s3-fs/src/prefetch/task.rs

@@ -17,15 +17,15 @@ pub struct RequestTask<Client: ObjectClient> {
     remaining: usize,
     range: RequestRange,
     part_queue: PartQueue<Client>,
-    backpressure_controller: BackpressureController<Client>,
+    backpressure_controller: BackpressureController,
 }
 
 impl<Client: ObjectClient> RequestTask<Client> {
     pub fn from_handle(
         task_handle: RemoteHandle<()>,
         range: RequestRange,
         part_queue: PartQueue<Client>,
-        backpressure_controller: BackpressureController<Client>,
+        backpressure_controller: BackpressureController,
     ) -> Self {
         Self {
             _task_handle: task_handle,

mountpoint-s3-fs/src/upload.rs

@@ -28,7 +28,7 @@ pub use incremental::AppendUploadRequest;
 pub struct Uploader<Client: ObjectClient> {
     client: Client,
     runtime: Runtime,
-    mem_limiter: Arc<MemoryLimiter<Client>>,
+    mem_limiter: Arc<MemoryLimiter>,
     storage_class: Option<String>,
     server_side_encryption: ServerSideEncryption,
     buffer_size: usize,
@@ -73,7 +73,7 @@ where
     pub fn new(
         client: Client,
         runtime: Runtime,
-        mem_limiter: Arc<MemoryLimiter<Client>>,
+        mem_limiter: Arc<MemoryLimiter>,
         storage_class: Option<String>,
         server_side_encryption: ServerSideEncryption,
         buffer_size: usize,