armdb 0.5.0

//! SQPOLL noise experiment matrix (research doc: armdb/docs/research/sqpoll-noise.md).
//!
//! Run: cargo run --release -p armdb --features bench-internals --example sqpoll_matrix
//!
//! Scenarios:
//!   1. single-collection ConstTree: IoBackend × threads × shard_count
//!   2. multiplier: 1 vs N independent ConstTree (Engines) — poller-thread count
//!   3. flusher submission-frequency: armour::Db, flusher on/off, with a low
//!      background write rate so each 1s flusher tick has dirty data to submit —
//!      does the flusher keep SQPOLL pollers awake (tick < sq_thread_idle)?
//!   4. ATTACH_WQ PoC: N independent SQPOLL rings vs N rings attached to one wq.
//!
//! Metrics: CPU x (process-CPU/wall), throughput (Mops/s), p99 latency, live
//! iou-sqp-* thread count. Reps: 1 warmup + 3 measured, median reported.
//! Note: there is no fsync axis — Config::enable_fsync is a no-op for Bitcask.

use std::time::{Duration, Instant};

use armdb::{Config, ConstTree, IoBackend};
use tempfile::TempDir;

const PREFILL: u64 = 1_000_000;
const OPS: u64 = 500_000;
const WARMUP_REPS: usize = 1;
const MEASURED_REPS: usize = 3;

/// splitmix64 — bijection on u64; well-distributed keys spread across all shards.
fn key64(i: u64) -> u64 {
    let mut z = i.wrapping_add(0x9E3779B97F4A7C15);
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
    z ^ (z >> 31)
}
fn key_bytes(i: u64) -> [u8; 8] {
    key64(i).to_be_bytes()
}

/// Process CPU time (user + sys) in seconds, from /proc/self/stat. Linux-only.
fn cpu_seconds() -> f64 {
    #[cfg(target_os = "linux")]
    {
        let stat = std::fs::read_to_string("/proc/self/stat").unwrap();
        let rest = stat.rsplit(')').next().unwrap();
        let fields: Vec<&str> = rest.split_whitespace().collect();
        let utime: f64 = fields[11].parse().unwrap();
        let stime: f64 = fields[12].parse().unwrap();
        (utime + stime) / 100.0 // CLK_TCK
    }
    #[cfg(not(target_os = "linux"))]
    {
        0.0
    }
}

/// Count live `iou-sqp-*` kernel threads belonging to this process.
fn count_iou_sqp() -> usize {
    let mut n = 0;
    if let Ok(rd) = std::fs::read_dir("/proc/self/task") {
        for e in rd.flatten() {
            if let Ok(comm) = std::fs::read_to_string(e.path().join("comm")) {
                if comm.trim_start().starts_with("iou-sqp") {
                    n += 1;
                }
            }
        }
    }
    n
}

/// Percentile from a pre-sorted slice of nanos.
fn pctl(sorted: &[u64], p: f64) -> u64 {
    if sorted.is_empty() {
        return 0;
    }
    let idx = (((sorted.len() as f64) * p).ceil() as usize).saturating_sub(1);
    sorted[idx.min(sorted.len() - 1)]
}

fn median_f64(mut xs: Vec<f64>) -> f64 {
    xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
    xs[xs.len() / 2]
}

fn governor() -> String {
    std::fs::read_to_string("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor")
        .map(|s| s.trim().to_string())
        .unwrap_or_else(|_| "unknown".into())
}

/// The four flush-path strategies under test.
fn io_variants() -> Vec<(&'static str, IoBackend)> {
    vec![
        (
            "sqpoll-2000",
            IoBackend::Uring {
                sqpoll_idle_ms: Some(2000),
            },
        ),
        (
            "sqpoll-50",
            IoBackend::Uring {
                sqpoll_idle_ms: Some(50),
            },
        ),
        (
            "no-sqpoll",
            IoBackend::Uring {
                sqpoll_idle_ms: None,
            },
        ),
        ("pwrite", IoBackend::Pwrite),
    ]
}

struct Cell {
    io: &'static str,
    shards: usize,
    threads: usize,
    mops: f64,
    cpu_x: f64,
    p99_us: f64,
    sqp_threads: usize,
}

