aria2-rust-pro/crates/aria2-rust-pro-tests/benches/rpc_pressure/bt_visibility.rs

#![expect(
    clippy::redundant_pub_crate,
    reason = "criterion bench entry points are re-exported only to the private bench root module"
)]

use super::support::{
    BT_TORRENT_FIXTURE, BenchmarkId, Criterion, InProcessRpcDispatcher, RpcMethod, RpcValue,
    RuntimeConfig, Throughput, TorrentPeerModel, TrackerPeerListModel, TrackerResponseModel,
    rpc_request,
};

#[derive(Clone, Copy, Debug)]
struct BtVisibilityPressureScenario {
    task_count: usize,
    rounds: usize,
}

fn seed_bt_visibility_dispatcher(
    scenario: BtVisibilityPressureScenario,
) -> (InProcessRpcDispatcher, Vec<String>) {
    let runtime = RuntimeConfig {
        allow_jsonrpc: true,
        allow_xmlrpc: true,
        split: 4,
        max_connections_per_server: 4,
        max_connection_per_server: 4,
        min_split_size: 1024,
        piece_length: 1024,
        ..RuntimeConfig::default()
    };
    let mut dispatcher = InProcessRpcDispatcher::with_runtime(runtime);
    let mut gids = Vec::with_capacity(scenario.task_count);
    for index in 0..scenario.task_count {
        let add_response = dispatcher.dispatch_json(rpc_request(
            RpcMethod::Aria2AddTorrent,
            vec![RpcValue::String(BT_TORRENT_FIXTURE.to_owned())],
        ));
        let gid = match add_response.result {
            Some(RpcValue::String(gid)) => gid,
            other => panic!("unexpected addTorrent result in visibility seed: {other:?}"),
        };
        dispatcher
            .apply_tracker_announce_result(
                &gid,
                &TrackerResponseModel {
                    peers: TrackerPeerListModel {
                        interval_sec: 900,
                        peers: vec![TorrentPeerModel {
                            ip: format!("198.51.100.{}", (index % 200) + 1),
                            port: 51413 + (index as u16 % 32),
                            peer_id: None,
                            client_name: None,
                            interested: false,
                            choked: false,
                        }],
                        min_interval_sec: None,
                        tracker_id: Some(format!("bench-{index}")),
                    },
                    scrape: None,
                },
            )
            .expect("tracker seed should populate visible bt peers");
        dispatcher
            .apply_bt_runtime_tick(
                &gid,
                8_192 + (index as u64 % 4) * 2_048,
                if index % 3 == 0 { 4_096 } else { 0 },
                512 + index as u64,
                128 + index as u64,
                if index % 3 == 0 { 5 } else { 0 },
                if index % 3 == 0 { 5 } else { 0 },
                index % 3 == 0,
                Some(2 + (index % 6) as u32),
            )
            .expect("runtime tick should seed visible bt progress");
        if index % 3 == 0 {
            dispatcher
                .set_bt_seeding_state(&gid, true, Some(1_000 + index as u64))
                .expect("seeded torrents should enter seeding");
            dispatcher
                .tick_bt_runtime_clock(&gid, 1_010 + index as u64, true)
                .expect("seeded torrents should advance share clocks");
        }
        gids.push(gid);
    }
    (dispatcher, gids)
}

fn run_bt_visibility_pressure(
    dispatcher: &mut InProcessRpcDispatcher,
    gids: &[String],
    scenario: BtVisibilityPressureScenario,
) -> usize {
    let mut calls = 0_usize;
    for round in 0..scenario.rounds {
        for (index, gid) in gids.iter().enumerate() {
            dispatcher
                .apply_bt_runtime_tick(
                    gid,
                    if round % 2 == 0 { 512 } else { 0 },
                    if round % 2 == 1 { 256 } else { 0 },
                    1_024 + round as u64 + index as u64,
                    256 + round as u64 + index as u64,
                    u64::from(index % 3 == 0),
                    u64::from(index % 3 == 0),
                    index % 3 == 0,
                    Some(2 + ((round + index) % 6) as u32),
                )
                .expect("pressure tick should keep bt runtime visible");

            let tell_status = dispatcher.dispatch_json(rpc_request(
                RpcMethod::Aria2TellStatus,
                vec![RpcValue::String(gid.clone())],
            ));
            match tell_status.result {
                Some(RpcValue::Object(payload)) => {
                    assert_eq!(payload.get("isBt"), Some(&RpcValue::Bool(true)));
                    assert!(payload.contains_key("announceList"));
                    assert!(payload.contains_key("bitfield"));
                    assert!(payload.contains_key("shareRatio"));
                    assert!(payload.contains_key("shareTime"));
                    assert!(payload.contains_key("files"));
                    assert!(payload.contains_key("seeder"));
                }
                other => panic!("unexpected tellStatus payload in bt visibility bench: {other:?}"),
            }
            calls += 1;

            let get_files = dispatcher.dispatch_json(rpc_request(
                RpcMethod::Aria2GetFiles,
                vec![RpcValue::String(gid.clone())],
            ));
            match get_files.result {
                Some(RpcValue::Array(entries)) => match entries.first() {
                    Some(RpcValue::Object(file)) => {
                        assert!(file.contains_key("selected"));
                        assert!(file.contains_key("completedLength"));
                        assert!(file.contains_key("bitfield"));
                    }
                    other => {
                        panic!("unexpected getFiles payload in bt visibility bench: {other:?}")
                    }
                },
                other => panic!("unexpected getFiles result in bt visibility bench: {other:?}"),
            }
            calls += 1;

            let get_peers = dispatcher.dispatch_json(rpc_request(
                RpcMethod::Aria2GetPeers,
                vec![RpcValue::String(gid.clone())],
            ));
            match get_peers.result {
                Some(RpcValue::Array(entries)) => match entries.first() {
                    Some(RpcValue::Object(peer)) => {
                        assert!(peer.contains_key("ip"));
                        assert!(peer.contains_key("port"));
                        assert!(peer.contains_key("peerChoking"));
                    }
                    other => {
                        panic!("unexpected getPeers payload in bt visibility bench: {other:?}")
                    }
                },
                other => panic!("unexpected getPeers result in bt visibility bench: {other:?}"),
            }
            calls += 1;
        }
    }
    calls
}

pub(super) fn bench_bt_visibility_pressure(c: &mut Criterion) {
    let mut group = c.benchmark_group("bt_visibility_pressure");
    for scenario in [
        BtVisibilityPressureScenario {
            task_count: 32,
            rounds: 4,
        },
        BtVisibilityPressureScenario {
            task_count: 64,
            rounds: 4,
        },
    ] {
        group.throughput(Throughput::Elements(
            (scenario.task_count * scenario.rounds * 3) as u64,
        ));
        group.bench_with_input(
            BenchmarkId::new("bt_visibility", scenario.task_count),
            &scenario,
            |b, &scenario| {
                b.iter_batched(
                    || seed_bt_visibility_dispatcher(scenario),
                    |(mut dispatcher, gids)| {
                        let calls = run_bt_visibility_pressure(&mut dispatcher, &gids, scenario);
                        assert_eq!(calls, scenario.task_count * scenario.rounds * 3);
                    },
                    criterion::BatchSize::SmallInput,
                );
            },
        );
    }
    group.finish();
}