memscope_rs/metadata/stack_trace/

capture.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

/// Represents a single frame in a stack trace
#[derive(Debug, Clone)]
pub struct StackFrame {
    /// Raw instruction pointer address
    pub instruction_pointer: usize,
    /// Symbol name if resolved
    pub symbol_name: Option<String>,
    /// Source filename if available
    pub filename: Option<String>,
    /// Line number in source file
    pub line_number: Option<u32>,
    /// Function or method name
    pub function_name: Option<String>,
}

/// Configuration for stack trace capture behavior
#[derive(Debug, Clone)]
pub struct CaptureConfig {
    /// Maximum number of frames to capture
    pub max_depth: usize,
    /// Number of top frames to skip (e.g., allocator internals)
    pub skip_frames: usize,
    /// Whether to resolve symbol information
    pub enable_symbols: bool,
    /// Whether to cache resolved symbols for performance
    pub cache_symbols: bool,
    /// Whether to filter out system/library frames
    pub filter_system_frames: bool,
}

impl Default for CaptureConfig {
    fn default() -> Self {
        Self {
            max_depth: 32,
            skip_frames: 2,
            enable_symbols: true,
            cache_symbols: true,
            filter_system_frames: true,
        }
    }
}

/// High-performance stack trace capture engine
pub struct StackTraceCapture {
    /// Capture configuration settings
    config: CaptureConfig,
    /// Whether capture is currently enabled
    enabled: AtomicBool,
    /// Total number of captures performed
    capture_count: AtomicUsize,
    /// Cache of resolved stack frames for performance
    frame_cache: HashMap<usize, StackFrame>,
}

impl StackTraceCapture {
    /// Create new stack trace capture instance with given configuration
    pub fn new(config: CaptureConfig) -> Self {
        Self {
            config,
            enabled: AtomicBool::new(true),
            capture_count: AtomicUsize::new(0),
            frame_cache: HashMap::new(),
        }
    }

    /// Capture full stack trace with symbol resolution
    /// Returns None if capture is disabled
    pub fn capture(&mut self) -> Option<Vec<StackFrame>> {
        if !self.enabled.load(Ordering::Relaxed) {
            return None;
        }

        self.capture_count.fetch_add(1, Ordering::Relaxed);

        let mut frames = Vec::with_capacity(self.config.max_depth);
        let mut frame_count = 0;

        #[cfg(feature = "backtrace")]
        {
            let bt = backtrace::Backtrace::new();

            for frame in bt.frames().iter().skip(self.config.skip_frames) {
                if frame_count >= self.config.max_depth {
                    break;
                }

                let ip = frame.ip() as usize;

                let stack_frame = if let Some(cached_frame) = self.frame_cache.get(&ip) {
                    cached_frame.clone()
                } else {
                    let new_frame = self.create_frame_from_backtrace(frame);
                    if self.config.cache_symbols {
                        self.frame_cache.insert(ip, new_frame.clone());
                    }
                    new_frame
                };

                if self.should_include_frame(&stack_frame) {
                    frames.push(stack_frame);
                    frame_count += 1;
                }
            }
        }

        #[cfg(not(feature = "backtrace"))]
        {
            let mut skip_count = 0;
            let mut current_ip = self.get_current_instruction_pointer();

            while frame_count < self.config.max_depth {
                if skip_count < self.config.skip_frames {
                    skip_count += 1;
                    current_ip = self.walk_stack_frame(current_ip)?;
                    continue;
                }

                let frame = if let Some(cached_frame) = self.frame_cache.get(&current_ip) {
                    cached_frame.clone()
                } else {
                    let new_frame = self.create_frame(current_ip);
                    if self.config.cache_symbols {
                        self.frame_cache.insert(current_ip, new_frame.clone());
                    }
                    new_frame
                };

                if self.should_include_frame(&frame) {
                    frames.push(frame);
                    frame_count += 1;
                }

                current_ip = self.walk_stack_frame(current_ip)?;
            }
        }

        Some(frames)
    }

