use std::ptr::NonNull;
use std::sync::Arc;
use crate::size_class::{BufferAllocation, SizeClassPool};
type TlsCacheEntry = (usize, NonNull<u8>, Arc<SizeClassPool>);
type TlsCache = std::cell::RefCell<Vec<TlsCacheEntry>>;
thread_local! {
static TLS_CACHE: TlsCache = const { std::cell::RefCell::new(Vec::new()) };
static TLS_CACHE_GUARD: TlsCacheGuard = const { TlsCacheGuard };
}
struct TlsCacheGuard;
pub(crate) fn take_tls_buffer(
min_bytes: usize,
) -> Option<(usize, NonNull<u8>, Arc<SizeClassPool>)> {
TLS_CACHE_GUARD.with(|_| {});
TLS_CACHE.with(|cache| {
let mut cache = cache.borrow_mut();
let index = cache
.iter()
.position(|(capacity, _, _)| *capacity >= min_bytes)?;
Some(cache.swap_remove(index))
})
}
pub(crate) fn store_tls_buffer(
capacity: usize,
ptr: NonNull<u8>,
owner: Arc<SizeClassPool>,
) -> bool {
TLS_CACHE_GUARD.with(|_| {});
TLS_CACHE.with(|cache| {
let mut cache = cache.borrow_mut();
if cache.len() >= 4 {
return false;
}
cache.push((capacity, ptr, owner));
true
})
}
impl Drop for TlsCacheGuard {
fn drop(&mut self) {
drain_tls_cache();
}
}
fn drain_tls_cache() {
let _ = TLS_CACHE.try_with(|cache| {
let mut cache = cache.borrow_mut();
while let Some((_, ptr, owner)) = cache.pop() {
let allocation = BufferAllocation { ptr };
owner.recycle_or_free(allocation);
}
});
}
#[doc(hidden)]
pub fn clear_tls_cache() {
drain_tls_cache();
}