core/test/performance/fixtures/http/concurrent.rs

//! 独立 Hyper 服务及三种语言共用的 HTTP/1.1 客户端，不链接 Dever 运行库。
use bytes::Bytes;
use http_body_util::{BodyExt, Full, Limited};
use hyper::{Request, Response, body::Incoming, service::service_fn};
use hyper_util::rt::{TokioIo, TokioTimer};
use std::{convert::Infallible, sync::Arc, time::Duration};
use tokio::{net::TcpListener, sync::{Barrier, Semaphore}, time::Instant};

async fn route(request: Request<Incoming>) -> Result<Response<Full<Bytes>>, Infallible> {
    let path = request.uri().path().to_owned();
    let body = match Limited::new(request.into_body(), 65536).collect().await {
        Ok(body) => body.to_bytes(),
        Err(_) => return Ok(Response::builder().status(400).body(Full::new(Bytes::new())).unwrap()),
    };
    let (status, kind, body) = match path.as_str() {
        "/json" => (200, "application/json", Bytes::from_static(b"{\"ok\":true}")),
        "/echo" => (200, "application/octet-stream", body),
        _ => (404, "text/plain", Bytes::new()),
    };
    Ok(Response::builder().status(status).header("content-type", kind).body(Full::new(body)).unwrap())
}

async fn serve(capacity: usize) -> Result<(), Box<dyn std::error::Error>> {
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let slots = Arc::new(Semaphore::new(capacity));
    println!("READY|{}", listener.local_addr()?.port());
    loop {
        let permit = slots.clone().acquire_owned().await?;
        let (stream, _) = listener.accept().await?;
        stream.set_nodelay(true)?;
        tokio::spawn(async move {
            let _permit = permit;
            let mut builder = hyper::server::conn::http1::Builder::new();
            builder.timer(TokioTimer::new()).header_read_timeout(Duration::from_secs(5)).max_buf_size(8192);
            if let Err(error) = builder.serve_connection(TokioIo::new(stream), service_fn(route)).await {
                if !error.is_incomplete_message() && !error.is_closed() {
                    eprintln!("connection: {error}");
                }
            }
        });
    }
}

async fn load(port: u16, path: String, concurrency: usize, seconds: f64) -> Result<(), Box<dyn std::error::Error>> {
    let payload = if path == "/echo" { Bytes::from(vec![b'x'; 65536]) } else { Bytes::new() };
    let expected = if path == "/echo" { payload.clone() } else { Bytes::from_static(b"{\"ok\":true}") };
    let kind = if path == "/echo" { "application/octet-stream" } else { "application/json" };
    let barrier = Arc::new(Barrier::new(concurrency + 1));
    let mut workers = Vec::new();
    for _ in 0..concurrency {
        let stream = tokio::net::TcpStream::connect(("127.0.0.1", port)).await?;
        stream.set_nodelay(true)?;
        let (mut sender, connection) = hyper::client::conn::http1::handshake(TokioIo::new(stream)).await?;
        let driver = tokio::spawn(connection);
        let barrier = barrier.clone();
        let payload = payload.clone();
        let expected = expected.clone();
        let path = path.clone();
        workers.push(tokio::spawn(async move {
            barrier.wait().await;
            let deadline = Instant::now() + Duration::from_secs_f64(seconds);
            let mut samples = Vec::new();
            let mut errors = 0_u64;
            while Instant::now() < deadline {
                let request = Request::builder().method(if path == "/echo" { "POST" } else { "GET" })
                    .uri(&path).header("host", "localhost").body(Full::new(payload.clone())).unwrap();
                let started = Instant::now();
                let exchange = async {
                    let response = sender.send_request(request).await.map_err(|error| error.to_string())?;
                    if response.status() != 200 || response.headers().get("content-type").and_then(|value| value.to_str().ok()) != Some(kind) {
                        return Err("status or content type differs".to_owned());
                    }
                    let body = response.into_body().collect().await.map_err(|error| error.to_string())?.to_bytes();
                    if body != expected { return Err("response body differs".to_owned()); }
                    Ok(())
                };
                match tokio::time::timeout(Duration::from_secs(5), exchange).await {
                    Ok(Ok(())) => samples.push(started.elapsed().as_micros() as u64),
                    result => {
                        eprintln!("request: {result:?}");
                        errors += 1;
                        break;
                    }
                }
            }
            drop(sender);
            driver.abort();
            let _ = driver.await;
            (samples, errors)
        }));
    }
    let started = Instant::now();
    barrier.wait().await;
    let mut samples = Vec::new();
    let mut errors = 0;
    for worker in workers {
        let (mut values, failed) = worker.await?;
        samples.append(&mut values);
        errors += failed;
    }
    let elapsed = started.elapsed().as_secs_f64();
    samples.sort_unstable();
    let percentile = |percent: usize| -> Option<f64> {
        if samples.is_empty() { None } else { Some(samples[(samples.len() * percent).div_ceil(100) - 1] as f64 / 1000.0) }
    };
    println!("{}", serde_json::json!({"concurrency":concurrency,"seconds":seconds,"elapsed_seconds":elapsed,
        "succeeded":samples.len(),"errors":errors,"qps":samples.len() as f64 / elapsed,
        "p50_ms":percentile(50),"p95_ms":percentile(95),"p99_ms":percentile(99)}));
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let arguments = std::env::args().skip(1).collect::<Vec<_>>();
    let server = arguments[0] == "serve";
    let runtime = tokio::runtime::Builder::new_multi_thread().worker_threads(if server { 2 } else { 4 }).enable_all().build()?;
    if server {
        runtime.block_on(serve(arguments[1].parse()?))
    } else {
        runtime.block_on(load(arguments[1].parse()?, arguments[2].clone(), arguments[3].parse()?, arguments[4].parse()?))
    }
}