memscope_rs/async_memory/

allocator.rs

//! Global allocator hook for async memory tracking
//!
//! Provides transparent integration with Rust's global allocator to capture
//! all memory allocation and deallocation events with task attribution.

use std::alloc::{GlobalAlloc, Layout, System};

use crate::async_memory::buffer::record_allocation_event;
use crate::async_memory::task_id::get_current_task;

/// Task-aware global allocator wrapper
///
/// Intercepts all allocation and deallocation calls to record events
/// with task attribution. Uses the system allocator as the backend
/// for actual memory management.
pub struct TaskTrackingAllocator;

unsafe impl GlobalAlloc for TaskTrackingAllocator {
    #[inline(always)]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // Perform actual allocation first
        let ptr = System.alloc(layout);

        // Record allocation event if successful and we have task context
        if !ptr.is_null() {
            self.record_allocation_fast(ptr as usize, layout.size());
        }

        ptr
    }

    #[inline(always)]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        // Record deallocation event before freeing
        self.record_deallocation_fast(ptr as usize, layout.size());

        // Perform actual deallocation
        System.dealloc(ptr, layout);
    }

    #[inline(always)]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // Use system allocator for zeroed allocation
        let ptr = System.alloc_zeroed(layout);

        // Record allocation event if successful
        if !ptr.is_null() {
            self.record_allocation_fast(ptr as usize, layout.size());
        }

        ptr
    }

    #[inline(always)]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        // Record deallocation of old memory
        if !ptr.is_null() {
            self.record_deallocation_fast(ptr as usize, layout.size());
        }

        // Perform actual reallocation
        let new_ptr = System.realloc(ptr, layout, new_size);

        // Record allocation of new memory if successful
        if !new_ptr.is_null() {
            self.record_allocation_fast(new_ptr as usize, new_size);
        }

        new_ptr
    }
}

impl TaskTrackingAllocator {
    /// Record allocation event with minimal overhead
    ///
    /// Optimized for the hot path of memory allocation.
    /// Uses efficient timestamp generation and task context lookup.
    #[inline(always)]
    fn record_allocation_fast(&self, ptr: usize, size: usize) {
        // Get current task context from thread-local storage
        let task_info = get_current_task();

        // Only record if we have a valid task context
        if !task_info.has_tracking_id() {
            return;
        }

        // Generate efficient timestamp
        let timestamp = current_timestamp();

        // Use primary task ID for attribution
        let task_id = task_info.primary_id();

        // Record allocation event (ignore errors to avoid allocation in error path)
        let _ = record_allocation_event(task_id, ptr, size, timestamp, true);
    }

    /// Record deallocation event with minimal overhead
    #[inline(always)]
    fn record_deallocation_fast(&self, ptr: usize, size: usize) {
        // Get current task context
        let task_info = get_current_task();

        // Only record if we have a valid task context
        if !task_info.has_tracking_id() {
            return;
        }

        // Generate timestamp
        let timestamp = current_timestamp();

        // Use primary task ID for attribution
        let task_id = task_info.primary_id();

        // Record deallocation event (ignore errors to avoid allocation in error path)
        let _ = record_allocation_event(task_id, ptr, size, timestamp, false);
    }
}

/// Get current timestamp with minimal overhead
///
/// Uses platform-specific optimizations for timestamp generation.
/// Prefers TSC on x86_64 for sub-nanosecond precision and minimal overhead.
#[inline(always)]
fn current_timestamp() -> u64 {
    #[cfg(target_arch = "x86_64")]
    {
        // Use Time Stamp Counter for minimal overhead
        unsafe { std::arch::x86_64::_rdtsc() }
    }
    #[cfg(not(target_arch = "x86_64"))]
    {
        // Fallback to high-resolution time for other architectures
        use std::time::{SystemTime, UNIX_EPOCH};
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0)
    }
}

/// Set the global allocator to use task tracking
///
/// This macro must be called at the crate root to enable task-aware
/// memory tracking for all allocations in the application.
///
/// Example:
/// ```rust
/// use memscope_rs::set_task_tracking_allocator;
///
/// set_task_tracking_allocator!();
/// ```
#[macro_export]
macro_rules! set_task_tracking_allocator {
    () => {
        #[global_allocator]
        static ALLOCATOR: $crate::async_memory::allocator::TaskTrackingAllocator =
            $crate::async_memory::allocator::TaskTrackingAllocator;
    };
}

