markbase/filetree-rocksdb/src/poc.rs

use anyhow::Result;
use std::time::{Duration, Instant};

fn main() -> Result<()> {
    println!("=== FileTree RocksDB POC Performance Test ===\n");

    let user_id = "test_rocksdb";

    println!("Step 1: Initialize RocksDB database...");
    let start = Instant::now();
    let tree = filetree_rocksdb::FileTreeRocksDB::init_user_db(user_id)?;
    let init_time = start.elapsed();
    println!("  ✓ Init time: {:?}", init_time);

    println!("\nStep 2: Insert 1,000 nodes (single insert)...");
    let nodes: Vec<filetree_rocksdb::FileNode> = (0..1000)
        .map(|i| filetree_rocksdb::FileTreeRocksDB::new_folder(&format!("folder_{}", i), None))
        .collect();

    let start = Instant::now();
    for node in &nodes {
        tree.insert_node(node)?;
    }
    let single_insert_time = start.elapsed();
    let single_throughput = nodes.len() as f64 / single_insert_time.as_secs_f64();
    println!("  ✓ Single insert: {:?}", single_insert_time);
    println!("  ✓ Throughput: {:.2} nodes/sec", single_throughput);

    println!("\nStep 3: Insert 10,000 nodes (batch insert)...");
    let nodes_batch: Vec<filetree_rocksdb::FileNode> = (0..10000)
        .map(|i| {
            filetree_rocksdb::FileTreeRocksDB::new_folder(&format!("batch_folder_{}", i), None)
        })
        .collect();

    let start = Instant::now();
    tree.insert_node_batch(&nodes_batch)?;
    let batch_insert_time = start.elapsed();
    let batch_throughput = nodes_batch.len() as f64 / batch_insert_time.as_secs_f64();
    println!("  ✓ Batch insert: {:?}", batch_insert_time);
    println!("  ✓ Throughput: {:.2} nodes/sec", batch_throughput);

    println!("\nStep 4: Query single node (10,000 iterations)...");
    let test_node_id = &nodes_batch[5000].node_id;
    let start = Instant::now();
    for _ in 0..10000 {
        let _ = tree.get_node(test_node_id)?;
    }
    let query_time = start.elapsed();
    let query_latency = query_time.as_nanos() as f64 / 10000.0;
    println!("  ✓ Total time: {:?}", query_time);
    println!("  ✓ Average latency: {:.2} ns", query_latency);

    println!("\nStep 5: Load all nodes...");
    let start = Instant::now();
    let all_nodes = tree.load_all()?;
    let load_time = start.elapsed();
    println!("  ✓ Load time: {:?}", load_time);
    println!("  ✓ Nodes loaded: {}", all_nodes.len());

    println!("\nStep 6: Concurrent reads (single process, simulated)...");
    let db_path = filetree_rocksdb::FileTreeRocksDB::user_db_path(user_id);
    let start = Instant::now();

    for i in 0..10000 {
        let node_id = format!("batch_folder_{}", i % 1000);
        let _ = tree.get_node(&node_id)?;
    }

    let concurrent_time = start.elapsed();
    let concurrent_ops = 10000;
    let concurrent_throughput = concurrent_ops as f64 / concurrent_time.as_secs_f64();
    println!("  ✓ Concurrent time: {:?}", concurrent_time);
    println!("  ✓ Total ops: {}", concurrent_ops);
    println!("  ✓ Throughput: {:.2} ops/sec", concurrent_throughput);

    println!("\nStep 7: Database size...");
    let db_size = get_directory_size(&db_path)?;
    println!(
        "  ✓ DB size: {} bytes ({:.2} MB)",
        db_size,
        db_size as f64 / 1024.0 / 1024.0
    );
    println!("  ✓ Nodes count: {}", tree.count_nodes()?);

    println!("\n=== Performance Summary ===");
    println!(
        "Single insert: {:?} ({:.2} nodes/sec)",
        single_insert_time, single_throughput
    );
    println!(
        "Batch insert: {:?} ({:.2} nodes/sec)",
        batch_insert_time, batch_throughput
    );
    println!("Query latency: {:.2} ns", query_latency);
    println!("Concurrent reads: {:.2} ops/sec", concurrent_throughput);
    println!("DB size: {:.2} MB", db_size as f64 / 1024.0 / 1024.0);

    println!("\nStep 8: Cleanup...");
    std::fs::remove_dir_all(&db_path)?;
    println!("  ✓ Test database removed");

    println!("\n✅ POC Test completed successfully!");

    Ok(())
}

fn get_directory_size(path: &str) -> Result<u64> {
    let mut total_size = 0;
    for entry in std::fs::read_dir(path)? {
        let entry = entry?;
        let metadata = entry.metadata()?;
        if metadata.is_file() {
            total_size += metadata.len();
        } else if metadata.is_dir() {
            total_size += get_directory_size(entry.path().to_str().unwrap())?;
        }
    }
    Ok(total_size)
}