zc2 0.0.26

//! Native async closed-loop load generator for the broker `/execute` path.
//!
//! `CONC` tokio tasks each fire requests back-to-back over a shared keep-alive
//! connection pool until a deadline — so exactly `CONC` requests are in flight
//! without spawning `CONC` OS threads. The runtime is pinned to a small number
//! of worker threads (`ZL_THREADS`, default 4) so the generator leaves the rest
//! of the box to the broker. This is the Rust replacement for the httpx
//! generator that capped at ~450 RPS on Python overhead.
//!
//! Run (from the zc crate):
//!   cargo run --release --example zload -- http://127.0.0.1:9000/execute 200 10
//! Args:  URL  CONCURRENCY  DURATION_SECS
//! Env:   ZL_THREADS (4)  ZL_WARMUP (3)  ZL_BODY (64)  ZL_USER (loaduser)

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

use reqwest::header::{HeaderMap, HeaderValue, CONTENT_TYPE};

struct TaskResult {
    ok: u64,
    err: u64,
    lat_ms: Vec<f32>,
}

async fn runner(
    client: Arc<reqwest::Client>,
    url: Arc<String>,
    body: Arc<Vec<u8>>,
    headers: HeaderMap,
    deadline: Instant,
    measure: bool,
) -> TaskResult {
    let mut ok = 0u64;
    let mut err = 0u64;
    let mut lat_ms = Vec::new();
    while Instant::now() < deadline {
        let t0 = Instant::now();
        let resp = client
            .post(url.as_str())
            .headers(headers.clone())
            .body(body.as_ref().clone())
            .send()
            .await;
        match resp {
            Ok(r) => {
                let status = r.status();
                // Drain the body so the connection returns to the pool for reuse.
                let _ = r.bytes().await;
                if status.is_success() {
                    ok += 1;
                } else {
                    err += 1;
                }
            }
            Err(_) => err += 1,
        }
        if measure {
            lat_ms.push(t0.elapsed().as_secs_f32() * 1000.0);
        }
    }
    TaskResult { ok, err, lat_ms }
}

async fn run_phase(
    client: Arc<reqwest::Client>,
    url: Arc<String>,
    body: Arc<Vec<u8>>,
    headers: &HeaderMap,
    conc: usize,
    dur: Duration,
    measure: bool,
) -> (u64, u64, Vec<f32>, f64) {
    let deadline = Instant::now() + dur;
    let mut tasks = Vec::with_capacity(conc);
    let start = Instant::now();
    for _ in 0..conc {
        tasks.push(tokio::spawn(runner(
            client.clone(),
            url.clone(),
            body.clone(),
            headers.clone(),
            deadline,
            measure,
        )));
    }
    let mut ok = 0u64;
    let mut err = 0u64;
    let mut lat: Vec<f32> = Vec::new();
    for t in tasks {
        if let Ok(r) = t.await {
            ok += r.ok;
            err += r.err;
            if measure {
                lat.extend(r.lat_ms);
            }
        }
    }
    let elapsed = start.elapsed().as_secs_f64();
    (ok, err, lat, elapsed)
}

fn pct(sorted: &[f32], p: f64) -> f32 {
    if sorted.is_empty() {
        return 0.0;
    }
    let idx = ((sorted.len() as f64) * p / 100.0) as usize;
    sorted[idx.min(sorted.len() - 1)]
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let url = args
        .get(1)
        .cloned()
        .unwrap_or_else(|| "http://127.0.0.1:9000/execute".into());
    let conc: usize = args.get(2).and_then(|v| v.parse().ok()).unwrap_or(200);
    let dur_secs: f64 = args.get(3).and_then(|v| v.parse().ok()).unwrap_or(10.0);

    let threads: usize = std::env::var("ZL_THREADS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4);
    let warmup: f64 = std::env::var("ZL_WARMUP")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(3.0);
    let body_size: usize = std::env::var("ZL_BODY")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(64);
    let user = std::env::var("ZL_USER").unwrap_or_else(|_| "loaduser".into());

    let mut headers = HeaderMap::new();
    headers.insert("X-Zakuro-User", HeaderValue::from_str(&user).unwrap());
    headers.insert(
        "X-Zakuro-Requirements",
        HeaderValue::from_static("{\"strategy\":\"round_robin\",\"estimated_duration_secs\":0.01}"),
    );
    headers.insert(
        CONTENT_TYPE,
        HeaderValue::from_static("application/octet-stream"),
    );

    let body = Arc::new(vec![b'x'; body_size]);
    let url = Arc::new(url);

    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(threads)
        .enable_all()
        .build()
        .unwrap();

    rt.block_on(async move {
        let client = Arc::new(
            reqwest::Client::builder()
                .pool_max_idle_per_host(conc + 16)
                .http1_only()
                .tcp_nodelay(true)
                .timeout(Duration::from_secs(30))
                .build()
                .unwrap(),
        );

        // Readiness: one request must succeed before we start (workers discovered).
        let mut ready = false;
        for _ in 0..60 {
            let r = client
                .post(url.as_str())
                .headers(headers.clone())
                .body(body.as_ref().clone())
                .send()
                .await;
            if let Ok(resp) = r {
                if resp.status().is_success() {
                    ready = true;
                    break;
                }
            }
            tokio::time::sleep(Duration::from_millis(500)).await;
        }
        if !ready {
            eprintln!(
                "[zload] WARNING: no successful warmup request — broker/workers may not be ready"
            );
        }

        if warmup > 0.0 {
            let _ = run_phase(
                client.clone(),
                url.clone(),
                body.clone(),
                &headers,
                conc,
                Duration::from_secs_f64(warmup),
                false,
            )
            .await;
        }

        let (ok, err, mut lat, elapsed) = run_phase(
            client.clone(),
            url.clone(),
            body.clone(),
            &headers,
            conc,
            Duration::from_secs_f64(dur_secs),
            true,
        )
        .await;
        lat.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let rps = if elapsed > 0.0 {
            ok as f64 / elapsed
        } else {
            0.0
        };
        println!(
            "conc={conc} threads={threads} dur={elapsed:.1}s  ok={ok} err={err}  rps={rps:.0}  \
             p50={:.2}ms p90={:.2}ms p99={:.2}ms p999={:.2}ms max={:.2}ms",
            pct(&lat, 50.0),
            pct(&lat, 90.0),
            pct(&lat, 99.0),
            pct(&lat, 99.9),
            pct(&lat, 100.0)
        );
    });
}