/// One configured cell: prefill, warmup, then MEASURED_REPS timed reps inserting
/// OPS unique new keys split across `threads`. Returns median throughput/CPU x and
/// pooled p99. (No fsync param — `enable_fsync` is a no-op for Bitcask; sustained
/// inserts already fire Write SQEs via natural 1MB write-buffer flushes.)
fn run_single(io: &'static str, backend: IoBackend, shards: usize, threads: usize) -> Cell {
    let dir = TempDir::new().unwrap();
    let mut config = Config::balanced().build();
    config.shard_count = shards;
    config.io_backend = backend;
    let tree = ConstTree::<[u8; 8], 16>::open(dir.path(), config).unwrap();

    for i in 0..PREFILL {
        tree.insert(&key_bytes(i), &[0u8; 16]).unwrap();
    }
    tree.flush().unwrap();

    let per_thread = OPS / threads as u64;
    let mut next_base = PREFILL;
    let mut mops_runs = Vec::new();
    let mut cpu_runs = Vec::new();
    let mut latencies: Vec<u64> = Vec::new();
    let mut sqp_threads = 0;

    for rep in 0..(WARMUP_REPS + MEASURED_REPS) {
        let base = next_base;
        next_base += OPS;
        let measured = rep >= WARMUP_REPS;
        let cpu0 = cpu_seconds();
        let t0 = Instant::now();
        let per_thread_lats: Vec<Vec<u64>> = std::thread::scope(|s| {
            let handles: Vec<_> = (0..threads)
                .map(|t| {
                    let tree = &tree;
                    let start = base + t as u64 * per_thread;
                    s.spawn(move || {
                        let mut lats =
                            Vec::with_capacity(if measured { per_thread as usize } else { 0 });
                        for i in start..start + per_thread {
                            let op = Instant::now();
                            tree.insert(&key_bytes(i), &[0u8; 16]).unwrap();
                            if measured {
                                lats.push(op.elapsed().as_nanos() as u64);
                            }
                        }
                        lats
                    })
                })
                .collect();
            handles.into_iter().map(|h| h.join().unwrap()).collect()
        });
        let wall = t0.elapsed().as_secs_f64();
        let cpu = cpu_seconds() - cpu0;
        if measured {
            mops_runs.push((per_thread * threads as u64) as f64 / wall / 1e6);
            cpu_runs.push(cpu / wall);
            for v in per_thread_lats {
                latencies.extend(v);
            }
            sqp_threads = sqp_threads.max(count_iou_sqp());
        }
    }
    latencies.sort_unstable();
    Cell {
        io,
        shards,
        threads,
        mops: median_f64(mops_runs),
        cpu_x: median_f64(cpu_runs),
        p99_us: pctl(&latencies, 0.99) as f64 / 1000.0,
        sqp_threads,
    }
}

fn print_cells(title: &str, cells: &[Cell]) {
    println!("\n=== {title} ===");
    println!(
        "{:>12} {:>7} {:>8} {:>9} {:>8} {:>9} {:>9}",
        "io", "shards", "threads", "Mops/s", "CPU x", "p99 us", "iou-sqp"
    );
    for c in cells {
        println!(
            "{:>12} {:>7} {:>8} {:>9.2} {:>8.1} {:>9.1} {:>9}",
            c.io, c.shards, c.threads, c.mops, c.cpu_x, c.p99_us, c.sqp_threads
        );
    }
}

struct MultCell {
    io: &'static str,
    collections: usize,
    shards_each: usize,
    mops: f64,
    cpu_x: f64,
    sqp_threads: usize,
}

/// Open `collections` independent ConstTree instances (each its own Engine →
/// own shards → own SQPOLL rings). Drive a write burst across all of them and
/// record CPU x and the live iou-sqp-* thread count.
fn run_multiplier(
    io: &'static str,
    backend: IoBackend,
    collections: usize,
    shards_each: usize,
) -> MultCell {
    let dirs: Vec<TempDir> = (0..collections).map(|_| TempDir::new().unwrap()).collect();
    let trees: Vec<ConstTree<[u8; 8], 16>> = dirs
        .iter()
        .map(|d| {
            let mut config = Config::balanced().build();
            config.shard_count = shards_each;
            config.io_backend = backend;
            ConstTree::<[u8; 8], 16>::open(d.path(), config).unwrap()
        })
        .collect();

    let threads = 8usize;
    let per_coll = OPS; // OPS new keys per collection
    let cpu0 = cpu_seconds();
    let t0 = Instant::now();
    std::thread::scope(|s| {
        for tree in &trees {
            for t in 0..threads {
                let tree = tree;
                let start = t as u64 * (per_coll / threads as u64);
                let end = start + per_coll / threads as u64;
                s.spawn(move || {
                    for i in start..end {
                        tree.insert(&key_bytes(i), &[0u8; 16]).unwrap();
                    }
                });
            }
        }
    });
    let wall = t0.elapsed().as_secs_f64();
    let cpu = cpu_seconds() - cpu0;
    let sqp_threads = count_iou_sqp();
    let total_ops = (per_coll * collections as u64) as f64;
    MultCell {
        io,
        collections,
        shards_each,
        mops: total_ops / wall / 1e6,
        cpu_x: cpu / wall,
        sqp_threads,
    }
}

