db-stresser/bin/redis.rs

use benchmark_oversdb::{StressTester, RedisClient};
use std::env;

#[tokio::main]
async fn main() -> Result<(), String> {
    let args: Vec<String> = std::env::args().collect();
    
    if args.len() < 2 {
        println!("Redis Stress Test & Benchmark Tool");
        println!("==================================");
        println!();
        println!("Usage: {} <test_type> [options]", args[0]);
        println!();
        println!("Test types:");
        println!("  quick      - Quick test (1000 ops, 10 clients)");
        println!("  standard   - Standard test (10000 ops, 50 clients)");
        println!("  intensive  - Intensive test (100000 ops, 100 clients)");
        println!("  extreme    - Extreme test (1000000 ops, 200 clients)");
        println!("  debug      - Debug test (100 ops, 5 clients)");
        println!("  custom <ops> <clients> <value_size> - Custom parameters");
        println!();
        println!("Environment variables:");
        println!("  REDIS_HOST     - Server host (default: 127.0.0.1)");
        println!("  REDIS_PORT     - Server port (default: 6379)");
        println!("  REDIS_PASSWORD - Password (optional)");
        println!();
        println!("Examples:");
        println!("  {} quick", args[0]);
        println!("  {} custom 50000 75 1024", args[0]);
        println!("  REDIS_PASSWORD=mypass {} standard", args[0]);
        return Ok(());
    }

    // Get connection details from environment or use defaults
    let host = env::var("REDIS_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
    let port = env::var("REDIS_PORT").unwrap_or_else(|_| "6379".to_string());
    let password = env::var("REDIS_PASSWORD").unwrap_or_else(|_| "".to_string());
    let addr = format!("{}:{}", host, port);

    let tester = StressTester::new(addr.clone(), "".to_string(), password.clone());
    
    println!("🔧 Connecting to Redis server at {}...", addr);
    if password != "" {
        println!("   Using password authentication");
    }
    
    match RedisClient::connect_with_auth(&addr, &password).await {
        Ok(mut client) => {
            println!("✓ Connection successful!");
            
            // Test connection with PING
            match client.ping().await {
                Ok(response) => println!("✓ Server responded: {}", response),
                Err(e) => {
                    println!("⚠ Warning: PING failed: {}", e);
                    println!("  Connection might still work for basic operations");
                }
            }
            
            // Optional: Clear database for clean benchmark
            if env::var("REDIS_FLUSH_DB").unwrap_or_else(|_| "false".to_string()) == "true" {
                println!("🧹 Flushing database...");
                match client.flush_db().await {
                    Ok(_) => println!("✓ Database flushed"),
                    Err(e) => println!("⚠ Warning: Failed to flush database: {}", e),
                }
            }
        },
        Err(e) => {
            println!("✗ Failed to connect: {}", e);
            println!("Make sure the Redis server is running on {}", addr);
            println!("Check your connection details and authentication:");
            println!("  REDIS_HOST=your_host REDIS_PORT=your_port {} <test_type>", args[0]);
            if password == "" {
                println!("  If authentication is required:");
                println!("  REDIS_PASSWORD=your_pass {} <test_type>", args[0]);
            }
            return Err(e);
        }
    }
    println!();

    let (ops, clients, value_size) = match args[1].as_str() {
        "quick" => (1_000, 10, 256),
        "standard" => (10_000, 50, 256),
        "intensive" => (100_000, 100, 256),
        "extreme" => (1_000_000, 200, 256),
        "debug" => (100, 5, 64),
        "custom" => {
            if args.len() < 5 {
                return Err("Custom test requires: custom <ops> <clients> <value_size>".to_string());
            }
            let ops = args[2].parse().map_err(|_| "Invalid operations count")?;
            let clients = args[3].parse().map_err(|_| "Invalid client count")?;
            let value_size = args[4].parse().map_err(|_| "Invalid value size")?;
            (ops, clients, value_size)
        },
        _ => return Err("Unknown test type. Use: quick, standard, intensive, extreme, debug, or custom".to_string()),
    };

    println!("🎯 Test Configuration:");
    println!("   Total operations: {}", ops);
    println!("   Concurrent clients: {}", clients);
    println!("   Value size: {} bytes", value_size);
    println!();

    // Prepare some data for read tests
    println!("📝 Preparing test data...");
    let mut prep_client = RedisClient::connect_with_auth(&addr, &password).await?;
    let prep_value = vec![0x42u8; value_size];
    
    for i in 0..10 {
        for j in 0..100 {
            let key = format!("worker{}:key{}", i, j);
            let _ = prep_client.set(&key, &prep_value).await;
        }
    }
    println!("✓ Test data prepared");
    println!();

    // Run benchmarks
    let mut results = Vec::new();

    // SET benchmark
    match tester.redis_benchmark_set(ops, clients, value_size).await {
        Ok(result) => {
            result.print();
            results.push(result);
        },
        Err(e) => println!("SET benchmark failed: {}", e),
    }

    // GET benchmark
    match tester.redis_benchmark_get(ops, clients).await {
        Ok(result) => {
            result.print();
            results.push(result);
        },
        Err(e) => println!("GET benchmark failed: {}", e),
    }

    // MIXED benchmark
    match tester.redis_benchmark_mixed(ops, clients, value_size).await {
        Ok(result) => {
            result.print();
            results.push(result);
        },
        Err(e) => println!("MIXED benchmark failed: {}", e),
    }

    // Summary
    println!("📈 SUMMARY");
    println!("==========");
    for result in &results {
        let error_rate = if result.total_operations + result.errors > 0 {
            format!("{:.1}%", (result.errors as f64 / (result.total_operations + result.errors) as f64) * 100.0)
        } else {
            "0.0%".to_string()
        };
        
        println!("{:8} | {:>10.0} ops/sec | {:>8.2}μs avg latency | {} errors ({})", 
            result.operation, 
            result.ops_per_second, 
            result.avg_latency_us,
            result.errors,
            error_rate
        );
    }
    
    if let Some(best) = results.iter().filter(|r| r.ops_per_second > 0.0).max_by(|a, b| a.ops_per_second.partial_cmp(&b.ops_per_second).unwrap()) {
        println!();
        println!("🏆 Best performance: {} with {:.0} ops/sec", best.operation, best.ops_per_second);
    }

    // Additional Redis-specific information
    println!();
    println!("📊 Redis Connection Info:");
    println!("   Server: {}", addr);
    println!("   Authentication: {}", if password != "" { "Yes" } else { "No" });
    println!();
    println!("💡 Tips for better Redis performance:");
    println!("   - Use Redis pipelining for batch operations");
    println!("   - Consider Redis Cluster for horizontal scaling");
    println!("   - Monitor memory usage with 'redis-cli info memory'");
    println!("   - Use appropriate data structures (strings, hashes, sets, etc.)");

    Ok(())
}