fluxmq 0.1.0

#![allow(dead_code)]
/// Protocol-specific arena allocator for high-frequency request/response processing
///
/// This module provides specialized arena memory allocation for Kafka protocol parsing
/// and serialization operations. Unlike storage-focused allocations, this targets
/// the rapid allocation/deallocation patterns in protocol processing.
///
/// Performance targets:
/// - 80-90% reduction in allocation overhead
/// - Zero-cost cleanup (arena drop)  
/// - Cache-friendly sequential allocation
use std::alloc::{alloc, dealloc, Layout};
use std::ptr::NonNull;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use tracing::debug;

/// Protocol processing arena optimized for request/response cycles
#[repr(C, align(64))] // Cache-line aligned for optimal performance
pub struct ProtocolArena {
    /// Raw memory region for the arena
    memory: NonNull<u8>,
    /// Total size of the arena
    size: usize,
    /// Current allocation position (atomic for thread safety)
    position: AtomicUsize,
    /// Arena identifier for debugging and metrics
    arena_id: u32,
    /// Number of allocations made from this arena
    allocation_count: AtomicUsize,
    /// Bytes allocated from this arena
    bytes_allocated: AtomicUsize,
}

impl ProtocolArena {
    /// Create a new protocol arena with specified size
    /// Typical size: 64KB for most Kafka requests, 256KB for batch operations
    pub fn new(size: usize, arena_id: u32) -> std::result::Result<Self, String> {
        if size == 0 || size > 16 * 1024 * 1024 {
            // Max 16MB per arena
            return Err(format!("Invalid arena size: {} bytes", size));
        }

        // Allocate aligned memory for the arena
        let layout = Layout::from_size_align(size, 64)
            .map_err(|e| format!("Failed to create layout: {}", e))?;

        let memory = unsafe { alloc(layout) };
        if memory.is_null() {
            return Err("Failed to allocate arena memory".to_string());
        }

        let arena = Self {
            memory: unsafe { NonNull::new_unchecked(memory) },
            size,
            position: AtomicUsize::new(0),
            arena_id,
            allocation_count: AtomicUsize::new(0),
            bytes_allocated: AtomicUsize::new(0),
        };

        debug!(
            "🏟️ PROTOCOL ARENA: Created arena {} with {}KB capacity",
            arena_id,
            size / 1024
        );

        Ok(arena)
    }

    /// Fast pointer-bump allocation with alignment
    pub fn allocate(&self, size: usize, align: usize) -> Option<NonNull<u8>> {
        if size == 0 {
            return None;
        }

        let mut current_pos = self.position.load(Ordering::Relaxed);

        loop {
            // Calculate aligned position
            let aligned_pos = (current_pos + align - 1) & !(align - 1);
            let new_pos = aligned_pos + size;

            // Check if allocation fits
            if new_pos > self.size {
                debug!(
                    "⚠️ PROTOCOL ARENA: Arena {} exhausted ({} / {} bytes used)",
                    self.arena_id, current_pos, self.size
                );
                return None;
            }

            // Try to claim this space atomically
            match self.position.compare_exchange_weak(
                current_pos,
                new_pos,
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => {
                    // Successfully allocated
                    self.allocation_count.fetch_add(1, Ordering::Relaxed);
                    self.bytes_allocated.fetch_add(size, Ordering::Relaxed);

                    let ptr = unsafe { self.memory.as_ptr().add(aligned_pos) };
                    return Some(unsafe { NonNull::new_unchecked(ptr) });
                }
                Err(actual) => {
                    // Another thread updated position, retry with new value
                    current_pos = actual;
                }
            }
        }
    }

    /// Allocate space for Vec with specific capacity
    pub fn allocate_vec<T>(&self, capacity: usize) -> Option<Vec<T>> {
        if capacity == 0 {
            return Some(Vec::new());
        }

        let size = capacity * std::mem::size_of::<T>();
        let align = std::mem::align_of::<T>();

        self.allocate(size, align).map(|ptr| {
            // Create Vec with arena-allocated memory
            // Safety: We've allocated sufficient aligned memory
            unsafe { Vec::from_raw_parts(ptr.as_ptr() as *mut T, 0, capacity) }
        })
    }

    /// Allocate String with specific capacity  
    pub fn allocate_string(&self, capacity: usize) -> Option<String> {
        if capacity == 0 {
            return Some(String::new());
        }

        self.allocate(capacity, 1).map(|ptr| {
            // Create String with arena-allocated memory
            // Safety: We've allocated sufficient memory for the string
            unsafe { String::from_raw_parts(ptr.as_ptr(), 0, capacity) }
        })
    }

