supermachine 0.1.1

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::OnceLock;
use std::time::Duration;

pub enum Stage {
    FrontendUnixRead,
    FrontendTcpRead,
    FrontendUnixToTcpWrite,
    FrontendTcpToUnixWrite,
    IoTcpRead,
    IoTcpWrite,
    IoCallback,
    RxqKick,
    GuestToTcpSend,
}

struct Counter {
    calls: AtomicU64,
    bytes: AtomicU64,
    us: AtomicU64,
}

impl Counter {
    const fn new() -> Self {
        Self {
            calls: AtomicU64::new(0),
            bytes: AtomicU64::new(0),
            us: AtomicU64::new(0),
        }
    }
}

static FRONTEND_UNIX_READ: Counter = Counter::new();
static FRONTEND_TCP_READ: Counter = Counter::new();
static FRONTEND_UNIX_TO_TCP_WRITE: Counter = Counter::new();
static FRONTEND_TCP_TO_UNIX_WRITE: Counter = Counter::new();
static IO_TCP_READ: Counter = Counter::new();
static IO_TCP_WRITE: Counter = Counter::new();
static IO_CALLBACK: Counter = Counter::new();
static RXQ_KICK: Counter = Counter::new();
static GUEST_TO_TCP_SEND: Counter = Counter::new();

pub fn enabled() -> bool {
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| {
        let on = matches!(
            std::env::var("SUPERMACHINE_MUX_PROFILE").as_deref(),
            Ok("1") | Ok("true") | Ok("yes") | Ok("on")
        );
        if on {
            start_reporter();
        }
        on
    })
}

pub fn record(stage: Stage, bytes: usize, us: u64) {
    if !enabled() {
        return;
    }
    let c = counter(stage);
    c.calls.fetch_add(1, Ordering::Relaxed);
    c.bytes.fetch_add(bytes as u64, Ordering::Relaxed);
    c.us.fetch_add(us, Ordering::Relaxed);
}

fn counter(stage: Stage) -> &'static Counter {
    match stage {
        Stage::FrontendUnixRead => &FRONTEND_UNIX_READ,
        Stage::FrontendTcpRead => &FRONTEND_TCP_READ,
        Stage::FrontendUnixToTcpWrite => &FRONTEND_UNIX_TO_TCP_WRITE,
        Stage::FrontendTcpToUnixWrite => &FRONTEND_TCP_TO_UNIX_WRITE,
        Stage::IoTcpRead => &IO_TCP_READ,
        Stage::IoTcpWrite => &IO_TCP_WRITE,
        Stage::IoCallback => &IO_CALLBACK,
        Stage::RxqKick => &RXQ_KICK,
        Stage::GuestToTcpSend => &GUEST_TO_TCP_SEND,
    }
}

fn start_reporter() {
    static STARTED: OnceLock<()> = OnceLock::new();
    STARTED.get_or_init(|| {
        std::thread::Builder::new()
            .name("mux-profile".into())
            .spawn(|| {
                let mut prev = snapshot();
                loop {
                    std::thread::sleep(Duration::from_secs(1));
                    let now = snapshot();
                    eprintln!("[mux-profile] {}", format_delta(&prev, &now));
                    prev = now;
                }
            })
            .ok();
    });
}

#[derive(Clone, Copy)]
struct Snap {
    calls: [u64; 9],
    bytes: [u64; 9],
    us: [u64; 9],
}

fn snapshot() -> Snap {
    let counters = [
        &FRONTEND_UNIX_READ,
        &FRONTEND_TCP_READ,
        &FRONTEND_UNIX_TO_TCP_WRITE,
        &FRONTEND_TCP_TO_UNIX_WRITE,
        &IO_TCP_READ,
        &IO_TCP_WRITE,
        &IO_CALLBACK,
        &RXQ_KICK,
        &GUEST_TO_TCP_SEND,
    ];
    let mut calls = [0; 9];
    let mut bytes = [0; 9];
    let mut us = [0; 9];
    for (i, c) in counters.iter().enumerate() {
        calls[i] = c.calls.load(Ordering::Relaxed);
        bytes[i] = c.bytes.load(Ordering::Relaxed);
        us[i] = c.us.load(Ordering::Relaxed);
    }
    Snap { calls, bytes, us }
}

fn format_delta(prev: &Snap, now: &Snap) -> String {
    const NAMES: [&str; 9] = [
        "front_unix_read",
        "front_tcp_read",
        "front_u2t_write",
        "front_t2u_write",
        "io_tcp_read",
        "io_tcp_write",
        "io_callback",
        "rxq_kick",
        "guest_to_tcp_send",
    ];
    let mut out = String::new();
    for i in 0..NAMES.len() {
        let calls = now.calls[i].saturating_sub(prev.calls[i]);
        if calls == 0 {
            continue;
        }
        let bytes = now.bytes[i].saturating_sub(prev.bytes[i]);
        let us = now.us[i].saturating_sub(prev.us[i]);
        if !out.is_empty() {
            out.push_str(" | ");
        }
        out.push_str(&format!(
            "{} calls={} bytes={} avg_us={:.2}",
            NAMES[i],
            calls,
            bytes,
            us as f64 / calls as f64
        ));
    }
    if out.is_empty() {
        "idle".to_owned()
    } else {
        out
    }
}