foxhunt/services/ml_training_service/tests/stress_memory_leak.rs

//! Memory Leak Stress Tests
//!
//! Long-running tests to detect memory leaks in the ML Training Service.
//! Validates memory stability over extended runtime and high allocation churn.

use anyhow::Result;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use sysinfo::{ProcessRefreshKind, RefreshKind, System};

use ml_training_service::orchestrator::TrainingJob;
use ml::training_pipeline::ProductionTrainingConfig;

/// Helper to create a minimal training config
fn create_test_config() -> ProductionTrainingConfig {
    ProductionTrainingConfig::default()
}

/// Helper to get current process memory usage in bytes
fn get_process_memory_usage() -> Option<u64> {
    let mut system = System::new_with_specifics(
        RefreshKind::new().with_processes(ProcessRefreshKind::everything())
    );
    system.refresh_processes();

    let pid = sysinfo::get_current_pid().ok()?;
    Some(system.process(pid)?.memory())
}

/// Test 1: Long-Running Service (1 Hour Runtime)
///
/// Monitors memory usage over 1 hour of continuous operation.
/// In stress test mode, runs for 5 minutes instead.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_long_running_service() -> Result<()> {
    println!("\n=== Test 1: Long-Running Service ===");
    println!("  NOTE: Running 5-minute version (1-hour version requires --features=long-stress)");

    let test_duration = Duration::from_secs(300); // 5 minutes
    let sample_interval = Duration::from_secs(10); // Sample every 10s

    let start = Instant::now();
    let initial_memory = get_process_memory_usage().unwrap_or(0);

    println!("  Initial memory: {} MB", initial_memory / 1_048_576);

    let mut memory_samples = Vec::new();
    memory_samples.push(initial_memory);

    // Run continuous operations for test duration
    let operations_count = Arc::new(AtomicU64::new(0));

    let ops_count = operations_count.clone();
    let worker = tokio::spawn(async move {
        loop {
            // Perform work (create jobs, simulate operations)
            let config = create_test_config();
            let _job = TrainingJob::new(
                "DQN".to_string(),
                config,
                "Long running test".to_string(),
                std::collections::HashMap::new(),
            );

            ops_count.fetch_add(1, Ordering::Relaxed);
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    });

    // Monitor memory usage
    while start.elapsed() < test_duration {
        tokio::time::sleep(sample_interval).await;

        if let Some(memory) = get_process_memory_usage() {
            memory_samples.push(memory);
            let elapsed = start.elapsed();
            println!("  [{:3}s] Memory: {} MB, Operations: {}",
                elapsed.as_secs(),
                memory / 1_048_576,
                operations_count.load(Ordering::Relaxed)
            );
        }
    }

    // Stop worker
    worker.abort();

    let final_memory = get_process_memory_usage().unwrap_or(0);
    let memory_growth = final_memory.saturating_sub(initial_memory);
    let growth_percent = (memory_growth as f64 / initial_memory as f64) * 100.0;

    let total_ops = operations_count.load(Ordering::Relaxed);

    println!("✓ Test 1 Results:");
    println!("  - Duration: {:?}", test_duration);
    println!("  - Initial memory: {} MB", initial_memory / 1_048_576);
    println!("  - Final memory: {} MB", final_memory / 1_048_576);
    println!("  - Memory growth: {} MB ({:.2}%)", memory_growth / 1_048_576, growth_percent);
    println!("  - Total operations: {}", total_ops);
    println!("  - Target: <1% growth per hour");

    // For 5-minute test, allow proportional growth: 1% / 12 ≈ 0.1%
    let max_allowed_growth = 0.2; // 0.2% for 5 minutes
    assert!(growth_percent < max_allowed_growth,
        "Memory growth {:.2}% exceeds {:.2}% target", growth_percent, max_allowed_growth);

    Ok(())
}

