photoncloud-monorepo/chainfire/crates/chainfire-server/benches/kv_bench.rs

use chainfire_client::ChainFireClient;
use chainfire_server::config::{ClusterConfig, NetworkConfig, NodeConfig, RaftConfig, ServerConfig, StorageConfig};
use chainfire_server::node::Node;
use chainfire_types::RaftRole;
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use std::time::Duration;
use tempfile::TempDir;
use tokio::runtime::Runtime;

const VALUE_SIZE: usize = 1024; // 1KB
const NUM_KEYS_THROUGHPUT: usize = 10_000; // 10K for throughput tests
const NUM_KEYS_LATENCY: usize = 100; // 100 for latency tests

fn create_test_node(temp_dir: &TempDir) -> (Node, Runtime) {
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()
        .unwrap();

    let config = ServerConfig {
        node: NodeConfig {
            id: 1,
            name: "benchmark-node".to_string(),
            role: "control_plane".to_string(),
        },
        cluster: ClusterConfig {
            id: 1,
            bootstrap: true,
            initial_members: vec![],
        },
        network: NetworkConfig {
            api_addr: "127.0.0.1:2379".parse().unwrap(),
            raft_addr: "127.0.0.1:2380".parse().unwrap(),
            gossip_addr: "127.0.0.1:2381".parse().unwrap(),
            tls: None,
        },
        storage: StorageConfig {
            data_dir: temp_dir.path().to_path_buf(),
        },
        raft: RaftConfig {
            role: RaftRole::Voter,
            tick_interval_ms: 100,
            election_timeout_ticks: 10,
            heartbeat_interval_ticks: 3,
            snapshot_interval_secs: 3600,
            max_applied_log_to_keep: 1000,
        },
    };

    let node = rt.block_on(async { Node::new(config).await.unwrap() });

    (node, rt)
}

fn bench_put_throughput(c: &mut Criterion) {
    let temp_dir = TempDir::new().unwrap();
    let (node, rt) = create_test_node(&temp_dir);

    // Start server
    let server_handle = rt.spawn(async move {
        node.run().await.unwrap();
    });

    // Give server time to start
    std::thread::sleep(Duration::from_millis(500));

    // Create client
    let mut client = rt.block_on(async {
        ChainFireClient::connect("http://127.0.0.1:2379")
            .await
            .unwrap()
    });

    let value = vec![b'x'; VALUE_SIZE];

    let mut group = c.benchmark_group("put_throughput");
    group.throughput(Throughput::Elements(NUM_KEYS_THROUGHPUT as u64));
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(30));

    group.bench_function(BenchmarkId::from_parameter(NUM_KEYS_THROUGHPUT), |b| {
        b.iter(|| {
            rt.block_on(async {
                for i in 0..NUM_KEYS_THROUGHPUT {
                    let key = format!("bench_key_{}", i);
                    client.put(black_box(&key), black_box(&value)).await.unwrap();
                }
            })
        });
    });

    group.finish();

    // Cleanup
    server_handle.abort();
    drop(rt);
}

fn bench_get_throughput(c: &mut Criterion) {
    let temp_dir = TempDir::new().unwrap();
    let (node, rt) = create_test_node(&temp_dir);

    // Start server
    let server_handle = rt.spawn(async move {
        node.run().await.unwrap();
    });

    // Give server time to start
    std::thread::sleep(Duration::from_millis(500));

    // Create client and populate data
    let mut client = rt.block_on(async {
        ChainFireClient::connect("http://127.0.0.1:2379")
            .await
            .unwrap()
    });

    let value = vec![b'x'; VALUE_SIZE];

    // Pre-populate keys
    rt.block_on(async {
        for i in 0..NUM_KEYS_THROUGHPUT {
            let key = format!("bench_key_{}", i);
            client.put(&key, &value).await.unwrap();
        }
    });

    let mut group = c.benchmark_group("get_throughput");
    group.throughput(Throughput::Elements(NUM_KEYS_THROUGHPUT as u64));
    group.sample_size(10);
    group.measurement_time(Duration::from_secs(30));

    group.bench_function(BenchmarkId::from_parameter(NUM_KEYS_THROUGHPUT), |b| {
        b.iter(|| {
            rt.block_on(async {
                for i in 0..NUM_KEYS_THROUGHPUT {
                    let key = format!("bench_key_{}", i);
                    let _ = client.get(black_box(&key)).await.unwrap();
                }
            })
        });
    });

    group.finish();

    // Cleanup
    server_handle.abort();
    drop(rt);
}

fn bench_put_latency(c: &mut Criterion) {
    let temp_dir = TempDir::new().unwrap();
    let (node, rt) = create_test_node(&temp_dir);

    // Start server
    let server_handle = rt.spawn(async move {
        node.run().await.unwrap();
    });

    // Give server time to start
    std::thread::sleep(Duration::from_millis(500));

    // Create client
    let mut client = rt.block_on(async {
        ChainFireClient::connect("http://127.0.0.1:2379")
            .await
            .unwrap()
    });

    let value = vec![b'x'; VALUE_SIZE];

    let mut group = c.benchmark_group("put_latency");
    group.sample_size(1000); // Larger sample for better p99/p999 estimates
    group.measurement_time(Duration::from_secs(60));

    group.bench_function("single_put", |b| {
        let mut key_counter = 0;
        b.iter(|| {
            let key = format!("latency_key_{}", key_counter);
            key_counter += 1;
            rt.block_on(async {
                client.put(black_box(&key), black_box(&value)).await.unwrap();
            })
        });
    });

    group.finish();

    // Cleanup
    server_handle.abort();
    drop(rt);
}

criterion_group!(benches, bench_put_throughput, bench_get_throughput, bench_put_latency);
criterion_main!(benches);