memscope_rs/core/

scope_tracker.rs

//! Scope tracking functionality for memory analysis

use crate::core::optimized_locks::OptimizedMutex;
use crate::core::sharded_locks::ShardedRwLock;
use crate::core::types::*;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// Global scope tracker instance
static GLOBAL_SCOPE_TRACKER: OnceLock<Arc<ScopeTracker>> = OnceLock::new();

/// Get the global scope tracker instance
pub fn get_global_scope_tracker() -> Arc<ScopeTracker> {
    GLOBAL_SCOPE_TRACKER
        .get_or_init(|| Arc::new(ScopeTracker::new()))
        .clone()
}

/// Unique identifier for scopes
pub type ScopeId = u64;

/// Core scope tracking functionality
pub struct ScopeTracker {
    /// Active scopes using sharded storage for better concurrency
    pub active_scopes: ShardedRwLock<ScopeId, ScopeInfo>,
    /// Completed scopes for analysis
    pub completed_scopes: OptimizedMutex<Vec<ScopeInfo>>,
    /// Scope hierarchy relationships
    pub scope_hierarchy: OptimizedMutex<ScopeHierarchy>,
    /// Next available scope ID using atomic counter
    next_scope_id: AtomicU64,
    /// Current scope stack per thread using sharded storage
    pub scope_stack: ShardedRwLock<String, Vec<ScopeId>>,
}

impl ScopeTracker {
    /// Create a new scope tracker
    pub fn new() -> Self {
        Self {
            active_scopes: ShardedRwLock::new(),
            completed_scopes: OptimizedMutex::new(Vec::new()),
            scope_hierarchy: OptimizedMutex::new(ScopeHierarchy {
                root_scopes: Vec::new(),
                scope_tree: HashMap::new(),
                max_depth: 0,
                total_scopes: 0,
                relationships: HashMap::new(),
                depth_map: HashMap::new(),
            }),
            next_scope_id: AtomicU64::new(1),
            scope_stack: ShardedRwLock::new(),
        }
    }

    /// Enter a new scope
    pub fn enter_scope(&self, name: String) -> TrackingResult<ScopeId> {
        let scope_id = self.allocate_scope_id();
        let thread_id = format!("{:?}", std::thread::current().id());
        let timestamp = current_timestamp();

        // Determine parent scope and depth
        let (parent_scope, depth) = {
            if let Some(thread_stack) = self.scope_stack.get(&thread_id) {
                if let Some(&parent_id) = thread_stack.last() {
                    if let Some(parent) = self.active_scopes.get(&parent_id) {
                        (Some(parent.name.clone()), parent.depth + 1)
                    } else {
                        (None, 0)
                    }
                } else {
                    (None, 0)
                }
            } else {
                (None, 0)
            }
        };

        // Create scope info
        let scope_info = ScopeInfo {
            name: name.clone(),
            parent: parent_scope.clone(),
            children: Vec::new(),
            depth,
            variables: Vec::new(),
            total_memory: 0,
            peak_memory: 0,
            allocation_count: 0,
            lifetime_start: Some(timestamp as u64),
            lifetime_end: None,
            is_active: true,
            start_time: timestamp as u64,
            end_time: None,
            memory_usage: 0,
            child_scopes: Vec::new(),
            parent_scope: parent_scope.clone(),
        };

        // Add to active scopes
        self.active_scopes.insert(scope_id, scope_info);

        // Update scope stack
        self.scope_stack.with_shard_write(&thread_id, |shard| {
            shard.entry(thread_id.clone()).or_default().push(scope_id);
        });

        // Update hierarchy
        if let Some(mut hierarchy) = self.scope_hierarchy.try_lock() {
            hierarchy.depth_map.insert(name.clone(), depth);

            if let Some(parent) = parent_scope.clone() {
                hierarchy
                    .relationships
                    .entry(parent)
                    .or_default()
                    .push(name.clone());
            } else {
                hierarchy.root_scopes.push(name);
            }
        }

        Ok(scope_id)
    }

