zc2 0.0.28

//! Native instant mock worker for broker throughput benchmarking.
//!
//! One process serves the broker's worker contract on a *range* of ports
//! (so the broker discovers N independent "workers" backed by a single shared
//! tokio runtime). Every POST returns immediately — no artificial delay — so
//! the broker frontend, not the worker, is the bottleneck. This is the worker
//! side of the native rig that replaces the Python uvicorn/FastAPI mock
//! (which capped throughput at a few k RPS per process).
//!
//! Run (from the zc crate):
//!   cargo run --release --example fastworker
//! Env:
//!   FW_PORTS   port range, inclusive, "3960-3967" (default)
//!   FW_THREADS tokio worker threads (default: 4)
//!   FW_DELAY_US  optional per-request delay in microseconds (default 0 = instant)

use std::time::Duration;

use axum::{
    response::IntoResponse,
    routing::{get, post},
    Json, Router,
};
use serde_json::json;

async fn health() -> impl IntoResponse {
    Json(json!({ "status": "ready" }))
}

async fn info() -> impl IntoResponse {
    // worker_type must be present for the broker's discovery to accept this.
    Json(json!({
        "name": "fastworker",
        "worker_type": "benchmark",
        "resources": {
            "cpus_total": 64.0, "cpus_available": 64.0,
            "memory_total": 64u64 * 1024 * 1024 * 1024,
            "memory_available": 64u64 * 1024 * 1024 * 1024,
            "gpus_total": 0, "gpus_available": 0
        },
        "pricing": { "price_per_hour": 0.0, "min_charge": 0.0 },
        "tags": []
    }))
}

async fn execute() -> impl IntoResponse {
    let delay_us: u64 = std::env::var("FW_DELAY_US")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(0);
    if delay_us > 0 {
        tokio::time::sleep(Duration::from_micros(delay_us)).await;
    }
    (
        [(axum::http::header::CONTENT_TYPE, "application/octet-stream")],
        b"ok".to_vec(),
    )
}

fn parse_ports(s: &str) -> Vec<u16> {
    if let Some((a, b)) = s.split_once('-') {
        let a: u16 = a.trim().parse().expect("bad start port");
        let b: u16 = b.trim().parse().expect("bad end port");
        (a..=b).collect()
    } else {
        s.split(',').filter_map(|p| p.trim().parse().ok()).collect()
    }
}

fn main() {
    let ports = parse_ports(&std::env::var("FW_PORTS").unwrap_or_else(|_| "3960-3967".into()));
    let threads: usize = std::env::var("FW_THREADS")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4);

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

    rt.block_on(async move {
        let mut handles = Vec::new();
        for port in ports.clone() {
            let app = Router::new()
                .route("/health", get(health))
                .route("/info", get(info))
                .route("/", post(execute))
                .route("/execute", post(execute));
            let addr = format!("127.0.0.1:{port}");
            let listener = tokio::net::TcpListener::bind(&addr)
                .await
                .unwrap_or_else(|e| panic!("bind {addr}: {e}"));
            handles.push(tokio::spawn(async move {
                axum::serve(listener, app).await.unwrap();
            }));
        }
        eprintln!(
            "[fastworker] serving {} ports {:?} on {} threads (FW_DELAY_US={})",
            ports.len(),
            (ports.first(), ports.last()),
            threads,
            std::env::var("FW_DELAY_US").unwrap_or_else(|_| "0".into())
        );
        for h in handles {
            let _ = h.await;
        }
    });
}