ferrix/crates/bench/examples/feed.rs

//! Feeds the synthetic Internet routing table (see the crate documentation)
//! to a running daemon over TCP, as one or more peers would, and holds the
//! sessions open afterwards. This is the load for profiling the daemon as
//! it takes in a full table; `scripts/profile/full-table.sh` runs it under
//! `perf`.
//!
//! Settings, from the environment:
//!
//! - `FEED_TO`: where the daemon listens (default `127.0.0.1:1790`).
//! - `FEED_PEERS`: how many peers send the table (default 1).
//! - `FEED_SCALE_DOWN`: as `BENCH_SCALE_DOWN`; 10 makes the table ten
//!   times smaller (default 1).
//! - `FEED_FROM`: the first three octets of an IPv4 /24, "192.0.2". Peer N
//!   connects from .N, which is also the next hop of its IPv4 routes.
//!   Without it, peers connect from wherever the system puts them.
//! - `FEED_ASN`: the first peer's AS number (default 64513); the others
//!   follow on from it.
//! - `FEED_HOLD_SECS`: how long the sessions stay open after the table is
//!   sent (default 30).
//! - `FEED_COLLECTOR=1`: before the peers, one more session, as a route
//!   collector would have: over IPv6, to `FEED_COLLECTOR_TO` (default
//!   `[::1]:1790`) from `FEED_COLLECTOR_FROM` if set, offering to receive
//!   several routes for a prefix (ADD-PATH) for both families and IPv4
//!   routes with IPv6 next hops, and reading and dropping whatever it is
//!   sent. Its AS number is 212232.
//!
//! What is printed: when the sessions are up and how long the writes took.
//! The writes end when the kernel's socket buffers have taken the last
//! bytes, which is before the daemon has read them: when the table is in
//! is for the daemon's API to say.

use std::net::{Ipv4Addr, SocketAddr};
use std::time::{Duration, Instant};

use bgp_bench::report::mib;
use bgp_bench::{Prefixes, Shape, env, table_stream};
use bgp_integration_tests::test_peer::{TestPeer, collector_capabilities};
use bgp_primitives::constants::AfiSafi;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpSocket;
use tokio::net::tcp::OwnedWriteHalf;

/// How often KEEPALIVEs go out while the sessions are held open: well
/// inside the 90 s hold time.
const KEEPALIVE_EVERY: Duration = Duration::from_secs(20);

#[tokio::main]
async fn main() {
    let to: SocketAddr = std::env::var("FEED_TO")
        .unwrap_or_else(|_| "127.0.0.1:1790".into())
        .parse()
        .expect("FEED_TO is an address and port");
    let peers = env("FEED_PEERS", 1) as u32;
    let shape = Shape::INTERNET.scaled_down(env("FEED_SCALE_DOWN", 1));
    let from = std::env::var("FEED_FROM").ok();
    let first_asn = env("FEED_ASN", 64513) as u32;
    let hold = Duration::from_secs(env("FEED_HOLD_SECS", 30) as u64);

    let started = Instant::now();
    let prefixes = Prefixes::generate(shape);
    let tables: Vec<Vec<u8>> = (1..=peers)
        .map(|peer| table_stream(&prefixes, shape, peer))
        .collect();
    let routes = shape.prefixes() * peers as usize;
    let bytes: usize = tables.iter().map(Vec::len).sum();
    println!(
        "table: {} IPv4 + {} IPv6 prefixes from each of {peers} peer(s): {routes} routes, {:.1} MiB on the wire, made in {:.2} s",
        shape.ipv4_prefixes,
        shape.ipv6_prefixes,
        mib(bytes),
        started.elapsed().as_secs_f64()
    );

    // The collector first, so that it is there as the table arrives, as a
    // router's collector session is after a restart.
    let collector = env("FEED_COLLECTOR", 0) == 1;
    let mut draining = Vec::new();
    let mut collector_writer = None;
    if collector {
        let to: SocketAddr = std::env::var("FEED_COLLECTOR_TO")
            .unwrap_or_else(|_| "[::1]:1790".into())
            .parse()
            .expect("FEED_COLLECTOR_TO is an IPv6 address and port");
        let socket = TcpSocket::new_v6().expect("a socket");
        if let Ok(from) = std::env::var("FEED_COLLECTOR_FROM") {
            let address: std::net::Ipv6Addr = from.parse().expect("an IPv6 address");
            socket
                .bind(SocketAddr::new(address.into(), 0))
                .expect("the collector's address is ours to connect from");
        }
        let stream = socket.connect(to).await.expect("the daemon listens");
        let mut session = TestPeer::from_stream(stream);
        session
            .handshake_with(
                212_232,
                Ipv4Addr::new(10, 0, 0, 200),
                90,
                collector_capabilities(
                    212_232,
                    &[AfiSafi::IPV4_UNICAST, AfiSafi::IPV6_UNICAST],
                    true,
                ),
            )
            .await
            .expect("the daemon takes the collector's session");
        let (writer, sink) = session.drain();
        draining.push(sink);
        println!("collector session established");
        collector_writer = Some(writer);
    }

    let mut sessions = Vec::new();
    for peer in 1..=peers {
        let socket = TcpSocket::new_v4().expect("a socket");
        if let Some(from) = &from {
            let address: Ipv4Addr = format!("{from}.{peer}")
                .parse()
                .expect("FEED_FROM is the first three octets of an IPv4 address");
            socket
                .bind(SocketAddr::new(address.into(), 0))
                .expect("the peer's address is ours to connect from");
        }
        let stream = socket.connect(to).await.expect("the daemon listens");
        let mut session = TestPeer::from_stream(stream);
        session
            .handshake(
                first_asn + peer - 1,
                Ipv4Addr::new(10, 0, 0, peer as u8),
                90,
            )
            .await
            .expect("the daemon takes the session");
        // Whatever the daemon sends back, KEEPALIVEs as a rule, is read
        // and dropped, so that the writer is never held up.
        let (writer, sink) = session.drain();
        draining.push(sink);
        sessions.push(writer);
    }
    println!("{peers} session(s) established; sending");

    let started = Instant::now();
    let senders: Vec<_> = sessions
        .into_iter()
        .zip(tables)
        .map(|(mut writer, table)| {
            tokio::spawn(async move {
                let started = Instant::now();
                writer.write_all(&table).await.expect("the daemon reads");
                (writer, started.elapsed())
            })
        })
        .collect();
    let mut writers: Vec<OwnedWriteHalf> = collector_writer.into_iter().collect();
    for (index, sender) in senders.into_iter().enumerate() {
        let (writer, took) = sender.await.expect("the sender runs");
        println!(
            "peer {}: table written in {:.3} s",
            index + 1,
            took.as_secs_f64()
        );
        writers.push(writer);
    }
    let elapsed = started.elapsed();
    println!(
        "all {routes} routes written in {:.3} s: {:.2} M routes/s into the sockets",
        elapsed.as_secs_f64(),
        routes as f64 / elapsed.as_secs_f64() / 1e6
    );

    let keepalive = TestPeer::keepalive_bytes();
    let deadline = Instant::now() + hold;
    while Instant::now() < deadline {
        tokio::time::sleep(KEEPALIVE_EVERY.min(deadline - Instant::now())).await;
        for writer in &mut writers {
            // A session the daemon closed is left alone.
            let _ = writer.write_all(&keepalive).await;
        }
    }
    println!("sessions closed after {} s", hold.as_secs());
    drop(draining);
}