    /// Exit a scope
    pub fn exit_scope(&self, scope_id: ScopeId) -> TrackingResult<()> {
        let thread_id = format!("{:?}", std::thread::current().id());
        let timestamp = current_timestamp();

        // Remove from active scopes and get scope info
        let mut scope_info = self.active_scopes.remove(&scope_id).ok_or_else(|| {
            TrackingError::InvalidPointer(format!("Invalid scope ID: {scope_id}"))
        })?;

        // Update end time
        scope_info.end_time = Some(timestamp as u64);
        scope_info.lifetime_end = Some(timestamp as u64);

        // Update scope stack
        self.scope_stack.with_shard_write(&thread_id, |shard| {
            if let Some(thread_stack) = shard.get_mut(&thread_id) {
                if let Some(pos) = thread_stack.iter().position(|&id| id == scope_id) {
                    thread_stack.remove(pos);
                }
            }
        });

        // Add to completed scopes
        let mut retries = 0;
        const MAX_RETRIES: u32 = 10;
        const RETRY_DELAY: Duration = Duration::from_millis(10);

        while retries < MAX_RETRIES {
            if let Some(mut completed_scopes) = self.completed_scopes.try_lock() {
                completed_scopes.push(scope_info);
                break;
            }

            // If we couldn't get the lock, wait a bit and try again
            retries += 1;
            if retries < MAX_RETRIES {
                std::thread::sleep(RETRY_DELAY);
            } else {
                // If we've exhausted our retries, log an error and drop the scope
                // This is better than silently losing the scope
                eprintln!(
                    "Failed to add scope {scope_id} to completed_scopes after {MAX_RETRIES} retries",
                );
            }
        }

        Ok(())
    }

    /// Associate a variable with the current scope
    pub fn associate_variable(
        &self,
        variable_name: String,
        memory_size: usize,
    ) -> TrackingResult<()> {
        let thread_id = format!("{:?}", std::thread::current().id());

        // Find current scope for this thread
        let current_scope_id = self
            .scope_stack
            .get(&thread_id)
            .and_then(|stack| stack.last().copied());

        if let Some(scope_id) = current_scope_id {
            self.active_scopes.with_shard_write(&scope_id, |shard| {
                if let Some(scope) = shard.get_mut(&scope_id) {
                    scope.variables.push(variable_name);
                    scope.memory_usage += memory_size;
                    scope.peak_memory = scope.peak_memory.max(scope.memory_usage);
                    scope.allocation_count += 1;
                }
            });
        }

        Ok(())
    }

    /// Get current scope analysis
    pub fn get_scope_analysis(&self) -> TrackingResult<ScopeAnalysis> {
        let mut all_scopes = Vec::new();

        // Collect active scopes from all shards
        for i in 0..16 {
            // Default shard count
            self.active_scopes.with_shard_read(&i, |shard| {
                all_scopes.extend(shard.values().cloned());
            });
        }

        // Add completed scopes
        if let Some(completed_scopes) = self.completed_scopes.try_lock() {
            all_scopes.extend(completed_scopes.iter().cloned());
        }

        let hierarchy = if let Some(hierarchy) = self.scope_hierarchy.try_lock() {
            hierarchy.clone()
        } else {
            // Return default hierarchy if lock fails
            ScopeHierarchy {
                root_scopes: Vec::new(),
                scope_tree: HashMap::new(),
                max_depth: 0,
                total_scopes: 0,
                relationships: HashMap::new(),
                depth_map: HashMap::new(),
            }
        };

        Ok(ScopeAnalysis {
            total_scopes: all_scopes.len(),
            active_scopes: all_scopes.iter().filter(|s| s.is_active).count(),
            max_depth: hierarchy.max_depth,
            average_lifetime: 1000.0, // Placeholder
            memory_efficiency: 0.8,   // Placeholder
            scopes: all_scopes,
            scope_hierarchy: hierarchy,
            cross_scope_references: Vec::new(), // Cross-scope reference tracking implementation
        })
    }

