mod class;
mod freelist;
use super::{IoBufMut, page_size};
use crate::{
iobuf::owner::PooledBuffer,
telemetry::metrics::{Counter, CounterFamily, EncodeLabelSet, GaugeFamily, Register, raw},
};
pub use class::BufferPoolThreadCache;
use class::SizeClassHandle;
pub(crate) use class::SizeClassLease;
use commonware_utils::{NZU32, NZUsize};
pub(super) use freelist::Freelist;
use std::{
collections::BTreeMap,
num::{NonZeroU32, NonZeroUsize},
sync::atomic::{AtomicUsize, Ordering},
};
use thiserror::Error;
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
use loom::sync::Arc;
} else {
use std::sync::Arc;
}
}
#[derive(Error, Debug, Clone, Copy, PartialEq, Eq)]
pub enum PoolError {
#[error("requested capacity exceeds maximum buffer size")]
Oversized,
#[error("pool exhausted for required size class")]
Exhausted,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum BufferPoolThreadCacheConfig {
Enabled(Option<NonZeroUsize>),
Disabled,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BufferPoolClassConfig {
pub size: NonZeroUsize,
pub max_buffers: NonZeroU32,
}
impl From<(NonZeroUsize, NonZeroU32)> for BufferPoolClassConfig {
fn from((size, max_buffers): (NonZeroUsize, NonZeroU32)) -> Self {
Self { size, max_buffers }
}
}
#[derive(Clone, Debug)]
pub struct BufferPoolConfig {
pool_min_size: usize,
class_limits: BTreeMap<NonZeroUsize, NonZeroU32>,
prefill: bool,
alignment: NonZeroUsize,
parallelism: NonZeroUsize,
pub(crate) thread_cache_config: BufferPoolThreadCacheConfig,
}
impl BufferPoolConfig {
pub fn for_network() -> Self {
Self {
pool_min_size: 0,
class_limits: BTreeMap::new(),
prefill: false,
alignment: NZUsize!(1),
parallelism: NZUsize!(1),
thread_cache_config: BufferPoolThreadCacheConfig::Enabled(None),
}
.with_size_class_range(NZUsize!(1024), NZUsize!(128 * 1024), NZU32!(4096))
}
pub fn for_storage() -> Self {
Self {
pool_min_size: 0,
class_limits: BTreeMap::new(),
prefill: false,
alignment: NZUsize!(1),
parallelism: NZUsize!(1),
thread_cache_config: BufferPoolThreadCacheConfig::Enabled(None),
}
.with_size_class_range(
NZUsize!(page_size()),
NZUsize!(8 * 1024 * 1024),
NZU32!(64),
)
}
const fn validate_class_size(size: NonZeroUsize) {
assert!(
size.get().is_power_of_two(),
"class size must be a power of two"
);
assert!(
size.get() <= isize::MAX as usize,
"class size must not exceed isize::MAX"
);
}
pub const fn with_pool_min_size(mut self, pool_min_size: usize) -> Self {
self.pool_min_size = pool_min_size;
self
}
pub fn with_size_class_range(
self,
min: NonZeroUsize,
max: NonZeroUsize,
max_buffers: NonZeroU32,
) -> Self {
Self::validate_class_size(min);
Self::validate_class_size(max);
assert!(max >= min, "max size must be >= min size");
self.with_size_classes(
(min.get().trailing_zeros()..=max.get().trailing_zeros())
.map(|exponent| (NZUsize!(1 << exponent), max_buffers)),
)
}
pub fn with_size_classes<I, C>(mut self, classes: I) -> Self
where
I: IntoIterator<Item = C>,
C: Into<BufferPoolClassConfig>,
{
let mut limits = BTreeMap::new();
for class in classes {
let class = class.into();
Self::validate_class_size(class.size);
assert!(
limits.insert(class.size, class.max_buffers).is_none(),
"duplicate class size {}",
class.size
);
}
assert!(
!limits.is_empty(),
"class layout must enable at least one class"
);
self.class_limits = limits;
self
}
pub fn with_size_class(mut self, size: NonZeroUsize, max_buffers: NonZeroU32) -> Self {
Self::validate_class_size(size);
self.class_limits.insert(size, max_buffers);
self
}
pub fn without_size_class(mut self, size: NonZeroUsize) -> Self {
Self::validate_class_size(size);
assert!(
self.class_limits.remove(&size).is_some(),
"cannot remove a class that is not enabled"
);
assert!(
!self.class_limits.is_empty(),
"cannot remove the final enabled class"
);
self
}
pub fn with_max_per_class(mut self, max_buffers: NonZeroU32) -> Self {
for limit in self.class_limits.values_mut() {
*limit = max_buffers;
}
self
}
pub fn with_bytes_per_class(mut self, bytes: NonZeroUsize) -> Self {
for (size, limit) in self.class_limits.iter_mut() {
let count = bytes.get() / size.get();
assert!(
count <= u32::MAX as usize,
"per-class byte weight derives a limit above u32::MAX"
);
*limit = NonZeroU32::new(count.max(1) as u32).expect("count is at least one");
}
self
}
pub const fn with_parallelism(mut self, parallelism: NonZeroUsize) -> Self {
self.parallelism = parallelism;
self
}
pub const fn with_max_thread_cache_capacity(mut self, capacity: NonZeroUsize) -> Self {
self.thread_cache_config = BufferPoolThreadCacheConfig::Enabled(Some(capacity));
self
}
pub const fn with_thread_cache_disabled(mut self) -> Self {
self.thread_cache_config = BufferPoolThreadCacheConfig::Disabled;
self
}
pub const fn with_prefill(mut self, prefill: bool) -> Self {
self.prefill = prefill;
self
}
pub const fn with_alignment(mut self, alignment: NonZeroUsize) -> Self {
self.alignment = alignment;
self
}
pub fn with_budget_bytes(mut self, budget: NonZeroUsize) -> Self {
let budget = budget.get() as u128;
let minimum: u128 = self
.class_limits
.keys()
.map(|size| size.get() as u128)
.sum();
assert!(
budget >= minimum,
"budget must cover at least one buffer from every enabled class"
);
const FRACTION_BITS: u32 = 64;
let scaled = |limit: NonZeroU32, scale: u128| -> u128 {
((limit.get() as u128).saturating_mul(scale) >> FRACTION_BITS).max(1)
};
let evaluate = |scale: u128| -> (u128, bool) {
self.class_limits
.iter()
.fold((0u128, true), |(total, fits), (&size, &limit)| {
let count = scaled(limit, scale);
(
total.saturating_add(count.saturating_mul(size.get() as u128)),
fits && count <= u32::MAX as u128,
)
})
};
let feasible = |scale: u128| -> bool {
let (total, fits) = evaluate(scale);
total <= budget && fits
};
let mut lo: u128 = 0;
let mut hi: u128 = (u32::MAX as u128 + 1) << FRACTION_BITS;
assert!(feasible(lo), "scale zero must be feasible");
while hi - lo > 1 {
let mid = lo + (hi - lo) / 2;
if feasible(mid) {
lo = mid;
} else {
hi = mid;
}
}
let (next_total, _) = evaluate(lo + 1);
assert!(
next_total > budget,
"budget requires scaling a class limit above u32::MAX"
);
for limit in self.class_limits.values_mut() {
let count = u32::try_from(scaled(*limit, lo)).expect("feasible count fits u32");
*limit = NonZeroU32::new(count).expect("count is at least one");
}
self
}
pub fn size_classes(&self) -> impl ExactSizeIterator<Item = BufferPoolClassConfig> + '_ {
self.class_limits
.iter()
.map(|(&size, &max_buffers)| BufferPoolClassConfig { size, max_buffers })
}
pub fn class_for(&self, size: usize) -> Option<BufferPoolClassConfig> {
self.size_classes().find(|class| class.size.get() >= size)
}
pub const fn pool_min_size(&self) -> usize {
self.pool_min_size
}
pub const fn prefill(&self) -> bool {
self.prefill
}
pub const fn alignment(&self) -> NonZeroUsize {
self.alignment
}
pub const fn parallelism(&self) -> NonZeroUsize {
self.parallelism
}
pub fn min_size(&self) -> NonZeroUsize {
*self
.class_limits
.first_key_value()
.expect("class layout must enable at least one class")
.0
}
pub fn max_size(&self) -> NonZeroUsize {
*self
.class_limits
.last_key_value()
.expect("class layout must enable at least one class")
.0
}
pub fn max_tracked_bytes(&self) -> usize {
self.class_limits
.iter()
.map(|(size, limit)| size.get().saturating_mul(limit.get() as usize))
.fold(0usize, usize::saturating_add)
}
fn validate(&self) {
assert!(
self.alignment.is_power_of_two(),
"alignment must be a power of two"
);
let min_size = self.min_size();
assert!(
min_size >= self.alignment,
"smallest class ({}) must be >= alignment ({})",
min_size,
self.alignment
);
assert!(
self.pool_min_size <= min_size.get(),
"pool_min_size ({}) must be <= smallest class ({})",
self.pool_min_size,
min_size
);
}
fn resolve_thread_cache_capacity(&self, class_limit: NonZeroU32) -> usize {
let class_limit = class_limit.get() as usize;
match self.thread_cache_config {
BufferPoolThreadCacheConfig::Enabled(None) => {
let effective_threads = self.parallelism.get().min(class_limit);
class_limit / effective_threads.saturating_mul(2)
}
BufferPoolThreadCacheConfig::Enabled(Some(capacity)) => capacity.get().min(class_limit),
BufferPoolThreadCacheConfig::Disabled => 0,
}
}
}
#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
struct SizeClassLabel {
size_class: u64,
}
struct PoolMetrics {
created: GaugeFamily<SizeClassLabel>,
exhausted_total: CounterFamily<SizeClassLabel>,
oversized_total: Counter,
}
impl PoolMetrics {
fn new(registry: &mut impl Register) -> Self {
Self {
created: registry.register(
"buffer_pool_created",
"Number of tracked buffers created for the pool",
raw::Family::default(),
),
exhausted_total: registry.register(
"buffer_pool_exhausted",
"Total number of failed allocations due to pool exhaustion",
raw::Family::default(),
),
oversized_total: registry.register(
"buffer_pool_oversized",
"Total number of allocation requests exceeding max buffer size",
raw::Counter::default(),
),
}
}
}
struct Allocation {
buffer: PooledBuffer,
is_new: bool,
}
pub(crate) struct BufferPoolInner {
config: BufferPoolConfig,
classes: Vec<SizeClassHandle>,
min_size: usize,
max_size: usize,
metrics: PoolMetrics,
}
impl Drop for BufferPoolInner {
fn drop(&mut self) {
self.classes.dedup_by(|a, b| a.same_class(b));
assert_eq!(self.classes.len(), self.config.size_classes().len());
for class in &self.classes {
class.drain_global();
}
}
}
impl BufferPoolInner {
#[inline(always)]
fn try_alloc(&self, class_index: usize, zero_on_new: bool) -> Option<Allocation> {
let class = &self.classes[class_index];
if let Some(buffer) = BufferPoolThreadCache::pop(class) {
return Some(Allocation {
buffer,
is_new: false,
});
}
self.try_alloc_new(class, zero_on_new)
}
#[inline(never)]
fn try_alloc_new(&self, class: &SizeClassHandle, zeroed: bool) -> Option<Allocation> {
let label = SizeClassLabel {
size_class: class.size() as u64,
};
let Some(buffer) = class.try_create(zeroed) else {
self.metrics.exhausted_total.get_or_create(&label).inc();
return None;
};
self.metrics.created.get_or_create(&label).inc();
Some(Allocation {
buffer,
is_new: true,
})
}
}
#[derive(Clone)]
pub struct BufferPool {
inner: Arc<BufferPoolInner>,
}
impl std::fmt::Debug for BufferPool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BufferPool")
.field("config", &self.inner.config)
.field("num_classes", &self.inner.config.size_classes().len())
.finish()
}
}
static NEXT_SIZE_CLASS_ID: AtomicUsize = AtomicUsize::new(0);
impl BufferPool {
pub(crate) fn new(config: BufferPoolConfig, registry: &mut impl Register) -> Self {
config.validate();
let metrics = PoolMetrics::new(registry);
let min_size = config.min_size().get();
let max_size = config.max_size().get();
let min_exponent = min_size.trailing_zeros() as usize;
let max_exponent = max_size.trailing_zeros() as usize;
let mut classes = Vec::with_capacity(max_exponent - min_exponent + 1);
for class_config in config.size_classes() {
let class_id = NEXT_SIZE_CLASS_ID.fetch_add(1, Ordering::Relaxed);
let handle = SizeClassHandle::new(
class_id,
class_config.size.get(),
config.alignment.get(),
class_config.max_buffers,
config.parallelism,
config.resolve_thread_cache_capacity(class_config.max_buffers),
config.prefill,
);
if config.prefill {
let label = SizeClassLabel {
size_class: class_config.size.get() as u64,
};
metrics
.created
.get_or_create(&label)
.set(class_config.max_buffers.get() as i64);
}
let index = class_config.size.get().trailing_zeros() as usize - min_exponent;
classes.resize(index + 1, handle);
}
Self {
inner: Arc::new(BufferPoolInner {
config,
classes,
min_size,
max_size,
metrics,
}),
}
}
#[inline(always)]
fn class_index(&self, size: usize) -> Option<usize> {
let min_size = self.inner.min_size;
let max_size = self.inner.max_size;
if size > max_size {
return None;
}
if size <= min_size {
return Some(0);
}
Some(
size.next_power_of_two()
.trailing_zeros()
.wrapping_sub(min_size.trailing_zeros()) as usize,
)
}
#[inline]
fn class_index_or_record_oversized(&self, capacity: usize) -> Option<usize> {
let class_index = self.class_index(capacity);
if class_index.is_none() {
self.inner.metrics.oversized_total.inc();
}
class_index
}
#[inline(always)]
pub fn try_alloc(&self, capacity: usize) -> Result<IoBufMut, PoolError> {
if capacity == 0 {
return Ok(IoBufMut::default());
}
if capacity < self.inner.config.pool_min_size {
return Ok(IoBufMut::with_alignment(
capacity,
self.inner.config.alignment,
));
}
let class_index = self
.class_index_or_record_oversized(capacity)
.ok_or(PoolError::Oversized)?;
let buffer = self
.inner
.try_alloc(class_index, false)
.map(|allocation| {
unsafe { IoBufMut::from_pooled_parts(allocation.buffer) }
})
.ok_or(PoolError::Exhausted)?;
Ok(buffer)
}
#[inline]
pub fn alloc(&self, capacity: usize) -> IoBufMut {
self.try_alloc(capacity).unwrap_or_else(|_| {
let size = capacity.max(1);
IoBufMut::with_alignment(size, self.inner.config.alignment)
})
}
pub unsafe fn alloc_len(&self, len: usize) -> IoBufMut {
let mut buf = self.alloc(len);
unsafe { buf.set_len(len) };
buf
}
pub fn try_alloc_zeroed(&self, len: usize) -> Result<IoBufMut, PoolError> {
if len == 0 {
return Ok(IoBufMut::default());
}
if len < self.inner.config.pool_min_size {
return Ok(IoBufMut::zeroed_with_alignment(
len,
self.inner.config.alignment,
));
}
let class_index = self
.class_index_or_record_oversized(len)
.ok_or(PoolError::Oversized)?;
let allocation = self
.inner
.try_alloc(class_index, true)
.ok_or(PoolError::Exhausted)?;
let mut buf = unsafe { IoBufMut::from_pooled_parts(allocation.buffer) };
if allocation.is_new {
unsafe { buf.set_len(len) };
} else {
unsafe {
std::ptr::write_bytes(buf.as_mut_ptr(), 0, len);
buf.set_len(len);
}
}
Ok(buf)
}
pub fn alloc_zeroed(&self, len: usize) -> IoBufMut {
self.try_alloc_zeroed(len).unwrap_or_else(|_| {
let size = len.max(1);
let mut buf = IoBufMut::zeroed_with_alignment(size, self.inner.config.alignment);
buf.truncate(len);
buf
})
}
pub fn config(&self) -> &BufferPoolConfig {
&self.inner.config
}
}
#[cfg(all(test, not(feature = "loom")))]
mod tests {
use super::{
class::tests::{
get_global_created, get_global_len, get_global_num_stripes, get_local_len,
get_thread_cache_capacity,
},
*,
};
use crate::{
iobuf::{IoBuf, cache_line_size},
telemetry::metrics::Registry,
};
use bytes::{Buf, BufMut};
use commonware_utils::NZU32;
use std::{
sync::{Arc, mpsc},
thread,
};
fn test_pool(config: BufferPoolConfig) -> BufferPool {
let mut registry = Registry::default();
BufferPool::new(config, &mut registry)
}
fn test_config(min_size: usize, max_size: usize, max_per_class: u32) -> BufferPoolConfig {
BufferPoolConfig::for_network()
.with_pool_min_size(0)
.with_size_class_range(
NZUsize!(min_size),
NZUsize!(max_size),
NZU32!(max_per_class),
)
.with_alignment(NZUsize!(page_size()))
}
fn sparse_config(classes: impl IntoIterator<Item = (usize, u32)>) -> BufferPoolConfig {
BufferPoolConfig::for_network()
.with_pool_min_size(0)
.with_size_classes(
classes
.into_iter()
.map(|(size, max_buffers)| (NZUsize!(size), NZU32!(max_buffers))),
)
.with_alignment(NZUsize!(page_size()))
}
fn classes_of(config: &BufferPoolConfig) -> Vec<(usize, u32)> {
config
.size_classes()
.map(|class| (class.size.get(), class.max_buffers.get()))
.collect()
}
fn get_allocated(pool: &BufferPool, size: usize) -> usize {
let class_index = pool.class_index(size).unwrap();
let class = &pool.inner.classes[class_index];
get_global_created(class) - get_global_len(class) - get_local_len(class)
}
fn get_available(pool: &BufferPool, size: usize) -> i64 {
let class_index = pool.class_index(size).unwrap();
let class = &pool.inner.classes[class_index];
(get_global_len(class) + get_local_len(class)) as i64
}
#[test]
fn test_page_size() {
let size = page_size();
assert!(size >= 4096);
assert!(size.is_power_of_two());
}
#[test]
fn test_config_validation() {
let page = page_size();
let config = test_config(page, page * 4, 10);
config.validate();
}
#[test]
fn test_explicit_thread_cache_capacity_clamps_to_class_limit() {
let page = page_size();
let config = test_config(page, page * 4, 10).with_max_thread_cache_capacity(NZUsize!(11));
config.validate();
let pool = test_pool(config);
let class_index = pool.class_index(page).unwrap();
assert_eq!(
get_thread_cache_capacity(&pool.inner.classes[class_index]),
10
);
let config = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(1024), NZU32!(4)), (NZUsize!(4096), NZU32!(64))])
.with_max_thread_cache_capacity(NZUsize!(16));
let pool = test_pool(config);
let small_index = pool.class_index(1024).unwrap();
let large_index = pool.class_index(4096).unwrap();
assert_eq!(
get_thread_cache_capacity(&pool.inner.classes[small_index]),
4
);
assert_eq!(
get_thread_cache_capacity(&pool.inner.classes[large_index]),
16
);
}
#[test]
#[should_panic(expected = "class size must be a power of two")]
fn test_config_invalid_min_size() {
let _ = BufferPoolConfig::for_network().with_size_class_range(
NZUsize!(3000),
NZUsize!(8192),
NZU32!(10),
);
}
#[test]
#[should_panic(expected = "class size must be a power of two")]
fn test_config_invalid_max_size() {
let _ = BufferPoolConfig::for_network().with_size_class_range(
NZUsize!(4096),
NZUsize!(12000),
NZU32!(10),
);
}
#[test]
#[should_panic(expected = "max size must be >= min size")]
fn test_config_range_rejects_max_below_min() {
let _ = BufferPoolConfig::for_network().with_size_class_range(
NZUsize!(8192),
NZUsize!(1024),
NZU32!(10),
);
}
#[test]
#[should_panic(expected = "class size must not exceed isize::MAX")]
fn test_config_rejects_class_size_above_isize_max() {
let _ = BufferPoolConfig::for_network()
.with_size_class(NZUsize!(1usize << (usize::BITS - 1)), NZU32!(1));
}
#[test]
#[should_panic(expected = "alignment must be a power of two")]
fn test_config_invalid_alignment() {
let page = page_size();
let config = test_config(page, page, 10).with_alignment(NZUsize!(page - 1));
config.validate();
}
#[test]
#[should_panic(expected = "must be >= alignment")]
fn test_config_min_size_below_alignment() {
let page = page_size();
let config = test_config(page, page, 10).with_alignment(NZUsize!(page * 2));
config.validate();
}
#[test]
#[should_panic(expected = "pool_min_size")]
fn test_config_pool_min_size_above_min_size() {
let page = page_size();
let config = test_config(page, page, 10).with_pool_min_size(page + 1);
config.validate();
}
#[test]
fn test_pool_class_index() {
let page = page_size();
let pool = test_pool(test_config(page, page * 8, 10));
assert_eq!(pool.inner.classes.len(), 4);
assert_eq!(pool.class_index(1), Some(0));
assert_eq!(pool.class_index(page), Some(0));
assert_eq!(pool.class_index(page + 1), Some(1));
assert_eq!(pool.class_index(page * 2), Some(1));
assert_eq!(pool.class_index(page * 4 + 1), Some(3));
assert_eq!(pool.class_index(page * 8 - 1), Some(3));
assert_eq!(pool.class_index(page * 8), Some(3));
assert_eq!(pool.class_index(page * 8 + 1), None);
}
#[test]
fn test_size_classes_replacement_normalizes_and_iterates() {
let config = BufferPoolConfig::for_network().with_size_classes([
(NZUsize!(1 << 20), NZU32!(16)),
(NZUsize!(4096), NZU32!(1024)),
(NZUsize!(65536), NZU32!(256)),
]);
assert_eq!(
classes_of(&config),
vec![(4096, 1024), (65536, 256), (1 << 20, 16)]
);
assert_eq!(config.size_classes().len(), 3);
assert_eq!(config.min_size().get(), 4096);
assert_eq!(config.max_size().get(), 1 << 20);
assert_eq!(
config.max_tracked_bytes(),
4096 * 1024 + 65536 * 256 + (1 << 20) * 16
);
let explicit = BufferPoolConfig::for_network().with_size_classes([BufferPoolClassConfig {
size: NZUsize!(512),
max_buffers: NZU32!(2),
}]);
assert_eq!(classes_of(&explicit), vec![(512, 2)]);
}
#[test]
#[should_panic(expected = "class layout must enable at least one class")]
fn test_size_classes_rejects_empty_input() {
let _ = BufferPoolConfig::for_network()
.with_size_classes(std::iter::empty::<BufferPoolClassConfig>());
}
#[test]
#[should_panic(expected = "duplicate class size 4096")]
fn test_size_classes_rejects_duplicates() {
let _ = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(4096), NZU32!(1)), (NZUsize!(4096), NZU32!(2))]);
}
#[test]
#[should_panic(expected = "class size must be a power of two")]
fn test_size_classes_rejects_non_power_of_two() {
let _ = BufferPoolConfig::for_network().with_size_classes([(NZUsize!(3000), NZU32!(1))]);
}
#[test]
fn test_size_class_upsert_and_removal() {
let base = BufferPoolConfig::for_network().with_size_class_range(
NZUsize!(1024),
NZUsize!(8192),
NZU32!(8),
);
let tuned = base.clone().with_size_class(NZUsize!(2048), NZU32!(64));
assert_eq!(
classes_of(&tuned),
vec![(1024, 8), (2048, 64), (4096, 8), (8192, 8)]
);
let extended = base.clone().with_size_class(NZUsize!(32768), NZU32!(2));
assert_eq!(extended.max_size().get(), 32768);
assert_eq!(classes_of(&extended).len(), 5);
let sparse = base.clone().without_size_class(NZUsize!(2048));
assert_eq!(classes_of(&sparse), vec![(1024, 8), (4096, 8), (8192, 8)]);
let narrowed = base.without_size_class(NZUsize!(1024));
assert_eq!(narrowed.min_size().get(), 2048);
let narrowed = narrowed.without_size_class(NZUsize!(8192));
assert_eq!(narrowed.max_size().get(), 4096);
let uniform = sparse.with_max_per_class(NZU32!(3));
assert_eq!(classes_of(&uniform), vec![(1024, 3), (4096, 3), (8192, 3)]);
}
#[test]
#[should_panic(expected = "cannot remove a class that is not enabled")]
fn test_without_size_class_rejects_absent_class() {
let _ = BufferPoolConfig::for_network().without_size_class(NZUsize!(1 << 30));
}
#[test]
#[should_panic(expected = "cannot remove the final enabled class")]
fn test_without_size_class_rejects_final_class() {
let _ = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(4096), NZU32!(1))])
.without_size_class(NZUsize!(4096));
}
#[test]
fn test_bytes_per_class_gives_equal_byte_weight() {
let config = BufferPoolConfig::for_network()
.with_size_classes([
(NZUsize!(1024), NZU32!(1)),
(NZUsize!(4096), NZU32!(1)),
(NZUsize!(1 << 20), NZU32!(1)),
])
.with_bytes_per_class(NZUsize!(64 * 1024));
assert_eq!(
classes_of(&config),
vec![(1024, 64), (4096, 16), (1 << 20, 1)]
);
let exact = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(4096), NZU32!(7))])
.with_bytes_per_class(NZUsize!(4096));
assert_eq!(classes_of(&exact), vec![(4096, 1)]);
}
#[test]
fn test_class_for_routes_to_smallest_fitting_class() {
let config = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(4096), NZU32!(4)), (NZUsize!(32768), NZU32!(2))]);
assert_eq!(config.class_for(0).unwrap().size.get(), 4096);
assert_eq!(config.class_for(4096).unwrap().size.get(), 4096);
assert_eq!(config.class_for(4097).unwrap().size.get(), 32768);
assert_eq!(config.class_for(16384).unwrap().size.get(), 32768);
assert_eq!(config.class_for(32768).unwrap().size.get(), 32768);
assert_eq!(config.class_for(32769), None);
}
#[test]
fn test_sparse_routing_allocates_next_enabled_class() {
let page = page_size();
let pool = test_pool(sparse_config([(page, 4), (page * 8, 4)]));
let buf = pool.try_alloc(1).unwrap();
assert_eq!(buf.capacity(), page);
let buf = pool.try_alloc(page).unwrap();
assert_eq!(buf.capacity(), page);
let buf = pool.try_alloc(page + 1).unwrap();
assert_eq!(buf.capacity(), page * 8);
let buf = pool.try_alloc(page * 4).unwrap();
assert_eq!(buf.capacity(), page * 8);
let buf = pool.try_alloc(page * 8).unwrap();
assert_eq!(buf.capacity(), page * 8);
assert_eq!(
pool.try_alloc(page * 8 + 1).unwrap_err(),
PoolError::Oversized
);
}
#[test]
fn test_sparse_routing_exhaustion_does_not_cascade() {
let page = page_size();
let pool = test_pool(sparse_config([(page, 1), (page * 8, 1)]));
let _small = pool.try_alloc(page).unwrap();
assert_eq!(pool.try_alloc(page).unwrap_err(), PoolError::Exhausted);
let _large = pool.try_alloc(page * 8).unwrap();
let fallback = pool.alloc(page);
assert!(!fallback.is_pooled());
assert_eq!(fallback.capacity(), page);
let small_fallback = pool.alloc(100);
assert!(!small_fallback.is_pooled());
assert!((100..108).contains(&small_fallback.capacity()));
}
#[test]
fn test_sparse_metrics_use_enabled_class_labels() {
let page = page_size();
let mut registry = Registry::default();
let pool = BufferPool::new(sparse_config([(page, 1), (page * 8, 1)]), &mut registry);
let _held = pool.try_alloc(page * 2).unwrap();
assert!(pool.try_alloc(page * 2).is_err());
assert!(pool.try_alloc(page * 16).is_err());
let encoded = registry.encode();
assert!(
encoded.contains(&format!(
"buffer_pool_created{{size_class=\"{}\"}} 1",
page * 8
)),
"metrics output: {encoded}"
);
assert!(
encoded.contains(&format!(
"buffer_pool_exhausted_total{{size_class=\"{}\"}} 1",
page * 8
)),
"metrics output: {encoded}"
);
assert!(
!encoded.contains(&format!("size_class=\"{}\"", page * 2)),
"metrics output: {encoded}"
);
assert!(
encoded.contains("buffer_pool_oversized_total 1"),
"metrics output: {encoded}"
);
}
#[test]
fn test_sparse_gap_allocations_share_one_class() {
let page = page_size();
let pool = test_pool(sparse_config([(page, 4), (page * 8, 4)]));
let direct = pool.class_index(page * 8).unwrap();
for size in [page + 1, page * 2, page * 4, page * 8] {
let index = pool.class_index(size).unwrap();
assert!(
pool.inner.classes[index].same_class(&pool.inner.classes[direct]),
"size {size} must alias the largest class"
);
}
let mut via_gap = pool.try_alloc(page * 2).unwrap();
let ptr = via_gap.as_mut_ptr();
drop(via_gap);
let mut direct_reuse = pool.try_alloc(page * 8).unwrap();
assert_eq!(direct_reuse.as_mut_ptr(), ptr);
}
#[test]
fn test_sparse_pool_drop_drains_each_unique_class_once() {
let page = page_size();
let pool =
test_pool(sparse_config([(page, 2), (page * 16, 2)]).with_thread_cache_disabled());
let class_index = pool.class_index(page * 16).unwrap();
let class = pool.inner.classes[class_index].clone();
let buf = pool.try_alloc(page * 2).unwrap();
drop(buf);
assert_eq!(get_global_len(&class), 1);
drop(pool);
assert_eq!(get_global_len(&class), 0);
assert_eq!(get_global_created(&class), 1);
}
#[test]
fn test_sparse_pool_debug_reports_unique_classes() {
let page = page_size();
let pool = test_pool(sparse_config([(page, 2), (page * 16, 2)]));
assert_eq!(pool.inner.classes.len(), 5);
let debug = format!("{pool:?}");
assert!(debug.contains("num_classes: 2"), "debug output: {debug}");
}
#[test]
fn test_sparse_prefill_creates_per_class_limits() {
let page = page_size();
let pool = test_pool(sparse_config([(page, 3), (page * 4, 1)]).with_prefill(true));
let small = &pool.inner.classes[pool.class_index(page).unwrap()];
let large = &pool.inner.classes[pool.class_index(page * 4).unwrap()];
assert_eq!(get_global_created(small), 3);
assert_eq!(get_global_len(small), 3);
assert_eq!(get_global_created(large), 1);
assert_eq!(get_global_len(large), 1);
let a = pool.try_alloc(page).unwrap();
let b = pool.try_alloc(page).unwrap();
let c = pool.try_alloc(page).unwrap();
assert!(pool.try_alloc(page).is_err());
drop((a, b, c));
let _gap = pool.try_alloc(page * 2).unwrap();
assert!(pool.try_alloc(page * 4).is_err());
}
#[test]
fn test_pool_alloc_and_return() {
let page = page_size();
let pool = test_pool(test_config(page, page * 4, 2));
let buf = pool.try_alloc(page).unwrap();
assert!(buf.capacity() >= page);
assert_eq!(buf.len(), 0);
drop(buf);
let buf2 = pool.try_alloc(page).unwrap();
assert!(buf2.capacity() >= page);
assert_eq!(buf2.len(), 0);
}
#[test]
fn test_alloc_len_sets_len() {
let page = page_size();
let pool = test_pool(test_config(page, page * 4, 2));
let mut buf = unsafe { pool.alloc_len(100) };
assert_eq!(buf.len(), 100);
buf.as_mut().fill(0xAB);
let frozen = buf.freeze();
assert_eq!(frozen.as_ref(), &[0xAB; 100]);
}
#[test]
fn test_alloc_zeroed_sets_len_and_zeros() {
let page = page_size();
let pool = test_pool(test_config(page, page * 4, 2));
let buf = pool.alloc_zeroed(100);
assert_eq!(buf.len(), 100);
assert!(buf.as_ref().iter().all(|&b| b == 0));
}
#[test]
fn test_try_alloc_zeroed_sets_len_and_zeros() {
let page = page_size();
let pool = test_pool(test_config(page, page * 4, 2));
let buf = pool.try_alloc_zeroed(page).unwrap();
assert!(buf.is_pooled());
assert_eq!(buf.len(), page);
assert!(buf.as_ref().iter().all(|&b| b == 0));
}
#[test]
fn test_alloc_zeroed_fallback_uses_untracked_zeroed_buffer() {
let page = page_size();
let pool = test_pool(test_config(page, page, 1));
let _pooled = pool.try_alloc(page).unwrap();
let buf = pool.alloc_zeroed(100);
assert!(!buf.is_pooled());
assert_eq!(buf.len(), 100);
assert!(buf.as_ref().iter().all(|&b| b == 0));
}
#[test]
fn test_alloc_zeroed_reuses_dirty_pooled_buffer() {
let page = page_size();
let pool = test_pool(test_config(page, page, 1));
let mut first = pool.alloc_zeroed(page);
assert!(first.is_pooled());
assert!(first.as_ref().iter().all(|&b| b == 0));
first.as_mut().fill(0xAB);
drop(first);
let second = pool.alloc_zeroed(page);
assert!(second.is_pooled());
assert_eq!(second.len(), page);
assert!(second.as_ref().iter().all(|&b| b == 0));
}
#[test]
fn test_requests_smaller_than_pool_min_size_bypass_pool() {
let pool = test_pool(
BufferPoolConfig::for_network()
.with_pool_min_size(512)
.with_size_class_range(NZUsize!(512), NZUsize!(1024), NZU32!(2))
.with_alignment(NZUsize!(128)),
);
let buf = pool.try_alloc(200).unwrap();
assert!(!buf.is_pooled());
assert_eq!(buf.capacity(), 200);
let zeroed = pool.try_alloc_zeroed(200).unwrap();
assert!(!zeroed.is_pooled());
assert_eq!(zeroed.len(), 200);
assert!(zeroed.as_ref().iter().all(|&b| b == 0));
let pooled = pool.try_alloc(512).unwrap();
assert!(pooled.is_pooled());
assert_eq!(pooled.capacity(), 512);
}
#[test]
fn test_zero_capacity_requests_bypass_pool() {
let page = page_size();
let pool = test_pool(test_config(page, page, 1));
let empty = pool.try_alloc(0).unwrap();
assert!(!empty.is_pooled());
assert_eq!(empty.capacity(), 0);
let zeroed = pool.try_alloc_zeroed(0).unwrap();
assert!(!zeroed.is_pooled());
assert_eq!(zeroed.len(), 0);
assert_eq!(zeroed.capacity(), 0);
assert_eq!(pool.alloc(0).capacity(), 0);
assert_eq!(pool.alloc_zeroed(0).len(), 0);
let real = pool.try_alloc(page).unwrap();
assert!(real.is_pooled());
assert_eq!(real.capacity(), page);
}
#[test]
fn test_pool_size_classes() {
let page = page_size();
let pool = test_pool(test_config(page, page * 4, 10));
let buf1 = pool.try_alloc(page).unwrap();
assert_eq!(buf1.capacity(), page);
let buf2 = pool.try_alloc(page + 1).unwrap();
assert_eq!(buf2.capacity(), page * 2);
let buf3 = pool.try_alloc(page * 3).unwrap();
assert_eq!(buf3.capacity(), page * 4);
}
#[test]
fn test_prefill() {
let page = NZUsize!(page_size());
let pool = test_pool(
BufferPoolConfig::for_network()
.with_pool_min_size(0)
.with_size_class_range(page, page, NZU32!(5))
.with_alignment(page)
.with_prefill(true),
);
let mut bufs = Vec::new();
for _ in 0..5 {
bufs.push(pool.try_alloc(page.get()).expect("alloc should succeed"));
}
assert!(pool.try_alloc(page.get()).is_err());
}
#[test]
fn test_config_for_network() {
let config = BufferPoolConfig::for_network();
config.validate();
assert_eq!(config.pool_min_size, 0);
assert_eq!(config.min_size().get(), 1024);
assert_eq!(config.max_size().get(), 128 * 1024);
let expected: Vec<(usize, u32)> = (10..=17).map(|e| (1usize << e, 4096)).collect();
assert_eq!(classes_of(&config), expected);
assert_eq!(config.parallelism, NZUsize!(1));
assert_eq!(
config.thread_cache_config,
BufferPoolThreadCacheConfig::Enabled(None)
);
assert!(!config.prefill);
assert_eq!(config.alignment.get(), 1);
}
#[test]
fn test_config_for_storage() {
let config = BufferPoolConfig::for_storage();
config.validate();
assert_eq!(config.pool_min_size, 0);
assert_eq!(config.min_size().get(), page_size());
assert_eq!(config.max_size().get(), 8 * 1024 * 1024);
let min_exponent = page_size().trailing_zeros();
let expected: Vec<(usize, u32)> = (min_exponent..=23).map(|e| (1usize << e, 64)).collect();
assert_eq!(classes_of(&config), expected);
assert_eq!(config.parallelism, NZUsize!(1));
assert_eq!(
config.thread_cache_config,
BufferPoolThreadCacheConfig::Enabled(None)
);
assert!(!config.prefill);
assert_eq!(config.alignment.get(), 1);
}
#[test]
fn test_storage_config_supports_default_allocations() {
let pool = test_pool(BufferPoolConfig::for_storage());
let buf = pool.try_alloc(8 * 1024 * 1024).unwrap();
assert_eq!(buf.capacity(), 8 * 1024 * 1024);
}
#[test]
fn test_config_builders() {
let page = NZUsize!(page_size());
let config = BufferPoolConfig::for_storage()
.with_pool_min_size(1024)
.with_parallelism(NZUsize!(4))
.with_max_thread_cache_capacity(NZUsize!(8))
.with_prefill(true)
.with_size_class_range(page, NZUsize!(128 * 1024), NZU32!(64));
config.validate();
assert_eq!(config.pool_min_size, 1024);
assert_eq!(config.min_size(), page);
assert_eq!(config.max_size().get(), 128 * 1024);
assert!(
config
.size_classes()
.all(|class| class.max_buffers.get() == 64)
);
assert_eq!(config.parallelism, NZUsize!(4));
assert_eq!(
config.thread_cache_config,
BufferPoolThreadCacheConfig::Enabled(Some(NZUsize!(8)))
);
assert!(config.prefill);
assert_eq!(config.alignment.get(), 1);
let aligned = BufferPoolConfig::for_network()
.with_pool_min_size(256)
.with_parallelism(NZUsize!(4))
.with_alignment(NZUsize!(256))
.with_size_class_range(NZUsize!(256), NZUsize!(128 * 1024), NZU32!(4096));
aligned.validate();
assert_eq!(aligned.parallelism, NZUsize!(4));
assert_eq!(
aligned.thread_cache_config,
BufferPoolThreadCacheConfig::Enabled(None)
);
assert_eq!(aligned.alignment.get(), 256);
assert_eq!(aligned.min_size().get(), 256);
}
#[test]
fn test_parallelism_policy_resolves_thread_cache_capacity() {
let page = page_size();
let pool = test_pool(test_config(page, page, 64).with_parallelism(NZUsize!(8)));
let class_index = pool.class_index(page).unwrap();
assert_eq!(
get_thread_cache_capacity(&pool.inner.classes[class_index]),
4
);
let pool = test_pool(test_config(page, page, 4096).with_parallelism(NZUsize!(8)));
let class_index = pool.class_index(page).unwrap();
assert_eq!(
get_thread_cache_capacity(&pool.inner.classes[class_index]),
256
);
}
#[test]
fn test_auto_thread_cache_disables_when_parallelism_exceeds_budget() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2).with_parallelism(NZUsize!(8)));
let class_index = pool.class_index(page).unwrap();
let class = &pool.inner.classes[class_index];
assert_eq!(get_thread_cache_capacity(class), 0);
let first = pool.try_alloc(page).expect("first tracked allocation");
let second = pool.try_alloc(page).expect("second tracked allocation");
let pool_for_thread = pool.clone();
let (returned_tx, returned_rx) = mpsc::channel();
let (release_tx, release_rx) = mpsc::channel();
let handle = thread::spawn(move || {
drop(first);
drop(second);
returned_tx.send(()).expect("signal returned buffers");
release_rx.recv().expect("release worker");
drop(pool_for_thread);
});
returned_rx.recv().expect("wait for returned buffers");
let _first = pool.try_alloc(page).expect("first global reuse");
let _second = pool.try_alloc(page).expect("second global reuse");
release_tx.send(()).expect("release worker");
handle.join().expect("worker should not panic");
}
#[test]
fn test_parallelism_policy_resolves_freelist_stripes() {
let page = page_size();
let pool = test_pool(test_config(page, page, 64).with_parallelism(NZUsize!(16)));
let class_index = pool.class_index(page).unwrap();
assert_eq!(get_global_num_stripes(&pool.inner.classes[class_index]), 16);
let pool = test_pool(test_config(page, page, 12).with_parallelism(NZUsize!(9)));
let class_index = pool.class_index(page).unwrap();
assert_eq!(get_global_num_stripes(&pool.inner.classes[class_index]), 8);
let pool = test_pool(
test_config(page, page, 64)
.with_parallelism(NZUsize!(16))
.with_thread_cache_disabled(),
);
let class_index = pool.class_index(page).unwrap();
assert_eq!(get_global_num_stripes(&pool.inner.classes[class_index]), 16);
}
#[test]
fn test_fixed_thread_cache_capacity_overrides_auto_capacity() {
let page = page_size();
let pool = test_pool(
test_config(page, page, 64)
.with_parallelism(NZUsize!(8))
.with_max_thread_cache_capacity(NZUsize!(7)),
);
let class_index = pool.class_index(page).unwrap();
assert_eq!(
get_thread_cache_capacity(&pool.inner.classes[class_index]),
7
);
assert_eq!(get_global_num_stripes(&pool.inner.classes[class_index]), 8);
}
#[test]
fn test_disabled_thread_cache_does_not_retain_buffers_locally() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2).with_thread_cache_disabled());
let class_index = pool.class_index(page).unwrap();
let class = &pool.inner.classes[class_index];
let tracked = pool.try_alloc(page).expect("tracked allocation");
drop(tracked);
assert_eq!(get_thread_cache_capacity(class), 0);
assert_eq!(get_local_len(class), 0);
assert_eq!(get_global_len(class), 1);
}
#[test]
fn test_config_with_budget_bytes() {
let base = BufferPoolConfig::for_network().with_size_class_range(
NZUsize!(4),
NZUsize!(16),
NZU32!(1),
);
let config = base.clone().with_budget_bytes(NZUsize!(280));
assert_eq!(classes_of(&config), vec![(4, 10), (8, 10), (16, 10)]);
assert_eq!(config.max_tracked_bytes(), 280);
let config = base.clone().with_budget_bytes(NZUsize!(279));
assert_eq!(classes_of(&config), vec![(4, 9), (8, 9), (16, 9)]);
let config = base.clone().with_budget_bytes(NZUsize!(28));
assert_eq!(classes_of(&config), vec![(4, 1), (8, 1), (16, 1)]);
let shaped_base = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(4), NZU32!(4)), (NZUsize!(16), NZU32!(1))]);
let shaped = shaped_base.clone().with_budget_bytes(NZUsize!(96));
assert_eq!(classes_of(&shaped), vec![(4, 12), (16, 3)]);
let shrunk = shaped_base.with_budget_bytes(NZUsize!(20));
assert_eq!(classes_of(&shrunk), vec![(4, 1), (16, 1)]);
let uneven = base.with_budget_bytes(NZUsize!(30));
assert_eq!(classes_of(&uneven), vec![(4, 1), (8, 1), (16, 1)]);
}
#[test]
fn test_config_with_budget_bytes_is_one_shot() {
let config = BufferPoolConfig::for_network()
.with_size_class_range(NZUsize!(4), NZUsize!(16), NZU32!(1))
.with_budget_bytes(NZUsize!(280));
assert_eq!(config.max_tracked_bytes(), 280);
let overridden = config.clone().with_max_per_class(NZU32!(100));
assert_eq!(overridden.max_tracked_bytes(), 2800);
let upserted = config.with_size_class(NZUsize!(32), NZU32!(100));
assert_eq!(upserted.max_tracked_bytes(), 280 + 32 * 100);
let base = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(1), NZU32!(3)), (NZUsize!(8), NZU32!(2))]);
let once = base.with_budget_bytes(NZUsize!(21));
assert_eq!(classes_of(&once), vec![(1, 4), (8, 2)]);
let twice = once.with_budget_bytes(NZUsize!(21));
assert_eq!(classes_of(&twice), vec![(1, 5), (8, 2)]);
}
#[test]
#[should_panic(expected = "budget must cover at least one buffer from every enabled class")]
fn test_config_with_budget_bytes_below_minimum() {
let _ = BufferPoolConfig::for_network()
.with_size_class_range(NZUsize!(4), NZUsize!(16), NZU32!(1))
.with_budget_bytes(NZUsize!(27));
}
#[test]
#[should_panic(expected = "budget requires scaling a class limit above u32::MAX")]
fn test_config_with_budget_bytes_above_u32() {
let _ = BufferPoolConfig::for_network()
.with_size_classes([(NZUsize!(1), NZU32!(1))])
.with_budget_bytes(NZUsize!(u32::MAX as usize + 2));
}
#[test]
fn test_config_with_budget_bytes_near_u32_breakpoints() {
cfg_if::cfg_if! {
if #[cfg(miri)] {
let budget = 10_000usize;
} else {
let budget = 1_000_000usize;
}
}
let a = u32::MAX;
let b = u32::MAX - 1;
let config = BufferPoolConfig::for_network()
.with_size_classes([
(NZUsize!(1), NonZeroU32::new(a).unwrap()),
(NZUsize!(2), NonZeroU32::new(b).unwrap()),
])
.with_budget_bytes(NonZeroUsize::new(budget).unwrap());
let expected = brute_force_budget(&[(1, a), (2, b)], budget as u128);
assert_eq!(
classes_of(&config)
.into_iter()
.map(|(_, limit)| limit)
.collect::<Vec<_>>(),
expected
);
}
fn brute_force_budget(shape: &[(usize, u32)], budget: u128) -> Vec<u32> {
let mut best: Option<Vec<u32>> = None;
let mut best_total = 0u128;
let mut candidates: Vec<(u128, u128)> = vec![(0, 1)];
for &(size, limit) in shape {
let max_count = (budget / size as u128).min(u32::MAX as u128);
for k in 1..=max_count {
candidates.push((k, limit as u128));
}
}
for (k, c) in candidates {
let counts: Vec<u128> = shape
.iter()
.map(|&(_, limit)| ((limit as u128 * k) / c).max(1))
.collect();
if counts.iter().any(|&count| count > u32::MAX as u128) {
continue;
}
let total: u128 = counts
.iter()
.zip(shape.iter())
.map(|(&count, &(size, _))| count * size as u128)
.sum();
if total <= budget && total >= best_total {
best_total = total;
best = Some(counts.iter().map(|&count| count as u32).collect());
}
}
best.expect("budget covers one buffer per class")
}
#[test]
fn test_pool_error_display() {
assert_eq!(
PoolError::Oversized.to_string(),
"requested capacity exceeds maximum buffer size"
);
assert_eq!(
PoolError::Exhausted.to_string(),
"pool exhausted for required size class"
);
}
#[test]
fn test_pool_debug_and_config_accessor() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let debug = format!("{pool:?}");
assert!(debug.contains("BufferPool"));
assert!(debug.contains("num_classes"));
assert_eq!(pool.config().min_size().get(), page);
}
#[test]
fn test_pooled_debug_and_empty_freeze_paths() {
let page = page_size();
let pool = test_pool(test_config(page, page, 3));
let pooled_mut = pool.try_alloc(page).expect("pooled allocation");
let pooled_mut_debug = format!("{pooled_mut:?}");
assert!(pooled_mut_debug.contains("IoBufMut"));
assert!(pooled_mut_debug.contains("cap"));
assert!(pooled_mut.is_pooled());
let empty = pool.try_alloc(page).expect("pooled allocation").freeze();
assert!(empty.is_empty());
assert!(!empty.is_pooled());
let mut non_empty = pool.try_alloc(page).expect("pooled allocation");
non_empty.put_slice(b"abc");
let pooled = non_empty.freeze();
let pooled_debug = format!("{pooled:?}");
assert!(pooled_debug.contains("IoBuf"));
assert!(pooled_debug.contains("pooled"));
assert!(pooled.is_pooled());
BufferPoolThreadCache::flush();
}
#[test]
fn test_freeze_returns_buffer_to_pool() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
assert_eq!(get_allocated(&pool, page), 0);
assert_eq!(get_available(&pool, page), 0);
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(b"x");
assert_eq!(get_allocated(&pool, page), 1);
assert_eq!(get_available(&pool, page), 0);
let iobuf = buf.freeze();
assert_eq!(get_allocated(&pool, page), 1);
drop(iobuf);
assert_eq!(get_allocated(&pool, page), 0);
assert_eq!(get_available(&pool, page), 1);
}
#[test]
fn test_refcount_and_copy_to_bytes_paths() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
{
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(&[0xAA; 100]);
let iobuf = buf.freeze();
let clone = iobuf.clone();
let slice = iobuf.slice(10..40);
let empty = iobuf.slice(10..10);
assert!(empty.is_empty());
drop(iobuf);
assert_eq!(get_allocated(&pool, page), 1);
drop(slice);
assert_eq!(get_allocated(&pool, page), 1);
drop(clone);
assert_eq!(get_allocated(&pool, page), 0);
}
{
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(&[0x42; 100]);
let mut iobuf = buf.freeze();
let zero = iobuf.copy_to_bytes(0);
assert!(zero.is_empty());
assert_eq!(iobuf.remaining(), 100);
let partial = iobuf.copy_to_bytes(30);
assert_eq!(&partial[..], &[0x42; 30]);
assert_eq!(iobuf.remaining(), 70);
let rest = iobuf.copy_to_bytes(70);
assert_eq!(&rest[..], &[0x42; 70]);
assert_eq!(iobuf.remaining(), 0);
let empty = iobuf.copy_to_bytes(0);
assert!(empty.is_empty());
drop(iobuf);
assert_eq!(get_allocated(&pool, page), 1);
drop(zero);
drop(partial);
assert_eq!(get_allocated(&pool, page), 1);
drop(rest);
assert_eq!(get_allocated(&pool, page), 0);
}
{
let buf = pool.try_alloc(page).unwrap();
let mut iobufmut = buf;
iobufmut.put_slice(&[0x7E; 100]);
let zero = iobufmut.copy_to_bytes(0);
assert!(zero.is_empty());
assert_eq!(iobufmut.remaining(), 100);
let partial = iobufmut.copy_to_bytes(30);
assert_eq!(&partial[..], &[0x7E; 30]);
assert_eq!(iobufmut.remaining(), 70);
let rest = iobufmut.copy_to_bytes(70);
assert_eq!(&rest[..], &[0x7E; 70]);
assert_eq!(iobufmut.remaining(), 0);
drop(iobufmut);
assert_eq!(get_allocated(&pool, page), 1);
drop(zero);
drop(partial);
assert_eq!(get_allocated(&pool, page), 1);
drop(rest);
assert_eq!(get_allocated(&pool, page), 0);
}
}
#[test]
fn test_iobuf_to_iobufmut_conversion_reuses_pool_for_non_full_unique_view() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(b"non-full");
assert_eq!(get_allocated(&pool, page), 1);
let iobuf = buf.freeze();
assert_eq!(iobuf.len(), 8);
assert_eq!(get_allocated(&pool, page), 1);
let iobufmut: IoBufMut = iobuf.into();
assert_eq!(iobufmut.as_ref(), b"non-full");
assert_eq!(
get_allocated(&pool, page),
1,
"pooled buffer should remain allocated after zero-copy IoBuf->IoBufMut conversion"
);
assert_eq!(get_available(&pool, page), 0);
drop(iobufmut);
assert_eq!(get_allocated(&pool, page), 0);
assert_eq!(get_available(&pool, page), 1);
}
#[test]
fn test_iobuf_try_into_mut_recycles_full_unique_view() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(&vec![0xAB; page]);
let iobuf = buf.freeze();
assert_eq!(get_allocated(&pool, page), 1);
let recycled = iobuf
.try_into_mut()
.expect("unique full-view pooled buffer should recycle");
assert_eq!(recycled.len(), page);
assert!(recycled.as_ref().iter().all(|&b| b == 0xAB));
assert_eq!(recycled.capacity(), page);
assert_eq!(get_allocated(&pool, page), 1);
drop(recycled);
assert_eq!(get_allocated(&pool, page), 0);
assert_eq!(get_available(&pool, page), 1);
}
#[test]
fn test_iobuf_try_into_mut_succeeds_for_unique_slice_and_fails_for_shared() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(&vec![0xCD; page]);
let iobuf = buf.freeze();
let sliced = iobuf.slice(1..page);
drop(iobuf);
let recycled = sliced
.try_into_mut()
.expect("unique sliced pooled buffer should recycle");
assert_eq!(recycled.len(), page - 1);
assert!(recycled.as_ref().iter().all(|&b| b == 0xCD));
assert_eq!(recycled.capacity(), page - 1);
assert_eq!(get_allocated(&pool, page), 1);
drop(recycled);
assert_eq!(get_allocated(&pool, page), 0);
assert_eq!(get_available(&pool, page), 1);
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(&vec![0xEF; page]);
let iobuf = buf.freeze();
let cloned = iobuf.clone();
let iobuf = iobuf
.try_into_mut()
.expect_err("shared pooled buffer must not convert to mutable");
drop(cloned);
drop(iobuf);
assert_eq!(get_allocated(&pool, page), 0);
assert!(get_available(&pool, page) >= 1);
}
#[test]
fn test_multithreaded_alloc_freeze_return() {
let page = page_size();
let pool = Arc::new(test_pool(test_config(page, page, 100)));
let mut handles = vec![];
cfg_if::cfg_if! {
if #[cfg(miri)] {
let iterations = 100;
} else {
let iterations = 1000;
}
}
for _ in 0..10 {
let pool = pool.clone();
let handle = thread::spawn(move || {
for _ in 0..iterations {
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(b"x");
let iobuf = buf.freeze();
let clones: Vec<_> = (0..5).map(|_| iobuf.clone()).collect();
drop(iobuf);
for clone in clones {
drop(clone);
}
}
});
handles.push(handle);
}
for handle in handles {
handle.join().unwrap();
}
let _buf = pool
.try_alloc(page)
.expect("pool should remain usable after multithreaded test");
}
#[test]
fn test_cross_thread_buffer_return() {
let page = page_size();
let pool = test_pool(test_config(page, page, 100));
let (tx, rx) = mpsc::channel();
for _ in 0..50 {
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(b"x");
let iobuf = buf.freeze();
tx.send(iobuf).unwrap();
}
drop(tx);
let handle = thread::spawn(move || {
while let Ok(iobuf) = rx.recv() {
drop(iobuf);
}
let class_index = pool
.class_index(page)
.expect("class exists for page-sized buffer");
assert_eq!(get_local_len(&pool.inner.classes[class_index]), 50);
assert_eq!(get_global_len(&pool.inner.classes[class_index]), 0);
for _ in 0..50 {
let _buf = pool
.try_alloc(page)
.expect("dropping thread should be able to reuse locally returned buffers");
}
});
handle.join().unwrap();
}
#[test]
fn test_pool_dropped_before_buffer() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let mut buf = pool.try_alloc(page).unwrap();
buf.put_slice(&[0u8; 100]);
let iobuf = buf.freeze();
drop(pool);
assert_eq!(iobuf.len(), 100);
drop(iobuf);
}
#[test]
fn test_pool_exhaustion_and_recovery() {
let page = page_size();
let pool = test_pool(test_config(page, page, 3));
let buf1 = pool.try_alloc(page).expect("first alloc");
let buf2 = pool.try_alloc(page).expect("second alloc");
let buf3 = pool.try_alloc(page).expect("third alloc");
assert!(pool.try_alloc(page).is_err(), "pool should be exhausted");
drop(buf1);
let buf4 = pool.try_alloc(page).expect("alloc after return");
assert!(pool.try_alloc(page).is_err(), "pool exhausted again");
drop(buf2);
drop(buf3);
drop(buf4);
assert_eq!(get_allocated(&pool, page), 0);
assert_eq!(get_available(&pool, page), 3);
let _buf5 = pool.try_alloc(page).expect("reuse from freelist");
assert_eq!(get_available(&pool, page), 2);
}
#[test]
fn test_try_alloc_errors() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let result = pool.try_alloc(page * 10);
assert_eq!(result.unwrap_err(), PoolError::Oversized);
let _buf1 = pool.try_alloc(page).unwrap();
let _buf2 = pool.try_alloc(page).unwrap();
let result = pool.try_alloc(page);
assert_eq!(result.unwrap_err(), PoolError::Exhausted);
}
#[test]
fn test_pool_metrics_track_created_exhausted_oversized() {
let page = page_size();
let mut registry = Registry::default();
let pool = BufferPool::new(test_config(page, page, 1), &mut registry);
let buf = pool.try_alloc(page).unwrap();
assert_eq!(pool.try_alloc(page).unwrap_err(), PoolError::Exhausted);
assert_eq!(pool.try_alloc(page * 2).unwrap_err(), PoolError::Oversized);
let encoded = registry.encode();
assert!(
encoded.contains(&format!("buffer_pool_created{{size_class=\"{page}\"}} 1")),
"created gauge missing: {encoded}"
);
assert!(
encoded.contains(&format!(
"buffer_pool_exhausted_total{{size_class=\"{page}\"}} 1"
)),
"exhausted counter missing: {encoded}"
);
assert!(
encoded.contains("buffer_pool_oversized_total 1"),
"oversized counter missing: {encoded}"
);
drop(buf);
}
#[test]
fn test_try_alloc_zeroed_errors() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let result = pool.try_alloc_zeroed(page * 10);
assert_eq!(result.unwrap_err(), PoolError::Oversized);
let _buf1 = pool.try_alloc_zeroed(page).unwrap();
let _buf2 = pool.try_alloc_zeroed(page).unwrap();
let result = pool.try_alloc_zeroed(page);
assert_eq!(result.unwrap_err(), PoolError::Exhausted);
}
#[test]
fn test_fallback_allocation() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let buf1 = pool.try_alloc(page).unwrap();
let buf2 = pool.try_alloc(page).unwrap();
assert!(buf1.is_pooled());
assert!(buf2.is_pooled());
let mut fallback_exhausted = pool.alloc(page);
assert!(!fallback_exhausted.is_pooled());
assert!((fallback_exhausted.as_mut_ptr() as usize).is_multiple_of(page));
assert_eq!(fallback_exhausted.capacity(), page);
let fallback_small = pool.alloc(100);
assert!(!fallback_small.is_pooled());
assert!((100..108).contains(&fallback_small.capacity()));
let mut fallback_oversized = pool.alloc(page * 10);
assert!(!fallback_oversized.is_pooled());
assert!((fallback_oversized.as_mut_ptr() as usize).is_multiple_of(page));
assert_eq!(fallback_oversized.capacity(), page * 10);
assert_eq!(get_allocated(&pool, page), 2);
drop(fallback_exhausted);
drop(fallback_oversized);
assert_eq!(get_allocated(&pool, page), 2);
drop(buf1);
drop(buf2);
assert_eq!(get_allocated(&pool, page), 0);
}
#[test]
fn test_is_pooled() {
let page = page_size();
let pool = test_pool(test_config(page, page, 10));
let pooled = pool.try_alloc(page).unwrap();
assert!(pooled.is_pooled());
let owned = IoBufMut::with_capacity(100);
assert!(!owned.is_pooled());
}
#[test]
fn test_iobuf_is_pooled() {
let page = page_size();
let pool = test_pool(test_config(page, page, 2));
let mut pooled = pool.try_alloc(page).unwrap();
pooled.put_slice(b"x");
let pooled = pooled.freeze();
assert!(pooled.is_pooled());
let fallback = pool.alloc(page * 10).freeze();
assert!(!fallback.is_pooled());
let bytes = IoBuf::copy_from_slice(b"hello");
assert!(!bytes.is_pooled());
}
#[test]
fn test_buffer_alignment() {
let page = page_size();
let cache_line = cache_line_size();
cfg_if::cfg_if! {
if #[cfg(miri)] {
let storage_config = BufferPoolConfig::for_storage()
.with_alignment(NZUsize!(page))
.with_max_per_class(NZU32!(32));
let network_config = BufferPoolConfig::for_network()
.with_alignment(NZUsize!(cache_line))
.with_max_per_class(NZU32!(32));
} else {
let storage_config =
BufferPoolConfig::for_storage().with_alignment(NZUsize!(page));
let network_config =
BufferPoolConfig::for_network().with_alignment(NZUsize!(cache_line));
}
}
let storage_buffer_pool = test_pool(storage_config);
let mut buf = storage_buffer_pool.try_alloc(100).unwrap();
assert_eq!(
buf.as_mut_ptr() as usize % page,
0,
"storage buffer not page-aligned"
);
let network_buffer_pool = test_pool(network_config);
let mut buf = network_buffer_pool.try_alloc(100).unwrap();
assert_eq!(
buf.as_mut_ptr() as usize % cache_line,
0,
"network buffer not cache-line aligned"
);
}
}
#[cfg(all(test, feature = "loom"))]
mod loom_tests {
use super::*;
use crate::telemetry::metrics::Registry;
use bytes::BufMut;
use loom::thread;
#[test]
fn freeze_clone_cross_thread_drop_then_reuse() {
loom::model(|| {
let mut registry = Registry::default();
let config = BufferPoolConfig::for_network()
.with_size_class_range(NZUsize!(64), NZUsize!(64), NZU32!(2))
.with_thread_cache_disabled();
let pool = BufferPool::new(config, &mut registry);
let mut buf = pool.alloc(64);
assert!(buf.is_pooled());
buf.put_slice(b"payload");
let frozen = buf.freeze();
let clone = frozen.clone();
let t = thread::spawn(move || {
assert_eq!(clone.as_ref(), b"payload");
drop(clone);
});
assert_eq!(frozen.as_ref(), b"payload");
drop(frozen);
t.join().unwrap();
let mut again = pool.alloc(64);
assert!(again.is_pooled());
again.put_slice(b"reuse");
assert_eq!(again.as_ref(), b"reuse");
});
}
#[test]
fn final_drop_races_pool_teardown() {
loom::model(|| {
let mut registry = Registry::default();
let config = BufferPoolConfig::for_network()
.with_size_class_range(NZUsize!(64), NZUsize!(64), NZU32!(1))
.with_thread_cache_disabled();
let pool = BufferPool::new(config, &mut registry);
let mut buf = pool.alloc(64);
assert!(buf.is_pooled());
buf.put_slice(b"x");
let frozen = buf.freeze();
let t = thread::spawn(move || drop(frozen));
drop(pool);
t.join().unwrap();
});
}
#[test]
fn final_drop_races_recheckout() {
loom::model(|| {
let mut registry = Registry::default();
let config = BufferPoolConfig::for_network()
.with_size_class_range(NZUsize!(64), NZUsize!(64), NZU32!(1))
.with_thread_cache_disabled();
let pool = BufferPool::new(config, &mut registry);
let mut buf = pool.alloc(64);
assert!(buf.is_pooled());
buf.put_slice(b"x");
let frozen = buf.freeze();
let clone = frozen.clone();
let t = thread::spawn(move || drop(clone));
drop(frozen);
if let Ok(mut again) = pool.try_alloc(64) {
assert!(again.is_pooled());
again.put_slice(b"y");
assert_eq!(again.as_ref(), b"y");
}
t.join().unwrap();
});
}
}