fn print_mult(cells: &[MultCell]) {
    println!("\n=== Scenario 2: multiplier (independent Engines, T=8) ===");
    println!(
        "{:>12} {:>12} {:>11} {:>9} {:>8} {:>9}",
        "io", "collections", "shards/coll", "Mops/s", "CPU x", "iou-sqp"
    );
    for c in cells {
        println!(
            "{:>12} {:>12} {:>11} {:>9.2} {:>8.1} {:>9}",
            c.io, c.collections, c.shards_each, c.mops, c.cpu_x, c.sqp_threads
        );
    }
}

use armdb::armour::Db;
use armdb::{CollectionMeta, NoHook, RapiraCodec};
use armour_core::{Fuid, GetType};
use rapira::Rapira;

// `hasher!` takes a literal key directly (no env var) so it works in release too;
// `const_hasher_or!`'s dev-key fallback is debug-only. The brand is irrelevant here
// (no schema FK layer in this benchmark).
armour_core::hasher!(bench_id, "sqpoll-matrix-dev-key");
type BenchId = Fuid<bench_id::Hasher>;

#[derive(Clone, Debug, PartialEq, GetType, Rapira)]
struct BenchRow {
    payload: u64,
}
impl CollectionMeta for BenchRow {
    type SelfId = BenchId;
    const NAME: &'static str = "bench_rows";
}

/// Submission-frequency probe. A LOW background write rate (10 small rows / 100ms)
/// keeps a little dirty data in the per-shard write buffer without ever filling it
/// (1 MB), so nothing is submitted to io_uring on its own. With the flusher ON
/// (1 s tick < 2 s `sq_thread_idle`) each tick flushes that dirty buffer → a Write
/// SQE → the SQPOLL poller is re-armed and never sleeps. With the flusher OFF
/// nothing is submitted, so the poller sleeps after ~2 s. We report CPU-per-second
/// and peak iou-sqp count over a 5 s window.
///
/// (This replaces the original burst-then-idle design: live `periodic_flush` →
/// `flush_buffers` returns immediately on an empty buffer — `armdb/src/shard.rs:681`
/// — so a pure idle window submits nothing and cannot keep pollers awake. R4.)
fn run_flusher(backend: IoBackend, flusher_on: bool) -> (f64, usize) {
    let dir = TempDir::new().unwrap();
    let flush = if flusher_on {
        Some(Duration::from_secs(1))
    } else {
        None
    };
    let db = std::sync::Arc::new(Db::open_with_intervals(dir.path(), None, flush).unwrap());

    let mut config = Config::balanced().build();
    config.shard_count = 8;
    config.io_backend = backend;
    let coll = db
        .open_typed_map::<BenchRow, RapiraCodec, _>(config, NoHook, &[])
        .unwrap();

    // Warm all 8 shards once (each gets an active file + buffer), then drain.
    for i in 0..10_000u64 {
        coll.put(&BenchId::from_u64(key64(i)), BenchRow { payload: i })
            .unwrap();
    }
    coll.flush_buffers().ok();

    // 5 s window: trickle a low background write rate; sample CPU + poller count.
    let cpu0 = cpu_seconds();
    let mut max_sqp = 0;
    let window = Duration::from_secs(5);
    let t0 = Instant::now();
    let mut n = 1_000_000u64;
    while t0.elapsed() < window {
        for _ in 0..10 {
            coll.put(&BenchId::from_u64(key64(n)), BenchRow { payload: n })
                .unwrap();
            n += 1;
        }
        max_sqp = max_sqp.max(count_iou_sqp());
        std::thread::sleep(Duration::from_millis(100));
    }
    let cpu = cpu_seconds() - cpu0;
    (cpu / window.as_secs_f64(), max_sqp) // CPU-seconds per wall-second over the window
}

