foxhunt/services/ml_training_service/tests/stress_streaming_load.rs

//! Streaming Load Stress Tests
//!
//! Tests gRPC streaming performance under extreme load conditions.
//! Validates concurrent stream handling, backpressure, and message throughput.

use anyhow::Result;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::broadcast;
use tokio::time::timeout;

use ml_training_service::orchestrator::TrainingStatusUpdate;

/// Helper to create status updates
fn create_status_update(job_num: u32, sequence: u32) -> TrainingStatusUpdate {
    TrainingStatusUpdate {
        job_id: uuid::Uuid::new_v4(),
        status: ml_training_service::orchestrator::JobStatus::Running,
        progress_percentage: (sequence % 100) as f32,
        current_epoch: sequence,
        total_epochs: 100,
        metrics: std::collections::HashMap::new(),
        message: format!("Job {} - Update {}", job_num, sequence),
        timestamp: chrono::Utc::now(),
        financial_metrics: None,
        resource_usage: ml_training_service::orchestrator::ResourceUsage {
            cpu_usage_percent: 50.0,
            memory_usage_gb: 1.0,
            gpu_usage_percent: Some(75.0),
            gpu_memory_usage_gb: Some(2.0),
            active_workers: 4,
        },
    }
}

/// Test 1: 100 Concurrent Watch Streams
///
/// Validates that the system can handle 100 concurrent streaming clients
/// without deadlocks, memory exhaustion, or stream corruption.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_100_concurrent_watch_streams() -> Result<()> {
    println!("\n=== Test 1: 100 Concurrent Watch Streams ===");

    let start = Instant::now();
    let (tx, _) = broadcast::channel(1000);

    let active_streams = Arc::new(AtomicU32::new(0));
    let messages_received = Arc::new(AtomicU64::new(0));

    let mut handles = Vec::new();

    // Spawn 100 concurrent stream consumers
    for i in 0..100 {
        let mut rx = tx.subscribe();
        let active = active_streams.clone();
        let received = messages_received.clone();

        let handle = tokio::spawn(async move {
            active.fetch_add(1, Ordering::Relaxed);

            let mut count = 0;
            let stream_timeout = Duration::from_secs(10);

            loop {
                match timeout(stream_timeout, rx.recv()).await {
                    Ok(Ok(_update)) => {
                        count += 1;
                        received.fetch_add(1, Ordering::Relaxed);

                        if count >= 100 {
                            break;
                        }
                    }
                    Ok(Err(_)) => break, // Channel closed
                    Err(_) => break, // Timeout
                }
            }

            active.fetch_sub(1, Ordering::Relaxed);
            count
        });

        handles.push(handle);
    }

    // Give streams time to subscribe
    tokio::time::sleep(Duration::from_millis(100)).await;

    // Send 100 status updates
    for seq in 0..100 {
        let update = create_status_update(0, seq);
        let _ = tx.send(update);
        tokio::time::sleep(Duration::from_millis(10)).await;
    }

    // Wait for all streams to complete
    let mut stream_counts = Vec::new();
    for handle in handles {
        if let Ok(count) = handle.await {
            stream_counts.push(count);
        }
    }

    let elapsed = start.elapsed();
    let total_received = messages_received.load(Ordering::Relaxed);
    let avg_per_stream = if !stream_counts.is_empty() {
        stream_counts.iter().sum::<u64>() / stream_counts.len() as u64
    } else {
        0
    };

    println!("✓ Test 1 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Active streams: {}/100", stream_counts.len());
    println!("  - Total messages received: {}", total_received);
    println!("  - Avg messages per stream: {}/100", avg_per_stream);
    println!("  - Expected total: {}", 100 * 100);

    assert_eq!(stream_counts.len(), 100, "Expected all 100 streams to complete");
    assert!(avg_per_stream >= 90, "Expected avg ≥90 messages per stream, got {}", avg_per_stream);

    Ok(())
}