    /// Get scope lifecycle metrics
    pub fn get_scope_lifecycle_metrics(&self) -> TrackingResult<Vec<ScopeLifecycleMetrics>> {
        let metrics = if let Some(completed_scopes) = self.completed_scopes.try_lock() {
            completed_scopes
                .iter()
                .map(|scope| {
                    let lifetime =
                        scope.end_time.unwrap_or(current_timestamp() as u64) - scope.start_time;
                    let efficiency_score = if scope.peak_memory > 0 {
                        scope.memory_usage as f64 / scope.peak_memory as f64
                    } else {
                        1.0
                    };

                    ScopeLifecycleMetrics {
                        scope_name: scope.name.clone(),
                        variable_count: scope.variables.len(),
                        average_lifetime_ms: lifetime as f64,
                        total_memory_usage: scope.memory_usage,
                        peak_memory_usage: scope.peak_memory,
                        allocation_frequency: 1.0, // Simplified
                        deallocation_efficiency: efficiency_score,
                        completed_allocations: scope.allocation_count,
                        memory_growth_events: 0,
                        peak_concurrent_variables: scope.variables.len(),
                        memory_efficiency_ratio: if scope.peak_memory > 0 {
                            scope.memory_usage as f64 / scope.peak_memory as f64
                        } else {
                            1.0
                        },
                        ownership_transfer_events: 0,
                        fragmentation_score: 0.0,
                        instant_allocations: 0,
                        short_term_allocations: 0,
                        medium_term_allocations: 0,
                        long_term_allocations: 0,
                        suspected_leaks: 0,
                        risk_distribution: crate::core::types::RiskDistribution::default(),
                        scope_metrics: Vec::new(),
                        type_lifecycle_patterns: Vec::new(),
                    }
                })
                .collect()
        } else {
            Vec::new()
        };

        Ok(metrics)
    }

    /// Get all scopes (active and completed) for data localization
    pub fn get_all_scopes(&self) -> Vec<ScopeInfo> {
        let mut all_scopes = Vec::new();

        // Add active scopes from all shards
        for i in 0..16 {
            // Default shard count
            self.active_scopes.with_shard_read(&i, |shard| {
                all_scopes.extend(shard.values().cloned());
            });
        }

        // Add completed scopes
        if let Some(completed_scopes) = self.completed_scopes.try_lock() {
            all_scopes.extend(completed_scopes.iter().cloned());
        }

        all_scopes
    }

    /// Allocate a new unique scope ID using atomic operations
    fn allocate_scope_id(&self) -> ScopeId {
        self.next_scope_id.fetch_add(1, Ordering::Relaxed)
    }
}

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

/// RAII scope guard for automatic scope management
pub struct ScopeGuard {
    scope_id: ScopeId,
    tracker: Arc<ScopeTracker>,
}

impl ScopeGuard {
    /// Enter a new scope with automatic cleanup
    pub fn enter(name: &str) -> TrackingResult<Self> {
        let tracker = get_global_scope_tracker();
        let scope_id = tracker.enter_scope(name.to_string())?;

        Ok(Self { scope_id, tracker })
    }

    /// Get the scope ID
    pub fn scope_id(&self) -> ScopeId {
        self.scope_id
    }
}

impl Drop for ScopeGuard {
    fn drop(&mut self) {
        let _ = self.tracker.exit_scope(self.scope_id);
    }
}
fn current_timestamp() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis())
        .unwrap_or(0)
}

/// Macro for tracking scopes with automatic cleanup
#[macro_export]
macro_rules! track_scope {
    ($scope_name:expr) => {
        let _scope_guard = $crate::scope_tracker::ScopeGuard::enter($scope_name)?;
    };
    ($scope_name:expr, $block:block) => {{
        let _scope_guard = $crate::scope_tracker::ScopeGuard::enter($scope_name)?;
        $block
    }};
}

