jvmrs 0.1.2

#![allow(dead_code)]
//! Thread-Local Allocation Buffers (TLAB) for fast object allocation
//!
//! TLAB allows each thread to allocate from its own buffer without
//! synchronization, reducing contention and improving allocation speed.

use std::cell::RefCell;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

/// TLAB configuration
#[derive(Debug, Clone)]
pub struct TlabConfig {
    /// Size of each TLAB in bytes
    pub tlab_size: usize,

    /// Maximum number of TLABs in the global pool
    pub max_tlabs: usize,

    /// Enable TLAB allocation
    pub enabled: bool,
}

impl Default for TlabConfig {
    fn default() -> Self {
        Self {
            tlab_size: 64 * 1024, // 64KB per TLAB
            max_tlabs: 16,
            enabled: true,
        }
    }
}

/// Thread-local allocation buffer
pub struct Tlab {
    /// Base pointer of the TLAB
    base: *mut u8,

    /// Current allocation pointer
    top: *mut u8,

    /// End pointer (base + size)
    end: *mut u8,

    /// Thread ID that owns this TLAB
    thread_id: u32,

    /// TLAB size
    size: usize,

    /// Number of allocations from this TLAB
    allocation_count: usize,

    /// Bytes allocated from this TLAB
    bytes_allocated: usize,
}

unsafe impl Send for Tlab {}

impl Tlab {
    /// Create a new TLAB with the given size
    pub fn new(size: usize, thread_id: u32) -> Self {
        let layout = std::alloc::Layout::from_size_align(size, 8)
            .expect("Invalid TLAB size or alignment");
        let base = unsafe { std::alloc::alloc(layout) };
        
        if base.is_null() {
            panic!("Failed to allocate TLAB of {} bytes", size);
        }

        Self {
            base,
            top: base,
            end: unsafe { base.add(size) },
            thread_id,
            size,
            allocation_count: 0,
            bytes_allocated: 0,
        }
    }

    /// Allocate an object from the TLAB
    pub fn allocate(&mut self, size: usize, alignment: usize) -> Option<*mut u8> {
        // Align current pointer
        let aligned_top = (self.top as usize + alignment - 1) & !(alignment - 1);
        let aligned_top = aligned_top as *mut u8;

        // Check if there's enough space
        if unsafe { aligned_top.add(size) } > self.end {
            return None; // TLAB is full
        }

        // Allocate
        let ptr = aligned_top;
        self.top = unsafe { aligned_top.add(size) };
        self.allocation_count += 1;
        self.bytes_allocated += size;

        Some(ptr)
    }

    /// Get remaining space in the TLAB
    pub fn remaining(&self) -> usize {
        self.end as usize - self.top as usize
    }

    /// Get the thread ID that owns this TLAB
    pub fn thread_id(&self) -> u32 {
        self.thread_id
    }

    /// Get allocation statistics
    pub fn stats(&self) -> TlabStats {
        TlabStats {
            allocation_count: self.allocation_count,
            bytes_allocated: self.bytes_allocated,
            remaining: self.remaining(),
            utilization: if self.size > 0 {
                (self.bytes_allocated as f64 / self.size as f64) * 100.0
            } else {
                0.0
            },
        }
    }

    /// Reset the TLAB for reuse
    pub fn reset(&mut self) {
        self.top = self.base;
        self.allocation_count = 0;
        self.bytes_allocated = 0;
    }
}

impl Drop for Tlab {
    fn drop(&mut self) {
        if !self.base.is_null() {
            unsafe {
                if let Ok(layout) = std::alloc::Layout::from_size_align(self.size, 8) {
                    std::alloc::dealloc(self.base, layout);
                }
                // If layout fails, we can't deallocate - this is a critical error
                // but we can't panic in drop, so we log and continue
                log::error!("Failed to create layout for TLAB deallocation");
            }
        }
    }
}

/// TLAB statistics
#[derive(Debug, Clone, Copy)]
pub struct TlabStats {
    pub allocation_count: usize,
    pub bytes_allocated: usize,
    pub remaining: usize,
    pub utilization: f64,
}