/// Test 2: Allocate/Deallocate 10K Jobs
///
/// Creates and destroys 10,000 jobs to test allocation churn.
/// Validates that memory returns to baseline after cleanup.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_allocate_deallocate_10k_jobs() -> Result<()> {
    println!("\n=== Test 2: Allocate/Deallocate 10K Jobs ===");

    let initial_memory = get_process_memory_usage().unwrap_or(0);
    println!("  Initial memory: {} MB", initial_memory / 1_048_576);

    // Force garbage collection baseline
    tokio::time::sleep(Duration::from_millis(100)).await;

    let allocation_start = Instant::now();
    let mut jobs = Vec::new();

    // Allocate 10,000 jobs
    println!("  Allocating 10,000 jobs...");
    for i in 0..10000 {
        let config = create_test_config();
        let job = TrainingJob::new(
            "DQN".to_string(),
            config,
            format!("Alloc test {}", i),
            std::collections::HashMap::new(),
        );
        jobs.push(job);

        if (i + 1) % 2000 == 0 {
            let mem = get_process_memory_usage().unwrap_or(0);
            println!("    {} jobs allocated, memory: {} MB", i + 1, mem / 1_048_576);
        }
    }

    let after_alloc_memory = get_process_memory_usage().unwrap_or(0);
    let allocation_time = allocation_start.elapsed();

    println!("  ✓ Allocation complete in {:?}", allocation_time);
    println!("  Memory after allocation: {} MB", after_alloc_memory / 1_048_576);

    // Deallocate all jobs
    println!("  Deallocating jobs...");
    let dealloc_start = Instant::now();
    jobs.clear();

    // Give time for cleanup
    tokio::time::sleep(Duration::from_millis(500)).await;

    let after_dealloc_memory = get_process_memory_usage().unwrap_or(0);
    let deallocation_time = dealloc_start.elapsed();

    println!("  ✓ Deallocation complete in {:?}", deallocation_time);
    println!("  Memory after deallocation: {} MB", after_dealloc_memory / 1_048_576);

    let memory_leaked = after_dealloc_memory.saturating_sub(initial_memory);
    let leak_percent = (memory_leaked as f64 / initial_memory as f64) * 100.0;

    println!("✓ Test 2 Results:");
    println!("  - Jobs allocated: 10,000");
    println!("  - Allocation time: {:?}", allocation_time);
    println!("  - Deallocation time: {:?}", deallocation_time);
    println!("  - Memory leaked: {} MB ({:.2}%)", memory_leaked / 1_048_576, leak_percent);
    println!("  - Target: <5% leak after cleanup");

    assert!(leak_percent < 5.0, "Memory leak {:.2}% exceeds 5% threshold", leak_percent);

    Ok(())
}

/// Test 3: gRPC Stream Cleanup Validation
///
/// Tests that gRPC streams are properly cleaned up and don't leak memory.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_grpc_stream_cleanup() -> Result<()> {
    println!("\n=== Test 3: gRPC Stream Cleanup Validation ===");

    let initial_memory = get_process_memory_usage().unwrap_or(0);
    println!("  Initial memory: {} MB", initial_memory / 1_048_576);

    // Create and destroy 1000 broadcast channels (simulating gRPC streams)
    println!("  Creating 1000 broadcast channels...");

    for batch in 0..10 {
        let mut channels = Vec::new();

        // Create 100 channels
        for _ in 0..100 {
            let (tx, _rx) = tokio::sync::broadcast::channel::<String>(100);
            channels.push(tx);
        }

        // Send messages
        for tx in &channels {
            for i in 0..10 {
                let _ = tx.send(format!("Message {}", i));
            }
        }

        // Drop channels
        channels.clear();

        tokio::time::sleep(Duration::from_millis(50)).await;

        if (batch + 1) % 2 == 0 {
            let mem = get_process_memory_usage().unwrap_or(0);
            println!("    Batch {}/10, memory: {} MB", batch + 1, mem / 1_048_576);
        }
    }

    tokio::time::sleep(Duration::from_millis(500)).await;

    let final_memory = get_process_memory_usage().unwrap_or(0);
    let memory_leaked = final_memory.saturating_sub(initial_memory);
    let leak_percent = (memory_leaked as f64 / initial_memory as f64) * 100.0;

    println!("✓ Test 3 Results:");
    println!("  - Channels created/destroyed: 1000");
    println!("  - Final memory: {} MB", final_memory / 1_048_576);
    println!("  - Memory leaked: {} MB ({:.2}%)", memory_leaked / 1_048_576, leak_percent);
    println!("  - Target: <3% leak");

    assert!(leak_percent < 3.0, "Stream leak {:.2}% exceeds 3% threshold", leak_percent);

    Ok(())
}

