mockforge-bench 0.3.214

//! QoS / DSCP traffic-class load generator (#933).
//!
//! Srikanth on #79 asked whether `mockforge bench` can emit different network
//! traffic classes (Voice, Video, Background, Best-Effort) in one run so the
//! path under test can be exercised against its QoS handling. k6 (the default
//! bench engine) is an application-layer HTTP client and exposes no L3/L4
//! knobs, so this is a NATIVE generator built on raw `socket2` sockets.
//!
//! DSCP is the top 6 bits of the IPv4 `IP_TOS` byte (`TOS = DSCP << 2`, ECN 0).
//! [`connect_marked`] creates the TCP socket, sets `IP_TOS` (and optionally
//! clamps `TCP_MAXSEG`) BEFORE connect, then hands the connected socket to
//! Tokio. Each request opens its own marked connection and sends a minimal
//! HTTP/1.1 request, so a single run can mix classes by weight.
//!
//! Jumbo frames (a NIC MTU property) and true IP fragmentation (kernel/path
//! controlled) are NOT socket knobs and are documented as OS-level operations
//! (`ip link set ... mtu 9000`, `tc`/`netem`) rather than built here. IPv6
//! traffic-class marking (`IPV6_TCLASS`) is a follow-up; v1 marks IPv4 IP_TOS.

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

use socket2::{Domain, Protocol, Socket, Type};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::sync::Mutex;

/// A single traffic class: a human name and its DSCP code point (0-63).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TrafficClass {
    /// Display name (`voice`, `video`, `best-effort`, `background`, `dscp46`).
    pub name: String,
    /// DSCP code point, 0-63.
    pub dscp: u8,
}

impl TrafficClass {
    /// The IPv4 `IP_TOS` byte for this class: `DSCP << 2` with ECN bits 0.
    pub fn tos_byte(&self) -> u8 {
        dscp_to_tos(self.dscp)
    }
}

/// Convert a DSCP code point (0-63) into the 8-bit IP_TOS/Traffic-Class byte.
/// DSCP occupies the high 6 bits; the low 2 (ECN) are left 0.
pub fn dscp_to_tos(dscp: u8) -> u8 {
    (dscp & 0x3f) << 2
}

/// Resolve a preset name to its DSCP code point.
fn preset_dscp(name: &str) -> Option<u8> {
    match name.to_ascii_lowercase().as_str() {
        // Expedited Forwarding — interactive voice.
        "voice" | "ef" => Some(46),
        // Assured Forwarding 41 — interactive video.
        "video" | "af41" => Some(34),
        // Default / Best Effort.
        "best-effort" | "be" | "default" => Some(0),
        // Class Selector 1 — scavenger / background (email, backups).
        "background" | "cs1" | "scavenger" => Some(8),
        _ => None,
    }
}

/// Parse a `--class` spec: `NAME[:WEIGHT]`, where NAME is a preset
/// (`voice`/`video`/`best-effort`/`background`) or `dscpNN` (NN = 0-63).
/// Weight defaults to 1. Returns the class plus its relative weight.
pub fn parse_class(spec: &str) -> Result<(TrafficClass, u32), String> {
    let (name_part, weight) = match spec.split_once(':') {
        Some((n, w)) => {
            let weight: u32 =
                w.parse().map_err(|_| format!("invalid weight in '{spec}' (want an integer)"))?;
            if weight == 0 {
                return Err(format!("weight must be >= 1 in '{spec}'"));
            }
            (n, weight)
        }
        None => (spec, 1),
    };
    let name_part = name_part.trim();
    if name_part.is_empty() {
        return Err("empty traffic class name".to_string());
    }

    let dscp = if let Some(dscp) = preset_dscp(name_part) {
        dscp
    } else if let Some(num) =
        name_part.strip_prefix("dscp").or_else(|| name_part.strip_prefix("DSCP"))
    {
        let n: u8 = num.parse().map_err(|_| format!("invalid DSCP number in '{spec}'"))?;
        if n > 63 {
            return Err(format!("DSCP must be 0-63, got {n}"));
        }
        n
    } else {
        return Err(format!(
            "unknown traffic class '{name_part}' (presets: voice, video, best-effort, background; or dscpNN)"
        ));
    };

    Ok((
        TrafficClass {
            name: name_part.to_ascii_lowercase(),
            dscp,
        },
        weight,
    ))
}