/// TLAB manager for the heap
pub struct TlabManager {
    /// Configuration
    config: TlabConfig,

    /// Global TLAB pool (for reuse)
    pool: RefCell<Vec<Option<Tlab>>>,

    /// Thread-local TLABs
    thread_tlabs: RefCell<Vec<Option<Tlab>>>,

    /// Statistics
    total_allocations: Arc<AtomicUsize>,
    tlab_hits: Arc<AtomicUsize>,
    tlab_misses: Arc<AtomicUsize>,
}

impl TlabManager {
    /// Create a new TLAB manager
    pub fn new(config: TlabConfig) -> Self {
        let pool = (0..config.max_tlabs).map(|_| None).collect();

        Self {
            config,
            pool: RefCell::new(pool),
            thread_tlabs: RefCell::new(Vec::new()),
            total_allocations: Arc::new(AtomicUsize::new(0)),
            tlab_hits: Arc::new(AtomicUsize::new(0)),
            tlab_misses: Arc::new(AtomicUsize::new(0)),
        }
    }

    /// Get or create a TLAB for the current thread
    pub fn get_tlab(&self, thread_id: u32) -> Option<*mut Tlab> {
        if !self.config.enabled {
            return None;
        }

        // Check if thread already has a TLAB
        {
            let mut tlabs = self.thread_tlabs.borrow_mut();
            if let Some(idx) = tlabs.iter().position(|t| {
                t.as_ref()
                    .map(|t| t.thread_id() == thread_id)
                    .unwrap_or(false)
            }) {
                return tlabs[idx].as_mut().map(|t| t as *mut Tlab);
            }
        }

        // Try to get a TLAB from the pool
        let mut pool = self.pool.borrow_mut();
        for tlab in pool.iter_mut() {
            if tlab.is_some() {
                let tlab_obj = tlab.take().unwrap();

                // Add to thread-local TLABs
                let mut tlabs = self.thread_tlabs.borrow_mut();
                tlabs.push(Some(tlab_obj));

                return tlabs
                    .last_mut()
                    .and_then(|t| t.as_mut())
                    .map(|t| t as *mut Tlab);
            }
        }

        // Create a new TLAB
        let tlab_obj = Tlab::new(self.config.tlab_size, thread_id);
        let mut tlabs = self.thread_tlabs.borrow_mut();
        tlabs.push(Some(tlab_obj));

        tlabs
            .last_mut()
            .and_then(|t| t.as_mut())
            .map(|t| t as *mut Tlab)
    }

    /// Allocate an object using TLAB
    pub fn allocate(&self, size: usize, alignment: usize, thread_id: u32) -> Option<*mut u8> {
        self.total_allocations.fetch_add(1, Ordering::Relaxed);

        if let Some(tlab_ptr) = self.get_tlab(thread_id) {
            unsafe {
                let tlab = &mut *tlab_ptr;
                if let Some(ptr) = tlab.allocate(size, alignment) {
                    self.tlab_hits.fetch_add(1, Ordering::Relaxed);
                    return Some(ptr);
                }

                // TLAB is full, return it to the pool
                self.tlab_misses.fetch_add(1, Ordering::Relaxed);

                // Find slot in pool
                let mut pool = self.pool.borrow_mut();
                for slot in pool.iter_mut() {
                    if slot.is_none() {
                        let mut new_tlab = Tlab::new(self.config.tlab_size, thread_id);
                        std::mem::swap(tlab, &mut new_tlab);
                        *slot = Some(new_tlab);
                        break;
                    }
                }
            }
        }

        // Fall back to heap allocation
        None
    }

    /// Get TLAB statistics
    pub fn stats(&self) -> TlabManagerStats {
        let total = self.total_allocations.load(Ordering::Relaxed);
        let hits = self.tlab_hits.load(Ordering::Relaxed);
        let misses = self.tlab_misses.load(Ordering::Relaxed);

        TlabManagerStats {
            total_allocations: total,
            tlab_hits: hits,
            tlab_misses: misses,
            hit_rate: if total > 0 {
                (hits as f64 / total as f64) * 100.0
            } else {
                0.0
            },
            active_tlabs: self.thread_tlabs.borrow().len(),
        }
    }