/// Does IORING_SETUP_ATTACH_WQ collapse N SQPOLL pollers to 1? Build N independent
/// SQPOLL rings and count iou-sqp-*, then build 1 SQPOLL ring + (N-1) rings
/// attached to its workqueue and count again. GPT-5 flagged that ATTACH_WQ shares
/// io-wq workers but may NOT remove the per-ring poller — this measures it.
#[cfg(target_os = "linux")]
fn attach_wq_poc(n: usize) {
    use rustix_uring::{IoUring, cqueue, squeue};
    use std::os::unix::io::AsRawFd;

    let independent: Vec<IoUring<squeue::Entry, cqueue::Entry>> = (0..n)
        .map(|_| IoUring::builder().setup_sqpoll(2000).build(256).unwrap())
        .collect();
    std::thread::sleep(Duration::from_millis(150));
    let indep_count = count_iou_sqp();
    drop(independent);
    std::thread::sleep(Duration::from_millis(300)); // let pollers exit

    let first: IoUring<squeue::Entry, cqueue::Entry> =
        IoUring::builder().setup_sqpoll(2000).build(256).unwrap();
    let wq_fd = first.as_raw_fd();
    let mut attached = vec![first];
    for _ in 1..n {
        let mut b = IoUring::builder();
        b.setup_sqpoll(2000);
        // SAFETY: wq_fd is a live io_uring fd owned by `first`, alive for this loop.
        unsafe {
            b.setup_attach_wq(wq_fd);
        }
        let ring: IoUring<squeue::Entry, cqueue::Entry> = b.build(256).unwrap();
        attached.push(ring);
    }
    std::thread::sleep(Duration::from_millis(150));
    let attached_count = count_iou_sqp();
    drop(attached);

    println!("\n=== Scenario 4: ATTACH_WQ PoC (N={n} rings) ===");
    println!("  independent SQPOLL rings : {indep_count} iou-sqp threads");
    println!("  SQPOLL + ATTACH_WQ rings : {attached_count} iou-sqp threads");
    println!(
        "  verdict: ATTACH_WQ {} pollers",
        if attached_count < indep_count {
            "COLLAPSES"
        } else {
            "does NOT collapse"
        }
    );
}

#[cfg(not(target_os = "linux"))]
fn attach_wq_poc(_n: usize) {
    println!("\n=== Scenario 4: ATTACH_WQ PoC — skipped (not Linux) ===");
}

fn main() {
    let cores = std::thread::available_parallelism().unwrap();
    println!(
        "sqpoll_matrix: {cores} logical cores, governor={}",
        governor()
    );
    println!("PREFILL={PREFILL} OPS={OPS} reps={MEASURED_REPS} (+{WARMUP_REPS} warmup)\n");
    // Scenario 1: single-collection grid at shard_count=16.
    let mut cells = Vec::new();
    for (name, backend) in io_variants() {
        for &threads in &[1usize, 8, 16] {
            cells.push(run_single(name, backend, 16, threads));
        }
    }
    // Headline poller-count-vs-cores at shard_count=8.
    for (name, backend) in io_variants().into_iter().filter(|(n, _)| *n != "sqpoll-50") {
        for &threads in &[1usize, 8] {
            cells.push(run_single(name, backend, 8, threads));
        }
    }
    print_cells("Scenario 1: single ConstTree (insert, end-to-end)", &cells);

    // Scenario 2: 1 vs N independent Engines — poller multiplier.
    let mut mult = Vec::new();
    for (name, backend) in io_variants().into_iter().filter(|(n, _)| *n != "sqpoll-50") {
        for &collections in &[1usize, 10] {
            mult.push(run_multiplier(name, backend, collections, 8));
        }
    }
    print_mult(&mult);

    // Scenario 3: flusher submission-frequency.
    println!(
        "\n=== Scenario 3: flusher submission-frequency (armour::Db, 8 shards, 5s window, low bg write rate) ==="
    );
    println!(
        "{:>12} {:>10} {:>14} {:>9}",
        "io", "flusher", "CPU/s", "iou-sqp"
    );
    for (name, backend) in io_variants()
        .into_iter()
        .filter(|(n, _)| *n == "sqpoll-2000" || *n == "pwrite")
    {
        for &flusher_on in &[true, false] {
            let (cpu_per_s, sqp) = run_flusher(backend, flusher_on);
            println!("{name:>12} {:>10} {cpu_per_s:>14.3} {sqp:>9}", flusher_on);
        }
    }
    // Scenario 4: ATTACH_WQ PoC.
    attach_wq_poc(8);
    attach_wq_poc(16);
}