Buffer Pool Trim

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/proxy/direct_relay.rs

@@ -276,6 +276,7 @@ where
     stats.increment_user_connects(user);
     let _direct_connection_lease = stats.acquire_direct_connection_lease();
 
+    let buffer_pool_trim = Arc::clone(&buffer_pool);
     let relay_result = relay_bidirectional(
         client_reader,
         client_writer,
@@ -321,6 +322,7 @@ where
         Err(e) => debug!(user = %user, error = %e, "Direct relay ended with error"),
     }
 
+    buffer_pool_trim.trim_to(buffer_pool_trim.max_buffers().min(64));
     relay_result
 }
 

src/proxy/middle_relay.rs

@@ -1317,6 +1317,7 @@ where
     );
     clear_relay_idle_candidate(conn_id);
     me_pool.registry().unregister(conn_id).await;
+    buffer_pool.trim_to(buffer_pool.max_buffers().min(64));
     result
 }
 

src/stream/buffer_pool.rs

@@ -148,6 +148,41 @@ impl BufferPool {
         self.buffer_size
     }
 
+    /// Maximum number of buffers the pool will retain.
+    pub fn max_buffers(&self) -> usize {
+        self.max_buffers
+    }
+
+    /// Current number of pooled buffers.
+    pub fn pooled(&self) -> usize {
+        self.buffers.len()
+    }
+
+    /// Total buffers allocated (pooled + checked out).
+    pub fn allocated(&self) -> usize {
+        self.allocated.load(Ordering::Relaxed)
+    }
+
+    /// Best-effort number of buffers currently checked out.
+    pub fn in_use(&self) -> usize {
+        self.allocated().saturating_sub(self.pooled())
+    }
+
+    /// Trim pooled buffers down to a target count.
+    pub fn trim_to(&self, target_pooled: usize) {
+        let target = target_pooled.min(self.max_buffers);
+        loop {
+            if self.buffers.len() <= target {
+                break;
+            }
+            if self.buffers.pop().is_some() {
+                self.decrement_allocated();
+            } else {
+                break;
+            }
+        }
+    }
+
     /// Preallocate buffers to fill the pool
     pub fn preallocate(&self, count: usize) {
         let to_alloc = count.min(self.max_buffers);