db/ourdb/examples/benchmark.rs

use ourdb::{OurDB, OurDBConfig, OurDBSetArgs};
use std::path::PathBuf;
use std::time::{Duration, Instant};

fn main() -> Result<(), ourdb::Error> {
    // Parse command line arguments
    let args: Vec<String> = std::env::args().collect();
    
    let (num_operations, record_size, incremental_mode, keysize) = parse_args(&args);
    
    println!("OurDB Benchmark");
    println!("===============");
    println!("Operations: {}", num_operations);
    println!("Record size: {} bytes", record_size);
    println!("Mode: {}", if incremental_mode { "Incremental" } else { "Key-Value" });
    println!("Key size: {} bytes", keysize);
    println!();
    
    // Create a temporary directory for the database
    let db_path = std::env::temp_dir().join(format!("ourdb_benchmark_{}", std::process::id()));
    std::fs::create_dir_all(&db_path)?;
    
    println!("Database path: {}", db_path.display());
    
    // Create a new database
    let config = OurDBConfig {
        path: db_path.clone(),
        incremental_mode,
        file_size: Some(1024 * 1024),
        keysize: Some(keysize),
    };
    
    let mut db = OurDB::new(config)?;
    
    // Prepare test data
    let test_data = vec![b'X'; record_size];
    
    // Benchmark write operations
    println!("\nBenchmarking writes...");
    let start = Instant::now();
    
    let mut ids = Vec::with_capacity(num_operations);
    for i in 0..num_operations {
        let id = if incremental_mode {
            db.set(OurDBSetArgs { id: None, data: &test_data })?
        } else {
            // In key-value mode, we provide explicit IDs
            let id = i as u32 + 1;
            db.set(OurDBSetArgs { id: Some(id), data: &test_data })?;
            id
        };
        ids.push(id);
    }
    
    let write_duration = start.elapsed();
    print_performance_stats("Write", num_operations, write_duration);
    
    // Benchmark read operations (sequential)
    println!("\nBenchmarking sequential reads...");
    let start = Instant::now();
    
    for &id in &ids {
        let _ = db.get(id)?;
    }
    
    let read_duration = start.elapsed();
    print_performance_stats("Sequential read", num_operations, read_duration);
    
    // Benchmark random reads
    println!("\nBenchmarking random reads...");
    let start = Instant::now();
    
    use std::collections::HashSet;
    let mut rng = rand::thread_rng();
    let mut random_indices = HashSet::new();
    
    // Select 20% of the IDs randomly for testing
    let sample_size = num_operations / 5;
    while random_indices.len() < sample_size {
        let idx = rand::Rng::gen_range(&mut rng, 0..ids.len());
        random_indices.insert(idx);
    }
    
    for idx in random_indices {
        let _ = db.get(ids[idx])?;
    }
    
    let random_read_duration = start.elapsed();
    print_performance_stats("Random read", sample_size, random_read_duration);
    
    // Benchmark update operations
    println!("\nBenchmarking updates...");
    let start = Instant::now();
    
    for &id in &ids[0..num_operations/2] {
        db.set(OurDBSetArgs { id: Some(id), data: &test_data })?;
    }
    
    let update_duration = start.elapsed();
    print_performance_stats("Update", num_operations/2, update_duration);
    
    // Benchmark history retrieval
    println!("\nBenchmarking history retrieval...");
    let start = Instant::now();
    
    for &id in &ids[0..num_operations/10] {
        let _ = db.get_history(id, 2)?;
    }
    
    let history_duration = start.elapsed();
    print_performance_stats("History retrieval", num_operations/10, history_duration);
    
    // Benchmark delete operations
    println!("\nBenchmarking deletes...");
    let start = Instant::now();
    
    for &id in &ids[0..num_operations/4] {
        db.delete(id)?;
    }
    
    let delete_duration = start.elapsed();
    print_performance_stats("Delete", num_operations/4, delete_duration);
    
    // Close and clean up
    db.close()?;
    std::fs::remove_dir_all(&db_path)?;
    
    println!("\nBenchmark completed successfully");
    
    Ok(())
}

fn parse_args(args: &[String]) -> (usize, usize, bool, u8) {
    let mut num_operations = 100000;
    let mut record_size = 100;
    let mut incremental_mode = true;
    let mut keysize = 4;
    
    for i in 1..args.len() {
        if args[i] == "--ops" && i + 1 < args.len() {
            if let Ok(n) = args[i + 1].parse() {
                num_operations = n;
            }
        } else if args[i] == "--size" && i + 1 < args.len() {
            if let Ok(n) = args[i + 1].parse() {
                record_size = n;
            }
        } else if args[i] == "--keyvalue" {
            incremental_mode = false;
        } else if args[i] == "--keysize" && i + 1 < args.len() {
            if let Ok(n) = args[i + 1].parse() {
                if [2, 3, 4, 6].contains(&n) {
                    keysize = n;
                }
            }
        } else if args[i] == "--help" {
            print_usage();
            std::process::exit(0);
        }
    }
    
    (num_operations, record_size, incremental_mode, keysize)
}

fn print_usage() {
    println!("OurDB Benchmark Tool");
    println!("Usage: cargo run --example benchmark [OPTIONS]");
    println!();
    println!("Options:");
    println!("  --ops N          Number of operations to perform (default: 100000)");
    println!("  --size N         Size of each record in bytes (default: 100)");
    println!("  --keyvalue       Use key-value mode instead of incremental mode");
    println!("  --keysize N      Key size in bytes (2, 3, 4, or 6) (default: 4)");
    println!("  --help           Print this help message");
}

fn print_performance_stats(operation: &str, count: usize, duration: Duration) {
    let ops_per_second = count as f64 / duration.as_secs_f64();
    let ms_per_op = duration.as_secs_f64() * 1000.0 / count as f64;
    
    println!("{} performance:", operation);
    println!("  Total time: {:.2} seconds", duration.as_secs_f64());
    println!("  Operations: {}", count);
    println!("  Speed: {:.2} ops/sec", ops_per_second);
    println!("  Average: {:.3} ms/op", ms_per_op);
}