db/tst/examples/performance.rs

use std::io::{self, Write};
use std::time::{Duration, Instant};
use tst::TST;

// Function to generate a test value of specified size
fn generate_test_value(index: usize, size: usize) -> Vec<u8> {
    let base_value = format!("val{:08}", index);
    let mut value = Vec::with_capacity(size);

    // Fill with repeating pattern to reach desired size
    while value.len() < size {
        value.extend_from_slice(base_value.as_bytes());
    }

    // Truncate to exact size
    value.truncate(size);

    value
}

// Number of records to insert
const TOTAL_RECORDS: usize = 100_000;
// How often to report progress (every X records)
const PROGRESS_INTERVAL: usize = 1_000;
// How many records to use for performance sampling
const PERFORMANCE_SAMPLE_SIZE: usize = 100;

fn main() -> Result<(), tst::Error> {
    // Create a temporary directory for the database
    let db_path = std::env::temp_dir().join("tst_performance_test");

    // Completely remove and recreate the directory to ensure a clean start
    if db_path.exists() {
        std::fs::remove_dir_all(&db_path)?;
    }
    std::fs::create_dir_all(&db_path)?;

    println!("Creating ternary search tree at: {}", db_path.display());
    println!("Will insert {} records and show progress...", TOTAL_RECORDS);

    // Create a new TST
    let mut tree = TST::new(db_path.to_str().unwrap(), true)?;

    // Track overall time
    let start_time = Instant::now();

    // Track performance metrics
    let mut insertion_times = Vec::with_capacity(TOTAL_RECORDS / PROGRESS_INTERVAL);
    let mut last_batch_time = Instant::now();
    let mut last_batch_records = 0;

    // Insert records and track progress
    for i in 0..TOTAL_RECORDS {
        let key = format!("key:{:08}", i);
        // Generate a 100-byte value
        let value = generate_test_value(i, 100);

        // Time the insertion of every Nth record for performance sampling
        if i % PERFORMANCE_SAMPLE_SIZE == 0 {
            let insert_start = Instant::now();
            tree.set(&key, value)?;
            let insert_duration = insert_start.elapsed();

            // Only print detailed timing for specific samples to avoid flooding output
            if i % (PERFORMANCE_SAMPLE_SIZE * 10) == 0 {
                println!("Record {}: Insertion took {:?}", i, insert_duration);
            }
        } else {
            tree.set(&key, value)?;
        }

        // Show progress at intervals
        if (i + 1) % PROGRESS_INTERVAL == 0 || i == TOTAL_RECORDS - 1 {
            let records_in_batch = i + 1 - last_batch_records;
            let batch_duration = last_batch_time.elapsed();
            let records_per_second = records_in_batch as f64 / batch_duration.as_secs_f64();

            insertion_times.push((i + 1, batch_duration));

            print!(
                "\rProgress: {}/{} records ({:.2}%) - {:.2} records/sec",
                i + 1,
                TOTAL_RECORDS,
                (i + 1) as f64 / TOTAL_RECORDS as f64 * 100.0,
                records_per_second
            );
            io::stdout().flush().unwrap();

            last_batch_time = Instant::now();
            last_batch_records = i + 1;
        }
    }

    let total_duration = start_time.elapsed();
    println!("\n\nPerformance Summary:");
    println!(
        "Total time to insert {} records: {:?}",
        TOTAL_RECORDS, total_duration
    );
    println!(
        "Average insertion rate: {:.2} records/second",
        TOTAL_RECORDS as f64 / total_duration.as_secs_f64()
    );

    // Show performance trend
    println!("\nPerformance Trend (records inserted vs. time per batch):");
    for (i, (record_count, duration)) in insertion_times.iter().enumerate() {
        if i % 10 == 0 || i == insertion_times.len() - 1 {
            // Only show every 10th point to avoid too much output
            println!(
                "  After {} records: {:?} for {} records ({:.2} records/sec)",
                record_count,
                duration,
                PROGRESS_INTERVAL,
                PROGRESS_INTERVAL as f64 / duration.as_secs_f64()
            );
        }
    }

    // Test access performance with distributed samples
    println!("\nTesting access performance with distributed samples...");
    let mut total_get_time = Duration::new(0, 0);
    let num_samples = 1000;

    // Use a simple distribution pattern instead of random
    for i in 0..num_samples {
        // Distribute samples across the entire range
        let sample_id = (i * (TOTAL_RECORDS / num_samples)) % TOTAL_RECORDS;
        let key = format!("key:{:08}", sample_id);

        let get_start = Instant::now();
        let _ = tree.get(&key)?;
        total_get_time += get_start.elapsed();
    }

    println!(
        "Average time to retrieve a record: {:?}",
        total_get_time / num_samples as u32
    );

    // Test prefix search performance
    println!("\nTesting prefix search performance...");
    let prefixes = ["key:0", "key:1", "key:5", "key:9"];

    for prefix in &prefixes {
        let list_start = Instant::now();
        let keys = tree.list(prefix)?;
        let list_duration = list_start.elapsed();

        println!(
            "Found {} keys with prefix '{}' in {:?}",
            keys.len(),
            prefix,
            list_duration
        );
    }

    // Clean up (optional)
    if std::env::var("KEEP_DB").is_err() {
        std::fs::remove_dir_all(&db_path)?;
        println!("\nCleaned up database directory");
    } else {
        println!("\nDatabase kept at: {}", db_path.display());
    }

    Ok(())
}