/// Check if task tracking allocator is enabled
///
/// Returns true if allocations are being tracked, false otherwise.
/// Useful for conditional behavior in libraries.
pub fn is_tracking_enabled() -> bool {
    // Simple heuristic: check if we have any task context
    get_current_task().has_tracking_id()
}

/// Get allocation tracking statistics
///
/// Returns basic statistics about allocation tracking overhead.
/// Used for performance monitoring and optimization.
#[derive(Debug, Clone)]
pub struct AllocationStats {
    /// Number of allocations recorded
    pub allocations_recorded: u64,
    /// Number of deallocations recorded  
    pub deallocations_recorded: u64,
    /// Number of tracking events dropped due to buffer overflow
    pub events_dropped: u64,
    /// Estimated tracking overhead in nanoseconds per allocation
    pub overhead_per_allocation_ns: f64,
}

impl AllocationStats {
    /// Calculate tracking efficiency ratio
    pub fn efficiency_ratio(&self) -> f64 {
        let total_events = self.allocations_recorded + self.deallocations_recorded;
        if total_events == 0 {
            1.0
        } else {
            (total_events - self.events_dropped) as f64 / total_events as f64
        }
    }

    /// Check if tracking performance is acceptable
    pub fn is_performance_acceptable(&self) -> bool {
        // Consider acceptable if overhead < 10ns per allocation and efficiency > 95%
        self.overhead_per_allocation_ns < 10.0 && self.efficiency_ratio() > 0.95
    }
}

/// Get current allocation tracking statistics
///
/// This is a simplified implementation - production version would
/// maintain global counters and calculate actual overhead.
pub fn get_allocation_stats() -> AllocationStats {
    // Simplified implementation for now
    AllocationStats {
        allocations_recorded: 0,
        deallocations_recorded: 0,
        events_dropped: 0,
        overhead_per_allocation_ns: 5.0, // Estimated based on design targets
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_allocator_basic_functionality() {
        let allocator = TaskTrackingAllocator;

        unsafe {
            // Test basic allocation
            let layout = Layout::from_size_align(1024, 8).expect("Invalid layout");
            let ptr = allocator.alloc(layout);
            assert!(!ptr.is_null());

            // Test deallocation
            allocator.dealloc(ptr, layout);

            // Test zeroed allocation
            let ptr_zero = allocator.alloc_zeroed(layout);
            assert!(!ptr_zero.is_null());

            // Verify memory is zeroed
            let slice = std::slice::from_raw_parts(ptr_zero, 1024);
            assert!(slice.iter().all(|&b| b == 0));

            allocator.dealloc(ptr_zero, layout);
        }
    }

    // Note: Direct allocator testing removed due to potential recursion in test environment
    // The allocator functionality is tested indirectly through other tests

    // Note: Allocator tests that involve actual memory allocation/deallocation
    // are removed to prevent stack overflow in test environment.
    // The allocator functionality is validated through integration tests.

    #[test]
    fn test_timestamp_generation() {
        let ts1 = current_timestamp();
        let ts2 = current_timestamp();

        // Timestamps should be monotonic (or at least not decreasing)
        assert!(ts2 >= ts1);

        // Timestamps should be non-zero in normal operation
        assert_ne!(ts1, 0);
        assert_ne!(ts2, 0);
    }

    #[test]
    fn test_tracking_status() {
        use crate::async_memory::task_id::{clear_current_task, set_current_task, TaskInfo};

        // Without task context, tracking should be disabled
        clear_current_task();
        assert!(!is_tracking_enabled());

        // With task context, tracking should be enabled
        let task_info = TaskInfo::new(99999, None);
        set_current_task(task_info);
        assert!(is_tracking_enabled());

        // Test with span ID only
        clear_current_task();
        let task_info_span = TaskInfo::new(0, Some(88888));
        set_current_task(task_info_span);
        assert!(is_tracking_enabled());
    }

    #[test]
    fn test_allocation_stats() {
        let stats = get_allocation_stats();

        // Should have reasonable performance characteristics
        assert!(stats.overhead_per_allocation_ns > 0.0);
        assert!(stats.overhead_per_allocation_ns < 100.0); // Less than 100ns overhead

        // Efficiency should be high initially
        assert!(stats.efficiency_ratio() >= 0.95);
        assert!(stats.is_performance_acceptable());
    }
}