    /// Reset all TLABs (e.g., after GC)
    pub fn reset_all(&self) {
        let mut tlabs = self.thread_tlabs.borrow_mut();
        for tlab in tlabs.iter_mut() {
            if let Some(t) = tlab {
                t.reset();
            }
        }
    }

    /// Clear all TLABs (e.g., for thread shutdown)
    pub fn clear_all(&self) {
        self.thread_tlabs.borrow_mut().clear();
    }
}

/// TLAB manager statistics
#[derive(Debug, Clone, Copy)]
pub struct TlabManagerStats {
    pub total_allocations: usize,
    pub tlab_hits: usize,
    pub tlab_misses: usize,
    pub hit_rate: f64,
    pub active_tlabs: usize,
}

impl Default for TlabManager {
    fn default() -> Self {
        Self::new(TlabConfig::default())
    }
}

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

    #[test]
    fn test_tlab_allocation() {
        let mut tlab = Tlab::new(1024, 1);

        // Allocate small object
        let ptr1 = tlab.allocate(16, 8);
        assert!(ptr1.is_some());
        assert_eq!(tlab.stats().allocation_count, 1);

        // Allocate another object
        let ptr2 = tlab.allocate(32, 8);
        assert!(ptr2.is_some());
        assert_eq!(tlab.stats().allocation_count, 2);

        // Pointers should be different
        assert_ne!(ptr1.unwrap(), ptr2.unwrap());
    }

    #[test]
    fn test_tlab_full() {
        let mut tlab = Tlab::new(100, 1);

        // Allocate 90 bytes
        let ptr1 = tlab.allocate(90, 8);
        assert!(ptr1.is_some());

        // Try to allocate 20 bytes (should fail)
        let ptr2 = tlab.allocate(20, 8);
        assert!(ptr2.is_none());
    }

    #[test]
    fn test_tlab_reset() {
        let mut tlab = Tlab::new(1024, 1);

        // Allocate some objects
        tlab.allocate(100, 8);
        tlab.allocate(200, 8);

        assert_eq!(tlab.stats().allocation_count, 2);

        // Reset TLAB
        tlab.reset();

        assert_eq!(tlab.stats().allocation_count, 0);
        assert_eq!(tlab.stats().bytes_allocated, 0);
    }

    #[test]
    fn test_tlab_manager() {
        let manager = TlabManager::new(TlabConfig {
            tlab_size: 1024,
            max_tlabs: 2,
            enabled: true,
        });

        // Allocate using TLAB
        let ptr = manager.allocate(100, 8, 1);
        assert!(ptr.is_some());

        let stats = manager.stats();
        assert_eq!(stats.total_allocations, 1);
        assert_eq!(stats.tlab_hits, 1);
    }

    #[test]
    fn test_tlab_stats() {
        let mut tlab = Tlab::new(1024, 1);

        tlab.allocate(100, 8);
        tlab.allocate(200, 8);

        let stats = tlab.stats();
        assert_eq!(stats.allocation_count, 2);
        assert_eq!(stats.bytes_allocated, 300);
        assert!(stats.utilization > 0.0);
    }

    #[test]
    fn test_tlab_alignment() {
        let mut tlab = Tlab::new(1024, 1);

        // Allocate with different alignments
        let ptr1 = tlab.allocate(16, 8);
        let ptr2 = tlab.allocate(32, 16);
        let ptr3 = tlab.allocate(64, 32);

        assert!(ptr1.is_some());
        assert!(ptr2.is_some());
        assert!(ptr3.is_some());

        // Check alignments
        assert_eq!(ptr1.unwrap() as usize % 8, 0);
        assert_eq!(ptr2.unwrap() as usize % 16, 0);
        assert_eq!(ptr3.unwrap() as usize % 32, 0);
    }
}