armature-core 0.8.2

//! Thread-Local `BytesMut` Buffer Pool
//!
//! This module provides a high-performance thread-local buffer pool for
//! reusing `BytesMut` buffers across HTTP request/response handling.
//! By avoiding repeated allocations, this can improve throughput by 4-5%.
//!
//! ## Why Thread-Local?
//!
//! - **No synchronization**: Thread-local access is lock-free
//! - **Cache locality**: Buffers stay in CPU cache
//! - **Zero contention**: Each thread has its own pool
//! - **Automatic cleanup**: Buffers released when thread exits
//!
//! ## Usage
//!
//! ```rust,ignore
//! use armature_core::buffer_pool::{acquire_buffer, PooledBuffer};
//!
//! // Acquire a buffer from the pool
//! let mut buf = acquire_buffer(BufferSize::Medium);
//!
//! // Use the buffer
//! buf.extend_from_slice(b"Hello, World!");
//!
//! // Buffer automatically returns to pool when dropped
//! drop(buf);
//! ```
//!
//! ## Buffer Sizes
//!
//! | Size | Capacity | Use Case |
//! |------|----------|----------|
//! | Tiny | 256B | Headers, small strings |
//! | Small | 4KB | JSON responses |
//! | Medium | 16KB | Typical request bodies |
//! | Large | 64KB | File uploads, bulk data |
//! | Huge | 256KB | Streaming, large payloads |
//!
//! ## Performance Impact
//!
//! - **Without pool**: ~200ns per allocation (malloc + memset)
//! - **With pool**: ~5ns per acquisition (stack operation)
//! - **Savings**: ~195ns per request * 100k req/s = significant

use bytes::BytesMut;
use std::cell::RefCell;
use std::ops::{Deref, DerefMut};
use std::sync::atomic::{AtomicU64, Ordering};

// ============================================================================
// Buffer Size Categories
// ============================================================================

/// Pre-defined buffer size categories for optimal pooling
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BufferSize {
    /// 256 bytes - headers, small strings
    Tiny,
    /// 4 KB - JSON responses, small bodies
    Small,
    /// 16 KB - typical request bodies
    Medium,
    /// 64 KB - larger bodies, file chunks
    Large,
    /// 256 KB - streaming, bulk data
    Huge,
    /// Custom size (will be rounded up to nearest power of 2)
    Custom(usize),
}

impl BufferSize {
    /// Get the actual byte capacity for this size
    #[inline]
    pub const fn capacity(self) -> usize {
        match self {
            Self::Tiny => 256,
            Self::Small => 4096,
            Self::Medium => 16384,
            Self::Large => 65536,
            Self::Huge => 262144,
            Self::Custom(n) => n.next_power_of_two(),
        }
    }

    /// Determine the best size category for a given byte count
    #[inline]
    pub const fn for_bytes(bytes: usize) -> Self {
        if bytes <= 256 {
            Self::Tiny
        } else if bytes <= 4096 {
            Self::Small
        } else if bytes <= 16384 {
            Self::Medium
        } else if bytes <= 65536 {
            Self::Large
        } else if bytes <= 262144 {
            Self::Huge
        } else {
            Self::Custom(bytes)
        }
    }

    /// Get the pool index for this size
    #[inline]
    const fn pool_index(self) -> usize {
        match self {
            Self::Tiny => 0,
            Self::Small => 1,
            Self::Medium => 2,
            Self::Large => 3,
            Self::Huge => 4,
            Self::Custom(_) => 5, // Custom sizes share a pool
        }
    }
}

// ============================================================================
// Pool Configuration
// ============================================================================

/// Configuration for the buffer pool
#[derive(Debug, Clone)]
pub struct PoolConfig {
    /// Maximum buffers per size category
    pub max_per_size: usize,
    /// Pre-allocate this many buffers on first use
    pub preallocate: usize,
    /// Maximum age (in acquisitions) before buffer is discarded
    pub max_age: u32,
    /// Enable statistics collection (small overhead)
    pub collect_stats: bool,
}

impl Default for PoolConfig {
    fn default() -> Self {
        Self {
            max_per_size: 64,
            preallocate: 4,
            max_age: 10000,
            collect_stats: true,
        }
    }
}

impl PoolConfig {
    /// High-performance configuration with larger pools
    pub fn high_performance() -> Self {
        Self {
            max_per_size: 128,
            preallocate: 16,
            max_age: 50000,
            collect_stats: false,
        }
    }