    /// Get current arena utilization statistics
    pub fn get_stats(&self) -> ArenaStats {
        let used = self.position.load(Ordering::Relaxed);
        let allocations = self.allocation_count.load(Ordering::Relaxed);
        let allocated_bytes = self.bytes_allocated.load(Ordering::Relaxed);

        ArenaStats {
            arena_id: self.arena_id,
            total_size: self.size,
            used_bytes: used,
            allocation_count: allocations,
            allocated_bytes,
            utilization: (used as f64 / self.size as f64) * 100.0,
        }
    }

    /// Check if arena has sufficient space for allocation
    pub fn can_allocate(&self, size: usize, align: usize) -> bool {
        let current_pos = self.position.load(Ordering::Relaxed);
        let aligned_pos = (current_pos + align - 1) & !(align - 1);
        aligned_pos + size <= self.size
    }

    /// Reset arena for reuse (clear all allocations)
    pub fn reset(&self) {
        self.position.store(0, Ordering::Relaxed);
        self.allocation_count.store(0, Ordering::Relaxed);
        self.bytes_allocated.store(0, Ordering::Relaxed);

        debug!(
            "🔄 PROTOCOL ARENA: Reset arena {} ({}KB ready for reuse)",
            self.arena_id,
            self.size / 1024
        );
    }
}

impl Drop for ProtocolArena {
    fn drop(&mut self) {
        let stats = self.get_stats();
        debug!(
            "🏁 PROTOCOL ARENA: Dropping arena {} - {} allocations, {:.1}% utilization",
            self.arena_id, stats.allocation_count, stats.utilization
        );

        // Deallocate arena memory
        let layout = Layout::from_size_align(self.size, 64).unwrap();
        unsafe {
            dealloc(self.memory.as_ptr(), layout);
        }
    }
}

// Arena is safe to send between threads
unsafe impl Send for ProtocolArena {}
unsafe impl Sync for ProtocolArena {}

/// Arena utilization statistics
#[derive(Debug, Clone)]
pub struct ArenaStats {
    pub arena_id: u32,
    pub total_size: usize,
    pub used_bytes: usize,
    pub allocation_count: usize,
    pub allocated_bytes: usize,
    pub utilization: f64,
}

impl ArenaStats {
    pub fn is_efficient(&self) -> bool {
        self.utilization > 75.0 && self.allocation_count > 5
    }

    pub fn report(&self) -> String {
        format!(
            "Arena {}: {:.1}% used ({} / {}KB), {} allocations",
            self.arena_id,
            self.utilization,
            self.used_bytes / 1024,
            self.total_size / 1024,
            self.allocation_count
        )
    }
}

/// Pool of protocol arenas for different request types
pub struct ProtocolArenaPool {
    /// Small arenas for simple requests (API versions, metadata)
    small_arenas: std::collections::VecDeque<Arc<ProtocolArena>>,
    /// Medium arenas for typical produce/fetch requests  
    medium_arenas: std::collections::VecDeque<Arc<ProtocolArena>>,
    /// Large arenas for batch operations
    large_arenas: std::collections::VecDeque<Arc<ProtocolArena>>,
    /// Next arena ID for debugging
    next_arena_id: AtomicUsize,
}

impl ProtocolArenaPool {
    pub fn new() -> Self {
        Self {
            small_arenas: std::collections::VecDeque::new(),
            medium_arenas: std::collections::VecDeque::new(),
            large_arenas: std::collections::VecDeque::new(),
            next_arena_id: AtomicUsize::new(1),
        }
    }

    /// Get arena suitable for request size
    pub fn get_arena(
        &mut self,
        estimated_size: usize,
    ) -> std::result::Result<Arc<ProtocolArena>, String> {
        let arena_size = match estimated_size {
            0..=8192 => 16384,      // 16KB for small requests
            8193..=65536 => 131072, // 128KB for medium requests
            _ => 524288,            // 512KB for large requests
        };

        let arena_pool = match arena_size {
            16384 => &mut self.small_arenas,
            131072 => &mut self.medium_arenas,
            _ => &mut self.large_arenas,
        };

        // Try to reuse existing arena
        if let Some(arena) = arena_pool.pop_front() {
            if arena.can_allocate(estimated_size, 8) {
                return Ok(arena);
            }
            // Arena too full, reset and try again
            arena.reset();
            if arena.can_allocate(estimated_size, 8) {
                return Ok(arena);
            }
        }

        // Create new arena
        let arena_id = self.next_arena_id.fetch_add(1, Ordering::Relaxed) as u32;
        let arena = Arc::new(ProtocolArena::new(arena_size, arena_id)?);

        debug!(
            "🆕 PROTOCOL ARENA: Created new {}KB arena {} for request size {}",
            arena_size / 1024,
            arena_id,
            estimated_size
        );

        Ok(arena)
    }

