memscope_rs/unified/strategies/

async_strategy.rs

// Async Memory Tracking Strategy
// Optimized implementation for async/await applications
// Uses task-local storage and async context awareness

use crate::unified::tracking_dispatcher::{
    MemoryTracker, TrackerConfig, TrackerError, TrackerStatistics, TrackerType,
};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tracing::{debug, info, warn};

/// Async memory tracking strategy
/// Specialized for async/await applications with task-local tracking
/// Provides context-aware tracking across async task boundaries
pub struct AsyncStrategy {
    /// Configuration for this strategy instance
    config: Option<TrackerConfig>,
    /// Global async tracking state
    global_state: Arc<AsyncGlobalState>,
    /// Task registry for managing async tasks
    task_registry: Arc<AsyncTaskRegistry>,
    /// Aggregated metrics across all tasks
    global_metrics: Arc<AsyncGlobalMetrics>,
}

/// Global state for async tracking coordination
/// Maintains task-level coordination without blocking async execution
#[derive(Debug)]
struct AsyncGlobalState {
    /// Whether async tracking is currently active
    is_active: AtomicU64, // 0 = inactive, 1 = active
    /// Total number of active tracking tasks
    active_tasks: AtomicUsize,
    /// Session start timestamp
    session_start_ns: AtomicU64,
    /// Next allocation ID (globally unique across tasks)
    next_allocation_id: AtomicU64,
    /// Next task ID for task identification
    next_task_id: AtomicU64,
}

/// Registry for async tasks participating in tracking
/// Manages task lifecycle and coordination
#[derive(Debug)]
struct AsyncTaskRegistry {
    /// Total tasks that have registered for tracking
    total_registered_tasks: AtomicUsize,
    /// Currently active tracking tasks
    active_tracking_tasks: AtomicUsize,
    /// Peak concurrent tasks
    peak_concurrent_tasks: AtomicUsize,
}

/// Global metrics aggregated from all async tasks
/// Provides system-wide view of async memory tracking
#[derive(Debug)]
struct AsyncGlobalMetrics {
    /// Total allocations across all tasks
    total_allocations: AtomicU64,
    /// Total bytes allocated across all tasks
    total_bytes_allocated: AtomicU64,
    /// Total task spawns tracked
    total_task_spawns: AtomicU64,
    /// Average task lifetime (nanoseconds)
    _avg_task_lifetime_ns: AtomicU64,
    /// Total tracking overhead in bytes
    total_overhead_bytes: AtomicUsize,
}

impl Default for AsyncGlobalState {
    /// Initialize async global state with inactive values
    fn default() -> Self {
        Self {
            is_active: AtomicU64::new(0),
            active_tasks: AtomicUsize::new(0),
            session_start_ns: AtomicU64::new(0),
            next_allocation_id: AtomicU64::new(1),
            next_task_id: AtomicU64::new(1),
        }
    }
}

impl Default for AsyncTaskRegistry {
    /// Initialize task registry with zero counts
    fn default() -> Self {
        Self {
            total_registered_tasks: AtomicUsize::new(0),
            active_tracking_tasks: AtomicUsize::new(0),
            peak_concurrent_tasks: AtomicUsize::new(0),
        }
    }
}

impl Default for AsyncGlobalMetrics {
    /// Initialize async metrics with zero values
    fn default() -> Self {
        Self {
            total_allocations: AtomicU64::new(0),
            total_bytes_allocated: AtomicU64::new(0),
            total_task_spawns: AtomicU64::new(0),
            _avg_task_lifetime_ns: AtomicU64::new(0),
            total_overhead_bytes: AtomicUsize::new(0),
        }
    }
}

impl AsyncStrategy {
    /// Create new async strategy instance
    /// Initializes async coordination structures
    pub fn new() -> Self {
        debug!("Creating new async strategy");

        Self {
            config: None,
            global_state: Arc::new(AsyncGlobalState::default()),
            task_registry: Arc::new(AsyncTaskRegistry::default()),
            global_metrics: Arc::new(AsyncGlobalMetrics::default()),
        }
    }

    /// Register current async task for tracking
    /// Should be called when entering async context
    pub fn register_current_task(&self) -> Result<u64, TrackerError> {
        let task_id = self
            .global_state
            .next_task_id
            .fetch_add(1, Ordering::Relaxed);

        debug!("Registering async task for tracking: id={}", task_id);

        // Update task registry
        self.task_registry
            .total_registered_tasks
            .fetch_add(1, Ordering::Relaxed);
        let current_active = self
            .task_registry
            .active_tracking_tasks
            .fetch_add(1, Ordering::Relaxed)
            + 1;

        // Update peak concurrent tasks
        let current_peak = self
            .task_registry
            .peak_concurrent_tasks
            .load(Ordering::Relaxed);
        if current_active > current_peak {
            self.task_registry
                .peak_concurrent_tasks
                .store(current_active, Ordering::Relaxed);
        }

        // Update global metrics
        self.global_metrics
            .total_task_spawns
            .fetch_add(1, Ordering::Relaxed);

        info!(
            "Async task registered: id={}, active_tasks={}",
            task_id, current_active
        );
        Ok(task_id)
    }

