memscope_rs/capture/platform/

allocator.rs

use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
use std::time::Instant;

/// Platform-specific allocator hooking system
pub struct PlatformAllocator {
    original_alloc: AtomicPtr<u8>,
    original_dealloc: AtomicPtr<u8>,
    config: HookConfig,
    stats: HookStats,
}

impl PlatformAllocator {
    pub fn original_alloc(&self) -> *mut u8 {
        self.original_alloc.load(Ordering::SeqCst)
    }
    pub fn original_dealloc(&self) -> *mut u8 {
        self.original_dealloc.load(Ordering::SeqCst)
    }
}

/// Configuration for allocation hooks
#[derive(Debug, Clone)]
pub struct HookConfig {
    /// Whether to track allocations
    pub track_allocations: bool,
    /// Whether to track deallocations
    pub track_deallocations: bool,
    /// Minimum allocation size to track
    pub min_tracked_size: usize,
    /// Maximum allocation size to track
    pub max_tracked_size: usize,
    /// Sample rate for tracking (0.0 to 1.0)
    pub sample_rate: f64,
}

/// Statistics for allocation hooks
#[derive(Debug)]
struct HookStats {
    /// Total allocations intercepted
    total_allocations: AtomicUsize,
    /// Total deallocations intercepted
    total_deallocations: AtomicUsize,
    /// Total bytes allocated
    total_bytes_allocated: AtomicUsize,
    /// Total bytes deallocated
    total_bytes_deallocated: AtomicUsize,
    /// Hook overhead time
    total_hook_time: AtomicUsize,
}

/// Result of allocation hook
#[derive(Debug, Clone)]
pub struct HookResult {
    /// Whether allocation should proceed
    pub should_proceed: bool,
    /// Whether to track this allocation
    pub should_track: bool,
    /// Additional metadata
    pub metadata: Option<AllocationMetadata>,
}

/// Information about an allocation
#[derive(Debug, Clone)]
pub struct AllocationInfo {
    /// Pointer to allocated memory
    pub ptr: *mut u8,
    /// Size of allocation
    pub size: usize,
    /// Alignment requirement
    pub align: usize,
    /// Timestamp of allocation
    pub timestamp: Instant,
    /// Thread ID that made allocation
    pub thread_id: ThreadId,
    /// Stack trace if captured
    pub stack_trace: Option<Vec<usize>>,
}

/// Metadata for allocation tracking
#[derive(Debug, Clone)]
pub struct AllocationMetadata {
    /// Type name if known
    pub type_name: Option<String>,
    /// Source location if available
    pub source_location: Option<SourceLocation>,
    /// Custom tags
    pub tags: Vec<String>,
}

/// Source code location information
#[derive(Debug, Clone)]
pub struct SourceLocation {
    /// File name
    pub file: String,
    /// Line number
    pub line: u32,
    /// Column number
    pub column: Option<u32>,
}

/// Thread identifier
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ThreadId(pub u64);

/// Allocation hook function type
pub type AllocationHook = fn(&AllocationInfo) -> HookResult;

/// Deallocation hook function type
pub type DeallocationHook = fn(*mut u8, usize) -> bool;

impl PlatformAllocator {
    /// Create new platform allocator
    pub fn new() -> Self {
        Self {
            original_alloc: AtomicPtr::new(std::ptr::null_mut()),
            original_dealloc: AtomicPtr::new(std::ptr::null_mut()),
            config: HookConfig::default(),
            stats: HookStats::new(),
        }
    }

    /// Install allocation hooks
    pub fn install_hooks(&mut self) -> Result<(), HookError> {
        #[cfg(target_os = "linux")]
        {
            self.install_linux_hooks()
        }

        #[cfg(target_os = "windows")]
        {
            self.install_windows_hooks()
        }

        #[cfg(target_os = "macos")]
        {
            self.install_macos_hooks()
        }

        #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos")))]
        {
            Err(HookError::UnsupportedPlatform)
        }
    }

    /// Remove allocation hooks
    pub fn remove_hooks(&mut self) -> Result<(), HookError> {
        #[cfg(target_os = "linux")]
        {
            self.remove_linux_hooks()
        }

        #[cfg(target_os = "windows")]
        {
            self.remove_windows_hooks()
        }

        #[cfg(target_os = "macos")]
        {
            self.remove_macos_hooks()
        }

        #[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "macos")))]
        {
            Err(HookError::UnsupportedPlatform)
        }
    }

    /// Get hook statistics
    pub fn get_statistics(&self) -> AllocationStatistics {
        AllocationStatistics {
            total_allocations: self.stats.total_allocations.load(Ordering::Relaxed),
            total_deallocations: self.stats.total_deallocations.load(Ordering::Relaxed),
            total_bytes_allocated: self.stats.total_bytes_allocated.load(Ordering::Relaxed),
            total_bytes_deallocated: self.stats.total_bytes_deallocated.load(Ordering::Relaxed),
            current_allocations: self
                .stats
                .total_allocations
                .load(Ordering::Relaxed)
                .saturating_sub(self.stats.total_deallocations.load(Ordering::Relaxed)),
            current_bytes: self
                .stats
                .total_bytes_allocated
                .load(Ordering::Relaxed)
                .saturating_sub(self.stats.total_bytes_deallocated.load(Ordering::Relaxed)),
            average_hook_overhead: self.calculate_average_overhead(),
        }
    }

