//! Buffer pool for efficient I/O operations. //! //! Provides pooled, aligned buffers that can be reused to reduce allocation //! overhead. Buffer alignment is configurable: use page alignment for storage I/O //! (required for direct I/O and DMA), or cache-line alignment for network I/O //! (reduces fragmentation). //! //! # Thread Safety //! //! [`BufferPool`] is `Send + Sync` and can be safely shared across threads. //! Allocation and deallocation use atomic counters together with a global //! freelist split across mutex-protected stripes, each with a fixed slot limit, //! plus per-thread caches. //! //! Global freelist operations use blocking mutexes. After a local cache miss or //! spill, an operation can wait for a stripe lock. A preempted lock holder can //! therefore delay the operation even when another stripe could provide or //! accept a buffer. //! //! # Pool Lifecycle //! //! Tracked buffers held by pooled views or cached in thread-local bins keep a //! strong reference to the originating size class. Buffers can outlive the //! public [`BufferPool`] handle and still return to their original size class. //! - Untracked fallback allocations store no class reference and deallocate //! directly when dropped. //! - Requests smaller than [`BufferPoolConfig::pool_min_size`] bypass pooling //! entirely and return untracked aligned allocations from both //! [`BufferPool::try_alloc`] and [`BufferPool::alloc`]. //! - Dropping [`BufferPool`] drains only the shared global freelists. Pooled //! views and buffers cached in a live thread's local cache can keep their //! size class alive until they are dropped or the thread exits. //! //! # Size Classes //! //! Buffers are organized into power-of-two size classes. The enabled classes //! do not need to be contiguous, and each class has its own tracked-buffer //! limit. For example, with enabled classes 4096, 8192, and 32768: //! - Class 0: 4096 bytes //! - Class 1: 8192 bytes //! - Class 2: 32768 bytes //! //! Allocation requests round up to the smallest enabled class that fits, so a //! 16000-byte request above is served by the 32768-byte class. Requests larger //! than the largest enabled class return [`PoolError::Oversized`] from //! [`BufferPool::try_alloc`], or fall back to an untracked aligned heap //! allocation from [`BufferPool::alloc`]. A request routed to an exhausted //! class returns [`PoolError::Exhausted`] without trying larger classes. //! //! # Cache Structure //! //! Each size class uses a two-level allocator: //! - a small per-thread local cache for steady-state same-thread reuse //! - a shared global freelist for refill and spill between threads //! //! When a local cache misses, the pool refills a small batch from the global //! freelist before attempting to create a new tracked buffer. Returned buffers //! first try to re-enter the dropping thread's local cache, spilling a bounded //! batch back to the global freelist if needed. 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; } } /// Error returned when buffer pool allocation fails. #[derive(Error, Debug, Clone, Copy, PartialEq, Eq)] pub enum PoolError { /// The requested capacity exceeds the maximum buffer size. #[error("requested capacity exceeds maximum buffer size")] Oversized, /// The pool is exhausted for the required size class. #[error("pool exhausted for required size class")] Exhausted, } /// Policy for sizing each thread's cache within a buffer pool size class. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum BufferPoolThreadCacheConfig { /// Enable thread-local caching. /// /// `None` derives the per-thread cache size from the pool's per-class /// capacity and expected parallelism, reserving about half of each class /// for the shared freelist. Small per-class budgets may resolve to zero, /// disabling thread-local caching so free buffers do not become stranded in /// other threads. /// /// `Some(n)` uses an explicit per-thread cache size, clamped independently /// to each size class's limit. Enabled(Option), /// Disable thread-local caching and route all reuse through the shared global freelist. Disabled, } /// Configuration for one enabled buffer pool size class. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct BufferPoolClassConfig { /// Buffer size for this class. Must be a power of two. pub size: NonZeroUsize, /// Maximum number of tracked buffers in this class. /// /// Size-class slots are identified by `u32`, so the per-class limit is /// capped by this type. pub max_buffers: NonZeroU32, } impl From<(NonZeroUsize, NonZeroU32)> for BufferPoolClassConfig { fn from((size, max_buffers): (NonZeroUsize, NonZeroU32)) -> Self { Self { size, max_buffers } } } /// Configuration for a buffer pool. /// /// The class layout is a set of power-of-two size classes, each with its own /// tracked-buffer limit. Enabled classes do not need to be contiguous, and /// requests route to the smallest enabled class that fits. /// /// Shape builders do not commute. Each builder applies to the layout produced /// by the previous one: replacement builders ([`Self::with_size_class_range`], /// [`Self::with_size_classes`]) discard the current layout, uniform builders /// ([`Self::with_max_per_class`], [`Self::with_bytes_per_class`]) overwrite /// every enabled limit, and [`Self::with_budget_bytes`] snapshots and rescales /// the shape that exists at that call. #[derive(Clone, Debug)] pub struct BufferPoolConfig { /// Minimum request size that should use pooled allocation. /// /// Requests smaller than this bypass the pool and use direct aligned /// allocation instead. A value of `0` means all eligible requests use the /// pool. pool_min_size: usize, /// Enabled size classes, keyed by power-of-two size. Sizes absent from the /// map are disabled. /// /// Builders maintain the invariant that at least one class is enabled. class_limits: BTreeMap, /// Whether to create every tracked buffer during pool construction. /// /// When enabled, each size class creates its configured limit of buffers /// and parks them in the class-global freelist before the pool is /// returned. This moves allocation cost to startup and makes the first /// reuse path avoid heap allocation. prefill: bool, /// Buffer alignment. Must be a power of two. alignment: NonZeroUsize, /// Expected number of threads concurrently accessing the pool. /// /// This sizes the shared global freelist stripes. It is also used to derive /// thread-cache capacity when the thread-cache policy is automatic, using /// approximately half of each class limit divided across expected threads. parallelism: NonZeroUsize, /// Policy for sizing the per-thread local cache in each size class. /// /// By default, thread-cache capacity is derived from [`Self::parallelism`] /// and each class limit. [`Self::with_max_thread_cache_capacity`] uses an /// explicit per-thread cache size clamped to each class limit. /// [`Self::with_thread_cache_disabled`] bypasses thread-local caches. pub(crate) thread_cache_config: BufferPoolThreadCacheConfig, } impl BufferPoolConfig { /// Network I/O preset: 1KB to 128KB buffers, 4096 per class, not prefilled. /// /// Network operations typically need multiple concurrent buffers per /// connection (message, encoding, encryption) so we allow 4096 buffers per /// size class. 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)) } /// Storage I/O preset: `page_size` (usually 4KB) to 8MB buffers, 64 per class, /// not prefilled. pub fn for_storage() -> Self { Self { pool_min_size: 0, class_limits: BTreeMap::new(), prefill: false, // TODO (#2960): this needs to be page/block aligned for O_DIRECT 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), ) } /// Validates a class size, panicking on invalid values. /// /// Sizes above `isize::MAX` are rejected here because `Layout` cannot /// represent them, which would otherwise surface as a misleading panic at /// pool construction. 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" ); } /// Returns a copy of this config with a new minimum request size that uses pooling. pub const fn with_pool_min_size(mut self, pool_min_size: usize) -> Self { self.pool_min_size = pool_min_size; self } /// Returns a copy of this config whose layout is the inclusive, contiguous /// power-of-two range from `min` to `max` with a uniform limit. /// /// This replaces the complete class layout. /// /// # Panics /// /// - `min` or `max` is not a power of two /// - `min` or `max` exceeds `isize::MAX` /// - `max < min` 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)), ) } /// Returns a copy of this config whose layout is exactly the given classes. /// /// This replaces the complete class layout. Input order does not matter, /// classes are normalized into ascending size order. /// /// # Panics /// /// - `classes` is empty /// - a class size is not a power of two /// - a class size exceeds `isize::MAX` /// - two classes have the same size pub fn with_size_classes(mut self, classes: I) -> Self where I: IntoIterator, C: Into, { 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 } /// Returns a copy of this config with the given class enabled, replacing /// its limit if it is already enabled. /// /// # Panics /// /// - `size` is not a power of two /// - `size` exceeds `isize::MAX` 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 } /// Returns a copy of this config with the given class removed. /// /// Requests that previously routed to the removed class route to the next /// larger enabled class. /// /// # Panics /// /// - `size` is not a power of two /// - `size` exceeds `isize::MAX` /// - no class with `size` is enabled /// - the class is the final enabled class 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 } /// Returns a copy of this config with the same limit on every enabled class. pub fn with_max_per_class(mut self, max_buffers: NonZeroU32) -> Self { for limit in self.class_limits.values_mut() { *limit = max_buffers; } self } /// Returns a copy of this config where every enabled class has /// approximately the same tracked-byte weight. /// /// Each enabled class's limit becomes `max(1, bytes / size)`, so limits /// halve as class sizes double. This is a one-shot count transformation, /// not a stored byte policy, and it never disables a class. /// /// # Panics /// /// Panics if a derived limit exceeds `u32::MAX`. 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 } /// Returns a copy of this config with a new expected parallelism. /// /// The global freelist derives its stripe count from this target and the /// class capacity. This value also controls thread-cache capacity when the /// thread-cache policy is automatic. The automatic policy reserves about /// half of each class for the global freelist and divides the remaining /// capacity across expected threads. pub const fn with_parallelism(mut self, parallelism: NonZeroUsize) -> Self { self.parallelism = parallelism; self } /// Returns a copy of this config with an explicit per-thread cache size. /// /// Each size class keeps a small per-thread cache of free buffers for /// same-thread reuse. By default its capacity is derived per class from /// the class limit and [`Self::parallelism`], reserving about half of the /// class for the shared global freelist. An explicit capacity replaces /// that derivation and may be larger or smaller than the derived value. /// /// The effective capacity for each class is `min(capacity, class limit)`. /// Clamping happens independently per class, so one small class cannot /// invalidate the configuration. /// /// Buffers held in a thread's cache are invisible to other threads until /// they spill to the global freelist or the thread exits, and each thread /// can retain up to the effective capacity of every class it touches. /// Larger values favor same-thread reuse while smaller values favor /// cross-thread visibility and a lower per-thread memory ceiling. /// /// Global-freelist striping is set separately by [`Self::with_parallelism`]. pub const fn with_max_thread_cache_capacity(mut self, capacity: NonZeroUsize) -> Self { self.thread_cache_config = BufferPoolThreadCacheConfig::Enabled(Some(capacity)); self } /// Returns a copy of this config with thread-local caching disabled. /// /// Global-freelist striping is set separately by [`Self::with_parallelism`]. pub const fn with_thread_cache_disabled(mut self) -> Self { self.thread_cache_config = BufferPoolThreadCacheConfig::Disabled; self } /// Returns a copy of this config with a new prefill setting. pub const fn with_prefill(mut self, prefill: bool) -> Self { self.prefill = prefill; self } /// Returns a copy of this config with a new alignment. pub const fn with_alignment(mut self, alignment: NonZeroUsize) -> Self { self.alignment = alignment; self } /// Returns a copy of this config with all class limits proportionally /// rescaled under a strict tracked-byte ceiling. /// /// This snapshots the currently enabled classes and their limits, then /// chooses the greatest common proportional scale for which the total /// tracked capacity `sum(size * scaled_limit)` stays within `budget`, /// where `scaled_limit = max(1, floor(limit * scale))`. Scaling may raise /// or lower limits, never disables a class, and may deliberately leave /// part of the budget unused rather than distort the requested shape. /// /// The budget covers tracked buffer payload capacity only. It does not /// include allocator metadata, alignment overhead, or pool bookkeeping. /// /// This is a one-shot transformation, not a stored policy. Later builder /// calls may change the resulting total, and calling this again rescales /// the already scaled limits rather than the shape they were derived from. /// /// # Panics /// /// - `budget` is smaller than one buffer from every enabled class /// - the budget would require scaling a limit above `u32::MAX` pub fn with_budget_bytes(mut self, budget: NonZeroUsize) -> Self { let budget = budget.get() as u128; // The smallest expressible footprint keeps one buffer per class. 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" ); // Scales are unsigned Q64.64 fixed-point. Since class limits fit in // u32, consecutive distinct count-change breakpoints k1/c1 and k2/c2 // are separated by at least 1/(c1*c2) > 2^-64, i.e. more than one // fixed-point step, so the maximal feasible count plateau always // contains a representable scale and the binary search below finds an // optimal count vector. const FRACTION_BITS: u32 = 64; // Scaled limit for one class. The u32 slot-identifier bound is // enforced by `evaluate` and the post-search assert below. Saturating // math keeps the evaluation monotonic for scales beyond that bound // instead of overflowing. let scaled = |limit: NonZeroU32, scale: u128| -> u128 { ((limit.get() as u128).saturating_mul(scale) >> FRACTION_BITS).max(1) }; // Saturating total tracked bytes at a scale, and whether every scaled // limit still fits u32 slot identifiers. Both constraints are // monotonically violated as the scale grows, so feasibility is a // prefix of the scale axis and binary search applies. 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 }; // Scale zero floors every class at one buffer, which the minimum // check above proved feasible. The upper bound exceeds any scale that // could keep the smallest possible limit within u32, so it is // infeasible and the search invariant holds at both ends. 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; } } // The next representable scale is infeasible. If its total would // still fit the budget, the binding constraint is the u32 limit // bound, which the caller must resolve instead of silently capping. let (next_total, _) = evaluate(lo + 1); assert!( next_total > budget, "budget requires scaling a class limit above u32::MAX" ); // Rescale each limit in place at the optimal feasible scale. 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 } /// Returns an iterator over enabled classes in ascending size order. pub fn size_classes(&self) -> impl ExactSizeIterator + '_ { self.class_limits .iter() .map(|(&size, &max_buffers)| BufferPoolClassConfig { size, max_buffers }) } /// Returns the enabled class that serves a pooled request of `size` bytes, /// or `None` if `size` exceeds the largest enabled class. /// /// Requests route to the smallest enabled class that fits, so in sparse /// layouts the returned class may be much larger than the request. This /// reports class shape only: zero-sized requests and requests below /// [`Self::pool_min_size`] bypass the pool, and oversized requests fall /// back to untracked aligned allocations with capacity at least as large /// as the request. pub fn class_for(&self, size: usize) -> Option { self.size_classes().find(|class| class.size.get() >= size) } /// Returns the minimum request size that uses pooled allocation. pub const fn pool_min_size(&self) -> usize { self.pool_min_size } /// Returns whether every tracked buffer is created during pool construction. pub const fn prefill(&self) -> bool { self.prefill } /// Returns the buffer alignment. pub const fn alignment(&self) -> NonZeroUsize { self.alignment } /// Returns the expected number of threads concurrently accessing the pool. pub const fn parallelism(&self) -> NonZeroUsize { self.parallelism } /// Returns the smallest enabled class size. pub fn min_size(&self) -> NonZeroUsize { *self .class_limits .first_key_value() .expect("class layout must enable at least one class") .0 } /// Returns the largest enabled class size. pub fn max_size(&self) -> NonZeroUsize { *self .class_limits .last_key_value() .expect("class layout must enable at least one class") .0 } /// Returns `sum(class size * class limit)`, saturating at `usize::MAX`. /// /// A saturated result means the configured maximum tracked capacity is at /// least that large. 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) } /// Validates cross-field constraints, panicking on invalid values. /// /// Layout-local mistakes panic at the builder call that introduces them. /// The constraints here span independently configured fields, so they are /// deferred to pool construction to keep builder order unrestricted. /// /// # Panics /// /// - `alignment` is not a power of two /// - the smallest enabled class is smaller than `alignment` /// - `pool_min_size` is larger than the smallest enabled class 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 ); } /// Resolves the effective per-thread cache size for one size class. /// /// Derived capacities divide half of the class limit across the expected /// parallelism so cross-thread reuse remains effective. Small class limits /// may resolve to zero. An explicit capacity replaces the derivation and /// clamps to the class limit. 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, } } } /// Label for buffer pool metrics, identifying the size class. #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)] struct SizeClassLabel { size_class: u64, } /// Metrics for the buffer pool. struct PoolMetrics { /// Number of tracked buffers created for the size class. created: GaugeFamily, /// Total number of failed allocations (pool exhausted). exhausted_total: CounterFamily, /// Total number of oversized allocation requests. 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(), ), // Counters are registered without the `_total` suffix because the // prometheus encoder appends it to counter names. 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(), ), } } } /// Internal allocation result for pooled allocations. struct Allocation { buffer: PooledBuffer, is_new: bool, } /// Internal state of the buffer pool. pub(crate) struct BufferPoolInner { config: BufferPoolConfig, /// Exponent-indexed routing vector. /// /// Entry `i` serves requests that round up to `min_size << i`. Disabled /// exponents hold cloned handles that alias the next enabled class, so /// the allocation path resolves any request with plain arithmetic and one /// vector index. Aliased entries form contiguous runs that end at the /// enabled class's own exponent. classes: Vec, /// Smallest enabled class size, cached off [`BufferPoolConfig`] so the /// allocation path never scans the configuration table. min_size: usize, /// Largest enabled class size, cached for the same reason. max_size: usize, metrics: PoolMetrics, } impl Drop for BufferPoolInner { fn drop(&mut self) { // The public pool is going away. Drain globally parked buffers while // the pool-owned class handles are still live. Pooled views and live // TLS cache entries own their own size-class references. If they // return later, they park their buffer and release the reference that // kept the class alive. // // Routing entries alias the next enabled class in contiguous runs, so // dropping consecutive duplicates leaves one live handle per unique // class and drains each class once. 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 { /// Try to allocate a buffer from the given size class. /// /// Uses a three-tier strategy: /// 1. **Thread-local cache** (fast path): no atomics, no contention. /// 2. **Global freelist**: striped pop, then batch-refill the local cache /// when the local bin is large enough to amortize shared-queue traffic. /// 3. **New allocation**: reserve a slot in the global freelist, then /// allocate from the heap. /// /// If `zero_on_new` is true, newly-created buffers are allocated with /// `alloc_zeroed`. Reused buffers are never re-zeroed here. #[inline(always)] fn try_alloc(&self, class_index: usize, zero_on_new: bool) -> Option { let class = &self.classes[class_index]; // Reuse path: try the thread-local cache first, then the global // freelist with batch refill when the local cache is large enough. if let Some(buffer) = BufferPoolThreadCache::pop(class) { return Some(Allocation { buffer, is_new: false, }); } // Slow path: create a new tracked buffer and update metrics. self.try_alloc_new(class, zero_on_new) } /// Creates a new tracked buffer after the reuse path fails. /// /// This is separate from [`Self::try_alloc`] so the steady-state allocation /// path can inline the TLS hit without also carrying slot reservation, /// metrics, and heap-allocation code. #[inline(never)] fn try_alloc_new(&self, class: &SizeClassHandle, zeroed: bool) -> Option { 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, }) } } /// A pool of reusable, aligned buffers. /// /// Buffers are organized into power-of-two size classes. When a buffer is /// requested, the smallest size class that fits is used. Pooled buffers are /// automatically returned when their final owning view is dropped. /// /// # Alignment /// /// Buffer alignment is guaranteed only at the allocation base, where a /// freshly allocated buffer's pointer starts. After [`bytes::Buf::advance`], /// the pointer returned by `as_mut_ptr()` may no longer be aligned. For /// direct I/O operations that require alignment, do not advance the buffer /// before use. /// /// # Thread-local caching /// /// Returned buffers are cached per thread for reuse. After the pool is /// dropped, buffers still cached on other threads are reclaimed when those /// threads exit or call [`BufferPoolThreadCache::flush`]. A long-lived /// thread that used a since-dropped pool retains its cached buffers until /// then. Processes that create and drop many pools should reuse threads' /// pools or flush explicitly. #[derive(Clone)] pub struct BufferPool { inner: Arc, } 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() } } /// Global allocator for size-class TLS registry ids. /// /// `class_id` is the key used by each thread's cache registry. It must be global, not /// pool-local, because the same thread-local registry serves every /// [`BufferPool`] touched by the thread. Without a global id, two different /// pools could share a class index and accidentally share one local cache. /// /// Ids are monotonic and never reused. Reuse would make stale per-thread cache /// state ambiguous after a pool is dropped and a later pool creates a new size /// class with the same id. Avoiding reuse means the hot path can index directly /// without generation checks, at the cost of possible holes in each thread's /// sparse registry. /// /// Relaxed ordering is sufficient: the atomic operation is only used to assign /// unique ids, not to publish any associated size-class state. static NEXT_SIZE_CLASS_ID: AtomicUsize = AtomicUsize::new(0); impl BufferPool { /// Creates a new buffer pool with the given configuration. /// /// # Panics /// /// Panics if the configuration is invalid. 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; // Create one allocator per enabled class and expand the exponent-indexed // routing vector up to it. Every exponent in the enabled span resolves // to the smallest enabled class at or above it, so the gap entries // `resize` fills below each class alias that class. Prefill happens // inside `SizeClassHandle::new`, once per unique class. 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, ); // Initialize created metrics after constructor 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, }), } } /// Returns the routing index for a given size, or `None` if `size` exceeds /// the largest enabled class. /// /// The routing vector is exponent-indexed, so this arithmetic is identical /// for contiguous and sparse layouts. Disabled exponents resolve to an /// aliased handle for the next enabled class. #[inline(always)] fn class_index(&self, size: usize) -> Option { 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); } // Pool construction guarantees `min_size` and `max_size` are powers of // two. Since `min_size < size <= max_size`, `next_power_of_two()` // resolves to a valid routing entry and its exponent must be greater // than `min_size`'s exponent. Use wrapping arithmetic to avoid a // release overflow-check branch in this hot helper. Some( size.next_power_of_two() .trailing_zeros() .wrapping_sub(min_size.trailing_zeros()) as usize, ) } /// Returns the size class index for `capacity`, recording oversized metrics on failure. #[inline] fn class_index_or_record_oversized(&self, capacity: usize) -> Option { let class_index = self.class_index(capacity); if class_index.is_none() { self.inner.metrics.oversized_total.inc(); } class_index } /// Attempts to allocate a buffer without falling back on pool miss. /// /// Unlike [`Self::alloc`], this method does not fall back to untracked /// allocation on exhaustion or oversized requests. Requests smaller than /// [`BufferPoolConfig::pool_min_size`] intentionally bypass pooling and /// return an untracked aligned allocation instead. /// /// The returned buffer has `len() == 0` and `capacity() >= capacity`. /// /// Zero-capacity requests return a detached empty buffer without touching /// the pool. /// /// # Initialization /// /// The returned buffer contains **uninitialized memory**. Do not read from /// it until data has been written. /// /// # Errors /// /// - [`PoolError::Oversized`]: `capacity` exceeds `max_size` /// - [`PoolError::Exhausted`]: pool exhausted for the required size class #[inline(always)] pub fn try_alloc(&self, capacity: usize) -> Result { 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| { // SAFETY: pooled allocations returned by the pool have an // initialized live lease. unsafe { IoBufMut::from_pooled_parts(allocation.buffer) } }) .ok_or(PoolError::Exhausted)?; Ok(buffer) } /// Allocates a buffer with capacity for at least `capacity` bytes. /// /// The returned buffer has `len() == 0` and `capacity() >= capacity`, /// matching the semantics of [`IoBufMut::with_capacity`] and /// [`bytes::BytesMut::with_capacity`]. Use [`bytes::BufMut::put_slice`] or /// other [`bytes::BufMut`] methods to write data to the buffer. /// /// Zero-capacity requests return a detached empty buffer without touching /// the pool. /// /// If the pool can provide a buffer (capacity within limits and pool not /// exhausted), this returns a pooled buffer that will be returned to the /// pool when dropped. Requests smaller than [`BufferPoolConfig::pool_min_size`] /// bypass pooling and return an untracked aligned allocation. Oversized or /// exhausted requests also fall back to an untracked aligned heap allocation /// that is deallocated when dropped. /// /// Use [`Self::try_alloc`] if eligible requests must fail instead of falling /// back to direct allocation. /// /// # Initialization /// /// The returned buffer contains **uninitialized memory**. Do not read from /// it until data has been written. #[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) }) } /// Allocates a buffer and sets its readable length to `len` without /// initializing bytes. /// /// Equivalent to [`Self::alloc`] followed by [`IoBufMut::set_len`]. /// /// # Safety /// /// Caller must ensure all bytes are initialized before any read operation. pub unsafe fn alloc_len(&self, len: usize) -> IoBufMut { let mut buf = self.alloc(len); // SAFETY: guaranteed by caller. unsafe { buf.set_len(len) }; buf } /// Attempts to allocate a zero-initialized buffer without falling back on /// pool miss. /// /// Unlike [`Self::alloc_zeroed`], this method does not fall back to /// untracked allocation on exhaustion or oversized requests. Requests /// smaller than [`BufferPoolConfig::pool_min_size`] intentionally bypass /// pooling and return an untracked aligned allocation instead. /// /// The returned buffer has `len() == len` and `capacity() >= len`. /// Zero-length requests return a detached empty buffer without touching /// the pool. /// /// # Initialization /// /// Bytes in `0..len` are initialized to zero. Bytes in `len..capacity` /// may be uninitialized. /// /// # Errors /// /// - [`PoolError::Oversized`]: `len` exceeds `max_size` /// - [`PoolError::Exhausted`]: pool exhausted for the required size class pub fn try_alloc_zeroed(&self, len: usize) -> Result { 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)?; // SAFETY: pooled allocations returned by the pool have an initialized // live lease. let mut buf = unsafe { IoBufMut::from_pooled_parts(allocation.buffer) }; if allocation.is_new { // SAFETY: newly allocated zeroed buffers have `capacity() >= len` // and the readable bytes are initialized by zeroed allocation. unsafe { buf.set_len(len) }; } else { // Reused buffers may contain old bytes, re-zero requested readable range. // SAFETY: `as_mut_ptr()` is valid for writes up to `capacity() >= len` bytes. unsafe { std::ptr::write_bytes(buf.as_mut_ptr(), 0, len); buf.set_len(len); } } Ok(buf) } /// Allocates a zero-initialized buffer with readable length `len`. /// /// The returned buffer has `len() == len` and `capacity() >= len`. /// Zero-length requests return a detached empty buffer without touching /// the pool. /// /// If the pool can provide a buffer (len within limits and pool not /// exhausted), this returns a pooled buffer that will be returned to the /// pool when dropped. Requests smaller than [`BufferPoolConfig::pool_min_size`] /// bypass pooling and return an untracked aligned allocation. Oversized or /// exhausted requests also fall back to an untracked aligned heap allocation /// that is deallocated when dropped. /// /// Use this for read APIs that require an initialized `&mut [u8]`. This /// avoids `unsafe set_len` at callsites. /// /// Use [`Self::try_alloc_zeroed`] if eligible requests must fail instead of /// falling back to direct allocation. /// /// # Initialization /// /// Bytes in `0..len` are initialized to zero. Bytes in `len..capacity` /// may be uninitialized. pub fn alloc_zeroed(&self, len: usize) -> IoBufMut { self.try_alloc_zeroed(len).unwrap_or_else(|_| { // Pool exhausted or oversized: allocate untracked zeroed memory. let size = len.max(1); let mut buf = IoBufMut::zeroed_with_alignment(size, self.inner.config.alignment); buf.truncate(len); buf }) } /// Returns the pool configuration. 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) } /// Creates a test config with page alignment. 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())) } /// Creates a page-aligned test config with exactly the given classes. fn sparse_config(classes: impl IntoIterator) -> 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())) } /// Collects the enabled classes as `(size, max_buffers)` pairs. fn classes_of(config: &BufferPoolConfig) -> Vec<(usize, u32)> { config .size_classes() .map(|class| (class.size.get(), class.max_buffers.get())) .collect() } /// Helper to get the number of caller-owned tracked buffers for a size class. /// /// With TLS enabled, tracked buffers can be free in either the shared /// freelist or the current thread's local cache. 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) } /// Helper to get the number of free buffers visible to the current thread. 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(); // An explicit capacity above a class limit clamps to that limit // instead of invalidating the configuration. 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 ); // Per-class limits clamp independently: a small class cannot lower a // larger class's explicit capacity. 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)); // Classes: page, page*2, page*4, page*8 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() { // Unsorted explicit input normalizes into ascending size order. 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 ); // BufferPoolClassConfig values work as inputs too. 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::()); } #[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), ); // Upsert replaces an enabled class's limit in place. 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)] ); // Upsert can also add a class outside the current span. 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); // Removing a middle class leaves a gap. let sparse = base.clone().without_size_class(NZUsize!(2048)); assert_eq!(classes_of(&sparse), vec![(1024, 8), (4096, 8), (8192, 8)]); // Removing an endpoint narrows the derived bounds. 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); // Uniform overwrite applies to every enabled class. 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)); // Classes smaller than the target get bytes/size buffers, classes // larger than the target floor at one buffer. assert_eq!( classes_of(&config), vec![(1024, 64), (4096, 16), (1 << 20, 1)] ); // A target equal to the class size derives exactly one buffer. 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))]); // Requests at or below the smallest class route to it. assert_eq!(config.class_for(0).unwrap().size.get(), 4096); assert_eq!(config.class_for(4096).unwrap().size.get(), 4096); // Requests in the gap route to the next enabled class, even when // their natural power-of-two exponent is disabled. 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); // Requests above the largest class have no serving class. assert_eq!(config.class_for(32769), None); } #[test] fn test_sparse_routing_allocates_next_enabled_class() { // Classes `page` and `8 * page` with the two exponents between them // disabled: requests in the gap route forward to the larger class. let page = page_size(); let pool = test_pool(sparse_config([(page, 4), (page * 8, 4)])); // Below the first class routes to it. let buf = pool.try_alloc(1).unwrap(); assert_eq!(buf.capacity(), page); // Exact fit for the first class. let buf = pool.try_alloc(page).unwrap(); assert_eq!(buf.capacity(), page); // One byte into the gap routes to the next enabled class. let buf = pool.try_alloc(page + 1).unwrap(); assert_eq!(buf.capacity(), page * 8); // A request whose natural class is disabled routes forward too. let buf = pool.try_alloc(page * 4).unwrap(); assert_eq!(buf.capacity(), page * 8); // Exact fit for the last class. let buf = pool.try_alloc(page * 8).unwrap(); assert_eq!(buf.capacity(), page * 8); // Above the last class is oversized. 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)])); // Exhaust the small class. A page-sized request must report // exhaustion even though the larger class still has capacity. let _small = pool.try_alloc(page).unwrap(); assert_eq!(pool.try_alloc(page).unwrap_err(), PoolError::Exhausted); // The larger class is unaffected. let _large = pool.try_alloc(page * 8).unwrap(); // The untracked fallback is sized from the request rather than a pool // class. The aligned owner may round its usable capacity up by at most // seven bytes. 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() { // Metrics attribute to the enabled class that served the request, so // a request routed through a gap lands on the larger class's label. let page = page_size(); let mut registry = Registry::default(); let pool = BufferPool::new(sparse_config([(page, 1), (page * 8, 1)]), &mut registry); // Allocate through the gap, then exhaust the routed class through it. let _held = pool.try_alloc(page * 2).unwrap(); assert!(pool.try_alloc(page * 2).is_err()); // Request above the largest class records an oversized attempt. assert!(pool.try_alloc(page * 16).is_err()); let encoded = registry.encode(); // Both created and exhausted count against the larger class, and the // disabled exponent of the natural request never appears as a label. 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() { // Requests routed through a gap and requests hitting the class // directly must share the same allocator, TLS cache, and metrics. let page = page_size(); let pool = test_pool(sparse_config([(page, 4), (page * 8, 4)])); // All gap exponents alias the same class as the direct exponent. 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" ); } // A buffer allocated through the gap returns to the aliased class and // is reusable through the direct route. 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() { // Dropping a sparse pool must reclaim globally parked buffers exactly // as the contiguous pool does, draining each unique class once even // though several routing entries alias it. 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(); // Keep a test-owned handle so the class remains inspectable after // the pool is dropped below. let class = pool.inner.classes[class_index].clone(); // Park one buffer allocated through the gap in the global freelist. 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)])); // Two enabled classes span five exponents, Debug must report two. 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)); // Each unique class prefilled exactly its own limit, aliases add none. 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); // Prefilled capacity is immediately allocatable and bounded. 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)); // Allocate a buffer - returns buffer with len=0, capacity >= requested let buf = pool.try_alloc(page).unwrap(); assert!(buf.capacity() >= page); assert_eq!(buf.len(), 0); // Drop returns to pool drop(buf); // Can allocate again 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)); // SAFETY: we immediately initialize all bytes before reading. 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)); // Exhaust pooled capacity for this class. 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)); // Dirty the buffer before returning it to the pool. 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() { // A single-slot pool: if a zero-capacity request claimed a buffer, // the follow-up real allocation would fail with Exhausted. 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)); // Small request gets smallest class let buf1 = pool.try_alloc(page).unwrap(); assert_eq!(buf1.capacity(), page); // Larger request gets appropriate class 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), ); // Should be able to allocate max_per_class buffers immediately let mut bufs = Vec::new(); for _ in 0..5 { bufs.push(pool.try_alloc(page.get()).expect("alloc should succeed")); } // Next allocation should fail 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() { // The storage preset's max_size (8 MB) should be allocatable out of the box. 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); // Alignment can be tuned explicitly as long as the smallest class is // also adjusted. 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(); // Half the class budget is divided across expected threads. 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 ); // Large classes scale past the previous eight-slot cap. 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(); // With only two buffers and eight expected threads, the auto policy's // per-thread share is zero: 2 / (2 * min(8, 2)) == 0. In that case the // pool should disable TLS instead of forcing every thread to retain at // least one buffer. 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); // Exhaust the size class so the only way the main thread can allocate // again is if the worker's returned buffers are globally visible. 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 || { // Return both buffers from another thread. The thread stays alive // after the drops, so any TLS entries it retained would remain // invisible to the main thread until `release_rx` fires. 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"); // Both allocations must succeed while the worker thread is still // alive. Before auto capacity could resolve to zero, one returned // buffer could remain stranded in the worker's TLS cache and this // second allocation would report exhaustion. 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); // When expected parallelism rounds above capacity, the freelist caps // stripes so every stripe can contain at least one slot. 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); // Disabling thread-local caches should not change global striping. 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(); // Fixed capacity should bypass the derived parallelism heuristic. 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); // Disabled thread caching still routes returns through the global // freelist, but should never retain buffers in the current thread. 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() { // Classes: 4, 8, 16 (sum = 28). Budget 280 scales the uniform shape // to exactly 10 buffers per class. 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); // The budget is a strict ceiling: 279 cannot afford the tenth round. let config = base.clone().with_budget_bytes(NZUsize!(279)); assert_eq!(classes_of(&config), vec![(4, 9), (8, 9), (16, 9)]); // The minimum footprint keeps one buffer per class. let config = base.clone().with_budget_bytes(NZUsize!(28)); assert_eq!(classes_of(&config), vec![(4, 1), (8, 1), (16, 1)]); // Scaling preserves a nonuniform shape proportionally: limits (4, 1) // over sizes (4, 16) cost 32 per round, so budget 96 affords a 3x // scale of the whole shape. 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)]); // Scaling can also shrink an existing shape. let shrunk = shaped_base.with_budget_bytes(NZUsize!(20)); assert_eq!(classes_of(&shrunk), vec![(4, 1), (16, 1)]); // Rounding never disables a class, so an uneven budget leaves an // intentionally unused remainder. 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() { // The budget is not a stored policy: later builder calls apply to the // scaled limits and may exceed the former budget. 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); // Rescaling applies to the already scaled limits, not the shape they // were derived from, so repeating a budget can compound the rounding // floors into a slightly different shape. 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() { // One-byte class: any budget beyond u32::MAX buffers must panic // instead of silently capping the limit. 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() { // Two classes with limits near u32::MAX exercise the tightest // count-change breakpoint separation the Q64.64 scale must resolve. // Reduce the budget under miri so the brute-force reference stays // fast in the interpreter. 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()); // Brute-force the optimal proportional vector along the scale axis. let expected = brute_force_budget(&[(1, a), (2, b)], budget as u128); assert_eq!( classes_of(&config) .into_iter() .map(|(_, limit)| limit) .collect::>(), expected ); } /// Reference implementation of proportional budget scaling. /// /// Walks the count vectors produced along the scale axis in breakpoint /// order and returns the last one whose total fits the budget. fn brute_force_budget(shape: &[(usize, u32)], budget: u128) -> Vec { // Collect candidate scales k/c for every class and every count k the // budget could possibly afford, then evaluate the count vector at each. let mut best: Option> = 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 { // counts_i = max(1, floor(c_i * k / c)) let counts: Vec = 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(); // Candidate vectors are componentwise ordered along the scale // axis, so the maximal fitting total identifies a unique vector. 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() { // Debug formatting and config accessor should be consistent. 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() { // Debug formatting for pooled mutable/immutable handles, and empty // freeze should detach without retaining the pool allocation. 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)); // Initially: 0 allocated, 0 available assert_eq!(get_allocated(&pool, page), 0); assert_eq!(get_available(&pool, page), 0); // Allocate, write, and freeze. Empty freeze deliberately detaches from // the pool, so this test keeps a non-empty immutable view alive. 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(); // Still allocated (held by IoBuf) assert_eq!(get_allocated(&pool, page), 1); // Drop the IoBuf - buffer should return to pool 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)); // Refcount behavior: // - clone/slice keep the pooled allocation alive // - empty slice does not keep ownership { 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); } // IoBuf::copy_to_bytes behavior: // - zero-length copy is empty and non-advancing // - partial copy advances while keeping ownership alive // - full drain transfers ownership out of source // - zero-length copy on already-empty source stays detached { 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); // Zero-length copy on empty should not transfer ownership. 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); } // IoBufMut::copy_to_bytes mirrors the immutable ownership semantics. { 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() { // IoBuf -> IoBufMut should recover pooled ownership for unique non-full views. 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"); // Conversion reused pooled storage instead of copying. 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); // Dropping returns the pooled buffer. 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() { // try_into_mut on a uniquely-owned full-view pooled IoBuf should recover // mutable ownership without copying, preserving data and pool tracking. 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); // Unique full view should recycle. 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)); // Unique sliced views can recover mutable ownership without copying. 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); // Shared views still cannot recover mutable ownership. 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![]; // Reduce iterations under miri (atomics are slow) cfg_if::cfg_if! { if #[cfg(miri)] { let iterations = 100; } else { let iterations = 1000; } } // Spawn multiple threads that allocate, freeze, clone, and drop 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(); // Write a byte so freeze produces a live pooled owner // (freezing an empty buffer detaches and releases it, // which would leave the refcount protocol unexercised). buf.put_slice(b"x"); let iobuf = buf.freeze(); // Clone a few times let clones: Vec<_> = (0..5).map(|_| iobuf.clone()).collect(); drop(iobuf); // Drop clones for clone in clones { drop(clone); } } }); handles.push(handle); } // Wait for all threads for handle in handles { handle.join().unwrap(); } // Worker threads may retain free buffers in their own local caches, so // the main thread cannot assert that all of them are visible here. // It should still be able to allocate successfully once the workers finish. let _buf = pool .try_alloc(page) .expect("pool should remain usable after multithreaded test"); } #[test] fn test_cross_thread_buffer_return() { // Allocate on one thread, freeze, send to another thread, drop there let page = page_size(); let pool = test_pool(test_config(page, page, 100)); let (tx, rx) = mpsc::channel(); // Allocate and freeze on main thread 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); // Receive and drop on another thread. Cross-thread returns initialize // the dropping thread's local cache, so the buffers remain local to that // thread instead of bouncing through the global freelist. 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() { // What happens if the pool is dropped while buffers are still in use? // The size class remains alive until the last tracked buffer is dropped. 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 the pool while buffer is still alive drop(pool); // Buffer should still be usable assert_eq!(iobuf.len(), 100); // Dropping the buffer should not panic and should return to the retained size class. drop(iobuf); // No assertion here - we just want to make sure it doesn't panic } #[test] fn test_pool_exhaustion_and_recovery() { // Test pool exhaustion and recovery. let page = page_size(); let pool = test_pool(test_config(page, page, 3)); // Exhaust the pool 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"); // Return one buffer drop(buf1); // Should be able to allocate again let buf4 = pool.try_alloc(page).expect("alloc after return"); assert!(pool.try_alloc(page).is_err(), "pool exhausted again"); // Return all and verify freelist reuse drop(buf2); drop(buf3); drop(buf4); assert_eq!(get_allocated(&pool, page), 0); assert_eq!(get_available(&pool, page), 3); // Allocate again - should reuse from freelist let _buf5 = pool.try_alloc(page).expect("reuse from freelist"); assert_eq!(get_available(&pool, page), 2); } #[test] fn test_try_alloc_errors() { // Test try_alloc error variants. let page = page_size(); let pool = test_pool(test_config(page, page, 2)); // Oversized request let result = pool.try_alloc(page * 10); assert_eq!(result.unwrap_err(), PoolError::Oversized); // Exhaust pool 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); // One created buffer, then exhaustion, then an oversized request. 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() { // try_alloc_zeroed should return the same error variants as try_alloc. let page = page_size(); let pool = test_pool(test_config(page, page, 2)); // Oversized request. let result = pool.try_alloc_zeroed(page * 10); assert_eq!(result.unwrap_err(), PoolError::Oversized); // Exhaust pool, then verify Exhausted error. 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() { // Test fallback allocation when pool is exhausted or oversized. let page = page_size(); let pool = test_pool(test_config(page, page, 2)); // Exhaust the pool let buf1 = pool.try_alloc(page).unwrap(); let buf2 = pool.try_alloc(page).unwrap(); assert!(buf1.is_pooled()); assert!(buf2.is_pooled()); // Fallback via alloc() when exhausted - still aligned, but untracked, // and sized from the requested capacity. 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())); // Fallback via alloc() when oversized - still aligned, but untracked. 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); // Verify pool counters unchanged by fallback allocations assert_eq!(get_allocated(&pool, page), 2); // Drop fallback buffers - should not affect pool counters drop(fallback_exhausted); drop(fallback_oversized); assert_eq!(get_allocated(&pool, page), 2); // Drop tracked buffers - counters should decrease drop(buf1); drop(buf2); assert_eq!(get_allocated(&pool, page), 0); } #[test] fn test_is_pooled() { // IoBufMut from the pool should report is_pooled, while heap-backed // buffers should not. 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()); // Oversized alloc uses untracked fallback allocation. 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(); // Reduce the class limits under miri (atomics are slow) 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)); } } // Storage preset - page aligned 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" ); // Network preset - cache-line 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; // Models the pooled buffer lifecycle across threads: checkout, freeze, // clone, cross-thread final drop, and reuse from the same pool. Whichever // thread drops last must return the buffer to the global freelist with // the refcount sentinel intact so the next checkout works without // reinitialization. The thread cache is disabled so the return path is // the loom-modeled global freelist rather than OS thread-local state, // which loom cannot reset between interleavings. #[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(); // The buffer returned through whichever drop was final. Checkout // must succeed and expose a writable buffer again. let mut again = pool.alloc(64); assert!(again.is_pooled()); again.put_slice(b"reuse"); assert_eq!(again.as_ref(), b"reuse"); }); } // Models the teardown edge: a pooled buffer's final drop (which parks the // buffer and then releases its size-class lease) racing the public pool's // drop (which drains the global freelist and then releases the pool-owned // class reference). Whichever release is last drops the SizeClass and its // freelist. Parking strictly before releasing is what keeps the freelist // alive until the return finishes publishing under its stripe lock. Loom // performs no liveness check when tracked state is dropped, so swapping // that order manifests as a use-after-free of the freed freelist state. // This corrupts Loom's internal object state and is caught by its internal // assertions in practice. The Loom-tracked class strong count still // verifies the release accounting itself. #[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(); }); } // Models a re-checkout racing the final drop of a shared pooled buffer. // The final drop leaves the refcount at the sentinel (the losing handle's // Release decrement lands on 1, or the race-final path re-stores it with // a Relaxed store) before the freelist's Release publication. A // successful concurrent take must observe that sentinel (asserted in // Freelist::claim under loom) before handing the slot to a new mutable // handle. #[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); // The single slot may still be checked out (Exhausted) or already // returned by whichever drop was final. A successful claim must // expose a writable buffer with the sentinel restored. 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(); }); } }