    /// Get high-precision timestamp in nanoseconds
    fn get_timestamp_ns() -> u64 {
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos() as u64)
            .unwrap_or(0)
    }

    /// Export collected data as JSON compatible with existing format
    fn export_as_json(&self) -> Result<String, TrackerError> {
        // Note: In a complete implementation, this would collect data from
        // task-local storage across all async tasks. For this demo, we return
        // basic structure with metadata.

        let mut output = serde_json::Map::new();
        output.insert("allocations".to_string(), serde_json::Value::Array(vec![]));
        output.insert("strategy_metadata".to_string(), serde_json::json!({
            "strategy_type": "async",
            "total_allocations": self.global_metrics.total_allocations.load(Ordering::Relaxed),
            "total_bytes": self.global_metrics.total_bytes_allocated.load(Ordering::Relaxed),
            "total_tasks": self.task_registry.total_registered_tasks.load(Ordering::Relaxed),
            "peak_concurrent_tasks": self.task_registry.peak_concurrent_tasks.load(Ordering::Relaxed),
            "overhead_bytes": self.global_metrics.total_overhead_bytes.load(Ordering::Relaxed)
        }));

        serde_json::to_string_pretty(&output).map_err(|e| TrackerError::DataCollectionFailed {
            reason: format!("JSON serialization failed: {}", e),
        })
    }
}

impl MemoryTracker for AsyncStrategy {
    /// Initialize strategy with provided configuration
    fn initialize(&mut self, config: TrackerConfig) -> Result<(), TrackerError> {
        debug!("Initializing async strategy with config: {:?}", config);

        // Validate configuration
        if config.sample_rate < 0.0 || config.sample_rate > 1.0 {
            return Err(TrackerError::InvalidConfiguration {
                reason: "Sample rate must be between 0.0 and 1.0".to_string(),
            });
        }

        // Store configuration
        self.config = Some(config);

        // Reset global state
        self.global_state.is_active.store(0, Ordering::Relaxed);
        self.global_state.active_tasks.store(0, Ordering::Relaxed);
        self.global_state
            .next_allocation_id
            .store(1, Ordering::Relaxed);
        self.global_state.next_task_id.store(1, Ordering::Relaxed);

        // Reset metrics
        self.global_metrics
            .total_allocations
            .store(0, Ordering::Relaxed);
        self.global_metrics
            .total_bytes_allocated
            .store(0, Ordering::Relaxed);
        self.global_metrics
            .total_task_spawns
            .store(0, Ordering::Relaxed);

        info!("Async strategy initialized successfully");
        Ok(())
    }

    /// Start active memory tracking for async contexts
    fn start_tracking(&mut self) -> Result<(), TrackerError> {
        debug!("Starting async tracking");

        let was_active = self.global_state.is_active.swap(1, Ordering::Relaxed);
        if was_active == 1 {
            warn!("Async tracking was already active");
            return Ok(());
        }

        // Record session start time
        self.global_state
            .session_start_ns
            .store(Self::get_timestamp_ns(), Ordering::Relaxed);

        info!("Async tracking started successfully");
        Ok(())
    }

    /// Stop tracking and collect data from all async tasks
    fn stop_tracking(&mut self) -> Result<Vec<u8>, TrackerError> {
        debug!("Stopping async tracking");

        let was_active = self.global_state.is_active.swap(0, Ordering::Relaxed);
        if was_active == 0 {
            warn!("Async tracking was not active");
        }

        // Export data
        let json_data = self.export_as_json()?;

        let total_allocations = self
            .global_metrics
            .total_allocations
            .load(Ordering::Relaxed);
        let total_tasks = self
            .task_registry
            .total_registered_tasks
            .load(Ordering::Relaxed);

        info!(
            "Async tracking stopped: {} allocations from {} tasks",
            total_allocations, total_tasks
        );

        Ok(json_data.into_bytes())
    }

    /// Get current tracking statistics
    fn get_statistics(&self) -> TrackerStatistics {
        TrackerStatistics {
            allocations_tracked: self
                .global_metrics
                .total_allocations
                .load(Ordering::Relaxed),
            memory_tracked_bytes: self
                .global_metrics
                .total_bytes_allocated
                .load(Ordering::Relaxed),
            overhead_bytes: self
                .global_metrics
                .total_overhead_bytes
                .load(Ordering::Relaxed) as u64,
            tracking_duration_ms: {
                let start_ns = self.global_state.session_start_ns.load(Ordering::Relaxed);
                if start_ns > 0 {
                    (Self::get_timestamp_ns() - start_ns) / 1_000_000
                } else {
                    0
                }
            },
        }
    }

    /// Check if strategy is currently active
    fn is_active(&self) -> bool {
        self.global_state.is_active.load(Ordering::Relaxed) == 1
    }

    /// Get strategy type identifier
    fn tracker_type(&self) -> TrackerType {
        TrackerType::AsyncTracker
    }
}

impl Default for AsyncStrategy {
    /// Create strategy with default configuration
    fn default() -> Self {
        Self::new()
    }
}