    /// Update hook configuration
    pub fn update_config(&mut self, config: HookConfig) {
        self.config = config;
    }

    #[cfg(target_os = "linux")]
    fn install_linux_hooks(&mut self) -> Result<(), HookError> {
        // Linux-specific implementation using LD_PRELOAD or similar
        //
        // NOTE: This is a placeholder implementation. A full implementation would:
        // 1. Use LD_PRELOAD to intercept malloc/free calls
        // 2. Or use dlsym to hook into libc allocation functions
        // 3. Register the GlobalAlloc implementation with the system
        //
        // Current implementation: No-op (hooks are already installed via GlobalAlloc trait)
        // Future implementation: Add LD_PRELOAD support for external libraries
        tracing::warn!("Linux hooks are not yet implemented. Using GlobalAlloc trait instead.");
        Ok(())
    }

    #[cfg(target_os = "linux")]
    fn remove_linux_hooks(&mut self) -> Result<(), HookError> {
        // Linux-specific cleanup
        // Restore original malloc/free if they were hooked
        tracing::info!("Linux hooks cleanup: No action needed (using GlobalAlloc trait)");
        Ok(())
    }

    #[cfg(target_os = "windows")]
    fn install_windows_hooks(&mut self) -> Result<(), HookError> {
        // Windows-specific implementation using detours or similar
        //
        // NOTE: This is a placeholder implementation. A full implementation would:
        // 1. Use Detours library to hook HeapAlloc/HeapFree
        // 2. Or use VirtualProtect to modify import address table
        // 3. Register the GlobalAlloc implementation with the system
        //
        // Current implementation: No-op (hooks are already installed via GlobalAlloc trait)
        // Future implementation: Add Detours support for external libraries
        tracing::warn!("Windows hooks are not yet implemented. Using GlobalAlloc trait instead.");
        Ok(())
    }

    #[cfg(target_os = "windows")]
    fn remove_windows_hooks(&mut self) -> Result<(), HookError> {
        // Windows-specific cleanup
        // Restore original HeapAlloc/HeapFree if they were hooked
        tracing::info!("Windows hooks cleanup: No action needed (using GlobalAlloc trait)");
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn install_macos_hooks(&mut self) -> Result<(), HookError> {
        // macOS-specific implementation using interpose or similar
        //
        // NOTE: This is a placeholder implementation. A full implementation would:
        // 1. Use DYLD_INSERT_LIBRARIES to intercept malloc/free
        // 2. Or use interpose to hook into system allocation functions
        // 3. Register the GlobalAlloc implementation with the system
        //
        // Current implementation: No-op (hooks are already installed via GlobalAlloc trait)
        // Future implementation: Add DYLD_INSERT_LIBRARIES support for external libraries
        tracing::warn!("macOS hooks are not yet implemented. Using GlobalAlloc trait instead.");
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn remove_macos_hooks(&mut self) -> Result<(), HookError> {
        // macOS-specific cleanup
        // Restore original malloc/free if they were hooked
        tracing::info!("macOS hooks cleanup: No action needed (using GlobalAlloc trait)");
        Ok(())
    }

    fn calculate_average_overhead(&self) -> f64 {
        let total_time = self.stats.total_hook_time.load(Ordering::Relaxed);
        let total_calls = self.stats.total_allocations.load(Ordering::Relaxed)
            + self.stats.total_deallocations.load(Ordering::Relaxed);

        if total_calls > 0 {
            total_time as f64 / total_calls as f64
        } else {
            0.0
        }
    }

    /// Handle allocation event
    pub fn handle_allocation(&self, info: &AllocationInfo) -> HookResult {
        let start_time = Instant::now();

        // Update statistics
        self.stats.total_allocations.fetch_add(1, Ordering::Relaxed);
        self.stats
            .total_bytes_allocated
            .fetch_add(info.size, Ordering::Relaxed);

        // Check if we should track this allocation
        let should_track = self.should_track_allocation(info);

        // Record hook overhead
        let overhead = start_time.elapsed().as_nanos() as usize;
        self.stats
            .total_hook_time
            .fetch_add(overhead, Ordering::Relaxed);

        HookResult {
            should_proceed: true,
            should_track,
            metadata: self.extract_metadata(info),
        }
    }

    /// Handle deallocation event
    pub fn handle_deallocation(&self, _ptr: *mut u8, size: usize) -> bool {
        let start_time = Instant::now();

        // Update statistics
        self.stats
            .total_deallocations
            .fetch_add(1, Ordering::Relaxed);
        self.stats
            .total_bytes_deallocated
            .fetch_add(size, Ordering::Relaxed);

        // Record hook overhead
        let overhead = start_time.elapsed().as_nanos() as usize;
        self.stats
            .total_hook_time
            .fetch_add(overhead, Ordering::Relaxed);

        true
    }

    fn should_track_allocation(&self, info: &AllocationInfo) -> bool {
        // Check size limits
        if info.size < self.config.min_tracked_size || info.size > self.config.max_tracked_size {
            return false;
        }

        // Apply sampling with random decision to avoid bias
        if self.config.sample_rate < 1.0 {
            let sample_decision = rand::random::<f64>();
            if sample_decision >= self.config.sample_rate {
                return false;
            }
        }

        true
    }

    fn extract_metadata(&self, _info: &AllocationInfo) -> Option<AllocationMetadata> {
        // Extract metadata from allocation context
        // Real implementation would use debug info, compiler hints, or source location tracking
        // This feature is not yet implemented
        None
    }
}

impl HookStats {
    fn new() -> Self {
        Self {
            total_allocations: AtomicUsize::new(0),
            total_deallocations: AtomicUsize::new(0),
            total_bytes_allocated: AtomicUsize::new(0),
            total_bytes_deallocated: AtomicUsize::new(0),
            total_hook_time: AtomicUsize::new(0),
        }
    }
}

/// Statistics about allocation hooks
#[derive(Debug, Clone)]
pub struct AllocationStatistics {
    /// Total number of allocations
    pub total_allocations: usize,
    /// Total number of deallocations
    pub total_deallocations: usize,
    /// Total bytes allocated
    pub total_bytes_allocated: usize,
    /// Total bytes deallocated
    pub total_bytes_deallocated: usize,
    /// Current active allocations
    pub current_allocations: usize,
    /// Current active bytes
    pub current_bytes: usize,
    /// Average hook overhead in nanoseconds
    pub average_hook_overhead: f64,
}

/// Errors that can occur during hook installation
#[derive(Debug, Clone, PartialEq)]
pub enum HookError {
    /// Platform not supported
    UnsupportedPlatform,
    /// Permission denied
    PermissionDenied,
    /// Hook already installed
    AlreadyInstalled,
    /// Hook not installed
    NotInstalled,
    /// System error
    SystemError(String),
}

impl Default for HookConfig {
    fn default() -> Self {
        Self {
            track_allocations: true,
            track_deallocations: true,
            min_tracked_size: 1,
            max_tracked_size: usize::MAX,
            sample_rate: 1.0,
        }
    }
}

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

impl std::fmt::Display for HookError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HookError::UnsupportedPlatform => {
                write!(f, "Platform not supported for allocation hooking")
            }
            HookError::PermissionDenied => write!(f, "Permission denied for hook installation"),
            HookError::AlreadyInstalled => write!(f, "Allocation hooks already installed"),
            HookError::NotInstalled => write!(f, "Allocation hooks not installed"),
            HookError::SystemError(msg) => write!(f, "System error: {}", msg),
        }
    }
}