/// Enhanced track_var macro that also associates with current scope
#[macro_export]
macro_rules! track_var_with_scope {
    ($var:ident) => {{
        // Track the variable normally
        let result = $crate::_track_var_impl(&$var, stringify!($var));

        // Also associate with current scope
        if result.is_ok() {
            if let Some(size) = $crate::Trackable::get_heap_ptr(&$var) {
                let scope_tracker = $crate::scope_tracker::get_global_scope_tracker();
                let _ = scope_tracker
                    .associate_variable(stringify!($var).to_string(), std::mem::size_of_val(&$var));
            }
        }

        result
    }};
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::{thread, time::Duration};

    #[test]
    fn test_scope_tracker_creation() {
        let tracker = ScopeTracker::new();
        assert_eq!(tracker.next_scope_id.load(Ordering::SeqCst), 1);
    }

    #[test]
    fn test_enter_and_exit_scope() {
        let tracker = ScopeTracker::new();
        let scope_id = tracker.enter_scope("test_scope".to_string()).unwrap();

        // Check scope is active
        assert!(tracker.active_scopes.get(&scope_id).is_some());

        // Exit scope
        tracker.exit_scope(scope_id).unwrap();

        // Check scope is no longer active
        assert!(tracker.active_scopes.get(&scope_id).is_none());
    }

    #[test]
    fn test_scope_hierarchy() {
        let tracker = ScopeTracker::new();

        // Enter parent scope
        let parent_id = tracker.enter_scope("parent".to_string()).unwrap();

        // Get the thread ID for checking the scope stack
        let thread_id = format!("{:?}", std::thread::current().id());

        // Check that parent scope is in the active scopes
        assert!(tracker.active_scopes.get(&parent_id).is_some());

        // Check that parent scope is in the thread's scope stack
        let scope_stack = tracker.scope_stack.get(&thread_id).unwrap();
        assert_eq!(scope_stack.last(), Some(&parent_id));

        // Enter child scope
        let child_id = tracker.enter_scope("child".to_string()).unwrap();

        // Check that child scope is in the active scopes
        assert!(tracker.active_scopes.get(&child_id).is_some());

        // Check that child scope is now at the top of the thread's scope stack
        let scope_stack = tracker.scope_stack.get(&thread_id).unwrap();
        assert_eq!(scope_stack.last(), Some(&child_id));

        // Check that parent scope is still in the stack
        assert!(scope_stack.contains(&parent_id));

        // Clean up - exit child scope first
        tracker.exit_scope(child_id).unwrap();

        // Check that parent scope is now at the top of the stack again
        let scope_stack = tracker.scope_stack.get(&thread_id).unwrap();
        assert_eq!(scope_stack.last(), Some(&parent_id));

        // Exit parent scope
        tracker.exit_scope(parent_id).unwrap();

        // Check that the stack is now empty
        let scope_stack = tracker.scope_stack.get(&thread_id);
        assert!(scope_stack.is_none() || scope_stack.unwrap().is_empty());
    }

    #[test]
    fn test_variable_association() {
        let tracker = ScopeTracker::new();
        let scope_id = tracker.enter_scope("test_scope".to_string()).unwrap();

        // Associate variable
        tracker
            .associate_variable("test_var".to_string(), 1024)
            .unwrap();

        // Check variable was associated
        let scope = tracker.active_scopes.get(&scope_id).unwrap();
        assert_eq!(scope.variables.len(), 1);
        assert_eq!(scope.variables[0], "test_var");
        assert_eq!(scope.memory_usage, 1024);

        tracker.exit_scope(scope_id).unwrap();
    }

    #[test]
    fn test_scope_guard() {
        // Get a fresh tracker for this test
        let tracker = Arc::new(ScopeTracker::new());

        // Set this tracker as the global instance for this test
        let _ = GLOBAL_SCOPE_TRACKER.set(Arc::clone(&tracker));

        let scope_id;
        let start_time;

        // Use a block to ensure the guard is dropped
        {
            let _guard = ScopeGuard::enter("guarded_scope").unwrap();
            scope_id = _guard.scope_id();

            // Scope should be active
            assert!(tracker.active_scopes.get(&scope_id).is_some());

            // Check the scope is marked as active in the ScopeInfo
            let scope = tracker.active_scopes.get(&scope_id).unwrap();
            assert!(scope.is_active);
            assert!(scope.end_time.is_none());

            // Save the start time for later comparison
            start_time = scope.start_time;

            // Associate variable
            tracker
                .associate_variable("guard_var".to_string(), 512)
                .unwrap();

            // Verify the variable was associated
            let scope = tracker.active_scopes.get(&scope_id).unwrap();
            assert_eq!(scope.variables.len(), 1);
            assert_eq!(scope.variables[0], "guard_var");
            assert!(scope.memory_usage >= 512);
        } // Guard drops here, scope should exit automatically

        // Scope should no longer be active
        assert!(tracker.active_scopes.get(&scope_id).is_none());

        // Check that the scope was moved to completed_scopes
        let completed_scopes = tracker.completed_scopes.lock();
        assert!(!completed_scopes.is_empty());

        // Find our completed scope
        let completed_scope = completed_scopes.last().expect("No completed scopes found");

        // Verify the completed scope has the correct properties
        assert_eq!(completed_scope.name, "guarded_scope");
        assert!(completed_scope.end_time.is_some(), "End time should be set");

        // Compare the end_time with the saved start_time
        assert!(
            completed_scope.end_time.unwrap() >= start_time,
            "End time should be after start time"
        );

        assert!(
            completed_scope.memory_usage >= 512,
            "Memory usage should be at least 512 bytes"
        );

        // Note: The is_active flag is not currently updated in exit_scope,
        // so we won't check it. The scope is considered completed based on end_time being set.
    }

    #[test]
    fn test_concurrent_scope_operations() {
        // Use a single tracker for all threads
        let tracker = Arc::new(ScopeTracker::new());
        let num_threads = 5;

        // We'll collect the scope names that were actually created and completed
        let created_scopes = Arc::new(std::sync::Mutex::new(Vec::new()));
        let completed_scopes_copy = Arc::new(std::sync::Mutex::new(Vec::new()));

        // Create a barrier to ensure all threads start at the same time
        let start_barrier = Arc::new(std::sync::Barrier::new(num_threads));

        // Spawn worker threads
        let handles = (0..num_threads)
            .map(|i| {
                let tracker = Arc::clone(&tracker);
                let start_barrier = start_barrier.clone();
                let created_scopes = Arc::clone(&created_scopes);
                let completed_scopes_copy = Arc::clone(&completed_scopes_copy);

                thread::spawn(move || {
                    // Wait for all threads to be ready
                    start_barrier.wait();

                    // Each thread creates its own scope
                    let scope_name = format!("thread_{i}");
                    let scope_id = tracker.enter_scope(scope_name.clone()).unwrap();

                    // Verify the scope was created
                    assert!(tracker.active_scopes.get(&scope_id).is_some());

                    // Store the created scope name
                    created_scopes.lock().unwrap().push(scope_name.clone());

                    // Simulate some work
                    thread::sleep(Duration::from_millis(10));

                    // Associate a variable
                    let var_name = format!("var_{i}");
                    tracker.associate_variable(var_name, 256).unwrap();

                    // Verify the variable was associated
                    let scope = tracker.active_scopes.get(&scope_id).unwrap();
                    assert_eq!(scope.variables.len(), 1);

                    // Exit the scope
                    tracker.exit_scope(scope_id).unwrap();

                    // Store the completed scope name
                    completed_scopes_copy.lock().unwrap().push(scope_name);

                    // Return the scope ID for verification
                    scope_id
                })
            })
            .collect::<Vec<_>>();

        // Wait for all threads to complete and collect their results
        let _: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();

        // Get the created and completed scopes
        let created = created_scopes.lock().unwrap();
        let completed = completed_scopes_copy.lock().unwrap();

        // Verify all scopes were created and completed
        assert_eq!(created.len(), num_threads);
        assert_eq!(completed.len(), num_threads);

        // Get the completed scopes from the tracker
        let tracker_scopes = tracker.completed_scopes.lock();

        // Verify we have a reasonable number of completed scopes in the tracker
        // We use >= to account for any potential race conditions where some scopes might be lost
        // due to lock contention, but we still want to ensure most scopes were processed
        assert!(
            !tracker_scopes.is_empty(),
            "No scopes were completed in the tracker"
        );

        // Log the number of completed scopes for debugging
        println!(
            "Completed {} out of {} scopes in the tracker",
            tracker_scopes.len(),
            num_threads
        );

        // Get the scope names from the tracker
        let mut tracker_scope_names: Vec<String> =
            tracker_scopes.iter().map(|s| s.name.clone()).collect();

        // Sort for consistent comparison
        let mut expected_names = created.clone();
        expected_names.sort();
        tracker_scope_names.sort();

        // Verify all expected scopes are in the tracker
        assert_eq!(
            expected_names, tracker_scope_names,
            "Expected scopes in tracker: {expected_names:?}, but found: {tracker_scope_names:?}",
        );

        // Verify all scopes in the tracker have proper end times and are not active
        for scope in tracker_scopes.iter() {
            assert!(
                scope.end_time.is_some(),
                "Scope {} has no end time",
                scope.name
            );
            assert!(
                scope.end_time.unwrap() >= scope.start_time,
                "Scope {} has end time before start time",
                scope.name
            );

            // Verify the scope is in our completed list
            assert!(
                completed.contains(&scope.name),
                "Scope {} not found in completed list",
                scope.name
            );
        }
    }
}