    /// Memory-efficient configuration
    pub fn memory_efficient() -> Self {
        Self {
            max_per_size: 16,
            preallocate: 2,
            max_age: 1000,
            collect_stats: true,
        }
    }
}

// ============================================================================
// Pool Statistics
// ============================================================================

/// Global statistics for buffer pool operations
#[derive(Debug, Default)]
pub struct PoolStats {
    /// Buffers acquired from pool (reused)
    hits: AtomicU64,
    /// Buffers allocated fresh (pool miss)
    misses: AtomicU64,
    /// Buffers returned to pool
    returns: AtomicU64,
    /// Buffers discarded (pool full or too old)
    discards: AtomicU64,
    /// Total bytes allocated
    bytes_allocated: AtomicU64,
    /// Current buffers in all pools
    pooled_count: AtomicU64,
}

impl PoolStats {
    /// Create new statistics
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a pool hit
    #[inline]
    pub fn record_hit(&self) {
        self.hits.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a pool miss
    #[inline]
    pub fn record_miss(&self, bytes: usize) {
        self.misses.fetch_add(1, Ordering::Relaxed);
        self.bytes_allocated
            .fetch_add(bytes as u64, Ordering::Relaxed);
    }

    /// Record a return to pool
    #[inline]
    pub fn record_return(&self) {
        self.returns.fetch_add(1, Ordering::Relaxed);
        self.pooled_count.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a discard
    #[inline]
    pub fn record_discard(&self) {
        self.discards.fetch_add(1, Ordering::Relaxed);
    }

    /// Record buffer taken from pool
    #[inline]
    pub fn record_taken(&self) {
        self.pooled_count.fetch_sub(1, Ordering::Relaxed);
    }

    /// Get hit count
    pub fn hits(&self) -> u64 {
        self.hits.load(Ordering::Relaxed)
    }

    /// Get miss count
    pub fn misses(&self) -> u64 {
        self.misses.load(Ordering::Relaxed)
    }

    /// Get return count
    pub fn returns(&self) -> u64 {
        self.returns.load(Ordering::Relaxed)
    }

    /// Get discard count
    pub fn discards(&self) -> u64 {
        self.discards.load(Ordering::Relaxed)
    }

    /// Get total bytes allocated
    pub fn bytes_allocated(&self) -> u64 {
        self.bytes_allocated.load(Ordering::Relaxed)
    }

    /// Get current pooled buffer count
    pub fn pooled_count(&self) -> u64 {
        self.pooled_count.load(Ordering::Relaxed)
    }

    /// Get hit rate as percentage
    pub fn hit_rate(&self) -> f64 {
        let hits = self.hits() as f64;
        let total = hits + self.misses() as f64;
        if total > 0.0 {
            (hits / total) * 100.0
        } else {
            0.0
        }
    }
}

/// Global statistics instance
static POOL_STATS: PoolStats = PoolStats {
    hits: AtomicU64::new(0),
    misses: AtomicU64::new(0),
    returns: AtomicU64::new(0),
    discards: AtomicU64::new(0),
    bytes_allocated: AtomicU64::new(0),
    pooled_count: AtomicU64::new(0),
};

/// Get the global pool statistics
pub fn pool_stats() -> &'static PoolStats {
    &POOL_STATS
}

// ============================================================================
// Thread-Local Buffer Pool
// ============================================================================

/// A single size category pool.
///
/// Each pooled buffer is stored alongside an `age` counter tracking how
/// many times it has been acquired from the pool; buffers whose age
/// reaches `PoolConfig::max_age` are discarded on release instead of
/// being kept for reuse.
struct SizePool {
    buffers: Vec<(BytesMut, u32)>,
    max_size: usize,
    capacity: usize,
}

impl SizePool {
    fn new(capacity: usize, max_size: usize, preallocate: usize) -> Self {
        let preallocate = preallocate.min(max_size);
        let mut buffers = Vec::with_capacity(max_size);
        for _ in 0..preallocate {
            buffers.push((BytesMut::with_capacity(capacity), 0));
        }
        Self {
            buffers,
            max_size,
            capacity,
        }
    }

    #[inline]
    fn acquire(&mut self, collect_stats: bool) -> Option<(BytesMut, u32)> {
        self.buffers.pop().map(|(mut buf, age)| {
            buf.clear();
            if collect_stats {
                POOL_STATS.record_hit();
                POOL_STATS.record_taken();
            }
            (buf, age)
        })
    }