impl std::error::Error for HookError {}

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

    #[test]
    fn test_platform_allocator_creation() {
        let allocator = PlatformAllocator::new();
        let stats = allocator.get_statistics();

        assert_eq!(stats.total_allocations, 0);
        assert_eq!(stats.total_deallocations, 0);
        assert_eq!(stats.current_allocations, 0);
    }

    #[test]
    fn test_hook_config() {
        let config = HookConfig::default();
        assert!(config.track_allocations);
        assert!(config.track_deallocations);
        assert_eq!(config.min_tracked_size, 1);
        assert_eq!(config.sample_rate, 1.0);
    }

    #[test]
    fn test_allocation_info() {
        let info = AllocationInfo {
            ptr: std::ptr::null_mut(),
            size: 1024,
            align: 8,
            timestamp: Instant::now(),
            thread_id: ThreadId(1),
            stack_trace: None,
        };

        assert_eq!(info.size, 1024);
        assert_eq!(info.align, 8);
    }

    #[test]
    fn test_hook_statistics() {
        let allocator = PlatformAllocator::new();

        let info = AllocationInfo {
            ptr: 0x1000 as *mut u8,
            size: 100,
            align: 8,
            timestamp: Instant::now(),
            thread_id: ThreadId(1),
            stack_trace: None,
        };

        let result = allocator.handle_allocation(&info);
        assert!(result.should_proceed);

        let stats = allocator.get_statistics();
        assert_eq!(stats.total_allocations, 1);
        assert_eq!(stats.total_bytes_allocated, 100);
    }

    #[test]
    fn test_sample_rate_filtering() {
        let mut allocator = PlatformAllocator::new();
        allocator.config.sample_rate = 0.5;

        // Test multiple allocations to check sampling
        let mut tracked_count = 0;
        for i in 0..1000 {
            let info = AllocationInfo {
                ptr: (0x1000 + i) as *mut u8,
                size: 64,
                align: 8,
                timestamp: Instant::now(),
                thread_id: ThreadId(1),
                stack_trace: None,
            };

            if allocator.should_track_allocation(&info) {
                tracked_count += 1;
            }
        }

        // Should track roughly 50% with some variance
        assert!(tracked_count > 400 && tracked_count < 600);
    }
}