/// Test 2: 10K Messages/Second Broadcast
///
/// Tests high-throughput message broadcasting to validate
/// the system can sustain 10,000 messages per second.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_10k_messages_per_second_broadcast() -> Result<()> {
    println!("\n=== Test 2: 10K Messages/Second Broadcast ===");

    let start = Instant::now();
    let (tx, _) = broadcast::channel(10000);

    let messages_sent = Arc::new(AtomicU64::new(0));
    let messages_received = Arc::new(AtomicU64::new(0));

    // Spawn 10 stream consumers
    let mut consumer_handles = Vec::new();
    for i in 0..10 {
        let mut rx = tx.subscribe();
        let received = messages_received.clone();

        let handle = tokio::spawn(async move {
            let mut count = 0;
            while let Ok(_update) = rx.recv().await {
                count += 1;
                received.fetch_add(1, Ordering::Relaxed);

                if count >= 10000 {
                    break;
                }
            }
            count
        });

        consumer_handles.push(handle);
    }

    // Give consumers time to subscribe
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Send messages as fast as possible
    let producer_start = Instant::now();
    let sent = messages_sent.clone();

    let producer_handle = tokio::spawn(async move {
        for seq in 0..10000 {
            let update = create_status_update(0, seq);
            if tx.send(update).is_ok() {
                sent.fetch_add(1, Ordering::Relaxed);
            }
        }
        producer_start.elapsed()
    });

    let send_duration = producer_handle.await?;

    // Wait for consumers to catch up
    tokio::time::sleep(Duration::from_millis(500)).await;

    let sent = messages_sent.load(Ordering::Relaxed);
    let received = messages_received.load(Ordering::Relaxed);
    let throughput = sent as f64 / send_duration.as_secs_f64();

    println!("✓ Test 2 Results:");
    println!("  - Send duration: {:?}", send_duration);
    println!("  - Messages sent: {}", sent);
    println!("  - Messages received: {} (across 10 consumers)", received);
    println!("  - Throughput: {:.0} msg/sec", throughput);
    println!("  - Target: >10,000 msg/sec");

    assert_eq!(sent, 10000, "Expected all 10,000 messages to be sent");
    assert!(throughput >= 10000.0, "Throughput {:.0} msg/sec below 10K target", throughput);

    Ok(())
}

/// Test 3: Slow Consumer (Backpressure Validation)
///
/// Tests that slow consumers don't block fast producers and that
/// backpressure mechanisms work correctly.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_slow_consumer_backpressure() -> Result<()> {
    println!("\n=== Test 3: Slow Consumer Backpressure ===");

    let start = Instant::now();
    let (tx, _) = broadcast::channel(1000);

    let fast_received = Arc::new(AtomicU64::new(0));
    let slow_received = Arc::new(AtomicU64::new(0));

    // Fast consumer
    let fast_rx = tx.subscribe();
    let fast_count = fast_received.clone();
    let fast_handle = tokio::spawn(async move {
        let mut rx = fast_rx;
        while let Ok(_update) = rx.recv().await {
            fast_count.fetch_add(1, Ordering::Relaxed);
        }
    });

    // Slow consumer (10ms delay per message)
    let slow_rx = tx.subscribe();
    let slow_count = slow_received.clone();
    let slow_handle = tokio::spawn(async move {
        let mut rx = slow_rx;
        while let Ok(_update) = rx.recv().await {
            tokio::time::sleep(Duration::from_millis(10)).await;
            slow_count.fetch_add(1, Ordering::Relaxed);
        }
    });

    // Give consumers time to subscribe
    tokio::time::sleep(Duration::from_millis(50)).await;

    // Send 1000 messages rapidly
    for seq in 0..1000 {
        let update = create_status_update(0, seq);
        let _ = tx.send(update);
    }

    // Wait for processing
    tokio::time::sleep(Duration::from_secs(2)).await;

    // Drop sender to close streams
    drop(tx);

    // Wait for consumers
    timeout(Duration::from_secs(15), async {
        fast_handle.await.ok();
        slow_handle.await.ok();
    }).await?;

    let elapsed = start.elapsed();
    let fast = fast_received.load(Ordering::Relaxed);
    let slow = slow_received.load(Ordering::Relaxed);

    println!("✓ Test 3 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Fast consumer received: {}/1000", fast);
    println!("  - Slow consumer received: {}/1000", slow);
    println!("  - Backpressure handled: {}", slow < fast);

    assert!(fast >= 900, "Fast consumer should receive most messages, got {}", fast);
    assert!(slow < fast, "Slow consumer should lag behind fast consumer");

    Ok(())
}

