//! Two routers of one AS, each taking in full tables from its own peers
//! and sending the other every path (ADD-PATH both ways, `all_paths`), as
//! a mesh of routers that each hold the network's whole view do. How
//! long until each router has the other's routes is what this measures:
//! the exporter reading the table's journal while sessions fill the
//! table, which nothing else exercises at this size.
//!
//! Then a new peer takes the first router's table, as a collector or a
//! router coming back does: how soon its first UPDATE arrives and how
//! long the table takes are measured too. It is a peer of the routers'
//! AS over IPv4, so what it gets is the first router's IPv4 routes from
//! its own peers: the other router's are not reflected to it, and IPv6
//! routes need an IPv6 session.
//!
//! Run with `cargo bench -p bgp-bench --bench mesh`. `BENCH_PEERS`
//! (default 4) feed the first router and `BENCH_PEERS_B` (default 2) the
//! second; `BENCH_SCALE_DOWN` shrinks the table. The routers peer over
//! IPv6 loopback and the feeders come over IPv4 loopback; the new peer
//! comes from 127.0.0.2, which Linux has and macOS does not, so that
//! part is skipped on macOS. With `BENCH_RECEIVER_EARLY=1` the new peer
//! comes as the feeders start instead, so that its exporter goes
//! through the table while the table is being filled, as a router's
//! do after a restart. `BENCH_COLLECTORS` (default 0) more peers of the
//! first router, from 127.0.1.0/24, are peers in the sense of RFC 9234
//! told of learned routes, and the feeders' routes are marked as having
//! come across from a peer, so none of them goes: as many exporters
//! following the table while it fills and passing nearly all of it
//! over, as a router with many peers at exchanges has.

use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use bgp_bench::{Prefixes, Shape, announcements, encoded, env};
use bgp_integration_tests::peer_config;
use bgp_integration_tests::test_peer::TestPeer;
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::attributes::PathAttribute;
use bgp_packet::open::{AddPathEntry, Capability, SendReceive};
use bgp_server::config::PeerConfig;
use bgp_server::family::SUPPORTED;
use bgp_server::policy::{ExportPolicy, ImportPolicy, PeerRole};
use bgp_server::session::SessionRegistry;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};

/// An AS number that is in none of the synthetic paths.
const ASN: u32 = 4_200_000_000;

/// Where the peer that takes the whole table connects from.
const RECEIVER: Ipv4Addr = Ipv4Addr::new(127, 0, 0, 2);

/// Where the collectors connect from: this and the addresses after it.
fn collector_address(n: u32) -> Ipv4Addr {
    Ipv4Addr::new(127, 0, 1, n as u8)
}

/// The collectors' AS numbers start here.
const COLLECTOR_ASN: u32 = 65000;

/// What the routers send each other, and the receiver: every path.
fn every_path() -> Arc<ExportPolicy> {
    Arc::new(ExportPolicy {
        learned_routes: true,
        all_paths: true,
        ..ExportPolicy::default()
    })
}

fn free_port() -> u16 {
    std::net::TcpListener::bind((Ipv6Addr::LOCALHOST, 0))
        .expect("a port")
        .local_addr()
        .expect("an address")
        .port()
}

/// A router of the mesh, listening for feeders on IPv4 loopback and for
/// the other router on IPv6 loopback at `port`, which it dials at
/// `other`.
async fn router(
    name: &'static str,
    id: u8,
    port: u16,
    other: u16,
    with_receiver: bool,
    collectors: u32,
) -> TestServer {
    TestServer::start_with(ASN, |c| {
        // Two routers of one AS must not share a BGP identifier.
        c.router_id = Ipv4Addr::new(192, 0, 2, id);
        c.cluster_id = c.router_id;
        c.listen_endpoints = vec![
            SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 0),
            SocketAddr::new(Ipv6Addr::LOCALHOST.into(), port),
        ];
        c.peers = vec![PeerConfig {
            add_path: Some(SendReceive::Both),
            import: ImportPolicy::permissive(),
            export: every_path(),
            ..peer_config(
                name,
                SocketAddr::new(Ipv6Addr::LOCALHOST.into(), other),
                ASN,
            )
        }];
        if with_receiver {
            c.peers.push(PeerConfig {
                add_path: Some(SendReceive::Send),
                export: every_path(),
                ..peer_config("receiver", SocketAddr::new(RECEIVER.into(), 179), ASN)
            });
        }
        for n in 1..=collectors {
            c.peers.push(PeerConfig {
                role: Some(PeerRole::Peer),
                export: Arc::new(ExportPolicy {
                    learned_routes: true,
                    all_paths: false,
                    ..ExportPolicy::default()
                }),
                ..peer_config(
                    "collector",
                    SocketAddr::new(collector_address(n).into(), 179),
                    COLLECTOR_ASN + n,
                )
            });
        }
        c.connect_retry = Duration::from_millis(200);
    })
    .await
    .expect("a router")
}