    /// Capture lightweight stack trace (instruction pointers only)
    /// Much faster than full capture, suitable for hot paths
    pub fn capture_lightweight(&self) -> Option<Vec<usize>> {
        if !self.enabled.load(Ordering::Relaxed) {
            return None;
        }

        let mut instruction_pointers = Vec::with_capacity(self.config.max_depth);

        #[cfg(feature = "backtrace")]
        {
            let bt = backtrace::Backtrace::new();

            for frame in bt.frames().iter().skip(self.config.skip_frames) {
                if instruction_pointers.len() >= self.config.max_depth {
                    break;
                }
                instruction_pointers.push(frame.ip() as usize);
            }
        }

        #[cfg(not(feature = "backtrace"))]
        {
            let mut current_ip = self.get_current_instruction_pointer();
            let mut skip_count = 0;

            for _ in 0..self.config.max_depth {
                if skip_count < self.config.skip_frames {
                    skip_count += 1;
                    current_ip = self.walk_stack_frame(current_ip)?;
                    continue;
                }

                instruction_pointers.push(current_ip);
                current_ip = self.walk_stack_frame(current_ip)?;
            }
        }

        Some(instruction_pointers)
    }

    /// Enable stack trace capture
    pub fn enable(&self) {
        self.enabled.store(true, Ordering::Relaxed);
    }

    /// Disable stack trace capture for performance
    pub fn disable(&self) {
        self.enabled.store(false, Ordering::Relaxed);
    }

    /// Check if capture is currently enabled
    pub fn is_enabled(&self) -> bool {
        self.enabled.load(Ordering::Relaxed)
    }

    /// Get total number of captures performed
    pub fn get_capture_count(&self) -> usize {
        self.capture_count.load(Ordering::Relaxed)
    }

    /// Clear the symbol resolution cache
    pub fn clear_cache(&mut self) {
        self.frame_cache.clear();
    }

    /// Get current size of symbol cache
    pub fn cache_size(&self) -> usize {
        self.frame_cache.len()
    }

    #[cfg(target_os = "macos")]
    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn get_current_instruction_pointer(&self) -> usize {
        // Platform-specific implementation would use __builtin_return_address or similar
        // Real stack trace is captured via backtrace crate in walk_stack
        0
    }

    #[cfg(target_os = "linux")]
    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn get_current_instruction_pointer(&self) -> usize {
        // Platform-specific implementation would use libc::backtrace or similar
        // Real stack trace is captured via backtrace crate in walk_stack
        0
    }

    #[cfg(target_os = "windows")]
    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn get_current_instruction_pointer(&self) -> usize {
        // Platform-specific implementation would use _ReturnAddress or similar
        // Real stack trace is captured via backtrace crate in walk_stack
        0
    }

    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn walk_stack_frame(&self, _current_ip: usize) -> Option<usize> {
        // Platform-specific stack walking implementation
        // This feature is not yet implemented
        None
    }

    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn create_frame(&self, ip: usize) -> StackFrame {
        let mut frame = StackFrame {
            instruction_pointer: ip,
            symbol_name: None,
            filename: None,
            line_number: None,
            function_name: None,
        };

        if self.config.enable_symbols {
            // Symbol resolution would happen here
            frame.function_name = self.resolve_function_name(ip);
            frame.filename = self.resolve_filename(ip);
            frame.line_number = self.resolve_line_number(ip);
        }

        frame
    }

    fn should_include_frame(&self, frame: &StackFrame) -> bool {
        if !self.config.filter_system_frames {
            return true;
        }

        // Filter out system/library frames
        if let Some(filename) = &frame.filename {
            if filename.contains("/usr/lib") || filename.contains("/lib64") {
                return false;
            }
        }

        if let Some(function_name) = &frame.function_name {
            if function_name.starts_with("__libc_") || function_name.starts_with("_start") {
                return false;
            }
        }

        true
    }

    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn resolve_function_name(&self, _ip: usize) -> Option<String> {
        // Real implementation would use debug symbols
        // This feature is not yet implemented
        None
    }

    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn resolve_filename(&self, _ip: usize) -> Option<String> {
        // Real implementation would use debug symbols
        // This feature is not yet implemented
        None
    }

