calltrace-rs 1.1.4

//! Thread-Safe Call Tree Management
//!
//! This module manages hierarchical call trees for multiple threads,
//! providing thread-safe access and efficient storage.

use crate::dwarf_analyzer::FunctionInfo;
use crate::error::Result;
use crate::register_reader::{ArgumentValue, CapturedArgument, RegisterContext};
use dashmap::DashMap;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};

/// Unique identifier for call tree nodes
pub type NodeId = u32;

/// Thread ID type
pub type ThreadId = u64;

/// Call tree node representing a single function call
#[derive(Debug)]
pub struct CallNode {
    pub id: NodeId,
    pub function_name: String,
    pub function_address: u64,
    pub call_site: u64,

    // Timing information
    pub start_time: u64,
    pub end_time: Option<u64>,
    pub duration_us: Option<u64>,

    // Thread information
    pub thread_id: ThreadId,
    pub call_depth: usize,

    // Function signature and arguments
    pub function_info: Option<FunctionInfo>,
    pub arguments: Vec<CapturedArgument>,
    pub register_context: Option<RegisterContext>,
    pub return_value: Option<ArgumentValue>,

    // Tree structure
    pub parent: Option<NodeId>,
    pub children: Vec<NodeId>,

    // Thread creation tracking
    pub is_pthread_create: bool,
    pub created_thread_id: Option<ThreadId>,

    // Completion status
    pub is_completed: bool,
}

/// Lightweight call tree node for fast path (minimal memory footprint)
#[derive(Debug)]
pub struct FastCallNode {
    pub id: NodeId,
    pub function_address: u64,
    pub call_site: u64,
    pub start_time: u64,
    pub end_time: Option<u64>,
    pub thread_id: ThreadId,
    pub call_depth: usize,
    pub parent: Option<NodeId>,
    pub children: Vec<NodeId>,
    pub is_completed: bool,
}

/// Per-thread call tree state
#[derive(Debug)]
pub struct ThreadCallTree {
    pub thread_id: ThreadId,
    pub thread_name: String,
    pub start_time: u64,
    pub end_time: Option<u64>,
    pub is_active: bool,

    // Thread relationship
    pub parent_thread_id: Option<ThreadId>,
    pub is_main_thread: bool,
    pub creation_function: Option<String>,

    // Call stack
    pub root_node: Option<NodeId>,
    pub current_node: Option<NodeId>,
    pub call_stack: Vec<NodeId>,
    pub max_depth: usize,

    // Statistics
    pub total_calls: u64,
    pub total_errors: u64,
}

/// Thread relationship information
#[derive(Debug, Clone)]
pub struct ThreadRelationship {
    pub parent_thread_id: ThreadId,
    pub child_thread_id: ThreadId,
    pub creation_node_id: NodeId,
    pub creation_function: String,
    pub creation_time: u64,
}

/// Global call tree manager
pub struct CallTreeManager {
    // Thread-safe storage
    nodes: DashMap<NodeId, Arc<RwLock<CallNode>>>,
    fast_nodes: DashMap<NodeId, Arc<RwLock<FastCallNode>>>,
    thread_trees: DashMap<ThreadId, Arc<RwLock<ThreadCallTree>>>,
    thread_relationships: DashMap<ThreadId, ThreadRelationship>,

    // ID generation
    next_node_id: AtomicU32,

    // Global statistics
    total_nodes: AtomicU64,
    total_threads: AtomicU64,
    session_start_time: u64,

    // Pending thread creation tracking
    pending_thread_creations: Arc<Mutex<Vec<PendingThreadCreation>>>,
}

/// Pending thread creation (for pthread_create tracking)
#[derive(Debug)]
struct PendingThreadCreation {
    parent_thread_id: ThreadId,
    creation_node_id: NodeId,
    creation_time: u64,
}

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

impl CallTreeManager {
    /// Create a new call tree manager
    pub fn new() -> Self {
        Self {
            nodes: DashMap::new(),
            fast_nodes: DashMap::new(),
            thread_trees: DashMap::new(),
            thread_relationships: DashMap::new(),
            next_node_id: AtomicU32::new(1),
            total_nodes: AtomicU64::new(0),
            total_threads: AtomicU64::new(0),
            session_start_time: current_timestamp_us(),
            pending_thread_creations: Arc::new(Mutex::new(Vec::new())),
        }
    }

