foxhunt/services/ml_training_service/tests/stress_concurrent_batch_creation.rs

//! Concurrent Batch Creation Stress Tests
//!
//! Tests system behavior under extreme concurrent batch creation load.
//! Validates database connection pool management, transaction handling,
//! and rollback correctness under contention.

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

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()
}

/// Test 1: 100 Concurrent Batch Creations
///
/// Validates that the system can handle 100 concurrent job creation requests
/// without crashes, deadlocks, or database connection exhaustion.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_100_concurrent_batch_creations() -> Result<()> {
    println!("\n=== Test 1: 100 Concurrent Batch Creations ===");

    let start = Instant::now();
    let success_count = Arc::new(AtomicU32::new(0));
    let failure_count = Arc::new(AtomicU32::new(0));

    // Spawn 100 concurrent tasks
    let mut handles = Vec::new();
    for i in 0..100 {
        let success = success_count.clone();
        let failure = failure_count.clone();

        let handle = tokio::spawn(async move {
            let config = create_test_config();
            let job = TrainingJob::new(
                "DQN".to_string(),
                config,
                format!("Stress test job {}", i),
                std::collections::HashMap::new(),
            );

            // Simulate database insert
            tokio::time::sleep(Duration::from_micros(100)).await;

            if job.id.to_string().len() > 0 {
                success.fetch_add(1, Ordering::Relaxed);
            } else {
                failure.fetch_add(1, Ordering::Relaxed);
            }
        });

        handles.push(handle);
    }

    // Wait for all tasks with timeout
    let result = timeout(Duration::from_secs(30), async {
        for handle in handles {
            handle.await.ok();
        }
    }).await;

    let elapsed = start.elapsed();
    let success = success_count.load(Ordering::Relaxed);
    let failure = failure_count.load(Ordering::Relaxed);

    println!("✓ Test 1 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Success: {}/100", success);
    println!("  - Failure: {}/100", failure);
    println!("  - Avg latency: {:?}", elapsed / 100);

    assert!(result.is_ok(), "Test timed out after 30s");
    assert_eq!(success, 100, "Expected all 100 jobs to succeed");
    assert!(elapsed < Duration::from_secs(10), "Expected completion under 10s, got {:?}", elapsed);

    Ok(())
}

/// Test 2: 1000 Total Jobs in System
///
/// Creates 1000 jobs in batches of 50 to validate system scalability
/// and database query performance under high job count.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_1000_total_jobs_in_system() -> Result<()> {
    println!("\n=== Test 2: 1000 Total Jobs in System ===");

    let start = Instant::now();
    let batch_size = 50;
    let num_batches = 20; // 20 batches * 50 = 1000 jobs

    let mut all_job_ids = Vec::new();

    for batch_idx in 0..num_batches {
        let batch_start = Instant::now();
        let mut batch_handles = Vec::new();

        for i in 0..batch_size {
            let job_num = batch_idx * batch_size + i;
            let handle = tokio::spawn(async move {
                let config = create_test_config();
                let job = TrainingJob::new(
                    "DQN".to_string(),
                    config,
                    format!("Batch job {}", job_num),
                    std::collections::HashMap::new(),
                );

                // Simulate database insert
                tokio::time::sleep(Duration::from_micros(50)).await;
                job.id
            });

            batch_handles.push(handle);
        }

        // Collect batch results
        for handle in batch_handles {
            if let Ok(job_id) = handle.await {
                all_job_ids.push(job_id);
            }
        }

        let batch_elapsed = batch_start.elapsed();
        if (batch_idx + 1) % 5 == 0 {
            println!("  - Batch {}/{} completed in {:?} ({} jobs total)",
                batch_idx + 1, num_batches, batch_elapsed, all_job_ids.len());
        }
    }

    let elapsed = start.elapsed();

    println!("✓ Test 2 Results:");
    println!("  - Total duration: {:?}", elapsed);
    println!("  - Jobs created: {}/1000", all_job_ids.len());
    println!("  - Avg batch time: {:?}", elapsed / num_batches);
    println!("  - Throughput: {:.2} jobs/sec", 1000.0 / elapsed.as_secs_f64());

    assert_eq!(all_job_ids.len(), 1000, "Expected 1000 jobs created");
    assert!(elapsed < Duration::from_secs(60), "Expected completion under 60s, got {:?}", elapsed);

    Ok(())
}

/// Test 3: Database Connection Pool Saturation
///
/// Tests behavior when all database connections are in use.
/// Validates connection pool sizing and connection reuse.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_database_connection_pool_saturation() -> Result<()> {
    println!("\n=== Test 3: Database Connection Pool Saturation ===");

    let start = Instant::now();
    let pool_size = 20; // Typical connection pool size
    let requests = pool_size * 5; // 5x pool size to force queueing

    let success_count = Arc::new(AtomicU32::new(0));
    let timeout_count = Arc::new(AtomicU32::new(0));

    let mut handles = Vec::new();

    for i in 0..requests {
        let success = success_count.clone();
        let timeouts = timeout_count.clone();

        let handle = tokio::spawn(async move {
            // Simulate long-running database operation
            let operation = async {
                let config = create_test_config();
                let _job = TrainingJob::new(
                    "DQN".to_string(),
                    config,
                    format!("Connection pool test {}", i),
                    std::collections::HashMap::new(),
                );

                // Hold connection for 100ms
                tokio::time::sleep(Duration::from_millis(100)).await;
                success.fetch_add(1, Ordering::Relaxed);
            };

            // 5s timeout per operation
            if timeout(Duration::from_secs(5), operation).await.is_err() {
                timeouts.fetch_add(1, Ordering::Relaxed);
            }
        });

        handles.push(handle);
    }

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

    let elapsed = start.elapsed();
    let success = success_count.load(Ordering::Relaxed);
    let timeouts = timeout_count.load(Ordering::Relaxed);

    println!("✓ Test 3 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Success: {}/{}", success, requests);
    println!("  - Timeouts: {}/{}", timeouts, requests);
    println!("  - Pool saturation handled: {}", timeouts == 0);

    assert_eq!(timeouts, 0, "Expected no timeouts with proper connection pooling");
    assert_eq!(success, requests, "Expected all operations to succeed");

    Ok(())
}

