fix(io-core): preserve pooled capacity and accounting (#9171)

* fix(io-core): preserve pooled capacity and accounting

Co-Authored-By: heihutu <heihutu@gmail.com>
Co-Authored-By: zhi22915 <qiuzgang@gmail.com>

* fix(io-core): preserve pool allocation admission order

Co-Authored-By: heihutu <heihutu@gmail.com>
Co-Authored-By: zhi22915 <qiuzgang@gmail.com>

* docs(io-core): clarify pool return locking

Co-Authored-By: heihutu <heihutu@gmail.com>

Co-Authored-By: zhi22915 <qiuzgang@gmail.com>

---------

Co-authored-by: heihutu <heihutu@gmail.com>
Co-authored-by: zhi22915 <qiuzgang@gmail.com>

Cargo.lock

@@ -10090,6 +10090,7 @@ version = "1.0.2"
 dependencies = [
  "bytes",
  "hotpath",
+ "metrics",
  "rustfs-io-metrics",
  "thiserror 2.0.21",
  "tokio",

crates/io-core/Cargo.toml

@@ -42,6 +42,7 @@ rustfs-io-metrics = { workspace = true }
 tracing = { workspace = true }
 
 [dev-dependencies]
+metrics = { workspace = true }
 tokio = { workspace = true, features = ["rt-multi-thread", "macros", "fs"] }
 
 [lib]

crates/io-core/src/pool.rs

@@ -83,7 +83,10 @@ struct PoolTier {
     /// Queue of available buffers for reuse
     available_buffers: Mutex<Vec<BytesMut>>,
     /// Metrics for tracking this tier
-    metrics: Mutex<Option<Arc<BytesPoolMetrics>>>,
+    metrics: Arc<BytesPoolMetrics>,
+    /// Serialize acquisitions before taking buffers so a slow recorder does not
+    /// accumulate allocated buffers behind hit-rate publication.
+    acquire_metrics_lock: Mutex<()>,
     /// Total acquisitions for this tier
     tier_total_acquires: AtomicU64,
     /// Total hits for this tier
@@ -185,16 +188,10 @@ impl BytesPool {
     /// ```
     pub fn with_config(config: BytesPoolConfig) -> Self {
         let metrics = Arc::new(BytesPoolMetrics::default());
-        let small_pool = Arc::new(PoolTier::new(config.small_size, config.small_max, "small"));
+        let small_pool = Arc::new(PoolTier::new(config.small_size, config.small_max, "small", Arc::clone(&metrics)));
-        let medium_pool = Arc::new(PoolTier::new(config.medium_size, config.medium_max, "medium"));
+        let medium_pool = Arc::new(PoolTier::new(config.medium_size, config.medium_max, "medium", Arc::clone(&metrics)));
-        let large_pool = Arc::new(PoolTier::new(config.large_size, config.large_max, "large"));
+        let large_pool = Arc::new(PoolTier::new(config.large_size, config.large_max, "large", Arc::clone(&metrics)));
-        let xlarge_pool = Arc::new(PoolTier::new(config.xlarge_size, config.xlarge_max, "xlarge"));
+        let xlarge_pool = Arc::new(PoolTier::new(config.xlarge_size, config.xlarge_max, "xlarge", Arc::clone(&metrics)));
-
-        // Set metrics reference in all tiers
-        small_pool.set_metrics(Arc::clone(&metrics));
-        medium_pool.set_metrics(Arc::clone(&metrics));
-        large_pool.set_metrics(Arc::clone(&metrics));
-        xlarge_pool.set_metrics(Arc::clone(&metrics));
 
         Self {
             small_pool,
@@ -225,7 +222,7 @@ impl BytesPool {
     /// ```
     pub async fn acquire_buffer(&self, size: usize) -> PooledBuffer {
         let tier = self.select_tier(size);
-        let mut buffer = tier.acquire_buffer(size, &self.metrics).await;
+        let mut buffer = tier.acquire_buffer(size).await;
         if buffer._permit.is_some() {
             buffer.tier = Some(Arc::clone(tier));
         }
@@ -252,7 +249,7 @@ impl BytesPool {
     /// ```
     pub fn try_acquire_buffer(&self, size: usize) -> Option<PooledBuffer> {
         let tier = self.select_tier(size);
-        let mut buffer = tier.try_acquire_buffer(size, &self.metrics)?;
+        let mut buffer = tier.try_acquire_buffer(size)?;
         // Set tier reference for return on drop
         buffer.tier = Some(Arc::clone(tier));
         Some(buffer)
@@ -290,43 +287,38 @@ impl BytesPool {
 }
 
 impl PoolTier {
-    fn new(buffer_size: usize, max_buffers: usize, name: &'static str) -> Self {
+    fn new(buffer_size: usize, max_buffers: usize, name: &'static str, metrics: Arc<BytesPoolMetrics>) -> Self {
         Self {
             buffer_size,
             max_buffers,
             semaphore: Arc::new(Semaphore::new(max_buffers)),
             name,
             available_buffers: Mutex::new(Vec::new()),
-            metrics: Mutex::new(None),
+            metrics,
+            acquire_metrics_lock: Mutex::new(()),
             tier_total_acquires: AtomicU64::new(0),
             tier_pool_hits: AtomicU64::new(0),
             tier_current_allocated_bytes: AtomicU64::new(0),
         }
     }
 
-    fn set_metrics(&self, metrics: Arc<BytesPoolMetrics>) {
+    fn take_or_allocate_buffer(&self, size: usize) -> (BytesMut, bool) {
-        *self.metrics.lock().unwrap_or_else(|e| e.into_inner()) = Some(metrics);
+        let pool_metrics = &self.metrics;
-    }
-
-    fn take_or_allocate_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> (BytesMut, bool) {
         let buffer_opt = {
             let mut available = self.available_buffers.lock().unwrap_or_else(|e| e.into_inner());
-            available.pop()
+            let buffer = available.pop();
+            if buffer.is_some() {
+                pool_metrics.available_buffers.fetch_sub(1, Ordering::Relaxed);
+            }
+            buffer
         };
         let was_reused = buffer_opt.is_some();
-        if was_reused {
-            // A reused buffer leaves the available pool. `return_buffer` does a
-            // matching `fetch_add(1)` on push, so without this the
-            // `available_buffers` gauge only ever grows and never reflects the real
-            // pool size (backlog#806).
-            pool_metrics.available_buffers.fetch_sub(1, Ordering::Relaxed);
-        }
 
         let buffer = if let Some(mut buf) = buffer_opt {
             let previous_capacity = buf.capacity();
             buf.clear();
             if previous_capacity < size {
-                buf.reserve(size - previous_capacity);
+                buf.reserve(size);
             }
             let current_capacity = buf.capacity();
             if current_capacity > previous_capacity {
@@ -353,7 +345,10 @@ impl PoolTier {
         (buffer, was_reused)
     }
 
-    fn record_acquire_metrics(&self, pool_metrics: &BytesPoolMetrics, buffer_capacity: usize, was_reused: bool) {
+    fn record_acquire_metrics(&self, buffer_capacity: usize, was_reused: bool) {
+        let pool_metrics = &self.metrics;
+        pool_metrics.total_acquires.fetch_add(1, Ordering::Relaxed);
+        self.tier_total_acquires.fetch_add(1, Ordering::Relaxed);
         rustfs_io_metrics::record_bytes_pool_acquire(self.name, buffer_capacity, was_reused);
 
         if was_reused {
@@ -374,7 +369,7 @@ impl PoolTier {
         rustfs_io_metrics::record_bytes_pool_allocated(self.name, self.tier_current_allocated_bytes.load(Ordering::Relaxed));
     }
 
-    async fn acquire_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> PooledBuffer {
+    async fn acquire_buffer(&self, size: usize) -> PooledBuffer {
         // Acquire semaphore permit (owned for storage in PooledBuffer)
         let permit = match Arc::clone(&self.semaphore).acquire_owned().await {
             Ok(p) => p,
@@ -388,17 +383,10 @@ impl PoolTier {
             }
         };
 
-        // Use the pool's shared metrics for recording
+        let _guard = self.acquire_metrics_lock.lock().unwrap_or_else(|e| e.into_inner());
-        let _metrics_lock = self.metrics.lock().unwrap_or_else(|e| e.into_inner());
+        let (buffer, was_reused) = self.take_or_allocate_buffer(size);
-        let _metrics = _metrics_lock.as_ref().expect("operation should succeed");
-
-        // Record acquisition
-        pool_metrics.total_acquires.fetch_add(1, Ordering::Relaxed);
-        self.tier_total_acquires.fetch_add(1, Ordering::Relaxed);
-
-        let (buffer, was_reused) = self.take_or_allocate_buffer(size, pool_metrics);
         let buffer_capacity = buffer.capacity();
-        self.record_acquire_metrics(pool_metrics, buffer_capacity, was_reused);
+        self.record_acquire_metrics(buffer_capacity, was_reused);
 
         PooledBuffer {
             buffer: ManuallyDrop::new(buffer),
@@ -407,21 +395,14 @@ impl PoolTier {
         }
     }
 
-    fn try_acquire_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> Option<PooledBuffer> {
+    fn try_acquire_buffer(&self, size: usize) -> Option<PooledBuffer> {
         // Try to acquire permit without blocking
         let permit = Arc::clone(&self.semaphore).try_acquire_owned().ok()?;
 
-        // Use the pool's shared metrics for recording
+        let _guard = self.acquire_metrics_lock.lock().unwrap_or_else(|e| e.into_inner());
-        let _metrics_lock = self.metrics.lock().unwrap_or_else(|e| e.into_inner());
+        let (buffer, was_reused) = self.take_or_allocate_buffer(size);
-        let _metrics = _metrics_lock.as_ref().expect("operation should succeed");
-
-        // Record acquisition
-        pool_metrics.total_acquires.fetch_add(1, Ordering::Relaxed);
-        self.tier_total_acquires.fetch_add(1, Ordering::Relaxed);
-
-        let (buffer, was_reused) = self.take_or_allocate_buffer(size, pool_metrics);
         let buffer_capacity = buffer.capacity();
-        self.record_acquire_metrics(pool_metrics, buffer_capacity, was_reused);
+        self.record_acquire_metrics(buffer_capacity, was_reused);
 
         Some(PooledBuffer {
             buffer: ManuallyDrop::new(buffer),
@@ -430,7 +411,7 @@ impl PoolTier {
         })
     }
 
-    /// Return a buffer to the pool for reuse without ever blocking the caller.
+    /// Return a buffer to the pool for reuse without waiting for the queue lock.
     fn return_buffer(&self, buffer: BytesMut) {
         let mut buffer = Some(buffer);
 
@@ -438,11 +419,7 @@ impl PoolTier {
             && available.len() < self.max_buffers
         {
             available.push(buffer.take().expect("buffer should be present until returned"));
-            if let Ok(metrics) = self.metrics.try_lock()
+            self.metrics.available_buffers.fetch_add(1, Ordering::Relaxed);
-                && let Some(metrics) = metrics.as_ref()
-            {
-                metrics.available_buffers.fetch_add(1, Ordering::Relaxed);
-            }
         }
 
         if let Some(buffer) = buffer {
@@ -452,16 +429,12 @@ impl PoolTier {
                     Some(current.saturating_sub(released_bytes))
                 })
                 .ok();
-            if let Ok(metrics) = self.metrics.try_lock()
+            self.metrics
-                && let Some(metrics) = metrics.as_ref()
+                .current_allocated_bytes
-            {
+                .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
-                metrics
+                    Some(current.saturating_sub(released_bytes))
-                    .current_allocated_bytes
+                })
-                    .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
+                .ok();
-                        Some(current.saturating_sub(released_bytes))
-                    })
-                    .ok();
-            }
         }
         rustfs_io_metrics::record_bytes_pool_allocated(self.name, self.tier_current_allocated_bytes.load(Ordering::Relaxed));
     }
@@ -551,6 +524,44 @@ mod tests {
         assert!(buffer.capacity() >= 2048);
     }
 
+    #[tokio::test]
+    async fn reused_buffer_satisfies_larger_async_request() {
+        let pool = BytesPool::with_config(BytesPoolConfig {
+            small_max: 1,
+            ..Default::default()
+        });
+        let mut first = pool.acquire_buffer(32 * 1024).await;
+        first.extend_from_slice(b"previous contents");
+        assert!(first.capacity() < 48 * 1024, "fixture must require buffer growth");
+        drop(first);
+
+        let reused = pool.acquire_buffer(48 * 1024).await;
+        assert!(reused.is_empty());
+        assert!(reused.capacity() >= 48 * 1024);
+        assert_eq!(pool.metrics().pool_hits.load(Ordering::Relaxed), 1);
+        assert_eq!(
+            pool.metrics().current_allocated_bytes.load(Ordering::Relaxed),
+            u64::try_from(reused.capacity()).expect("buffer capacity fits in metrics")
+        );
+    }
+
+    #[test]
+    fn reused_buffer_satisfies_larger_try_request() {
+        let pool = BytesPool::with_config(BytesPoolConfig {
+            small_max: 1,
+            ..Default::default()
+        });
+        let mut first = pool.try_acquire_buffer(32 * 1024).expect("first buffer is available");
+        first.extend_from_slice(b"previous contents");
+        assert!(first.capacity() < 48 * 1024, "fixture must require buffer growth");
+        drop(first);
+
+        let reused = pool.try_acquire_buffer(48 * 1024).expect("returned buffer is available");
+        assert!(reused.is_empty());
+        assert!(reused.capacity() >= 48 * 1024);
+        assert_eq!(pool.metrics().pool_hits.load(Ordering::Relaxed), 1);
+    }
+
     #[tokio::test]
     async fn available_buffers_gauge_decrements_on_reuse() {
         // Regression (backlog#806): acquiring a pooled (reused) buffer must

crates/io-core/tests/bytes_pool_metrics_concurrency.rs

@@ -0,0 +1,121 @@
+// Copyright 2024 RustFS Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::{Mutex, mpsc};
+use std::thread;
+use std::time::Duration;
+
+use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit};
+use rustfs_io_core::{BytesPool, BytesPoolConfig};
+
+const WAIT: Duration = Duration::from_secs(5);
+const GATE_WAIT: Duration = Duration::from_secs(20);
+
+struct AcquisitionGate {
+    entered: mpsc::Sender<()>,
+    release: Mutex<mpsc::Receiver<()>>,
+}
+
+impl Recorder for AcquisitionGate {
+    fn describe_counter(&self, _key: KeyName, _unit: Option<Unit>, _description: SharedString) {}
+
+    fn describe_gauge(&self, _key: KeyName, _unit: Option<Unit>, _description: SharedString) {}
+
+    fn describe_histogram(&self, _key: KeyName, _unit: Option<Unit>, _description: SharedString) {}
+
+    fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter {
+        if key.name() == "rustfs_bytes_pool_acquisitions_total" {
+            self.entered
+                .send(())
+                .expect("notify test that acquisition reached the recorder");
+            self.release
+                .lock()
+                .expect("lock acquisition gate")
+                .recv_timeout(GATE_WAIT)
+                .expect("release the acquisition recorder gate");
+        }
+        Counter::noop()
+    }
+
+    fn register_gauge(&self, _key: &Key, _metadata: &Metadata<'_>) -> Gauge {
+        Gauge::noop()
+    }
+
+    fn register_histogram(&self, _key: &Key, _metadata: &Metadata<'_>) -> Histogram {
+        Histogram::noop()
+    }
+}
+
+#[test]
+fn returned_buffer_is_counted_while_an_acquisition_recorder_is_busy() {
+    let pool = BytesPool::with_config(BytesPoolConfig {
+        small_max: 2,
+        ..Default::default()
+    });
+    let held = pool.try_acquire_buffer(8 * 1024).expect("acquire the buffer to return");
+    assert_eq!(pool.available_buffers(), 0);
+
+    // This integration test is a dedicated test process: the recorder is local
+    // to the worker, while the metrics enablement flag is process-wide.
+    rustfs_io_metrics::set_metrics_enabled(true);
+    let (entered_tx, entered_rx) = mpsc::channel();
+    let (release_tx, release_rx) = mpsc::channel();
+    let worker_pool = pool.clone();
+    let acquirer = thread::spawn(move || {
+        let gate = AcquisitionGate {
+            entered: entered_tx,
+            release: Mutex::new(release_rx),
+        };
+        metrics::with_local_recorder(&gate, || {
+            worker_pool
+                .try_acquire_buffer(8 * 1024)
+                .expect("acquire a second same-tier buffer")
+        })
+    });
+
+    let entered = entered_rx.recv_timeout(WAIT);
+    if entered.is_err() {
+        let _ = release_tx.send(());
+        let _ = acquirer.join();
+        panic!("acquisition did not reach the recorder gate: {entered:?}");
+    }
+
+    let (dropped_tx, dropped_rx) = mpsc::channel();
+    let returner = thread::spawn(move || {
+        drop(held);
+        let _ = dropped_tx.send(());
+    });
+    let returned_while_blocked = dropped_rx.recv_timeout(WAIT);
+
+    // Always release before joining so a blocking Drop cannot deadlock the test.
+    let _ = release_tx.send(());
+    let returned = returner.join();
+    let acquired = acquirer.join();
+    rustfs_io_metrics::set_metrics_enabled(false);
+
+    assert!(returned_while_blocked.is_ok(), "returning a buffer must not wait on the recorder");
+    returned.expect("return worker should finish");
+    let acquired = acquired.expect("acquisition worker should finish");
+    let after_return = pool.available_buffers();
+    let reused = pool.try_acquire_buffer(8 * 1024).expect("reuse the returned buffer");
+    let after_reuse = pool.available_buffers();
+    drop(reused);
+    drop(acquired);
+    let after_both_return = pool.available_buffers();
+    assert_eq!(
+        (after_return, after_reuse, after_both_return),
+        (1, 0, 2),
+        "returned buffers must be counted, and reuse must not underflow the available gauge"
+    );
+}