    #[cfg_attr(feature = "backtrace", allow(dead_code))]
    fn resolve_line_number(&self, _ip: usize) -> Option<u32> {
        // Real implementation would use debug symbols
        // This feature is not yet implemented
        None
    }

    #[cfg(feature = "backtrace")]
    fn create_frame_from_backtrace(&self, frame: &backtrace::BacktraceFrame) -> StackFrame {
        let ip = frame.ip() as usize;

        let symbol_name = frame
            .symbols()
            .first()
            .and_then(|sym| sym.name())
            .map(|name| name.to_string());

        let filename = frame
            .symbols()
            .first()
            .and_then(|sym| sym.filename())
            .map(|path| path.display().to_string());

        let line_number = frame.symbols().first().and_then(|sym| sym.lineno());

        let function_name = symbol_name.clone();

        StackFrame {
            instruction_pointer: ip,
            symbol_name,
            filename,
            line_number,
            function_name,
        }
    }
}

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

impl StackFrame {
    pub fn new(ip: usize) -> Self {
        Self {
            instruction_pointer: ip,
            symbol_name: None,
            filename: None,
            line_number: None,
            function_name: None,
        }
    }

    pub fn with_symbols(
        ip: usize,
        function_name: Option<String>,
        filename: Option<String>,
        line_number: Option<u32>,
    ) -> Self {
        Self {
            instruction_pointer: ip,
            symbol_name: function_name.clone(),
            filename,
            line_number,
            function_name,
        }
    }

    pub fn is_resolved(&self) -> bool {
        self.function_name.is_some() || self.filename.is_some()
    }

    pub fn display_name(&self) -> String {
        if let Some(func) = &self.function_name {
            if let (Some(file), Some(line)) = (&self.filename, self.line_number) {
                format!("{}() at {}:{}", func, file, line)
            } else {
                format!("{}()", func)
            }
        } else {
            format!("0x{:x}", self.instruction_pointer)
        }
    }
}

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

    #[test]
    fn test_enable_disable() {
        let capture = StackTraceCapture::default();

        assert!(capture.is_enabled());

        capture.disable();
        assert!(!capture.is_enabled());

        capture.enable();
        assert!(capture.is_enabled());
    }

    #[test]
    fn test_frame_creation() {
        let frame = StackFrame::new(0x1234);
        assert_eq!(frame.instruction_pointer, 0x1234);
        assert!(!frame.is_resolved());

        let resolved_frame = StackFrame::with_symbols(
            0x5678,
            Some("test_func".to_string()),
            Some("test.rs".to_string()),
            Some(42),
        );
        assert!(resolved_frame.is_resolved());
        assert_eq!(resolved_frame.display_name(), "test_func() at test.rs:42");
    }

    #[test]
    fn test_capture_count() {
        let mut capture = StackTraceCapture::default();

        assert_eq!(capture.get_capture_count(), 0);

        capture.capture();
        assert_eq!(capture.get_capture_count(), 1);

        capture.capture();
        assert_eq!(capture.get_capture_count(), 2);
    }

    #[test]
    fn test_custom_config() {
        let config = CaptureConfig {
            max_depth: 10,
            skip_frames: 1,
            enable_symbols: false,
            cache_symbols: false,
            filter_system_frames: false,
        };

        let mut capture = StackTraceCapture::new(config);

        if let Some(frames) = capture.capture() {
            assert!(frames.len() <= 10);
            // With symbols disabled, frames should not be resolved
            for frame in &frames {
                assert!(frame.function_name.is_none() || !frame.is_resolved());
            }
        }
    }
}