/// Connects `count` feeders to `router`, numbered from `first`, and
/// returns their sessions with the bytes each is to send.
async fn feeders(
    router: &TestServer,
    prefixes: &Prefixes,
    shape: Shape,
    first: u32,
    count: u32,
    marked: bool,
) -> Vec<(TestPeer, Vec<u8>)> {
    let mut sessions = Vec::new();
    for peer in first..first + count {
        let mut updates = announcements(prefixes, shape, peer);
        if marked {
            // As having come across from a peer (RFC 9234): not for the
            // collectors, which are peers too.
            for update in &mut updates {
                update
                    .path_attributes
                    .push(PathAttribute::OnlyToCustomer(64512 + peer));
            }
        }
        let stream: Vec<u8> = encoded(&updates)
            .iter()
            .flat_map(|m| m.iter().copied())
            .collect();
        let mut session = TestPeer::connect(router.addr).await.expect("a connection");
        // No hold timer: a feeder has nothing to say after its table.
        session
            .handshake(64512 + peer, Ipv4Addr::new(10, 0, 0, peer as u8), 0)
            .await
            .expect("a session");
        sessions.push((session, stream));
    }
    sessions
}

fn main() {
    // The routers' logs, with `BENCH_LOG` as a `RUST_LOG` filter.
    if let Ok(filter) = std::env::var("BENCH_LOG") {
        tracing_subscriber::fmt()
            .with_env_filter(filter)
            .with_writer(std::io::stderr)
            .init();
    }
    let peers_a = env("BENCH_PEERS", 4) as u32;
    let peers_b = env("BENCH_PEERS_B", 2) as u32;
    let shape = Shape::INTERNET.scaled_down(env("BENCH_SCALE_DOWN", 1));
    let per_peer = shape.prefixes();
    println!(
        "== Two routers of AS{ASN}, every path both ways: A fed by {peers_a} peers, B by {peers_b}, {} prefixes each",
        per_peer
    );
    let prefixes = Prefixes::generate(shape);

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("a runtime");
    runtime.block_on(async {
        let (port_a, port_b) = (free_port(), free_port());
        let collectors = env("BENCH_COLLECTORS", 0) as u32;
        let a = router("b", 1, port_a, port_b, true, collectors).await;
        let b = router("a", 2, port_b, port_a, false, 0).await;
        a.wait_for_sessions(1).await.expect("the mesh session at A");
        b.wait_for_sessions(1).await.expect("the mesh session at B");
        // The collectors are there before the feeders, as a router's
        // peers are before its tables arrive after a restart.
        let mut draining = Vec::new();
        for n in 1..=collectors {
            let socket = tokio::net::TcpSocket::new_v4().expect("a socket");
            if let Err(e) = socket.bind(SocketAddr::new(collector_address(n).into(), 0)) {
                println!(
                    "collectors: skipped, {} is not ours ({e})",
                    collector_address(n)
                );
                break;
            }
            let stream = socket.connect(a.addr).await.expect("A listens");
            let mut session = TestPeer::from_stream(stream);
            session
                .handshake(COLLECTOR_ASN + n, Ipv4Addr::new(192, 0, 3, n as u8), 0)
                .await
                .expect("A takes a collector's session");
            let (mut reader, writer) = session.into_stream().into_split();
            draining.push((
                writer,
                tokio::spawn(async move {
                    let mut sink = vec![0u8; 1 << 16];
                    while matches!(reader.read(&mut sink).await, Ok(n) if n > 0) {}
                }),
            ));
        }
        if collectors > 0 {
            a.wait_for_sessions(1 + draining.len())
                .await
                .expect("the collectors' sessions at A");
        }

        let feed_a = feeders(&a, &prefixes, shape, 1, peers_a, collectors > 0).await;
        let feed_b = feeders(&b, &prefixes, shape, 1 + peers_a, peers_b, false).await;
        let own_a = per_peer * peers_a as usize;
        let own_b = per_peer * peers_b as usize;

        let receiver_routes = shape.ipv4_prefixes * peers_a as usize;
        let started = Instant::now();
        // Early: as the feeders start, or (2) once A has its own peers'
        // routes and B's are still coming, as a router's sessions come
        // up after a restart while its peers' tables arrive.
        let when = env("BENCH_RECEIVER_EARLY", 0);
        let early_receiver = (when > 0).then(|| {
            let (addr, registry, rib) = (a.addr, a.registry.clone(), Arc::clone(&a.rib));
            tokio::spawn(async move {
                while when == 2 && rib.len() < own_a {
                    tokio::time::sleep(Duration::from_millis(5)).await;
                }
                resync(addr, registry, receiver_routes).await
            })
        });
        let mut sending = Vec::new();
        for (mut session, stream) in feed_a.into_iter().chain(feed_b) {
            sending.push(tokio::spawn(async move {
                session.send_raw(&stream).await.expect("the router reads");
                session
            }));
        }
        let mut a_own = None;
        let mut b_own = None;
        let mut a_all = None;
        let mut b_all = None;
        while a_all.is_none() || b_all.is_none() {
            tokio::time::sleep(Duration::from_millis(5)).await;
            let (at_a, at_b) = (a.rib.len(), b.rib.len());
            let now = started.elapsed();
            if a_own.is_none() && at_a >= own_a {
                a_own = Some(now);
            }
            if b_own.is_none() && at_b >= own_b {
                b_own = Some(now);
            }
            if a_all.is_none() && at_a >= own_a + own_b {
                a_all = Some(now);
            }
            if b_all.is_none() && at_b >= own_a + own_b {
                b_all = Some(now);
            }
            if now > Duration::from_secs(env("BENCH_DEADLINE", 600) as u64) {
                let sessions = |router: &TestServer| -> Vec<String> {
                    router
                        .registry
                        .list()
                        .iter()
                        .map(|s| {
                            format!(
                                "{} routes {} rejected {}",
                                s.peer,
                                s.counters.routes(),
                                s.counters.rejected()
                            )
                        })
                        .collect()
                };
                panic!(
                    "not all routes arrived: A has {at_a} {:?}, B has {at_b} {:?}",
                    sessions(&a),
                    sessions(&b)
                );
            }
        }
        let row = |what: &str, when: Option<Duration>, routes: usize| {
            let when = when.unwrap_or_default();
            println!(
                "{what:<52} {:>8.3} s   {:.2} M routes/s",
                when.as_secs_f64(),
                routes as f64 / when.as_secs_f64() / 1e6
            );
        };
        row("A has its own peers' routes", a_own, own_a);
        row("B has its own peers' routes", b_own, own_b);
        row(
            "A has B's routes too (B's exporter, A's session)",
            a_all,
            own_b,
        );
        row(
            "B has A's routes too (A's exporter, B's session)",
            b_all,
            own_a,
        );
        // The feeders' sessions stay up until the end: a feeder going
        // would have its routes withdrawn from both tables meanwhile.
        let mut feeders = Vec::new();
        for session in sending {
            feeders.push(session.await.expect("a feeder is done"));
        }
        match early_receiver {
            Some(receiver) => receiver.await.expect("the receiver is done"),
            None => resync(a.addr, a.registry.clone(), receiver_routes).await,
        }
        drop(feeders);
        drop(draining);
    });
}

