mirror of
https://github.com/rustfs/rustfs.git
synced 2026-11-05 02:23:56 +00:00
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>
This commit is contained in:
Hauser Zhi
Generated
+1
@@ -10090,6 +10090,7 @@ version = "1.0.2"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"hotpath",
|
||||
"metrics",
|
||||
"rustfs-io-metrics",
|
||||
"thiserror 2.0.21",
|
||||
"tokio",
|
||||
|
||||
@@ -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]
|
||||
|
||||
+79
-68
@@ -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 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<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>) {
|
||||
*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<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 _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
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user