From babfeaed00fc82d287c2389df57b876f82f6e1c8 Mon Sep 17 00:00:00 2001 From: Hauser Zhi Date: Wed, 4 Nov 2026 22:07:58 +0000 Subject: [PATCH] fix(io-core): preserve pooled capacity and accounting (#9171) * fix(io-core): preserve pooled capacity and accounting Co-Authored-By: heihutu Co-Authored-By: zhi22915 * fix(io-core): preserve pool allocation admission order Co-Authored-By: heihutu Co-Authored-By: zhi22915 * docs(io-core): clarify pool return locking Co-Authored-By: heihutu Co-Authored-By: zhi22915 --------- Co-authored-by: heihutu Co-authored-by: zhi22915 --- Cargo.lock | 1 + crates/io-core/Cargo.toml | 1 + crates/io-core/src/pool.rs | 147 ++++++++++-------- .../tests/bytes_pool_metrics_concurrency.rs | 121 ++++++++++++++ 4 files changed, 202 insertions(+), 68 deletions(-) create mode 100644 crates/io-core/tests/bytes_pool_metrics_concurrency.rs diff --git a/Cargo.lock b/Cargo.lock index f7d6a9029..65af81ec0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10090,6 +10090,7 @@ version = "1.0.2" dependencies = [ "bytes", "hotpath", + "metrics", "rustfs-io-metrics", "thiserror 2.0.21", "tokio", diff --git a/crates/io-core/Cargo.toml b/crates/io-core/Cargo.toml index f3118c19a..5a51a2ea6 100644 --- a/crates/io-core/Cargo.toml +++ b/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] diff --git a/crates/io-core/src/pool.rs b/crates/io-core/src/pool.rs index 733f29b8a..e7cd7eeca 100644 --- a/crates/io-core/src/pool.rs +++ b/crates/io-core/src/pool.rs @@ -83,7 +83,10 @@ struct PoolTier { /// Queue of available buffers for reuse available_buffers: Mutex>, /// Metrics for tracking this tier - metrics: Mutex>>, + metrics: Arc, + /// 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 medium_pool = Arc::new(PoolTier::new(config.medium_size, config.medium_max, "medium")); - let large_pool = Arc::new(PoolTier::new(config.large_size, config.large_max, "large")); - let xlarge_pool = Arc::new(PoolTier::new(config.xlarge_size, config.xlarge_max, "xlarge")); - - // 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)); + 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", Arc::clone(&metrics))); + 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", 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 { 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) -> 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) { - *self.metrics.lock().unwrap_or_else(|e| e.into_inner()) = Some(metrics); - } - - fn take_or_allocate_buffer(&self, size: usize, pool_metrics: &BytesPoolMetrics) -> (BytesMut, bool) { + fn take_or_allocate_buffer(&self, size: usize) -> (BytesMut, bool) { + let pool_metrics = &self.metrics; 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 _metrics_lock = self.metrics.lock().unwrap_or_else(|e| e.into_inner()); - 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 _guard = self.acquire_metrics_lock.lock().unwrap_or_else(|e| e.into_inner()); + let (buffer, was_reused) = self.take_or_allocate_buffer(size); 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 { + fn try_acquire_buffer(&self, size: usize) -> Option { // 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 _metrics_lock = self.metrics.lock().unwrap_or_else(|e| e.into_inner()); - 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 _guard = self.acquire_metrics_lock.lock().unwrap_or_else(|e| e.into_inner()); + let (buffer, was_reused) = self.take_or_allocate_buffer(size); 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() - && let Some(metrics) = metrics.as_ref() - { - metrics.available_buffers.fetch_add(1, Ordering::Relaxed); - } + self.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() - && let Some(metrics) = metrics.as_ref() - { - metrics - .current_allocated_bytes - .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| { - Some(current.saturating_sub(released_bytes)) - }) - .ok(); - } + self.metrics + .current_allocated_bytes + .try_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + 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 diff --git a/crates/io-core/tests/bytes_pool_metrics_concurrency.rs b/crates/io-core/tests/bytes_pool_metrics_concurrency.rs new file mode 100644 index 000000000..68ebc9c69 --- /dev/null +++ b/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>, +} + +impl Recorder for AcquisitionGate { + fn describe_counter(&self, _key: KeyName, _unit: Option, _description: SharedString) {} + + fn describe_gauge(&self, _key: KeyName, _unit: Option, _description: SharedString) {} + + fn describe_histogram(&self, _key: KeyName, _unit: Option, _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" + ); +}