    /// Handle function entry
    pub fn function_enter(
        &self,
        function_address: u64,
        call_site: u64,
        function_info: Option<FunctionInfo>,
        arguments: Vec<CapturedArgument>,
        register_context: Option<RegisterContext>,
    ) -> Result<NodeId> {
        let thread_id = current_thread_id();

        // Ensure thread tree exists
        self.ensure_thread_tree(thread_id);

        // Create new call node
        let node_id = self.next_node_id.fetch_add(1, Ordering::Relaxed);

        let function_name = function_info
            .as_ref()
            .map(|info| crate::intern_string(&info.name))
            .unwrap_or_else(|| crate::intern_string(&crate::format_address(function_address)));

        let node = CallNode {
            id: node_id,
            function_name,
            function_address,
            call_site,
            start_time: current_timestamp_us(),
            end_time: None,
            duration_us: None,
            thread_id,
            call_depth: 0, // Will be set below
            function_info,
            arguments,
            register_context,
            return_value: None,
            parent: None,
            children: Vec::with_capacity(4), // Pre-allocate for typical function call tree
            is_pthread_create: false,
            created_thread_id: None,
            is_completed: false,
        };

        // Insert node
        self.nodes.insert(node_id, Arc::new(RwLock::new(node)));
        self.total_nodes.fetch_add(1, Ordering::Relaxed);

        // Update thread tree
        if let Some(tree_ref) = self.thread_trees.get(&thread_id) {
            let mut tree = tree_ref.write().unwrap();

            // Set call depth
            let call_depth = tree.call_stack.len();
            if let Some(node_ref) = self.nodes.get(&node_id) {
                node_ref.write().unwrap().call_depth = call_depth;
            }

            // Link to parent if exists
            if let Some(parent_id) = tree.current_node {
                if let Some(node_ref) = self.nodes.get(&node_id) {
                    node_ref.write().unwrap().parent = Some(parent_id);
                }

                // Add as child to parent
                if let Some(parent_ref) = self.nodes.get(&parent_id) {
                    parent_ref.write().unwrap().children.push(node_id);
                }
            } else {
                // This is the root node
                tree.root_node = Some(node_id);
            }

            // Update current node and call stack
            tree.current_node = Some(node_id);
            tree.call_stack.push(node_id);
            tree.total_calls += 1;

            if tree.call_stack.len() > tree.max_depth {
                tree.max_depth = tree.call_stack.len();
            }
        }

        Ok(node_id)
    }

    /// Handle function entry - fast path (minimal overhead)
    /// This is optimized for the case where argument capture is disabled
    #[inline]
    pub fn function_enter_fast_path(
        &self,
        function_address: u64,
        call_site: u64,
    ) -> Result<NodeId> {
        let thread_id = current_thread_id();

        // Ensure thread tree exists
        self.ensure_thread_tree(thread_id);

        // Create new fast call node
        let node_id = self.next_node_id.fetch_add(1, Ordering::Relaxed);

        let node = FastCallNode {
            id: node_id,
            function_address,
            call_site,
            start_time: current_timestamp_us(),
            end_time: None,
            thread_id,
            call_depth: 0, // Will be set below
            parent: None,
            children: Vec::with_capacity(4), // Pre-allocate for typical function call tree
            is_completed: false,
        };

        // Insert fast node
        self.fast_nodes.insert(node_id, Arc::new(RwLock::new(node)));
        self.total_nodes.fetch_add(1, Ordering::Relaxed);

        // Update thread tree with minimal locking
        if let Some(tree_ref) = self.thread_trees.get(&thread_id) {
            let mut tree = tree_ref.write().unwrap();

            // Set call depth
            let call_depth = tree.call_stack.len();
            if let Some(node_ref) = self.fast_nodes.get(&node_id) {
                node_ref.write().unwrap().call_depth = call_depth;
            }

            // Link to parent if exists
            if let Some(parent_id) = tree.current_node {
                if let Some(node_ref) = self.fast_nodes.get(&node_id) {
                    node_ref.write().unwrap().parent = Some(parent_id);
                }

                // Add as child to parent (check both fast and full nodes)
                if let Some(parent_ref) = self.fast_nodes.get(&parent_id) {
                    parent_ref.write().unwrap().children.push(node_id);
                } else if let Some(parent_ref) = self.nodes.get(&parent_id) {
                    parent_ref.write().unwrap().children.push(node_id);
                }
            } else {
                // This is the root node
                tree.root_node = Some(node_id);
            }

            // Update current node and call stack
            tree.current_node = Some(node_id);
            tree.call_stack.push(node_id);
            tree.total_calls += 1;

            if tree.call_stack.len() > tree.max_depth {
                tree.max_depth = tree.call_stack.len();
            }
        }

        Ok(node_id)
    }