    #[inline]
    fn release(&mut self, mut buf: BytesMut, age: u32, max_age: u32, collect_stats: bool) {
        let age = age.saturating_add(1);
        // Discard buffers that have aged past the configured limit; otherwise
        // keep them if the pool isn't full and the buffer isn't oversized.
        if age < max_age
            && self.buffers.len() < self.max_size
            && buf.capacity() <= self.capacity * 2
        {
            buf.clear();
            self.buffers.push((buf, age));
            if collect_stats {
                POOL_STATS.record_return();
            }
        } else if collect_stats {
            POOL_STATS.record_discard();
        }
    }

    #[allow(dead_code)] // Reserved for future direct allocation path
    fn allocate(&self) -> BytesMut {
        POOL_STATS.record_miss(self.capacity);
        BytesMut::with_capacity(self.capacity)
    }
}

/// Thread-local buffer pool with multiple size categories
struct ThreadLocalPool {
    /// Pools for each size category [Tiny, Small, Medium, Large, Huge, Custom]
    pools: [SizePool; 6],
    config: PoolConfig,
}

impl ThreadLocalPool {
    fn new(config: PoolConfig) -> Self {
        Self {
            pools: [
                SizePool::new(
                    BufferSize::Tiny.capacity(),
                    config.max_per_size,
                    config.preallocate,
                ),
                SizePool::new(
                    BufferSize::Small.capacity(),
                    config.max_per_size,
                    config.preallocate,
                ),
                SizePool::new(
                    BufferSize::Medium.capacity(),
                    config.max_per_size,
                    config.preallocate,
                ),
                SizePool::new(
                    BufferSize::Large.capacity(),
                    config.max_per_size,
                    config.preallocate,
                ),
                SizePool::new(
                    BufferSize::Huge.capacity(),
                    config.max_per_size,
                    config.preallocate,
                ),
                SizePool::new(
                    BufferSize::Custom(1024 * 1024).capacity(),
                    config.max_per_size / 4,
                    0, // Custom-size buffers are never pre-allocated
                ),
            ],
            config,
        }
    }

    #[inline]
    fn acquire(&mut self, size: BufferSize) -> (BytesMut, u32) {
        let idx = size.pool_index();
        let pool = &mut self.pools[idx];
        let collect_stats = self.config.collect_stats;

        pool.acquire(collect_stats).unwrap_or_else(|| {
            let capacity = if matches!(size, BufferSize::Custom(n) if n > 0) {
                size.capacity()
            } else {
                pool.capacity
            };
            if collect_stats {
                POOL_STATS.record_miss(capacity);
            }
            (BytesMut::with_capacity(capacity), 0)
        })
    }

    #[inline]
    fn release(&mut self, buf: BytesMut, age: u32, size: BufferSize) {
        let idx = size.pool_index();
        self.pools[idx].release(buf, age, self.config.max_age, self.config.collect_stats);
    }
}

thread_local! {
    /// Thread-local buffer pool
    static BUFFER_POOL: RefCell<ThreadLocalPool> = RefCell::new(
        ThreadLocalPool::new(PoolConfig::default())
    );
}

// ============================================================================
// Pooled Buffer (RAII Guard)
// ============================================================================

/// A buffer acquired from the pool that returns automatically when dropped.
///
/// This provides RAII semantics for buffer management - the buffer is
/// automatically returned to the pool when it goes out of scope.
///
/// # Example
///
/// ```rust,ignore
/// let mut buf = acquire_buffer(BufferSize::Medium);
/// buf.extend_from_slice(b"data");
/// // buf automatically returns to pool here
/// ```
pub struct PooledBuffer {
    inner: Option<BytesMut>,
    size: BufferSize,
    /// Number of times this buffer has been acquired from the pool.
    /// Used to discard buffers older than `PoolConfig::max_age` on release.
    age: u32,
}

impl PooledBuffer {
    /// Create a new pooled buffer
    fn new(buf: BytesMut, size: BufferSize, age: u32) -> Self {
        Self {
            inner: Some(buf),
            size,
            age,
        }
    }

    /// Take ownership of the inner buffer without returning to pool
    ///
    /// Use this when you need to pass the buffer to something that
    /// will own it (like converting to Bytes).
    #[inline]
    pub fn take(mut self) -> BytesMut {
        self.inner.take().expect("buffer already taken")
    }

    /// Freeze the buffer into immutable Bytes
    #[inline]
    pub fn freeze(mut self) -> bytes::Bytes {
        self.inner.take().expect("buffer already taken").freeze()
    }

    /// Get the buffer size category
    #[inline]
    pub fn size_category(&self) -> BufferSize {
        self.size
    }