/// Configuration for a QoS traffic-class bench run.
#[derive(Debug, Clone)]
pub struct QosBenchConfig {
    /// Target URL (`http://host:port/path`). HTTPS is not supported by this
    /// native generator (the point is raw socket control, not TLS); use a
    /// plain-HTTP endpoint or terminate TLS in front.
    pub target_url: String,
    /// HTTP method (GET/HEAD/POST/...). GET by default; body is empty.
    pub method: String,
    /// Traffic classes to mix, each with a relative weight.
    pub classes: Vec<(TrafficClass, u32)>,
    /// Concurrent workers.
    pub concurrency: u32,
    /// Total run duration.
    pub duration: Duration,
    /// Optional TCP_MAXSEG clamp (bytes) applied to every connection.
    pub mss: Option<u32>,
}

/// Per-class outcome from a QoS bench run.
#[derive(Debug, Clone)]
pub struct ClassStats {
    pub name: String,
    pub dscp: u8,
    pub tos_byte: u8,
    pub requests: u64,
    pub ok: u64,
    pub failed: u64,
    pub p50_ms: u64,
    pub p95_ms: u64,
}

/// Aggregate result from a QoS bench run.
#[derive(Debug, Clone)]
pub struct QosBenchResult {
    pub total_requests: u64,
    pub successful: u64,
    pub failed: u64,
    pub elapsed: Duration,
    pub req_per_sec: f64,
    pub per_class: Vec<ClassStats>,
    /// Set when the platform silently ignored IP_TOS (so the operator knows
    /// the marking may not have reached the wire).
    pub marking_unsupported: bool,
}

/// Create a fresh TCP socket for `addr` with its IPv4 `IP_TOS` set to `tos`
/// (and, when `mss` is given, `TCP_MAXSEG` clamped), BEFORE any connect, so
/// the very first SYN carries the DSCP marking. `tos_applied`, when given, is
/// set to `true` iff the kernel accepted the IP_TOS setsockopt. Kept separate
/// from the connect so a test can `getsockopt` the TOS back off the exact
/// socket the generator uses (see `connect_marked_socket_carries_tos`).
fn marked_socket(
    addr: SocketAddr,
    tos: u8,
    mss: Option<u32>,
    tos_applied: Option<&Arc<std::sync::atomic::AtomicBool>>,
) -> std::io::Result<Socket> {
    let domain = if addr.is_ipv4() {
        Domain::IPV4
    } else {
        Domain::IPV6
    };
    let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?;

    // DSCP marking. IPv4 IP_TOS only in v1; IPv6 IPV6_TCLASS is a follow-up.
    if addr.is_ipv4() {
        match socket.set_tos(u32::from(tos)) {
            Ok(()) => {
                if let Some(flag) = tos_applied {
                    flag.store(true, Ordering::Relaxed);
                }
            }
            Err(e) => {
                // Non-fatal: still generate load, just unmarked.
                tracing::debug!("set_tos({tos}) failed: {e}");
            }
        }
    }
    if let Some(m) = mss {
        // Best-effort; TCP_MAXSEG isn't settable on every platform.
        let _ = socket.set_mss(m);
    }
    Ok(socket)
}

/// Create a TCP connection to `addr` with its IPv4 `IP_TOS` set to `tos`
/// (and, when `mss` is given, `TCP_MAXSEG` clamped) BEFORE connect, so the
/// very first SYN carries the DSCP marking. Returns the connected Tokio
/// stream. `tos_applied` reports whether the kernel accepted the IP_TOS
/// setsockopt (some platforms reject it for non-privileged users).
async fn connect_marked(
    addr: SocketAddr,
    tos: u8,
    mss: Option<u32>,
    tos_applied: &Arc<std::sync::atomic::AtomicBool>,
) -> std::io::Result<TcpStream> {
    let socket = marked_socket(addr, tos, mss, Some(tos_applied))?;

    // Non-blocking connect: EINPROGRESS is expected; the socket becomes
    // writable when connect resolves, and take_error() reports the outcome.
    socket.set_nonblocking(true)?;
    match socket.connect(&addr.into()) {
        Ok(()) => {}
        Err(e) if e.raw_os_error() == Some(libc::EINPROGRESS) => {}
        Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {}
        Err(e) => return Err(e),
    }

    let std_stream: std::net::TcpStream = socket.into();
    let stream = TcpStream::from_std(std_stream)?;
    stream.writable().await?;
    if let Some(err) = stream.take_error()? {
        return Err(err);
    }
    Ok(stream)
}