/// A new peer takes the table of the router at `addr`, whose sessions
/// are in `registry`: `routes` routes with every path. How long until
/// the first UPDATE, and until all of it is sent.
async fn resync(addr: SocketAddr, registry: SessionRegistry, routes: usize) {
    let socket = tokio::net::TcpSocket::new_v4().expect("a socket");
    if let Err(e) = socket.bind(SocketAddr::new(RECEIVER.into(), 0)) {
        println!("A new peer taking A's routes: skipped, {RECEIVER} is not ours ({e})");
        return;
    }
    let stream = socket.connect(addr).await.expect("A listens");
    let mut session = TestPeer::from_stream(stream);
    let mut capabilities = vec![Capability::FourByteAsn { asn: ASN }];
    capabilities.extend(SUPPORTED.into_iter().map(Capability::multiprotocol));
    capabilities.push(Capability::AddPath {
        entries: SUPPORTED
            .iter()
            .map(|family| AddPathEntry {
                afi: family.afi,
                safi: family.safi,
                send_recv: SendReceive::Receive,
            })
            .collect(),
    });
    session
        .handshake_with(ASN, Ipv4Addr::new(192, 0, 2, 9), 90, capabilities)
        .await
        .expect("A takes the receiver's session");
    let started = Instant::now();
    let (reader, mut writer) = session.into_stream().into_split();
    // KEEPALIVEs, so that the session outlives the hold time.
    let keepalives = tokio::spawn(async move {
        let mut keepalive = [0xffu8; 19];
        keepalive[16..19].copy_from_slice(&[0, 19, 4]);
        loop {
            tokio::time::sleep(Duration::from_secs(20)).await;
            if writer.write_all(&keepalive).await.is_err() {
                break;
            }
        }
    });
    // When the first UPDATE arrived, in microseconds; 0 until it has.
    let first = Arc::new(AtomicU64::new(0));
    let updates = Arc::new(AtomicU64::new(0));
    let counting = tokio::spawn({
        let (first, updates) = (Arc::clone(&first), Arc::clone(&updates));
        async move {
            let mut reader = BufReader::with_capacity(1 << 20, reader);
            let mut header = [0u8; 19];
            let mut body = vec![0u8; 4096];
            while reader.read_exact(&mut header).await.is_ok() {
                let length = usize::from(u16::from_be_bytes([header[16], header[17]]));
                if reader.read_exact(&mut body[..length - 19]).await.is_err() {
                    break;
                }
                match header[18] {
                    // An UPDATE (RFC 4271 §4.1).
                    2 => {
                        let _ = first.compare_exchange(
                            0,
                            started.elapsed().as_micros() as u64,
                            Ordering::Relaxed,
                            Ordering::Relaxed,
                        );
                        let n = updates.fetch_add(1, Ordering::Relaxed) + 1;
                        if n % 1_000_000 == 0 {
                            println!(
                                "  the receiver has read {n} UPDATEs after {:.2} s",
                                started.elapsed().as_secs_f64()
                            );
                        }
                    }
                    3 => println!(
                        "the receiver was sent a NOTIFICATION, error {} subcode {}, after {} UPDATEs",
                        body[0],
                        body[1],
                        updates.load(Ordering::Relaxed)
                    ),
                    _ => {}
                }
            }
            println!(
                "the receiver's connection closed after {} UPDATEs",
                updates.load(Ordering::Relaxed)
            );
        }
    });
    let deadline = Duration::from_secs(env("BENCH_DEADLINE", 600) as u64);
    let advertised = loop {
        tokio::time::sleep(Duration::from_millis(5)).await;
        let advertised = registry
            .list()
            .iter()
            .find(|s| s.peer.addr.ip() == RECEIVER)
            .map_or(0, |s| s.counters.advertised());
        if advertised as usize >= routes {
            break advertised;
        }
        if started.elapsed() > deadline {
            let sessions: Vec<String> = registry
                .list()
                .iter()
                .map(|s| {
                    format!(
                        "{} {:?} export {:?} add_path_send {:?} advertised {}",
                        s.peer,
                        s.name,
                        s.export,
                        s.add_path_send,
                        s.counters.advertised()
                    )
                })
                .collect();
            panic!(
                "the receiver has {advertised} of {routes} routes and read {} UPDATEs; sessions at A: {sessions:#?}",
                updates.load(Ordering::Relaxed)
            );
        }
    };
    let took = started.elapsed();
    counting.abort();
    keepalives.abort();
    let first = Duration::from_micros(first.load(Ordering::Relaxed));
    println!(
        "{:<52} {:>8.3} s   {:.2} M routes/s   first UPDATE after {:.0} ms, {} UPDATEs",
        "A new peer has A's IPv4 routes (A's exporter)",
        took.as_secs_f64(),
        advertised as f64 / took.as_secs_f64() / 1e6,
        first.as_secs_f64() * 1e3,
        updates.load(Ordering::Relaxed),
    );
}