//! 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::report::row;
use bgp_bench::{Prefixes, Shape, announcements, encoded, env, table_stream};
use bgp_integration_tests::peer_config;
use bgp_integration_tests::test_peer::{TestPeer, collector_capabilities};
use bgp_integration_tests::test_server::TestServer;
use bgp_packet::attributes::PathAttribute;
use bgp_packet::message::msg_type;
use bgp_packet::open::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};
use tokio::net::tcp::OwnedWriteHalf;
use tokio::task::JoinHandle;

/// 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;

/// How often the tables are looked at while they fill.
const POLL: Duration = Duration::from_millis(5);

/// 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_port`, knowing it as `other_name`.
async fn router(
    other_name: &'static str,
    id: u8,
    port: u16,
    other_port: 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(
                other_name,
                SocketAddr::new(Ipv6Addr::LOCALHOST.into(), other_port),
                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 stream = if marked {
            // As having come across from a peer (RFC 9234): not for the
            // collectors, which are peers too.
            let mut updates = announcements(prefixes, shape, peer);
            for update in &mut updates {
                update
                    .path_attributes
                    .push(PathAttribute::OnlyToCustomer(64512 + peer));
            }
            encoded(&updates)
                .iter()
                .flat_map(|m| m.iter().copied())
                .collect()
        } else {
            table_stream(prefixes, shape, peer)
        };
        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
}

/// Every session of `registry`, one line each, for the panic when a
/// deadline passes: where the routes have got to, and what each session
/// was to be sent.
fn sessions_report(registry: &SessionRegistry) -> Vec<String> {
    registry
        .list()
        .iter()
        .map(|s| {
            format!(
                "{} {:?} export {:?} add_path_send {:?} routes {} rejected {} advertised {}",
                s.peer,
                s.name,
                s.export,
                s.add_path_send,
                s.counters.routes(),
                s.counters.rejected(),
                s.counters.advertised()
            )
        })
        .collect()
}

/// How big a run is, from the environment.
struct Settings {
    peers_a: u32,
    peers_b: u32,
    shape: Shape,
    collectors: u32,
}

impl Settings {
    fn from_env() -> Self {
        Self {
            peers_a: env("BENCH_PEERS", 4) as u32,
            peers_b: env("BENCH_PEERS_B", 2) as u32,
            shape: Shape::INTERNET.scaled_down(env("BENCH_SCALE_DOWN", 1)),
            collectors: env("BENCH_COLLECTORS", 0) as u32,
        }
    }

    /// How many routes A's own feeders bring, and B's.
    fn own_routes(&self) -> (usize, usize) {
        let per_peer = self.shape.prefixes();
        (
            per_peer * self.peers_a as usize,
            per_peer * self.peers_b as usize,
        )
    }

    /// How many routes the receiver is to be sent: A's IPv4 routes from
    /// its own peers.
    fn receiver_routes(&self) -> usize {
        self.shape.ipv4_prefixes * self.peers_a as usize
    }
}

/// The two routers, up and peered, with the collectors' sessions once
/// they are connected.
struct Mesh {
    settings: Settings,
    prefixes: Prefixes,
    a: TestServer,
    b: TestServer,
    /// The collectors' connections, drained, held until the end.
    collectors: Vec<(OwnedWriteHalf, JoinHandle<()>)>,
}

/// The feeders sending, from when they began.
struct Feeding {
    started: Instant,
    /// Each hands its session back once its table is written.
    sending: Vec<JoinHandle<TestPeer>>,
}

/// When each router had what: its own peers' routes, and the other's too.
struct Tables {
    a_own: Duration,
    b_own: Duration,
    a_all: Duration,
    b_all: Duration,
}

impl Mesh {
    /// Starts the routers, with the collectors and receiver configured at
    /// A, and waits for them to peer.
    async fn start(settings: Settings, prefixes: Prefixes) -> Self {
        let (port_a, port_b) = (free_port(), free_port());
        let a = router("b", 1, port_a, port_b, true, settings.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");
        Self {
            settings,
            prefixes,
            a,
            b,
            collectors: Vec::new(),
        }
    }

    /// Connects the collectors to A. They are there before the feeders, as
    /// a router's peers are before its tables arrive after a restart. An
    /// address that is not this machine's ends the connecting, with a
    /// word.
    async fn connect_collectors(&mut self) {
        for n in 1..=self.settings.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(self.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");
            self.collectors.push(session.drain());
        }
        if !self.collectors.is_empty() {
            self.a
                .wait_for_sessions(1 + self.collectors.len())
                .await
                .expect("the collectors' sessions at A");
        }
    }

    /// Connects the feeders to both routers and starts them sending.
    async fn feed(&self) -> Feeding {
        let Settings {
            peers_a,
            peers_b,
            shape,
            ..
        } = self.settings;
        let marked = !self.collectors.is_empty();
        let feed_a = feeders(&self.a, &self.prefixes, shape, 1, peers_a, marked).await;
        let feed_b = feeders(&self.b, &self.prefixes, shape, 1 + peers_a, peers_b, false).await;
        let started = Instant::now();
        let sending = feed_a
            .into_iter()
            .chain(feed_b)
            .map(|(mut session, stream)| {
                tokio::spawn(async move {
                    session.send_raw(&stream).await.expect("the router reads");
                    session
                })
            })
            .collect();
        Feeding { started, sending }
    }

    /// The receiver taking A's table while it fills, if `BENCH_RECEIVER_EARLY`
    /// asks for that: as the feeders start (1), or once A has its own
    /// peers' routes and B's are still coming (2), as a router's sessions
    /// come up after a restart while its peers' tables arrive.
    fn early_receiver(&self, deadline: Duration) -> Option<JoinHandle<()>> {
        let when = env("BENCH_RECEIVER_EARLY", 0);
        if when == 0 {
            return None;
        }
        let (own_a, _) = self.settings.own_routes();
        let routes = self.settings.receiver_routes();
        let (addr, registry, rib) = (
            self.a.addr,
            self.a.registry.clone(),
            Arc::clone(&self.a.rib),
        );
        Some(tokio::spawn(async move {
            while when == 2 && rib.len() < own_a {
                tokio::time::sleep(POLL).await;
            }
            resync(addr, registry, routes, deadline).await
        }))
    }

    /// Waits until both routers have every route, noting when each had
    /// what; or panics once `deadline` has passed since `started`.
    async fn wait_for_tables(&self, started: Instant, deadline: Duration) -> Tables {
        let (own_a, own_b) = self.settings.own_routes();
        let mut a_own = None;
        let mut b_own = None;
        let mut a_all = None;
        let mut b_all = None;
        let (a_all, b_all) = loop {
            tokio::time::sleep(POLL).await;
            let (at_a, at_b) = (self.a.rib.len(), self.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 let (Some(a_all), Some(b_all)) = (a_all, b_all) {
                break (a_all, b_all);
            }
            if now > deadline {
                panic!(
                    "not all routes arrived: A has {at_a} {:?}, B has {at_b} {:?}",
                    sessions_report(&self.a.registry),
                    sessions_report(&self.b.registry)
                );
            }
        };
        Tables {
            a_own: a_own.unwrap_or_default(),
            b_own: b_own.unwrap_or_default(),
            a_all,
            b_all,
        }
    }

    /// One row for each of the four moments of `tables`, with the rate the
    /// routes that made the moment came at.
    fn report(&self, tables: &Tables) {
        let (own_a, own_b) = self.settings.own_routes();
        let rate = |when: Duration, routes: usize| {
            format!("{:.2} M routes/s", routes as f64 / when.as_secs_f64() / 1e6)
        };
        row(
            "A has its own peers' routes",
            tables.a_own,
            rate(tables.a_own, own_a),
        );
        row(
            "B has its own peers' routes",
            tables.b_own,
            rate(tables.b_own, own_b),
        );
        row(
            "A has B's routes too (B's exporter, A's session)",
            tables.a_all,
            rate(tables.a_all, own_b),
        );
        row(
            "B has A's routes too (A's exporter, B's session)",
            tables.b_all,
            rate(tables.b_all, own_a),
        );
    }
}

impl Feeding {
    /// Waits for every feeder to have written its table, and returns the
    /// sessions, which are to stay up until the end: a feeder going would
    /// have its routes withdrawn from both tables meanwhile.
    async fn finished(self) -> Vec<TestPeer> {
        let mut feeders = Vec::new();
        for session in self.sending {
            feeders.push(session.await.expect("a feeder is done"));
        }
        feeders
    }
}

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 settings = Settings::from_env();
    let deadline = Duration::from_secs(env("BENCH_DEADLINE", 600) as u64);
    println!(
        "== Two routers of AS{ASN}, every path both ways: A fed by {} peers, B by {}, {} prefixes each",
        settings.peers_a,
        settings.peers_b,
        settings.shape.prefixes()
    );
    let prefixes = Prefixes::generate(settings.shape);

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("a runtime");
    runtime.block_on(async {
        let mut mesh = Mesh::start(settings, prefixes).await;
        mesh.connect_collectors().await;
        let early_receiver = mesh.early_receiver(deadline);
        let feeding = mesh.feed().await;
        let tables = mesh.wait_for_tables(feeding.started, deadline).await;
        mesh.report(&tables);
        let feeders = feeding.finished().await;
        match early_receiver {
            Some(receiver) => receiver.await.expect("the receiver is done"),
            None => {
                resync(
                    mesh.a.addr,
                    mesh.a.registry.clone(),
                    mesh.settings.receiver_routes(),
                    deadline,
                )
                .await
            }
        }
        drop(feeders);
        drop(mesh);
    });
}

/// The new peer taking a router's table: counting the UPDATEs it is sent
/// while keeping the session alive.
struct Receiver {
    started: Instant,
    /// When the first UPDATE arrived, in microseconds after `started`; 0
    /// until it has.
    first: Arc<AtomicU64>,
    updates: Arc<AtomicU64>,
    counting: JoinHandle<()>,
    keepalives: JoinHandle<()>,
}

impl Receiver {
    /// A session with the router at `addr` from [`RECEIVER`], as a
    /// collector of every path; or `None`, with a word, where that address
    /// is not this machine's to connect from.
    async fn connect(addr: SocketAddr) -> Option<TestPeer> {
        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 None;
        }
        let stream = socket.connect(addr).await.expect("A listens");
        let mut session = TestPeer::from_stream(stream);
        session
            .handshake_with(
                ASN,
                Ipv4Addr::new(192, 0, 2, 9),
                90,
                collector_capabilities(ASN, &SUPPORTED, false),
            )
            .await
            .expect("A takes the receiver's session");
        Some(session)
    }

    /// Starts counting what `session` is sent, from now. Every millionth
    /// UPDATE is noted, as is a NOTIFICATION or the connection closing:
    /// what went wrong, if the table never all arrives.
    fn count_updates(session: TestPeer) -> Self {
        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 keepalive = TestPeer::keepalive_bytes();
            loop {
                tokio::time::sleep(Duration::from_secs(20)).await;
                if writer.write_all(&keepalive).await.is_err() {
                    break;
                }
            }
        });
        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] {
                        msg_type::UPDATE => {
                            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()
                                );
                            }
                        }
                        msg_type::NOTIFICATION => 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)
                );
            }
        });
        Self {
            started,
            first,
            updates,
            counting,
            keepalives,
        }
    }

    /// Waits until the router whose sessions are in `registry` has
    /// advertised `routes` routes to the receiver, and returns how many it
    /// has; or panics once `deadline` has passed.
    async fn wait_for_advertised(
        &self,
        registry: &SessionRegistry,
        routes: usize,
        deadline: Duration,
    ) -> u64 {
        loop {
            tokio::time::sleep(POLL).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 {
                return advertised;
            }
            if self.started.elapsed() > deadline {
                panic!(
                    "the receiver has {advertised} of {routes} routes and read {} UPDATEs; sessions at A: {:#?}",
                    self.updates.load(Ordering::Relaxed),
                    sessions_report(registry)
                );
            }
        }
    }

    /// Prints the row for the table being all there, `advertised` routes
    /// of it, and stops counting.
    fn report(self, advertised: u64) {
        let took = self.started.elapsed();
        self.counting.abort();
        self.keepalives.abort();
        let first = Duration::from_micros(self.first.load(Ordering::Relaxed));
        row(
            "A new peer has A's IPv4 routes (A's exporter)",
            took,
            format!(
                "{:.2} M routes/s   first UPDATE after {:.0} ms, {} UPDATEs",
                advertised as f64 / took.as_secs_f64() / 1e6,
                first.as_secs_f64() * 1e3,
                self.updates.load(Ordering::Relaxed),
            ),
        );
    }
}

/// 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, deadline: Duration) {
    let Some(session) = Receiver::connect(addr).await else {
        return;
    };
    let receiver = Receiver::count_updates(session);
    let advertised = receiver
        .wait_for_advertised(&registry, routes, deadline)
        .await;
    receiver.report(advertised);
}