    /// Get the buffer capacity
    #[inline]
    pub fn capacity(&self) -> usize {
        self.inner.as_ref().map(|b| b.capacity()).unwrap_or(0)
    }
}

impl Deref for PooledBuffer {
    type Target = BytesMut;

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.inner.as_ref().expect("buffer already taken")
    }
}

impl DerefMut for PooledBuffer {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.inner.as_mut().expect("buffer already taken")
    }
}

impl Drop for PooledBuffer {
    fn drop(&mut self) {
        if let Some(buf) = self.inner.take() {
            BUFFER_POOL.with(|pool| {
                pool.borrow_mut().release(buf, self.age, self.size);
            });
        }
    }
}

impl std::fmt::Debug for PooledBuffer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PooledBuffer")
            .field("size", &self.size)
            .field("len", &self.inner.as_ref().map(|b| b.len()))
            .field("capacity", &self.inner.as_ref().map(|b| b.capacity()))
            .finish()
    }
}

// ============================================================================
// Public API
// ============================================================================

/// Acquire a buffer from the thread-local pool.
///
/// If the pool is empty, a new buffer is allocated.
/// The buffer is automatically returned to the pool when dropped.
///
/// # Example
///
/// ```rust,ignore
/// use armature_core::buffer_pool::{acquire_buffer, BufferSize};
///
/// let mut buf = acquire_buffer(BufferSize::Medium);
/// buf.extend_from_slice(b"Hello");
/// // buf returns to pool when dropped
/// ```
#[inline]
pub fn acquire_buffer(size: BufferSize) -> PooledBuffer {
    BUFFER_POOL.with(|pool| {
        let (buf, age) = pool.borrow_mut().acquire(size);
        PooledBuffer::new(buf, size, age)
    })
}

/// Acquire a buffer sized for the given byte count.
///
/// Automatically selects the appropriate size category.
///
/// # Example
///
/// ```rust,ignore
/// let mut buf = acquire_buffer_for_bytes(1024);
/// // Gets a Small (4KB) buffer
/// ```
#[inline]
pub fn acquire_buffer_for_bytes(bytes: usize) -> PooledBuffer {
    acquire_buffer(BufferSize::for_bytes(bytes))
}

/// Acquire a buffer and immediately write data to it.
///
/// # Example
///
/// ```rust,ignore
/// let buf = acquire_with_data(b"Hello, World!");
/// assert_eq!(&buf[..], b"Hello, World!");
/// ```
#[inline]
pub fn acquire_with_data(data: &[u8]) -> PooledBuffer {
    let mut buf = acquire_buffer(BufferSize::for_bytes(data.len()));
    buf.extend_from_slice(data);
    buf
}

/// Acquire a buffer for JSON serialization.
///
/// Pre-sized for typical JSON responses.
#[inline]
pub fn acquire_json_buffer() -> PooledBuffer {
    acquire_buffer(BufferSize::Small)
}

/// Acquire a buffer for request body reading.
///
/// Pre-sized for typical request bodies.
#[inline]
pub fn acquire_body_buffer() -> PooledBuffer {
    acquire_buffer(BufferSize::Medium)
}

/// Acquire a buffer for response building.
///
/// Pre-sized for typical HTTP responses.
#[inline]
pub fn acquire_response_buffer() -> PooledBuffer {
    acquire_buffer(BufferSize::Small)
}

/// Acquire a buffer for streaming data.
///
/// Larger buffer for bulk operations.
#[inline]
pub fn acquire_streaming_buffer() -> PooledBuffer {
    acquire_buffer(BufferSize::Large)
}

// ============================================================================
// Scoped Buffer Usage
// ============================================================================

/// Execute a closure with a pooled buffer.
///
/// The buffer is automatically returned to the pool after the closure.
///
/// # Example
///
/// ```rust,ignore
/// let result = with_buffer(BufferSize::Small, |buf| {
///     buf.extend_from_slice(b"data");
///     buf.len()
/// });
/// assert_eq!(result, 4);
/// ```
#[inline]
pub fn with_buffer<F, R>(size: BufferSize, f: F) -> R
where
    F: FnOnce(&mut BytesMut) -> R,
{
    let mut buf = acquire_buffer(size);
    f(&mut buf)
}