/// Send one minimal HTTP/1.1 request over an already-marked connection and
/// return the response status code. `Connection: close` so the server hangs
/// up and we can read to EOF without parsing Content-Length.
async fn send_request(
    mut stream: TcpStream,
    method: &str,
    host: &str,
    path: &str,
) -> std::io::Result<u16> {
    let req = format!(
        "{method} {path} HTTP/1.1\r\nHost: {host}\r\nUser-Agent: mockforge-bench-qos\r\nConnection: close\r\n\r\n"
    );
    stream.write_all(req.as_bytes()).await?;
    stream.flush().await?;

    // Read just enough to see the status line.
    let mut buf = [0u8; 256];
    let n = stream.read(&mut buf).await?;
    let head = String::from_utf8_lossy(&buf[..n]);
    // "HTTP/1.1 200 OK" -> 200
    let status = head
        .split_whitespace()
        .nth(1)
        .and_then(|s| s.parse::<u16>().ok())
        .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidData, "no status line"))?;
    // Drain the rest so the peer's close is clean (best effort).
    let mut sink = [0u8; 4096];
    while let Ok(r) = stream.read(&mut sink).await {
        if r == 0 {
            break;
        }
    }
    Ok(status)
}

/// Split a target URL into (host, port, path) for the raw generator.
/// Only plain `http://` is supported.
fn parse_http_target(url: &str) -> Result<(String, u16, String), String> {
    let rest = url
        .strip_prefix("http://")
        .ok_or_else(|| format!("QoS bench needs a plain http:// target, got '{url}'"))?;
    let (authority, path) = match rest.find('/') {
        Some(i) => (&rest[..i], &rest[i..]),
        None => (rest, "/"),
    };
    let (host, port) = match authority.rsplit_once(':') {
        Some((h, p)) => {
            let port: u16 = p.parse().map_err(|_| format!("invalid port in '{url}'"))?;
            (h.to_string(), port)
        }
        None => (authority.to_string(), 80),
    };
    if host.is_empty() {
        return Err(format!("empty host in '{url}'"));
    }
    let path = if path.is_empty() {
        "/".to_string()
    } else {
        path.to_string()
    };
    Ok((host, port, path))
}

/// Build the weighted selection table: class index repeated `weight` times.
fn weighted_indices(classes: &[(TrafficClass, u32)]) -> Vec<usize> {
    let mut table = Vec::new();
    for (i, (_, w)) in classes.iter().enumerate() {
        for _ in 0..*w {
            table.push(i);
        }
    }
    table
}

fn percentile(sorted: &[u64], p: f64) -> u64 {
    if sorted.is_empty() {
        return 0;
    }
    let idx = ((sorted.len() as f64 - 1.0) * p).round() as usize;
    sorted[idx.min(sorted.len() - 1)]
}

