use async_stream::stream;
use futures::StreamExt;
use std::{
collections::HashMap,
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
time::{Duration, Instant},
};
use tokio::sync::Mutex;
#[derive(Debug, Clone)]
pub struct ChunkedBenchConfig {
pub target_url: String,
pub method: reqwest::Method,
pub concurrency: u32,
pub duration: Duration,
pub chunk_size_bytes: usize,
pub total_size_bytes: usize,
pub chunk_interval_ms: u64,
pub headers: HashMap<String, String>,
pub skip_tls_verify: bool,
pub rps: Option<u32>,
pub no_keep_alive: bool,
pub body: Option<Arc<Vec<u8>>>,
pub request_timeout: Option<Duration>,
}
const MAX_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
const REQUEST_TIMEOUT_MARGIN: Duration = Duration::from_secs(60);
impl ChunkedBenchConfig {
fn body_len(&self) -> usize {
self.body.as_ref().map_or(self.total_size_bytes, |b| b.len())
}
pub fn min_upload_time(&self) -> Duration {
let chunks = self.body_len().div_ceil(self.chunk_size_bytes.max(1)) as u64;
Duration::from_millis(chunks.saturating_sub(1) * self.chunk_interval_ms)
}
pub fn auto_request_timeout(&self) -> Duration {
self.min_upload_time() * 2 + REQUEST_TIMEOUT_MARGIN
}
}
pub const JSON_PADDING_FIELD: &str = "_padding";
pub fn is_json_content_type(headers: &HashMap<String, String>) -> bool {
headers.iter().any(|(k, v)| {
if !k.eq_ignore_ascii_case("content-type") {
return false;
}
let media = v.split(';').next().unwrap_or("").trim().to_ascii_lowercase();
media == "application/json" || media.ends_with("+json")
})
}
pub fn build_json_body(seed: Option<&serde_json::Value>, total: usize) -> Vec<u8> {
let mut obj = match seed {
Some(serde_json::Value::Object(m)) => m.clone(),
Some(other) => {
let mut m = serde_json::Map::new();
m.insert("data".to_string(), other.clone());
m
}
None => serde_json::Map::new(),
};
obj.insert(JSON_PADDING_FIELD.to_string(), serde_json::Value::String(String::new()));
let base = serde_json::to_vec(&obj).unwrap_or_default().len();
let pad = total.saturating_sub(base);
obj.insert(JSON_PADDING_FIELD.to_string(), serde_json::Value::String("X".repeat(pad)));
serde_json::to_vec(&obj).unwrap_or_default()
}
struct Pacer {
interval: Duration,
next: Mutex<Instant>,
}
impl Pacer {
fn new(rps: u32) -> Self {
Self {
interval: Duration::from_secs_f64(1.0 / f64::from(rps)),
next: Mutex::new(Instant::now()),
}
}
async fn reserve(&self) -> Instant {
let mut next = self.next.lock().await;
let slot = (*next).max(Instant::now());
*next = slot + self.interval;
slot
}
}
#[derive(Debug, Clone)]
pub struct ErrorSample {
pub status: u16,
pub server_header: Option<String>,
pub body_excerpt: String,
}
#[derive(Debug, Clone)]
pub struct ChunkedBenchResult {
pub total_requests: u64,
pub successful: u64,
pub failed: u64,
pub bytes_sent: u64,
pub elapsed: Duration,
pub req_per_sec: f64,
pub load_window: Duration,
pub start_rate_per_sec: f64,
pub upload_mbps: f64,
pub latencies_ms: Vec<u64>,
pub avg_latency_ms: f64,
pub p50_ms: u64,
pub p95_ms: u64,
pub p99_ms: u64,
pub status_counts: HashMap<u16, u64>,
pub timed_out: u64,
pub request_timeout: Option<Duration>,
pub error_samples: Vec<ErrorSample>,
}
const MAX_ERROR_SAMPLES: usize = 5;
const ERROR_BODY_EXCERPT_BYTES: usize = 256;
pub async fn run(cfg: ChunkedBenchConfig) -> anyhow::Result<ChunkedBenchResult> {
if cfg.chunk_size_bytes == 0 {
anyhow::bail!("chunk_size_bytes must be > 0");
}
if cfg.total_size_bytes == 0 {
anyhow::bail!("total_size_bytes must be > 0");
}
if cfg.concurrency == 0 {
anyhow::bail!("concurrency must be >= 1");
}
if cfg.rps == Some(0) {
anyhow::bail!("rps must be >= 1");
}
let mut builder = reqwest::Client::builder().danger_accept_invalid_certs(cfg.skip_tls_verify);
if cfg.no_keep_alive {
builder = builder.pool_max_idle_per_host(0);
}
if let Some(timeout) = cfg.request_timeout {
builder = builder.connect_timeout(timeout.min(MAX_CONNECT_TIMEOUT));
}
let client = builder.build()?;
let pacer = cfg.rps.map(|n| Arc::new(Pacer::new(n)));
let total_requests = Arc::new(AtomicU64::new(0));
let successful = Arc::new(AtomicU64::new(0));
let failed = Arc::new(AtomicU64::new(0));
let timed_out = Arc::new(AtomicU64::new(0));
let bytes_sent = Arc::new(AtomicU64::new(0));
let latencies: Arc<Mutex<Vec<u64>>> = Arc::new(Mutex::new(Vec::with_capacity(8192)));
let status_counts: Arc<Mutex<HashMap<u16, u64>>> = Arc::new(Mutex::new(HashMap::new()));
let error_samples: Arc<Mutex<Vec<ErrorSample>>> = Arc::new(Mutex::new(Vec::new()));
let deadline = Instant::now() + cfg.duration;
let started_at = Instant::now();
let mut workers = Vec::with_capacity(cfg.concurrency as usize);
for _ in 0..cfg.concurrency {
let cfg = cfg.clone();
let client = client.clone();
let total_requests = total_requests.clone();
let successful = successful.clone();
let failed = failed.clone();
let timed_out = timed_out.clone();
let bytes_sent = bytes_sent.clone();
let latencies = latencies.clone();
let status_counts = status_counts.clone();
let error_samples = error_samples.clone();
let pacer = pacer.clone();
workers.push(tokio::spawn(async move {
while Instant::now() < deadline {
if let Some(pacer) = &pacer {
let slot = pacer.reserve().await;
if slot >= deadline {
break;
}
tokio::time::sleep_until(slot.into()).await;
}
let req_started = Instant::now();
match send_one_chunked_request(&client, &cfg).await {
Ok(SendResult { status, sample }) => {
successful.fetch_add(1, Ordering::Relaxed);
bytes_sent.fetch_add(cfg.body_len() as u64, Ordering::Relaxed);
let elapsed_ms = req_started.elapsed().as_millis() as u64;
latencies.lock().await.push(elapsed_ms);
*status_counts.lock().await.entry(status).or_insert(0) += 1;
if let Some(s) = sample {
let mut g = error_samples.lock().await;
if g.len() < MAX_ERROR_SAMPLES {
g.push(s);
}
}
}
Err(e) => {
failed.fetch_add(1, Ordering::Relaxed);
if e.downcast_ref::<reqwest::Error>()
.is_some_and(reqwest::Error::is_timeout)
{
timed_out.fetch_add(1, Ordering::Relaxed);
}
}
}
total_requests.fetch_add(1, Ordering::Relaxed);
}
}));
}
for w in workers {
let _ = w.await;
}
let elapsed = started_at.elapsed();
let load_window = cfg.duration;
let total = total_requests.load(Ordering::Relaxed);
let sent = bytes_sent.load(Ordering::Relaxed);
let mut samples: Vec<u64> = {
let mut g = latencies.lock().await;
std::mem::take(&mut *g)
};
let final_status_counts: HashMap<u16, u64> = {
let mut g = status_counts.lock().await;
std::mem::take(&mut *g)
};
let final_error_samples: Vec<ErrorSample> = {
let mut g = error_samples.lock().await;
std::mem::take(&mut *g)
};
samples.sort_unstable();
let avg = if samples.is_empty() {
0.0
} else {
samples.iter().copied().sum::<u64>() as f64 / samples.len() as f64
};
let p = |q: f64| -> u64 {
if samples.is_empty() {
return 0;
}
let idx = ((samples.len() as f64 - 1.0) * q).round() as usize;
samples[idx]
};
Ok(ChunkedBenchResult {
total_requests: total,
successful: successful.load(Ordering::Relaxed),
failed: failed.load(Ordering::Relaxed),
bytes_sent: sent,
elapsed,
req_per_sec: per_second(total as f64, elapsed),
load_window,
start_rate_per_sec: per_second(total as f64, load_window),
upload_mbps: per_second(sent as f64 * 8.0, elapsed) / 1_000_000.0,
avg_latency_ms: avg,
p50_ms: p(0.50),
p95_ms: p(0.95),
p99_ms: p(0.99),
latencies_ms: samples,
status_counts: final_status_counts,
timed_out: timed_out.load(Ordering::Relaxed),
request_timeout: cfg.request_timeout,
error_samples: final_error_samples,
})
}
pub(crate) fn per_second(amount: f64, window: Duration) -> f64 {
if window.is_zero() {
0.0
} else {
amount / window.as_secs_f64()
}
}
struct SendResult {
status: u16,
sample: Option<ErrorSample>,
}
async fn send_one_chunked_request(
client: &reqwest::Client,
cfg: &ChunkedBenchConfig,
) -> anyhow::Result<SendResult> {
let chunk_size = cfg.chunk_size_bytes;
let total = cfg.body_len();
let interval_ms = cfg.chunk_interval_ms;
let prebuilt = cfg.body.clone();
let body_stream = stream! {
let mut sent: usize = 0;
let filler = vec![b'X'; chunk_size];
while sent < total {
if interval_ms > 0 && sent > 0 {
tokio::time::sleep(Duration::from_millis(interval_ms)).await;
}
let next = std::cmp::min(chunk_size, total - sent);
let chunk = match &prebuilt {
Some(b) => b[sent..sent + next].to_vec(),
None => filler[..next].to_vec(),
};
sent += next;
yield Ok::<_, std::io::Error>(chunk);
}
};
let body = reqwest::Body::wrap_stream(body_stream.boxed());
let mut req = client.request(cfg.method.clone(), &cfg.target_url).body(body);
if let Some(timeout) = cfg.request_timeout {
req = req.timeout(timeout);
}
for (k, v) in &cfg.headers {
req = req.header(k, v);
}
let resp = req.send().await?;
let status = resp.status().as_u16();
let sample = if !(200..300).contains(&status) {
let server_header = resp
.headers()
.get(reqwest::header::SERVER)
.and_then(|v| v.to_str().ok())
.map(str::to_owned);
let bytes = resp.bytes().await.unwrap_or_default();
let take = std::cmp::min(bytes.len(), ERROR_BODY_EXCERPT_BYTES);
let body_excerpt = String::from_utf8_lossy(&bytes[..take]).trim().to_owned();
Some(ErrorSample {
status,
server_header,
body_excerpt,
})
} else {
None
};
Ok(SendResult { status, sample })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn error_sample_struct_holds_diagnostic_fields() {
let s = ErrorSample {
status: 503,
server_header: Some("nginx/1.21.0".into()),
body_excerpt: "upstream timed out".into(),
};
assert_eq!(s.status, 503);
assert_eq!(s.server_header.as_deref(), Some("nginx/1.21.0"));
assert_eq!(s.body_excerpt, "upstream timed out");
}
#[tokio::test]
async fn rejects_zero_concurrency() {
let cfg = ChunkedBenchConfig {
target_url: "http://127.0.0.1:1".into(),
method: reqwest::Method::POST,
concurrency: 0,
duration: Duration::from_millis(10),
chunk_size_bytes: 1024,
total_size_bytes: 4096,
chunk_interval_ms: 0,
headers: HashMap::new(),
skip_tls_verify: false,
rps: None,
no_keep_alive: false,
body: None,
request_timeout: None,
};
assert!(run(cfg).await.is_err());
}
#[tokio::test]
async fn rejects_zero_chunk_size() {
let cfg = ChunkedBenchConfig {
target_url: "http://127.0.0.1:1".into(),
method: reqwest::Method::POST,
concurrency: 1,
duration: Duration::from_millis(10),
chunk_size_bytes: 0,
total_size_bytes: 4096,
chunk_interval_ms: 0,
headers: HashMap::new(),
skip_tls_verify: false,
rps: None,
no_keep_alive: false,
body: None,
request_timeout: None,
};
assert!(run(cfg).await.is_err());
}
#[tokio::test]
async fn pacer_spaces_slots_evenly() {
let pacer = Pacer::new(10);
let a = pacer.reserve().await;
let b = pacer.reserve().await;
let c = pacer.reserve().await;
assert_eq!(b - a, Duration::from_millis(100));
assert_eq!(c - b, Duration::from_millis(100));
}
#[tokio::test]
async fn rps_caps_request_starts_across_workers() {
let mut server = mockito::Server::new_async().await;
let _m = server
.mock("POST", "/upload")
.with_status(200)
.expect_at_least(1)
.create_async()
.await;
let cfg = ChunkedBenchConfig {
target_url: format!("{}/upload", server.url()),
method: reqwest::Method::POST,
concurrency: 8,
duration: Duration::from_millis(1000),
chunk_size_bytes: 64,
total_size_bytes: 256,
chunk_interval_ms: 0,
headers: HashMap::new(),
skip_tls_verify: false,
rps: Some(5),
no_keep_alive: true,
body: None,
request_timeout: None,
};
let r = run(cfg).await.unwrap();
assert!((4..=6).contains(&r.total_requests), "got {} requests", r.total_requests);
assert_eq!(r.failed, 0);
assert_eq!(r.load_window, Duration::from_millis(1000));
assert!((r.start_rate_per_sec - r.total_requests as f64).abs() < 1e-9);
assert!(r.start_rate_per_sec <= 5.0 + 1e-9);
}
fn stalled_upload_cfg(url: String, request_timeout: Option<Duration>) -> ChunkedBenchConfig {
ChunkedBenchConfig {
target_url: url,
method: reqwest::Method::POST,
concurrency: 1,
duration: Duration::from_millis(200),
chunk_size_bytes: 64,
total_size_bytes: 256,
chunk_interval_ms: 0,
headers: HashMap::new(),
skip_tls_verify: false,
rps: None,
no_keep_alive: true,
body: None,
request_timeout,
}
}
#[tokio::test]
async fn request_timeout_frees_a_worker_stuck_on_a_silent_server() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let url = format!("http://{}/upload", listener.local_addr().unwrap());
tokio::spawn(async move {
loop {
let (mut sock, _) = listener.accept().await.unwrap();
tokio::spawn(async move {
let _ = tokio::io::copy(&mut sock, &mut tokio::io::sink()).await;
});
}
});
let started = Instant::now();
let r = run(stalled_upload_cfg(url, Some(Duration::from_millis(300)))).await.unwrap();
assert!(started.elapsed() < Duration::from_secs(5), "took {:?}", started.elapsed());
assert_eq!(r.total_requests, 1);
assert_eq!(r.failed, 1);
assert_eq!(r.timed_out, 1);
}
#[test]
fn auto_request_timeout_covers_the_minimum_upload_time() {
let mut cfg = stalled_upload_cfg(String::new(), None);
cfg.total_size_bytes = 1_048_576;
cfg.chunk_size_bytes = 4096;
cfg.chunk_interval_ms = 50;
assert_eq!(cfg.min_upload_time(), Duration::from_millis(12_750));
assert_eq!(cfg.auto_request_timeout(), Duration::from_millis(85_500));
cfg.chunk_interval_ms = 0;
assert_eq!(cfg.auto_request_timeout(), Duration::from_secs(60));
}
#[tokio::test]
async fn rejects_zero_rps() {
let cfg = ChunkedBenchConfig {
target_url: "http://127.0.0.1:1".into(),
method: reqwest::Method::POST,
concurrency: 1,
duration: Duration::from_millis(10),
chunk_size_bytes: 1024,
total_size_bytes: 4096,
chunk_interval_ms: 0,
headers: HashMap::new(),
skip_tls_verify: false,
rps: Some(0),
no_keep_alive: false,
body: None,
request_timeout: None,
};
assert!(run(cfg).await.is_err());
}
#[test]
fn json_body_is_valid_and_exact_size() {
let seed = serde_json::json!({"name": "widget", "count": 3});
for total in [64usize, 4096, 1_048_576] {
let body = build_json_body(Some(&seed), total);
assert_eq!(body.len(), total);
let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(v["name"], "widget");
assert_eq!(v["count"], 3);
}
}
#[test]
fn json_body_without_seed_or_non_object_seed() {
let body = build_json_body(None, 100);
assert_eq!(body.len(), 100);
assert!(serde_json::from_slice::<serde_json::Value>(&body).unwrap().is_object());
let body = build_json_body(Some(&serde_json::json!([1, 2])), 100);
assert_eq!(body.len(), 100);
let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(v["data"], serde_json::json!([1, 2]));
}
#[test]
fn json_body_smaller_than_minimum_stays_valid() {
let body = build_json_body(None, 1);
assert_eq!(body, br#"{"_padding":""}"#);
}
#[test]
fn detects_json_content_types() {
let h = |v: &str| HashMap::from([("content-type".to_string(), v.to_string())]);
assert!(is_json_content_type(&h("application/json")));
assert!(is_json_content_type(&h("Application/JSON; charset=utf-8")));
assert!(is_json_content_type(&h("application/vnd.api+json")));
assert!(!is_json_content_type(&h("application/octet-stream")));
assert!(!is_json_content_type(&HashMap::new()));
let upper = HashMap::from([("Content-Type".to_string(), "application/json".to_string())]);
assert!(is_json_content_type(&upper));
}
#[tokio::test]
async fn streams_prebuilt_json_body_intact() {
let mut server = mockito::Server::new_async().await;
let m = server
.mock("POST", "/items")
.match_body(mockito::Matcher::PartialJson(serde_json::json!({"name": "widget"})))
.with_status(201)
.expect_at_least(1)
.create_async()
.await;
let body = build_json_body(Some(&serde_json::json!({"name": "widget"})), 10_000);
let cfg = ChunkedBenchConfig {
target_url: format!("{}/items", server.url()),
method: reqwest::Method::POST,
concurrency: 1,
duration: Duration::from_millis(300),
chunk_size_bytes: 777,
total_size_bytes: 10_000,
chunk_interval_ms: 0,
headers: HashMap::from([("Content-Type".into(), "application/json".into())]),
skip_tls_verify: false,
rps: None,
no_keep_alive: false,
body: Some(Arc::new(body)),
request_timeout: None,
};
let r = run(cfg).await.unwrap();
assert_eq!(r.failed, 0);
assert!(r.status_counts.contains_key(&201), "statuses: {:?}", r.status_counts);
assert_eq!(r.bytes_sent, r.successful * 10_000);
m.assert_async().await;
}
}