/// Execute a closure that produces Bytes from a pooled buffer.
///
/// Optimized path for JSON serialization and similar operations.
///
/// # Example
///
/// ```rust,ignore
/// let bytes = buffer_to_bytes(BufferSize::Small, |buf| {
///     serde_json::to_writer(buf.writer(), &data)?;
///     Ok(())
/// })?;
/// ```
#[inline]
pub fn buffer_to_bytes<F, E>(size: BufferSize, f: F) -> Result<bytes::Bytes, E>
where
    F: FnOnce(&mut BytesMut) -> Result<(), E>,
{
    let mut buf = acquire_buffer(size);
    f(&mut buf)?;
    Ok(buf.freeze())
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_buffer_sizes() {
        assert_eq!(BufferSize::Tiny.capacity(), 256);
        assert_eq!(BufferSize::Small.capacity(), 4096);
        assert_eq!(BufferSize::Medium.capacity(), 16384);
        assert_eq!(BufferSize::Large.capacity(), 65536);
        assert_eq!(BufferSize::Huge.capacity(), 262144);
        assert_eq!(BufferSize::Custom(1000).capacity(), 1024); // Rounded up
    }

    #[test]
    fn test_size_for_bytes() {
        assert_eq!(BufferSize::for_bytes(100), BufferSize::Tiny);
        assert_eq!(BufferSize::for_bytes(1000), BufferSize::Small);
        assert_eq!(BufferSize::for_bytes(10000), BufferSize::Medium);
        assert_eq!(BufferSize::for_bytes(50000), BufferSize::Large);
        assert_eq!(BufferSize::for_bytes(200000), BufferSize::Huge);
    }

    #[test]
    fn test_acquire_and_release() {
        // Acquire a buffer
        let mut buf = acquire_buffer(BufferSize::Small);
        assert!(buf.capacity() >= BufferSize::Small.capacity());

        // Write to it
        buf.extend_from_slice(b"Hello, World!");
        assert_eq!(&buf[..], b"Hello, World!");

        // Drop returns to pool
        drop(buf);

        // Next acquire should hit pool
        let buf2 = acquire_buffer(BufferSize::Small);
        assert!(buf2.is_empty()); // Should be cleared
    }

    #[test]
    fn test_acquire_with_data() {
        let buf = acquire_with_data(b"test data");
        assert_eq!(&buf[..], b"test data");
    }

    #[test]
    fn test_pooled_buffer_take() {
        let buf = acquire_buffer(BufferSize::Tiny);
        let inner = buf.take();
        assert_eq!(inner.capacity(), BufferSize::Tiny.capacity());
    }

    #[test]
    fn test_pooled_buffer_freeze() {
        let mut buf = acquire_buffer(BufferSize::Tiny);
        buf.extend_from_slice(b"frozen");
        let bytes = buf.freeze();
        assert_eq!(&bytes[..], b"frozen");
    }

    #[test]
    fn test_with_buffer() {
        let len = with_buffer(BufferSize::Small, |buf| {
            buf.extend_from_slice(b"test");
            buf.len()
        });
        assert_eq!(len, 4);
    }

    #[test]
    fn test_buffer_to_bytes() {
        let bytes: Result<bytes::Bytes, std::convert::Infallible> =
            buffer_to_bytes(BufferSize::Tiny, |buf| {
                buf.extend_from_slice(b"bytes");
                Ok(())
            });
        assert_eq!(&bytes.unwrap()[..], b"bytes");
    }

    #[test]
    fn test_pool_stats() {
        // Reset-like behavior (stats are global, so just check they work)
        let stats = pool_stats();
        let initial_hits = stats.hits();

        // Acquire and release to generate stats
        let _buf = acquire_buffer(BufferSize::Tiny);
        drop(_buf);
        let _buf = acquire_buffer(BufferSize::Tiny);

        // Should have at least one hit or miss
        assert!(stats.hits() + stats.misses() > initial_hits);
    }

    #[test]
    fn test_specialized_acquires() {
        let json_buf = acquire_json_buffer();
        assert!(json_buf.capacity() >= BufferSize::Small.capacity());

        let body_buf = acquire_body_buffer();
        assert!(body_buf.capacity() >= BufferSize::Medium.capacity());

        let response_buf = acquire_response_buffer();
        assert!(response_buf.capacity() >= BufferSize::Small.capacity());

        let streaming_buf = acquire_streaming_buffer();
        assert!(streaming_buf.capacity() >= BufferSize::Large.capacity());
    }

    #[test]
    fn test_pool_config() {
        let config = PoolConfig::high_performance();
        assert_eq!(config.max_per_size, 128);
        assert!(!config.collect_stats);

        let config = PoolConfig::memory_efficient();
        assert_eq!(config.max_per_size, 16);
        assert!(config.collect_stats);
    }
}