/// Run the QoS traffic-class bench. Spawns `concurrency` workers that, until
/// `duration` elapses, pick a class by weight, open a fresh DSCP-marked
/// connection, and send one request, then aggregates per-class stats.
pub async fn run(cfg: QosBenchConfig) -> anyhow::Result<QosBenchResult> {
    if cfg.concurrency == 0 {
        anyhow::bail!("concurrency must be >= 1");
    }
    if cfg.classes.is_empty() {
        anyhow::bail!("at least one --class is required");
    }
    let (host, port, path) = parse_http_target(&cfg.target_url).map_err(|e| anyhow::anyhow!(e))?;
    let addr: SocketAddr = format!("{host}:{port}")
        .parse()
        .or_else(|_| {
            // Resolve a hostname to its first address.
            use std::net::ToSocketAddrs;
            (host.as_str(), port)
                .to_socket_addrs()
                .ok()
                .and_then(|mut it| it.next())
                .ok_or(())
        })
        .map_err(|_| anyhow::anyhow!("could not resolve target host '{host}:{port}'"))?;

    let table = Arc::new(weighted_indices(&cfg.classes));
    let classes = Arc::new(cfg.classes.clone());
    // Per-class counters + latency buckets.
    let n = cfg.classes.len();
    let requests: Arc<Vec<AtomicU64>> = Arc::new((0..n).map(|_| AtomicU64::new(0)).collect());
    let oks: Arc<Vec<AtomicU64>> = Arc::new((0..n).map(|_| AtomicU64::new(0)).collect());
    let fails: Arc<Vec<AtomicU64>> = Arc::new((0..n).map(|_| AtomicU64::new(0)).collect());
    let lats: Arc<Vec<Mutex<Vec<u64>>>> =
        Arc::new((0..n).map(|_| Mutex::new(Vec::new())).collect());
    let tos_applied = Arc::new(std::sync::atomic::AtomicBool::new(false));

    let deadline = Instant::now() + cfg.duration;
    let started = Instant::now();
    let host = Arc::new(host);
    let path = Arc::new(path);
    let method = Arc::new(cfg.method.to_uppercase());

    let mut workers = Vec::with_capacity(cfg.concurrency as usize);
    for w in 0..cfg.concurrency {
        let table = table.clone();
        let classes = classes.clone();
        let requests = requests.clone();
        let oks = oks.clone();
        let fails = fails.clone();
        let lats = lats.clone();
        let tos_applied = tos_applied.clone();
        let host = host.clone();
        let path = path.clone();
        let method = method.clone();
        let mss = cfg.mss;
        // Deterministic per-worker starting offset so workers don't all pick
        // the same class first; avoids Math.random-style nondeterminism.
        let mut cursor = w as usize;

        workers.push(tokio::spawn(async move {
            while Instant::now() < deadline {
                let idx = table[cursor % table.len()];
                cursor = cursor.wrapping_add(1);
                let tos = classes[idx].0.tos_byte();
                let started_req = Instant::now();
                requests[idx].fetch_add(1, Ordering::Relaxed);
                let outcome = async {
                    let stream = connect_marked(addr, tos, mss, &tos_applied).await?;
                    send_request(stream, &method, &host, &path).await
                }
                .await;
                match outcome {
                    Ok(status) if (200..400).contains(&status) => {
                        oks[idx].fetch_add(1, Ordering::Relaxed);
                        lats[idx].lock().await.push(started_req.elapsed().as_millis() as u64);
                    }
                    Ok(_) => {
                        fails[idx].fetch_add(1, Ordering::Relaxed);
                    }
                    Err(_) => {
                        fails[idx].fetch_add(1, Ordering::Relaxed);
                    }
                }
            }
        }));
    }

    for h in workers {
        let _ = h.await;
    }
    let elapsed = started.elapsed();

    let mut per_class = Vec::with_capacity(n);
    let mut total = 0u64;
    let mut total_ok = 0u64;
    let mut total_fail = 0u64;
    for (i, (class, _)) in cfg.classes.iter().enumerate() {
        let req = requests[i].load(Ordering::Relaxed);
        let ok = oks[i].load(Ordering::Relaxed);
        let fail = fails[i].load(Ordering::Relaxed);
        total += req;
        total_ok += ok;
        total_fail += fail;
        let mut l = lats[i].lock().await.clone();
        l.sort_unstable();
        per_class.push(ClassStats {
            name: class.name.clone(),
            dscp: class.dscp,
            tos_byte: class.tos_byte(),
            requests: req,
            ok,
            failed: fail,
            p50_ms: percentile(&l, 0.50),
            p95_ms: percentile(&l, 0.95),
        });
    }

    let secs = elapsed.as_secs_f64().max(f64::MIN_POSITIVE);
    Ok(QosBenchResult {
        total_requests: total,
        successful: total_ok,
        failed: total_fail,
        elapsed,
        req_per_sec: total as f64 / secs,
        per_class,
        marking_unsupported: !tos_applied.load(Ordering::Relaxed),
    })
}