    /// Return arena to pool for reuse
    pub fn return_arena(&mut self, arena: Arc<ProtocolArena>) {
        let stats = arena.get_stats();

        if Arc::strong_count(&arena) == 1 && stats.utilization < 90.0 {
            // Arena can be reused
            let arena_pool = match stats.total_size {
                16384 => &mut self.small_arenas,
                131072 => &mut self.medium_arenas,
                _ => &mut self.large_arenas,
            };

            arena_pool.push_back(arena);
        }
        // Otherwise arena will be dropped when Arc count reaches 0
    }

    /// Get pool statistics
    pub fn get_pool_stats(&self) -> PoolStats {
        PoolStats {
            small_arenas: self.small_arenas.len(),
            medium_arenas: self.medium_arenas.len(),
            large_arenas: self.large_arenas.len(),
            total_arenas: self.small_arenas.len()
                + self.medium_arenas.len()
                + self.large_arenas.len(),
            next_arena_id: self.next_arena_id.load(Ordering::Relaxed),
        }
    }
}

#[derive(Debug)]
pub struct PoolStats {
    pub small_arenas: usize,
    pub medium_arenas: usize,
    pub large_arenas: usize,
    pub total_arenas: usize,
    pub next_arena_id: usize,
}

impl PoolStats {
    pub fn report(&self) -> String {
        format!(
            "Protocol Arena Pool: {} small, {} medium, {} large (total: {}, next_id: {})",
            self.small_arenas,
            self.medium_arenas,
            self.large_arenas,
            self.total_arenas,
            self.next_arena_id
        )
    }
}

/// RAII wrapper for protocol arena allocation
pub struct ProtocolArenaGuard {
    arena: Arc<ProtocolArena>,
    initial_position: usize,
}

impl ProtocolArenaGuard {
    pub fn new(arena: Arc<ProtocolArena>) -> Self {
        let initial_position = arena.position.load(Ordering::Relaxed);
        Self {
            arena,
            initial_position,
        }
    }

    /// Access the underlying arena
    pub fn arena(&self) -> &ProtocolArena {
        &self.arena
    }

    /// Get allocation statistics for this guard's usage
    pub fn usage_stats(&self) -> (usize, usize) {
        let current_pos = self.arena.position.load(Ordering::Relaxed);
        let bytes_used = current_pos.saturating_sub(self.initial_position);
        let allocations = self.arena.allocation_count.load(Ordering::Relaxed);
        (bytes_used, allocations)
    }
}

impl Drop for ProtocolArenaGuard {
    fn drop(&mut self) {
        let (bytes_used, _) = self.usage_stats();
        if bytes_used > 0 {
            debug!(
                "📊 PROTOCOL ARENA: Guard released, used {} bytes from arena {}",
                bytes_used, self.arena.arena_id
            );
        }
    }
}

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

    #[test]
    fn test_protocol_arena_basic_allocation() {
        let arena = ProtocolArena::new(1024, 1).unwrap();

        // Test basic allocation
        let ptr = arena.allocate(100, 1).unwrap();
        assert!(!ptr.as_ptr().is_null());

        // Test stats
        let stats = arena.get_stats();
        assert_eq!(stats.allocation_count, 1);
        assert!(stats.used_bytes >= 100);
    }

    #[test]
    fn test_protocol_arena_vec_allocation() {
        let arena = ProtocolArena::new(8192, 2).unwrap();

        // Allocate Vec<u32> with capacity 100
        let vec: Vec<u32> = arena.allocate_vec(100).unwrap();
        assert_eq!(vec.capacity(), 100);
        assert_eq!(vec.len(), 0);

        let stats = arena.get_stats();
        assert_eq!(stats.allocation_count, 1);
    }

    #[test]
    fn test_protocol_arena_exhaustion() {
        let arena = ProtocolArena::new(100, 3).unwrap();

        // First allocation should succeed
        let ptr1 = arena.allocate(50, 1);
        assert!(ptr1.is_some());

        // Second allocation that exceeds capacity should fail
        let ptr2 = arena.allocate(60, 1);
        assert!(ptr2.is_none());
    }

    #[test]
    fn test_protocol_arena_reset() {
        let arena = ProtocolArena::new(1024, 4).unwrap();

        // Make some allocations
        arena.allocate(100, 1).unwrap();
        arena.allocate(200, 1).unwrap();

        let stats_before = arena.get_stats();
        assert!(stats_before.used_bytes > 0);
        assert_eq!(stats_before.allocation_count, 2);

        // Reset arena
        arena.reset();

        let stats_after = arena.get_stats();
        assert_eq!(stats_after.used_bytes, 0);
        assert_eq!(stats_after.allocation_count, 0);
    }
}