/// Test 4: Network Latency Simulation (100ms RTT)
///
/// Simulates high network latency to validate stream resilience.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_network_latency_simulation() -> Result<()> {
    println!("\n=== Test 4: Network Latency Simulation (100ms RTT) ===");

    let start = Instant::now();
    let (tx, _) = broadcast::channel(1000);

    let messages_received = Arc::new(AtomicU64::new(0));
    let mut handles = Vec::new();

    // Spawn 10 consumers with simulated network delay
    for _ in 0..10 {
        let mut rx = tx.subscribe();
        let received = messages_received.clone();

        let handle = tokio::spawn(async move {
            let mut count = 0;
            while let Ok(_update) = rx.recv().await {
                // Simulate 100ms network round-trip
                tokio::time::sleep(Duration::from_millis(100)).await;
                count += 1;
                received.fetch_add(1, Ordering::Relaxed);

                if count >= 50 {
                    break;
                }
            }
            count
        });

        handles.push(handle);
    }

    // Give consumers time to subscribe
    tokio::time::sleep(Duration::from_millis(100)).await;

    // Send 50 messages
    for seq in 0..50 {
        let update = create_status_update(0, seq);
        let _ = tx.send(update);
        tokio::time::sleep(Duration::from_millis(20)).await; // Paced sending
    }

    // Wait for all consumers
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let total_received = messages_received.load(Ordering::Relaxed);

    println!("✓ Test 4 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Messages received: {} (across 10 consumers)", total_received);
    println!("  - Expected: {}", 50 * 10);
    println!("  - Network latency: 100ms RTT");

    assert!(total_received >= 450, "Expected at least 450 messages, got {}", total_received);

    Ok(())
}

/// Test 5: Stream Multiplexing (16 Jobs Per Stream)
///
/// Tests multiplexing multiple job updates over single streams.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_stream_multiplexing() -> Result<()> {
    println!("\n=== Test 5: Stream Multiplexing (16 Jobs Per Stream) ===");

    let start = Instant::now();
    let (tx, _) = broadcast::channel(2000);

    let messages_received = Arc::new(AtomicU64::new(0));
    let unique_jobs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));

    // Spawn 10 stream consumers
    let mut handles = Vec::new();
    for _ in 0..10 {
        let mut rx = tx.subscribe();
        let received = messages_received.clone();
        let jobs = unique_jobs.clone();

        let handle = tokio::spawn(async move {
            let mut local_jobs = std::collections::HashSet::new();

            while let Ok(update) = rx.recv().await {
                received.fetch_add(1, Ordering::Relaxed);
                local_jobs.insert(update.job_id);

                if local_jobs.len() >= 16 {
                    break;
                }
            }

            // Merge into global set
            let mut global = jobs.lock().expect("INVARIANT: Lock should not be poisoned");
            global.extend(local_jobs.iter());
            local_jobs.len()
        });

        handles.push(handle);
    }

    // Give consumers time to subscribe
    tokio::time::sleep(Duration::from_millis(100)).await;

    // Send updates for 16 different jobs
    for job_num in 0..16 {
        for seq in 0..10 {
            let update = create_status_update(job_num, seq);
            let _ = tx.send(update);
        }
    }

    // Wait for consumers
    for handle in handles {
        handle.await.ok();
    }

    let elapsed = start.elapsed();
    let total_received = messages_received.load(Ordering::Relaxed);
    let unique_job_count = unique_jobs.lock().expect("INVARIANT: Lock should not be poisoned").len();

    println!("✓ Test 5 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Messages received: {}", total_received);
    println!("  - Unique jobs tracked: {}/16", unique_job_count);
    println!("  - Expected messages: {}", 16 * 10 * 10); // 16 jobs * 10 updates * 10 consumers

    assert_eq!(unique_job_count, 16, "Expected all 16 unique jobs to be tracked");
    assert!(total_received >= 1400, "Expected at least 1400 messages, got {}", total_received);

    Ok(())
}

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

    /// Benchmark streaming throughput
    #[tokio::test]
    #[ignore = "Stress test - run with --ignored"]
    async fn bench_streaming_throughput() -> Result<()> {
        println!("\n=== Benchmark: Streaming Throughput ===");

        let (tx, mut rx) = broadcast::channel(10000);

        let received_count = Arc::new(AtomicU64::new(0));
        let count = received_count.clone();

        // Consumer
        let consumer = tokio::spawn(async move {
            while rx.recv().await.is_ok() {
                count.fetch_add(1, Ordering::Relaxed);
            }
        });

        // Producer
        let start = Instant::now();
        for seq in 0..10000 {
            let update = create_status_update(0, seq);
            if tx.send(update).is_err() {
                break;
            }
        }
        let send_duration = start.elapsed();

        drop(tx);
        consumer.await.ok();

        let received = received_count.load(Ordering::Relaxed);
        let throughput = received as f64 / send_duration.as_secs_f64();

        println!("✓ Throughput Results:");
        println!("  - Messages: {}/10000", received);
        println!("  - Duration: {:?}", send_duration);
        println!("  - Throughput: {:.0} msg/sec", throughput);
        println!("  - Target: >1,000 msg/sec per stream");

        assert!(throughput >= 1000.0, "Throughput below 1K msg/sec target");

        Ok(())
    }
}