/// Render the result as a compact table for the CLI.
pub fn render_report(res: &QosBenchResult) -> String {
    let mut out = String::new();
    out.push_str(&format!(
        "QoS bench: {} requests in {:.1}s ({:.0} req/s), {} ok / {} failed\n",
        res.total_requests,
        res.elapsed.as_secs_f64(),
        res.req_per_sec,
        res.successful,
        res.failed
    ));
    if res.marking_unsupported {
        out.push_str(
            "  WARNING: IP_TOS/DSCP marking was not accepted by the kernel on this run; \
             traffic went out unmarked (needs a Unix host and, on some platforms, privileges).\n",
        );
    }
    out.push_str(&format!(
        "  {:<14} {:>5} {:>7} {:>8} {:>7} {:>8} {:>8}\n",
        "class", "dscp", "tos", "reqs", "ok", "p50(ms)", "p95(ms)"
    ));
    for c in &res.per_class {
        out.push_str(&format!(
            "  {:<14} {:>5} {:>7} {:>8} {:>7} {:>8} {:>8}\n",
            c.name,
            c.dscp,
            format!("0x{:02x}", c.tos_byte),
            c.requests,
            c.ok,
            c.p50_ms,
            c.p95_ms
        ));
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn dscp_maps_to_tos_high_six_bits() {
        // EF=46 -> 0xB8, AF41=34 -> 0x88, CS1=8 -> 0x20, BE=0 -> 0x00.
        assert_eq!(dscp_to_tos(46), 0xB8);
        assert_eq!(dscp_to_tos(34), 0x88);
        assert_eq!(dscp_to_tos(8), 0x20);
        assert_eq!(dscp_to_tos(0), 0x00);
        // ECN bits (low 2) always clear.
        assert_eq!(dscp_to_tos(63) & 0x03, 0);
    }

    #[test]
    fn parse_presets_and_custom_and_weights() {
        assert_eq!(parse_class("voice").unwrap().0.dscp, 46);
        assert_eq!(parse_class("video").unwrap().0.dscp, 34);
        assert_eq!(parse_class("best-effort").unwrap().0.dscp, 0);
        assert_eq!(parse_class("background").unwrap().0.dscp, 8);
        assert_eq!(parse_class("dscp46").unwrap().0.dscp, 46);
        // Weight parsing.
        let (c, w) = parse_class("voice:40").unwrap();
        assert_eq!((c.dscp, w), (46, 40));
        // Errors.
        assert!(parse_class("bogus").is_err());
        assert!(parse_class("dscp99").is_err()); // > 63
        assert!(parse_class("voice:0").is_err()); // weight 0
        assert!(parse_class("voice:x").is_err());
    }

    #[test]
    fn weighted_table_repeats_by_weight() {
        let classes = vec![
            (
                TrafficClass {
                    name: "a".into(),
                    dscp: 0,
                },
                3,
            ),
            (
                TrafficClass {
                    name: "b".into(),
                    dscp: 8,
                },
                1,
            ),
        ];
        let table = weighted_indices(&classes);
        assert_eq!(table.len(), 4);
        assert_eq!(table.iter().filter(|&&i| i == 0).count(), 3);
        assert_eq!(table.iter().filter(|&&i| i == 1).count(), 1);
    }

    #[test]
    fn parse_http_target_splits_host_port_path() {
        assert_eq!(
            parse_http_target("http://host:3000/up").unwrap(),
            ("host".to_string(), 3000, "/up".to_string())
        );
        assert_eq!(
            parse_http_target("http://host/").unwrap(),
            ("host".to_string(), 80, "/".to_string())
        );
        assert!(parse_http_target("https://host/").is_err());
    }

    /// The whole point of this feature: prove the socket the generator
    /// actually uses carries the DSCP marking. `marked_socket` is the exact
    /// helper `connect_marked` calls, so reading `IP_TOS` back off it with
    /// getsockopt proves every outbound SYN/segment carries `DSCP << 2`. The
    /// kernel stamps a socket's IP_TOS into the IP header of its packets, so
    /// this is the sender-side wire guarantee.
    #[tokio::test]
    async fn marked_socket_carries_dscp_and_connect_works() {
        let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();

        for dscp in [46u8, 34, 8, 0] {
            let tos = dscp_to_tos(dscp);
            let applied = Arc::new(std::sync::atomic::AtomicBool::new(false));
            let socket = marked_socket(addr, tos, None, Some(&applied)).unwrap();
            assert!(applied.load(Ordering::Relaxed), "set_tos must succeed for dscp {dscp}");
            let read_back = socket.tos().unwrap() as u8;
            assert_eq!(
                read_back, tos,
                "getsockopt IP_TOS must equal DSCP<<2 (dscp={dscp}, want {tos:#04x})"
            );
        }

        // And `connect_marked` (which uses `marked_socket`) reaches a live peer.
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let laddr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            let _ = listener.accept().await;
        });
        let applied = Arc::new(std::sync::atomic::AtomicBool::new(false));
        let stream = connect_marked(laddr, dscp_to_tos(46), None, &applied).await.unwrap();
        assert!(stream.peer_addr().is_ok());
        assert!(applied.load(Ordering::Relaxed), "connect_marked applied the TOS");
    }

    /// The MSS clamp is accepted by the socket (TCP_MAXSEG) without breaking
    /// the connect path.
    #[tokio::test]
    async fn mss_clamp_is_accepted() {
        let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
        // Should not error; getsockopt of MSS varies by platform so we only
        // assert the socket is still usable for a connect afterwards.
        let socket = marked_socket(addr, dscp_to_tos(0), Some(536), None).unwrap();
        drop(socket);
    }
}