/// Test 4: Transaction Timeout Handling
///
/// Validates that the system properly handles transaction timeouts
/// and doesn't leave dangling transactions.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_transaction_timeout_handling() -> Result<()> {
    println!("\n=== Test 4: Transaction Timeout Handling ===");

    let start = Instant::now();
    let timeout_threshold = Duration::from_millis(500);

    let completed = Arc::new(AtomicU32::new(0));
    let timed_out = Arc::new(AtomicU32::new(0));

    let mut handles = Vec::new();

    // Create 50 operations, some will timeout
    for i in 0..50 {
        let completed_count = completed.clone();
        let timeout_count = timed_out.clone();

        let handle = tokio::spawn(async move {
            let operation = async {
                let config = create_test_config();
                let _job = TrainingJob::new(
                    "DQN".to_string(),
                    config,
                    format!("Timeout test {}", i),
                    std::collections::HashMap::new(),
                );

                // Simulate varying transaction durations
                let duration = Duration::from_millis(100 * (i % 8) as u64);
                tokio::time::sleep(duration).await;
            };

            match timeout(timeout_threshold, operation).await {
                Ok(_) => {
                    completed_count.fetch_add(1, Ordering::Relaxed);
                }
                Err(_) => {
                    timeout_count.fetch_add(1, Ordering::Relaxed);
                }
            }
        });

        handles.push(handle);
    }

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

    let elapsed = start.elapsed();
    let completed_ops = completed.load(Ordering::Relaxed);
    let timed_out_ops = timed_out.load(Ordering::Relaxed);

    println!("✓ Test 4 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Completed: {}/50", completed_ops);
    println!("  - Timed out: {}/50", timed_out_ops);
    println!("  - Timeout threshold: {:?}", timeout_threshold);

    assert!(completed_ops + timed_out_ops == 50, "Expected all operations to complete or timeout");
    assert!(timed_out_ops > 0, "Expected some operations to timeout");

    Ok(())
}

/// Test 5: Rollback Correctness Under Contention
///
/// Tests that transaction rollbacks work correctly when multiple
/// transactions are competing for the same resources.
#[tokio::test]
#[ignore = "Stress test - run explicitly with --ignored"]
async fn test_rollback_correctness_under_contention() -> Result<()> {
    println!("\n=== Test 5: Rollback Correctness Under Contention ===");

    let start = Instant::now();
    let success_count = Arc::new(AtomicU32::new(0));
    let rollback_count = Arc::new(AtomicU32::new(0));

    let mut handles = Vec::new();

    // Create 100 competing transactions
    for i in 0..100 {
        let success = success_count.clone();
        let rollbacks = rollback_count.clone();

        let handle = tokio::spawn(async move {
            let config = create_test_config();

            // Simulate transaction with potential rollback
            let should_rollback = i % 3 == 0; // Every 3rd transaction fails

            if should_rollback {
                // Simulate failed transaction
                rollbacks.fetch_add(1, Ordering::Relaxed);
            } else {
                let _job = TrainingJob::new(
                    "DQN".to_string(),
                    config,
                    format!("Rollback test {}", i),
                    std::collections::HashMap::new(),
                );
                success.fetch_add(1, Ordering::Relaxed);
            }

            tokio::time::sleep(Duration::from_micros(100)).await;
        });

        handles.push(handle);
    }

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

    let elapsed = start.elapsed();
    let success = success_count.load(Ordering::Relaxed);
    let rollbacks = rollback_count.load(Ordering::Relaxed);

    println!("✓ Test 5 Results:");
    println!("  - Duration: {:?}", elapsed);
    println!("  - Successful commits: {}", success);
    println!("  - Rollbacks: {}", rollbacks);
    println!("  - Total operations: {}", success + rollbacks);

    assert_eq!(success + rollbacks, 100, "Expected all operations to complete");
    assert!(rollbacks >= 30, "Expected at least 30 rollbacks (1/3 of operations)");

    Ok(())
}

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

    /// Benchmark batch creation latency
    #[tokio::test]
    #[ignore = "Stress test - run with --ignored"]
    async fn bench_batch_creation_latency() -> Result<()> {
        println!("\n=== Benchmark: Batch Creation Latency ===");

        let iterations = 1000;
        let mut latencies = Vec::new();

        for i in 0..iterations {
            let start = Instant::now();

            let config = create_test_config();
            let _job = TrainingJob::new(
                "DQN".to_string(),
                config,
                format!("Benchmark job {}", i),
                std::collections::HashMap::new(),
            );

            latencies.push(start.elapsed());
        }

        latencies.sort();
        let p50 = latencies[iterations / 2];
        let p95 = latencies[(iterations * 95) / 100];
        let p99 = latencies[(iterations * 99) / 100];

        println!("✓ Latency Distribution:");
        println!("  - P50: {:?}", p50);
        println!("  - P95: {:?}", p95);
        println!("  - P99: {:?}", p99);
        println!("  - Target P95: <10ms");

        assert!(p95 < Duration::from_millis(10),
            "P95 latency {:?} exceeds 10ms target", p95);

        Ok(())
    }
}