tidaldb/tidal-net/benches/transport_throughput.rs

// Benchmark harness: unwrap on known-good setup is idiomatic here.
#![allow(clippy::unwrap_used)]
//! Benchmark comparing `InProcessTransport` vs `GrpcTransport` throughput.
//!
//! Measures segments/sec for both transports to quantify gRPC overhead.

use std::{collections::HashMap, net::SocketAddr, thread, time::Duration};

use criterion::{Criterion, Throughput, criterion_group, criterion_main};
use tidal_net::{GrpcTransport, config::GrpcTransportConfig};
use tidaldb::replication::{
    WalSegmentId,
    in_process::InProcessTransportFactory,
    shard::{RegionId, ShardId},
    transport::{Transport, WalSegmentPayload},
};

const SEGMENT_SIZE: usize = 1024; // 1 KB payload per segment

fn make_payload(seqno: u64) -> WalSegmentPayload {
    WalSegmentPayload {
        id: WalSegmentId::new(RegionId::SINGLE, ShardId(0), seqno),
        bytes: vec![0xAB; SEGMENT_SIZE],
        event_count: 10,
        leader_last_seq: seqno,
        stream_baseline: 0,
        term: 0,
        leader_region: 0,
    }
}

fn free_addr() -> SocketAddr {
    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
    listener.local_addr().unwrap()
}

fn bench_transports(c: &mut Criterion) {
    let mut group = c.benchmark_group("transport_throughput");
    group.throughput(Throughput::Elements(1));

    // ── InProcessTransport ──────────────────────────────────────────────

    {
        let shards = vec![ShardId(0), ShardId(1)];
        let mut transports = InProcessTransportFactory::new(&shards).build();
        let t0 = transports.remove(&ShardId(0)).unwrap();
        let t1 = transports.remove(&ShardId(1)).unwrap();

        // Drain receiver in background.
        let drain = thread::spawn(move || while t1.recv_segment().is_some() {});

        let mut seq = 0u64;
        group.bench_function("in_process", |b| {
            b.iter(|| {
                seq += 1;
                t0.send_segment(ShardId(1), make_payload(seq)).unwrap();
            });
        });

        drop(t0);
        let _ = drain.join();
    }

    // ── GrpcTransport ───────────────────────────────────────────────────

    {
        let addr0 = free_addr();
        let addr1 = free_addr();

        let config0 = GrpcTransportConfig {
            local_shard: ShardId(0),
            listen_addr: addr0,
            peers: HashMap::from([(ShardId(1), addr1)]),
            insecure: true,
            ..Default::default()
        };
        let config1 = GrpcTransportConfig {
            local_shard: ShardId(1),
            listen_addr: addr1,
            peers: HashMap::from([(ShardId(0), addr0)]),
            insecure: true,
            ..Default::default()
        };

        let t0 = GrpcTransport::new(config0).expect("grpc transport 0");
        let t1 = GrpcTransport::new(config1).expect("grpc transport 1");
        thread::sleep(Duration::from_millis(200));

        // Drain receiver in background.
        let drain = thread::spawn(move || while t1.recv_segment().is_some() {});

        let mut seq = 0u64;
        group.bench_function("grpc_localhost", |b| {
            b.iter(|| {
                seq += 1;
                t0.send_segment(ShardId(1), make_payload(seq)).unwrap();
            });
        });

        drop(t0);
        let _ = drain.join();
    }

    group.finish();
}

criterion_group!(benches, bench_transports);
criterion_main!(benches);