/// Test 4: Database Connection Leak Detection
///
/// Validates that database connections are properly returned to the pool.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_database_connection_leak() -> Result<()> {
    println!("\n=== Test 4: Database Connection Leak Detection ===");

    let initial_memory = get_process_memory_usage().unwrap_or(0);
    println!("  Initial memory: {} MB", initial_memory / 1_048_576);

    // Simulate 1000 database operations
    println!("  Simulating 1000 database operations...");

    for i in 0..1000 {
        // Simulate connection acquisition and query
        let config = create_test_config();
        let _job = TrainingJob::new(
            "DQN".to_string(),
            config,
            format!("DB test {}", i),
            std::collections::HashMap::new(),
        );

        // Simulate query execution time
        tokio::time::sleep(Duration::from_micros(100)).await;

        if (i + 1) % 200 == 0 {
            let mem = get_process_memory_usage().unwrap_or(0);
            println!("    {} operations, memory: {} MB", i + 1, mem / 1_048_576);
        }
    }

    tokio::time::sleep(Duration::from_millis(500)).await;

    let final_memory = get_process_memory_usage().unwrap_or(0);
    let memory_leaked = final_memory.saturating_sub(initial_memory);
    let leak_percent = (memory_leaked as f64 / initial_memory as f64) * 100.0;

    println!("✓ Test 4 Results:");
    println!("  - Database operations: 1000");
    println!("  - Final memory: {} MB", final_memory / 1_048_576);
    println!("  - Memory leaked: {} MB ({:.2}%)", memory_leaked / 1_048_576, leak_percent);
    println!("  - Target: <2% leak");

    assert!(leak_percent < 2.0, "Connection leak {:.2}% exceeds 2% threshold", leak_percent);

    Ok(())
}

/// Test 5: Monitor RSS Growth (<1% Per Hour)
///
/// Continuous monitoring of Resident Set Size to detect gradual leaks.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_monitor_rss_growth() -> Result<()> {
    println!("\n=== Test 5: Monitor RSS Growth ===");
    println!("  NOTE: Running 2-minute version for CI compatibility");

    let test_duration = Duration::from_secs(120); // 2 minutes
    let sample_interval = Duration::from_secs(5);

    let initial_memory = get_process_memory_usage().unwrap_or(0);
    println!("  Initial RSS: {} MB", initial_memory / 1_048_576);

    let mut samples = Vec::new();
    samples.push((0, initial_memory));

    let start = Instant::now();

    // Continuous workload
    let operations = Arc::new(AtomicU64::new(0));
    let ops = operations.clone();

    let worker = tokio::spawn(async move {
        loop {
            let config = create_test_config();
            let _job = TrainingJob::new(
                "DQN".to_string(),
                config,
                "RSS monitor".to_string(),
                std::collections::HashMap::new(),
            );
            ops.fetch_add(1, Ordering::Relaxed);
            tokio::time::sleep(Duration::from_millis(50)).await;
        }
    });

    // Monitor RSS
    while start.elapsed() < test_duration {
        tokio::time::sleep(sample_interval).await;

        if let Some(memory) = get_process_memory_usage() {
            let elapsed_secs = start.elapsed().as_secs();
            samples.push((elapsed_secs, memory));
            println!("  [{:3}s] RSS: {} MB", elapsed_secs, memory / 1_048_576);
        }
    }

    worker.abort();

    // Calculate growth rate
    let final_memory = samples.last().expect("INVARIANT: Collection should be non-empty").1;
    let total_growth = final_memory.saturating_sub(initial_memory);
    let growth_percent = (total_growth as f64 / initial_memory as f64) * 100.0;

    // Extrapolate to hourly rate (2 min -> 60 min)
    let hourly_growth_percent = growth_percent * 30.0;

    println!("✓ Test 5 Results:");
    println!("  - Duration: {:?}", test_duration);
    println!("  - Samples collected: {}", samples.len());
    println!("  - Total operations: {}", operations.load(Ordering::Relaxed));
    println!("  - Growth (2 min): {:.3}%", growth_percent);
    println!("  - Extrapolated hourly: {:.3}%", hourly_growth_percent);
    println!("  - Target: <1% per hour");

    assert!(hourly_growth_percent < 1.0,
        "RSS growth {:.3}%/hour exceeds 1% target", hourly_growth_percent);

    Ok(())
}