    /// Handle function exit
    pub fn function_exit(&self, function_address: u64) -> Result<()> {
        self.function_exit_with_return_value(function_address, None)
    }

    /// Handle function exit with return value
    pub fn function_exit_with_return_value(
        &self,
        _function_address: u64,
        return_value: Option<ArgumentValue>,
    ) -> Result<()> {
        let thread_id = current_thread_id();

        if let Some(tree_ref) = self.thread_trees.get(&thread_id) {
            let mut tree = tree_ref.write().unwrap();

            if let Some(current_id) = tree.current_node {
                // Complete the current node
                if let Some(node_ref) = self.nodes.get(&current_id) {
                    let mut node = node_ref.write().unwrap();
                    let now = current_timestamp_us();
                    node.end_time = Some(now);
                    node.duration_us = Some(now.saturating_sub(node.start_time));
                    node.return_value = return_value;
                    node.is_completed = true;
                }

                // Pop from call stack
                tree.call_stack.pop();

                // Update current node to parent
                if let Some(parent_id) = tree.call_stack.last() {
                    tree.current_node = Some(*parent_id);
                } else {
                    tree.current_node = None;
                }
            }
        }

        Ok(())
    }

    /// Handle function exit - fast path (minimal overhead)
    /// This is optimized for the case where argument capture is disabled
    #[inline]
    pub fn function_exit_fast_path(&self, _function_address: u64) -> Result<()> {
        let thread_id = current_thread_id();

        if let Some(tree_ref) = self.thread_trees.get(&thread_id) {
            let mut tree = tree_ref.write().unwrap();

            if let Some(current_id) = tree.current_node {
                // Complete the current node (check both fast and full nodes)
                if let Some(node_ref) = self.fast_nodes.get(&current_id) {
                    let mut node = node_ref.write().unwrap();
                    let now = current_timestamp_us();
                    node.end_time = Some(now);
                    node.is_completed = true;
                } else if let Some(node_ref) = self.nodes.get(&current_id) {
                    let mut node = node_ref.write().unwrap();
                    let now = current_timestamp_us();
                    node.end_time = Some(now);
                    node.duration_us = Some(now.saturating_sub(node.start_time));
                    node.is_completed = true;
                }

                // Pop from call stack
                tree.call_stack.pop();

                // Update current node to parent
                if let Some(parent_id) = tree.call_stack.last() {
                    tree.current_node = Some(*parent_id);
                } else {
                    tree.current_node = None;
                }
            }
        }

        Ok(())
    }

    /// Mark a node as pthread_create and register pending thread creation
    pub fn mark_pthread_create(&self, node_id: NodeId) -> Result<()> {
        if let Some(node_ref) = self.nodes.get(&node_id) {
            let mut node = node_ref.write().unwrap();
            node.is_pthread_create = true;

            // Add to pending thread creations
            let pending = PendingThreadCreation {
                parent_thread_id: node.thread_id,
                creation_node_id: node_id,
                creation_time: current_timestamp_us(),
            };

            self.pending_thread_creations.lock().unwrap().push(pending);
        }

        Ok(())
    }

    /// Handle new thread start (consume pending thread creation)
    pub fn thread_started(&self, new_thread_id: ThreadId) -> Result<()> {
        let mut pending = self.pending_thread_creations.lock().unwrap();

        if let Some(pos) = pending.iter().position(|_| true) {
            let creation = pending.remove(pos);

            // Create thread relationship
            let relationship = ThreadRelationship {
                parent_thread_id: creation.parent_thread_id,
                child_thread_id: new_thread_id,
                creation_node_id: creation.creation_node_id,
                creation_function: crate::string_constants::PTHREAD_CREATE.to_string(),
                creation_time: creation.creation_time,
            };

            self.thread_relationships
                .insert(new_thread_id, relationship);

            // Update the pthread_create node with the actual thread ID
            if let Some(node_ref) = self.nodes.get(&creation.creation_node_id) {
                node_ref.write().unwrap().created_thread_id = Some(new_thread_id);
            }
        }

        Ok(())
    }

