use std::collections::BTreeMap;
use std::env;
use std::ffi::OsString;
use std::fmt;
use std::mem::{size_of, ManuallyDrop, MaybeUninit};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
use tenferro_tensor::HostBufferRecycler;
use num_complex::{Complex32, Complex64};
use crate::CacheStats;
#[derive(Debug)]
pub(crate) enum UninitCheckoutToken {
Fresh { actual_capacity: usize },
Reused { actual_capacity: usize },
}
pub const BUFFER_POOL_MAX_RETAINED_BYTES_ENV: &str = "TENFERRO_BUFFER_POOL_MAX_RETAINED_BYTES";
pub const DEFAULT_MAX_RETAINED_CAPACITY_BYTES: usize = 100 * 1024 * 1024;
static DEFAULT_MAX_RETAINED_CAPACITY_FROM_ENV: OnceLock<usize> = OnceLock::new();
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct BufferPoolStats {
pub buffers: usize,
pub capacity_bytes: usize,
}
pub struct BufferPool {
state: Arc<SharedPool>,
}
#[derive(Debug)]
struct SharedPool(Mutex<PoolState>);
#[derive(Debug)]
struct PoolState {
f64_pool: BTreeMap<usize, Vec<Vec<f64>>>,
f32_pool: BTreeMap<usize, Vec<Vec<f32>>>,
i32_pool: BTreeMap<usize, Vec<Vec<i32>>>,
i64_pool: BTreeMap<usize, Vec<Vec<i64>>>,
bool_pool: BTreeMap<usize, Vec<Vec<bool>>>,
c64_pool: BTreeMap<usize, Vec<Vec<Complex64>>>,
c32_pool: BTreeMap<usize, Vec<Vec<Complex32>>>,
retained_capacity_bytes: usize,
max_retained_capacity_bytes: usize,
}
impl fmt::Debug for BufferPool {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BufferPool")
.field("stats", &self.stats())
.field(
"max_retained_capacity_bytes",
&self.max_retained_capacity_bytes(),
)
.finish_non_exhaustive()
}
}
pub trait PoolScalar:
Copy + Sized + Send + Sync + tenferro_tensor::TensorScalar + private::Sealed
{
fn pool_zero() -> Self;
fn pool_acquire_zeroed(pool: &mut BufferPool, len: usize) -> Vec<Self>;
fn pool_release(pool: &mut BufferPool, buf: Vec<Self>);
}
pub(crate) mod private {
use std::mem::MaybeUninit;
#[allow(private_interfaces)]
pub trait Sealed {
fn pool_acquire_uninit_tracked(
pool: &mut super::BufferPool,
len: usize,
) -> crate::Result<(Vec<MaybeUninit<Self>>, super::UninitCheckoutToken)>
where
Self: Sized;
fn pool_finish_recycled(
pool: &mut super::BufferPool,
checkout: super::UninitCheckoutToken,
) -> std::sync::Weak<dyn tenferro_tensor::HostBufferRecycler<Self>>
where
Self: Sized;
}
}
const EMPTIED_BIN_SCAN_BUDGET: usize = 16;
fn take_best_fit<T>(pool: &mut BTreeMap<usize, Vec<Vec<T>>>, len: usize) -> Option<Vec<T>> {
let mut cursor = len;
let mut emptied = 0usize;
let mut stale: Vec<usize> = Vec::new();
let serving = loop {
match pool
.range(cursor..)
.next()
.map(|(&key, vecs)| (key, vecs.is_empty()))
{
None => break None,
Some((key, false)) => break Some(key),
Some((key, true)) => {
emptied += 1;
if emptied > EMPTIED_BIN_SCAN_BUDGET {
stale.push(key);
}
cursor = key.saturating_add(1);
}
}
};
for key in stale {
pool.remove(&key);
}
let key = serving?;
let vecs = pool.get_mut(&key)?;
vecs.pop()
}
fn pool_len<T>(pool: &BTreeMap<usize, Vec<Vec<T>>>) -> usize {
pool.values().map(Vec::len).sum()
}
#[cfg(test)]
fn empty_bins<T>(pool: &BTreeMap<usize, Vec<Vec<T>>>) -> usize {
pool.values().filter(|vecs| vecs.is_empty()).count()
}
fn evict_one_from_pool<T>(pool: &mut BTreeMap<usize, Vec<Vec<T>>>) -> Option<usize> {
let key = *pool.iter().find(|(_, vecs)| !vecs.is_empty())?.0;
let vecs = pool.get_mut(&key)?;
let _ = vecs.pop()?;
if vecs.is_empty() {
pool.remove(&key);
}
Some(key.saturating_mul(size_of::<T>()))
}
#[derive(Clone, Copy)]
enum TypedPoolKind {
F64,
F32,
I32,
I64,
Bool,
C64,
C32,
}
fn smallest_pool_candidate<T>(
pool: &BTreeMap<usize, Vec<Vec<T>>>,
kind: TypedPoolKind,
) -> Option<(usize, TypedPoolKind)> {
pool.iter()
.find(|(_, vecs)| !vecs.is_empty())
.map(|(&capacity, _)| (capacity.saturating_mul(size_of::<T>()), kind))
}
macro_rules! impl_pool_scalar {
($ty:ty, $field:ident, $zero:expr) => {
#[allow(private_interfaces)]
impl private::Sealed for $ty {
fn pool_acquire_uninit_tracked(
pool: &mut BufferPool,
len: usize,
) -> crate::Result<(Vec<MaybeUninit<Self>>, UninitCheckoutToken)> {
let mut state = lock_pool(&pool.state);
let pool = &mut *state;
match take_best_fit(&mut pool.$field, len) {
Some(buf) => {
let cap = buf.capacity();
let bytes = cap.checked_mul(size_of::<Self>()).ok_or_else(|| {
crate::Error::invalid_argument(
"pooled_uninit_output",
"length",
"pool capacity byte length overflow",
)
})?;
pool.retained_capacity_bytes -= bytes;
let mut buf = ManuallyDrop::new(buf);
let buf = unsafe { Vec::from_raw_parts(buf.as_mut_ptr().cast(), len, cap) };
Ok((
buf,
UninitCheckoutToken::Reused {
actual_capacity: cap,
},
))
}
None => {
let mut buf = Vec::new();
buf.try_reserve_exact(len).map_err(|err| {
crate::Error::backend_source("pooled_uninit_output", err)
})?;
unsafe { buf.set_len(len) };
let actual_capacity = buf.capacity();
Ok((buf, UninitCheckoutToken::Fresh { actual_capacity }))
}
}
}
fn pool_finish_recycled(
pool: &mut BufferPool,
checkout: UninitCheckoutToken,
) -> Weak<dyn HostBufferRecycler<Self>> {
let _ = checkout;
let owner: Arc<dyn HostBufferRecycler<Self>> = pool.state.clone();
Arc::downgrade(&owner)
}
}
impl PoolScalar for $ty {
fn pool_zero() -> Self {
$zero
}
fn pool_acquire_zeroed(pool: &mut BufferPool, len: usize) -> Vec<Self> {
let mut state = lock_pool(&pool.state);
let pool = &mut *state;
match take_best_fit(&mut pool.$field, len) {
Some(mut buf) => {
pool.retained_capacity_bytes = pool
.retained_capacity_bytes
.saturating_sub(buf.capacity().saturating_mul(size_of::<Self>()));
buf.resize(len, Self::pool_zero());
buf.fill(Self::pool_zero());
buf
}
None => vec![Self::pool_zero(); len],
}
}
fn pool_release(pool: &mut BufferPool, buf: Vec<Self>) {
pool.state.recycle(buf);
}
}
impl HostBufferRecycler<$ty> for SharedPool {
fn recycle(&self, buf: Vec<$ty>) {
let mut state = lock_pool(self);
let pool = &mut *state;
let cap = buf.capacity();
if cap > 0 {
pool.retained_capacity_bytes = pool
.retained_capacity_bytes
.saturating_add(cap.saturating_mul(size_of::<$ty>()));
pool.$field.entry(cap).or_default().push(buf);
pool.enforce_retention_limit();
}
}
}
};
}
impl_pool_scalar!(f64, f64_pool, 0.0);
impl_pool_scalar!(f32, f32_pool, 0.0);
impl_pool_scalar!(i32, i32_pool, 0);
impl_pool_scalar!(i64, i64_pool, 0);
impl_pool_scalar!(bool, bool_pool, false);
impl_pool_scalar!(Complex64, c64_pool, Complex64::new(0.0, 0.0));
impl_pool_scalar!(Complex32, c32_pool, Complex32::new(0.0, 0.0));
impl BufferPool {
pub(crate) fn checkout_handle(&self) -> Self {
Self {
state: Arc::clone(&self.state),
}
}
#[cfg(test)]
pub(crate) fn bin_count(&self) -> usize {
let state = lock_pool(&self.state);
state.f64_pool.len()
+ state.f32_pool.len()
+ state.i32_pool.len()
+ state.i64_pool.len()
+ state.bool_pool.len()
+ state.c64_pool.len()
+ state.c32_pool.len()
}
#[cfg(test)]
pub(crate) fn empty_bin_count(&self) -> usize {
let state = lock_pool(&self.state);
empty_bins(&state.f64_pool)
+ empty_bins(&state.f32_pool)
+ empty_bins(&state.i32_pool)
+ empty_bins(&state.i64_pool)
+ empty_bins(&state.bool_pool)
+ empty_bins(&state.c64_pool)
+ empty_bins(&state.c32_pool)
}
pub fn new() -> Self {
Self::with_max_retained_capacity_bytes(default_max_retained_capacity_bytes())
}
pub fn with_max_retained_capacity_bytes(max_retained_capacity_bytes: usize) -> Self {
Self {
state: Arc::new(SharedPool(Mutex::new(PoolState {
f64_pool: BTreeMap::new(),
f32_pool: BTreeMap::new(),
i32_pool: BTreeMap::new(),
i64_pool: BTreeMap::new(),
bool_pool: BTreeMap::new(),
c64_pool: BTreeMap::new(),
c32_pool: BTreeMap::new(),
retained_capacity_bytes: 0,
max_retained_capacity_bytes,
}))),
}
}
pub fn unbounded() -> Self {
Self::with_max_retained_capacity_bytes(usize::MAX)
}
pub fn max_retained_capacity_bytes(&self) -> usize {
lock_pool(&self.state).max_retained_capacity_bytes
}
pub fn set_max_retained_capacity_bytes(&mut self, max_retained_capacity_bytes: usize) {
let mut state = lock_pool(&self.state);
state.max_retained_capacity_bytes = max_retained_capacity_bytes;
state.enforce_retention_limit();
}
pub fn len(&self) -> usize {
self.stats().buffers
}
pub fn retained_capacity_bytes(&self) -> usize {
self.stats().capacity_bytes
}
pub fn stats(&self) -> BufferPoolStats {
let state = lock_pool(&self.state);
BufferPoolStats {
buffers: pool_len(&state.f64_pool)
+ pool_len(&state.f32_pool)
+ pool_len(&state.i32_pool)
+ pool_len(&state.i64_pool)
+ pool_len(&state.bool_pool)
+ pool_len(&state.c64_pool)
+ pool_len(&state.c32_pool),
capacity_bytes: state.retained_capacity_bytes,
}
}
pub fn cache_stats(&self) -> CacheStats {
let stats = self.stats();
CacheStats {
entries: stats.buffers,
retained_bytes: stats.capacity_bytes,
hits: 0,
misses: 0,
evictions: 0,
clears: 0,
}
}
pub fn acquire_with_capacity<T: PoolScalar>(&mut self, cap: usize) -> Vec<T> {
if cap == 0 {
return Vec::new();
}
let (data, _checkout) = <T as private::Sealed>::pool_acquire_uninit_tracked(self, cap)
.expect("validated typed pool capacity must be acquirable");
let mut data = ManuallyDrop::new(data);
let ptr = data.as_mut_ptr().cast::<T>();
let capacity = data.capacity();
unsafe { Vec::from_raw_parts(ptr, 0, capacity) }
}
pub fn acquire_zeroed<T: PoolScalar>(&mut self, len: usize) -> Vec<T> {
T::pool_acquire_zeroed(self, len)
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn clear(&mut self) {
let mut state = lock_pool(&self.state);
state.clear();
}
}
fn lock_pool(state: &SharedPool) -> MutexGuard<'_, PoolState> {
state
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
impl PoolState {
fn clear(&mut self) {
self.f64_pool.clear();
self.f32_pool.clear();
self.i32_pool.clear();
self.i64_pool.clear();
self.bool_pool.clear();
self.c64_pool.clear();
self.c32_pool.clear();
self.retained_capacity_bytes = 0;
}
fn enforce_retention_limit(&mut self) {
while self.retained_capacity_bytes > self.max_retained_capacity_bytes {
let Some(evicted_bytes) = self.evict_smallest_retained_buffer() else {
self.retained_capacity_bytes = 0;
return;
};
if evicted_bytes == 0 {
if self.retained_capacity_bytes == 0 {
self.retained_capacity_bytes = 0;
return;
}
continue;
}
self.retained_capacity_bytes =
self.retained_capacity_bytes.saturating_sub(evicted_bytes);
}
}
fn evict_smallest_retained_buffer(&mut self) -> Option<usize> {
let candidates = [
smallest_pool_candidate(&self.f64_pool, TypedPoolKind::F64),
smallest_pool_candidate(&self.f32_pool, TypedPoolKind::F32),
smallest_pool_candidate(&self.i32_pool, TypedPoolKind::I32),
smallest_pool_candidate(&self.i64_pool, TypedPoolKind::I64),
smallest_pool_candidate(&self.bool_pool, TypedPoolKind::Bool),
smallest_pool_candidate(&self.c64_pool, TypedPoolKind::C64),
smallest_pool_candidate(&self.c32_pool, TypedPoolKind::C32),
];
let (_, kind) = candidates
.into_iter()
.flatten()
.min_by_key(|(bytes, _)| *bytes)?;
match kind {
TypedPoolKind::F64 => evict_one_from_pool(&mut self.f64_pool),
TypedPoolKind::F32 => evict_one_from_pool(&mut self.f32_pool),
TypedPoolKind::I32 => evict_one_from_pool(&mut self.i32_pool),
TypedPoolKind::I64 => evict_one_from_pool(&mut self.i64_pool),
TypedPoolKind::Bool => evict_one_from_pool(&mut self.bool_pool),
TypedPoolKind::C64 => evict_one_from_pool(&mut self.c64_pool),
TypedPoolKind::C32 => evict_one_from_pool(&mut self.c32_pool),
}
}
}
fn default_max_retained_capacity_bytes() -> usize {
*DEFAULT_MAX_RETAINED_CAPACITY_FROM_ENV.get_or_init(|| {
parse_default_max_retained_capacity_bytes(env::var_os(BUFFER_POOL_MAX_RETAINED_BYTES_ENV))
})
}
fn parse_default_max_retained_capacity_bytes(value: Option<OsString>) -> usize {
value
.and_then(|value| value.into_string().ok())
.and_then(|value| value.parse().ok())
.unwrap_or(DEFAULT_MAX_RETAINED_CAPACITY_BYTES)
}
impl Default for BufferPool {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests;