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::{Prefixes, Shape, announcements, encoded, env};
use bgp_integration_tests::test_peer::TestPeer;
use bgp_packet::open::{AddPathEntry, Capability, ExtendedNextHopEntry, SendReceive};
use bgp_primitives::constants::{Afi, AfiSafi, Safi};
use bgp_server::family::SUPPORTED;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpSocket;
use tokio::net::tcp::OwnedWriteHalf;

/// A KEEPALIVE message: the marker, the length (19) and the type (4).
const KEEPALIVE: [u8; 19] = {
    let mut bytes = [0xff; 19];
    bytes[16] = 0;
    bytes[17] = 19;
    bytes[18] = 4;
    bytes
};

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

fn mib(bytes: usize) -> f64 {
    bytes as f64 / (1024.0 * 1024.0)
}

#[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| {
            encoded(&announcements(&prefixes, shape, peer))
                .iter()
                .flat_map(|message| message.iter().copied())
                .collect()
        })
        .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);
        let mut capabilities = vec![Capability::FourByteAsn { asn: 212_232 }];
        capabilities.extend(SUPPORTED.into_iter().map(Capability::multiprotocol));
        capabilities.push(Capability::AddPath {
            entries: [AfiSafi::IPV4_UNICAST, AfiSafi::IPV6_UNICAST]
                .into_iter()
                .map(|family| AddPathEntry {
                    afi: family.afi,
                    safi: family.safi,
                    send_recv: SendReceive::Receive,
                })
                .collect(),
        });
        capabilities.push(Capability::ExtendedNextHop {
            entries: vec![ExtendedNextHopEntry {
                afi: Afi::IPV4,
                safi: Safi::UNICAST,
                nh_afi: Afi::IPV6,
            }],
        });
        session
            .handshake_with(212_232, Ipv4Addr::new(10, 0, 0, 200), 90, capabilities)
            .await
            .expect("the daemon takes the collector's session");
        let (mut reader, writer) = session.into_stream().into_split();
        draining.push(tokio::spawn(async move {
            let mut sink = [0u8; 65536];
            while let Ok(n) = reader.read(&mut sink).await {
                if n == 0 {
                    break;
                }
            }
        }));
        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");
        sessions.push(session.into_stream());
    }
    println!("{peers} session(s) established; sending");

    let started = Instant::now();
    let senders: Vec<_> = sessions
        .into_iter()
        .zip(tables)
        .map(|(stream, table)| {
            tokio::spawn(async move {
                let (mut reader, mut writer) = stream.into_split();
                // Whatever the daemon sends back, KEEPALIVEs as a rule, is
                // read and dropped, so that its writer is never held up.
                tokio::spawn(async move {
                    let mut sink = [0u8; 65536];
                    while let Ok(n) = reader.read(&mut sink).await {
                        if n == 0 {
                            break;
                        }
                    }
                });
                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 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);
}