    /// Ensure thread tree exists for the given thread
    fn ensure_thread_tree(&self, thread_id: ThreadId) {
        if !self.thread_trees.contains_key(&thread_id) {
            let tree = ThreadCallTree {
                thread_id,
                thread_name: {
                    // Use optimized formatting for thread name with safe fallback
                    if let Ok(result) = std::panic::catch_unwind(|| {
                        crate::FORMAT_BUFFER.with(|buffer| {
                            let mut buf = buffer.borrow_mut();
                            buf.clear();
                            use std::fmt::Write;
                            write!(buf, "Thread-{}", thread_id).unwrap();
                            buf.clone()
                        })
                    }) {
                        result
                    } else {
                        format!("Thread-{}", thread_id)
                    }
                },
                start_time: current_timestamp_us(),
                end_time: None,
                is_active: true,
                parent_thread_id: None,
                is_main_thread: false, // Will be determined later
                creation_function: None,
                root_node: None,
                current_node: None,
                call_stack: Vec::with_capacity(32), // Pre-allocate for typical call stack depth
                max_depth: 0,
                total_calls: 0,
                total_errors: 0,
            };

            self.thread_trees
                .insert(thread_id, Arc::new(RwLock::new(tree)));
            self.total_threads.fetch_add(1, Ordering::Relaxed);

            // Check for pending thread creation
            let _ = self.thread_started(thread_id);
        }
    }

    /// Get statistics
    pub fn get_statistics(&self) -> CallTreeStatistics {
        CallTreeStatistics {
            total_nodes: self.total_nodes.load(Ordering::Relaxed),
            total_threads: self.total_threads.load(Ordering::Relaxed),
            active_threads: self.thread_trees.len() as u64,
            session_duration_us: current_timestamp_us().saturating_sub(self.session_start_time),
        }
    }

    /// Get all thread IDs
    pub fn get_thread_ids(&self) -> Vec<ThreadId> {
        self.thread_trees.iter().map(|entry| *entry.key()).collect()
    }

    /// Get thread tree for a specific thread
    pub fn get_thread_tree(&self, thread_id: ThreadId) -> Option<Arc<RwLock<ThreadCallTree>>> {
        self.thread_trees
            .get(&thread_id)
            .map(|entry| entry.value().clone())
    }

    /// Get node by ID
    pub fn get_node(&self, node_id: NodeId) -> Option<Arc<RwLock<CallNode>>> {
        self.nodes.get(&node_id).map(|entry| entry.value().clone())
    }

    /// Get any node by ID (checks both full nodes and fast nodes)
    /// If found in fast_nodes, converts to full CallNode format for JSON output
    pub fn get_any_node(&self, node_id: NodeId) -> Option<Arc<RwLock<CallNode>>> {
        // First check full nodes
        if let Some(node) = self.nodes.get(&node_id) {
            return Some(node.value().clone());
        }

        // Then check fast nodes and convert
        if let Some(fast_node_ref) = self.fast_nodes.get(&node_id) {
            let fast_node = fast_node_ref.read().unwrap();

            // Convert FastCallNode to CallNode for JSON output
            let full_node = CallNode {
                id: fast_node.id,
                function_name: crate::format_address(fast_node.function_address), // Address-based name
                function_address: fast_node.function_address,
                call_site: fast_node.call_site,
                start_time: fast_node.start_time,
                end_time: fast_node.end_time,
                duration_us: fast_node
                    .end_time
                    .map(|end| end.saturating_sub(fast_node.start_time)),
                thread_id: fast_node.thread_id,
                call_depth: fast_node.call_depth,
                function_info: None,    // Fast nodes don't store function info
                arguments: Vec::new(),  // Fast nodes don't store arguments
                register_context: None, // Fast nodes don't store register context
                return_value: None,     // Fast nodes don't store return values
                parent: fast_node.parent,
                children: fast_node.children.clone(),
                is_pthread_create: false, // Fast nodes don't track pthread
                created_thread_id: None,
                is_completed: fast_node.is_completed,
            };

            return Some(Arc::new(RwLock::new(full_node)));
        }

        None
    }
}

/// Call tree statistics
#[derive(Debug, Clone)]
pub struct CallTreeStatistics {
    pub total_nodes: u64,
    pub total_threads: u64,
    pub active_threads: u64,
    pub session_duration_us: u64,
}

/// Get current timestamp in microseconds
fn current_timestamp_us() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_micros() as u64
}

/// Get current thread ID
fn current_thread_id() -> ThreadId {
    // Use platform-specific thread ID
    #[cfg(target_os = "linux")]
    {
        unsafe { libc::syscall(libc::SYS_gettid) as u64 }
    }

    #[cfg(not(target_os = "linux"))]
    {
        // Fallback to Rust thread ID hash
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        thread::current().id().hash(&mut hasher);
        hasher.finish()
    }
}