santh-bufpool 0.1.0

Typed buffer recycling with fixed size classes and lock-free checkout/return
Documentation
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(|_| {});
    // Note: Buffer is zeroed in PoolBuffer::drop before this function
    // is called. No need to zero again here.
    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() {
    // Note: try_with can fail if TLS has been destroyed (e.g., during thread panic).
    // In that rare case, buffers may leak. This is acceptable because:
    // 1. Thread panics during buffer drop are extremely rare
    // 2. The OS reclaims all memory on process exit anyway
    // 3. Normal thread exit always has TLS available and drains properly
    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);
        }
    });
}

/// Clear the thread-local cache and return all buffers to their pools.
///
/// This is exposed for tests that need deterministic pool state.
#[doc(hidden)]
pub fn clear_tls_cache() {
    drain_tls_cache();
}