db-stresser/bin/skytable.rs

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

#[tokio::main]
async fn main() -> Result<(), String> {
    let args: Vec<String> = std::env::args().collect();
    
    if args.len() < 2 {
        println!("Skytable 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!("  SKYTABLE_HOST     - Server host (default: 127.0.0.1)");
        println!("  SKYTABLE_USERNAME - Username (default: root)");
        println!("  SKYTABLE_PASSWORD - Password (default: password)");
        println!();
        println!("Examples:");
        println!("  {} quick", args[0]);
        println!("  {} custom 50000 75 1024", args[0]);
        println!("  SKYTABLE_PASSWORD=mypass {} standard", args[0]);
        return Ok(());
    }

    // Get connection details from environment or use defaults
    let host = env::var("SKYTABLE_HOST").unwrap_or_else(|_| "127.0.0.1".to_string());
    let username = env::var("SKYTABLE_USERNAME").unwrap_or_else(|_| "root".to_string());
    let password = env::var("SKYTABLE_PASSWORD").unwrap_or_else(|_| "password".to_string());

    let tester = StressTester::new(host.clone(), username.clone(), password.clone());
    
    println!("🔧 Connecting to Skytable server at {}:2003...", host);
    println!("   Using credentials: {}/{}", username, password);
    
    match SkytableClient::connect_with_auth(&host, &username, &password).await {
        Ok(mut client) => {
            println!("✓ Connection successful!");
            // Setup benchmark space
            match client.setup_benchmark_space().await {
                Ok(_) => println!("✓ Benchmark space setup complete!"),
                Err(e) => {
                    println!("⚠ Warning: Benchmark space setup failed: {}", e);
                    println!("  This might be normal if using basic key-value mode");
                }
            }
        },
        Err(e) => {
            println!("✗ Failed to connect: {}", e);
            println!("Make sure the Skytable server is running on {}:2003", host);
            println!("and that you have the correct credentials");
            println!("You can set credentials using environment variables:");
            println!("  SKYTABLE_USERNAME=your_user SKYTABLE_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 = SkytableClient::connect_with_auth(&host, &username, &password).await?;
    prep_client.setup_benchmark_space().await.ok(); // Ignore errors
    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.skytable_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.skytable_benchmark_get(ops, clients).await {
        Ok(result) => {
            result.print();
            results.push(result);
        },
        Err(e) => println!("GET benchmark failed: {}", e),
    }

    // MIXED benchmark
    match tester.skytable_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);
    }

    Ok(())
}