From e594d6f07911b289fb606b51252320a7da54c936 Mon Sep 17 00:00:00 2001
From: Alexey <247128645+axkurcom@users.noreply.github.com>
Date: Tue, 31 Mar 2026 19:22:36 +0300
Subject: [PATCH] Buffer Pool Trim

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
---
 src/proxy/direct_relay.rs |  2 ++
 src/proxy/middle_relay.rs |  1 +
 src/stream/buffer_pool.rs | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/src/proxy/direct_relay.rs b/src/proxy/direct_relay.rs
index 5d9c450..32f2370 100644
--- a/src/proxy/direct_relay.rs
+++ b/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
 }
 
diff --git a/src/proxy/middle_relay.rs b/src/proxy/middle_relay.rs
index d2f37a6..aebc628 100644
--- a/src/proxy/middle_relay.rs
+++ b/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
 }
 
diff --git a/src/stream/buffer_pool.rs b/src/stream/buffer_pool.rs
index 0040fa5..e5a231f 100644
--- a/src/stream/buffer_pool.rs
+++ b/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);