1use async_stream::stream;
39use futures::StreamExt;
40use std::{
41 collections::HashMap,
42 sync::{
43 atomic::{AtomicU64, Ordering},
44 Arc,
45 },
46 time::{Duration, Instant},
47};
48use tokio::sync::Mutex;
49
50#[derive(Debug, Clone)]
52pub struct ChunkedBenchConfig {
53 pub target_url: String,
55 pub method: reqwest::Method,
57 pub concurrency: u32,
59 pub duration: Duration,
61 pub chunk_size_bytes: usize,
63 pub total_size_bytes: usize,
65 pub chunk_interval_ms: u64,
67 pub headers: HashMap<String, String>,
70 pub skip_tls_verify: bool,
72 pub rps: Option<u32>,
76 pub no_keep_alive: bool,
79 pub body: Option<Arc<Vec<u8>>>,
83 pub request_timeout: Option<Duration>,
88}
89
90const MAX_CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
92const REQUEST_TIMEOUT_MARGIN: Duration = Duration::from_secs(60);
94
95impl ChunkedBenchConfig {
96 fn body_len(&self) -> usize {
98 self.body.as_ref().map_or(self.total_size_bytes, |b| b.len())
99 }
100
101 pub fn min_upload_time(&self) -> Duration {
104 let chunks = self.body_len().div_ceil(self.chunk_size_bytes.max(1)) as u64;
105 Duration::from_millis(chunks.saturating_sub(1) * self.chunk_interval_ms)
106 }
107
108 pub fn auto_request_timeout(&self) -> Duration {
111 self.min_upload_time() * 2 + REQUEST_TIMEOUT_MARGIN
112 }
113}
114
115pub const JSON_PADDING_FIELD: &str = "_padding";
117
118pub fn is_json_content_type(headers: &HashMap<String, String>) -> bool {
121 headers.iter().any(|(k, v)| {
122 if !k.eq_ignore_ascii_case("content-type") {
123 return false;
124 }
125 let media = v.split(';').next().unwrap_or("").trim().to_ascii_lowercase();
126 media == "application/json" || media.ends_with("+json")
127 })
128}
129
130pub fn build_json_body(seed: Option<&serde_json::Value>, total: usize) -> Vec<u8> {
138 let mut obj = match seed {
139 Some(serde_json::Value::Object(m)) => m.clone(),
140 Some(other) => {
141 let mut m = serde_json::Map::new();
142 m.insert("data".to_string(), other.clone());
143 m
144 }
145 None => serde_json::Map::new(),
146 };
147 obj.insert(JSON_PADDING_FIELD.to_string(), serde_json::Value::String(String::new()));
148 let base = serde_json::to_vec(&obj).unwrap_or_default().len();
149 let pad = total.saturating_sub(base);
151 obj.insert(JSON_PADDING_FIELD.to_string(), serde_json::Value::String("X".repeat(pad)));
152 serde_json::to_vec(&obj).unwrap_or_default()
153}
154
155struct Pacer {
157 interval: Duration,
158 next: Mutex<Instant>,
159}
160
161impl Pacer {
162 fn new(rps: u32) -> Self {
163 Self {
164 interval: Duration::from_secs_f64(1.0 / f64::from(rps)),
165 next: Mutex::new(Instant::now()),
166 }
167 }
168
169 async fn reserve(&self) -> Instant {
172 let mut next = self.next.lock().await;
173 let slot = (*next).max(Instant::now());
174 *next = slot + self.interval;
175 slot
176 }
177}
178
179#[derive(Debug, Clone)]
184pub struct ErrorSample {
185 pub status: u16,
186 pub server_header: Option<String>,
189 pub body_excerpt: String,
191}
192
193#[derive(Debug, Clone)]
195pub struct ChunkedBenchResult {
196 pub total_requests: u64,
197 pub successful: u64,
198 pub failed: u64,
199 pub bytes_sent: u64,
200 pub elapsed: Duration,
201 pub req_per_sec: f64,
204 pub load_window: Duration,
206 pub start_rate_per_sec: f64,
210 pub upload_mbps: f64,
212 pub latencies_ms: Vec<u64>,
213 pub avg_latency_ms: f64,
214 pub p50_ms: u64,
215 pub p95_ms: u64,
216 pub p99_ms: u64,
217 pub status_counts: HashMap<u16, u64>,
218 pub timed_out: u64,
220 pub request_timeout: Option<Duration>,
222 pub error_samples: Vec<ErrorSample>,
225}
226
227const MAX_ERROR_SAMPLES: usize = 5;
229const ERROR_BODY_EXCERPT_BYTES: usize = 256;
231
232pub async fn run(cfg: ChunkedBenchConfig) -> anyhow::Result<ChunkedBenchResult> {
235 if cfg.chunk_size_bytes == 0 {
236 anyhow::bail!("chunk_size_bytes must be > 0");
237 }
238 if cfg.total_size_bytes == 0 {
239 anyhow::bail!("total_size_bytes must be > 0");
240 }
241 if cfg.concurrency == 0 {
242 anyhow::bail!("concurrency must be >= 1");
243 }
244 if cfg.rps == Some(0) {
245 anyhow::bail!("rps must be >= 1");
246 }
247
248 let mut builder = reqwest::Client::builder().danger_accept_invalid_certs(cfg.skip_tls_verify);
249 if cfg.no_keep_alive {
250 builder = builder.pool_max_idle_per_host(0);
251 }
252 if let Some(timeout) = cfg.request_timeout {
253 builder = builder.connect_timeout(timeout.min(MAX_CONNECT_TIMEOUT));
254 }
255 let client = builder.build()?;
256 let pacer = cfg.rps.map(|n| Arc::new(Pacer::new(n)));
257
258 let total_requests = Arc::new(AtomicU64::new(0));
259 let successful = Arc::new(AtomicU64::new(0));
260 let failed = Arc::new(AtomicU64::new(0));
261 let timed_out = Arc::new(AtomicU64::new(0));
262 let bytes_sent = Arc::new(AtomicU64::new(0));
263 let latencies: Arc<Mutex<Vec<u64>>> = Arc::new(Mutex::new(Vec::with_capacity(8192)));
264 let status_counts: Arc<Mutex<HashMap<u16, u64>>> = Arc::new(Mutex::new(HashMap::new()));
265 let error_samples: Arc<Mutex<Vec<ErrorSample>>> = Arc::new(Mutex::new(Vec::new()));
266
267 let deadline = Instant::now() + cfg.duration;
268 let started_at = Instant::now();
269
270 let mut workers = Vec::with_capacity(cfg.concurrency as usize);
271 for _ in 0..cfg.concurrency {
272 let cfg = cfg.clone();
273 let client = client.clone();
274 let total_requests = total_requests.clone();
275 let successful = successful.clone();
276 let failed = failed.clone();
277 let timed_out = timed_out.clone();
278 let bytes_sent = bytes_sent.clone();
279 let latencies = latencies.clone();
280 let status_counts = status_counts.clone();
281 let error_samples = error_samples.clone();
282 let pacer = pacer.clone();
283
284 workers.push(tokio::spawn(async move {
285 while Instant::now() < deadline {
286 if let Some(pacer) = &pacer {
287 let slot = pacer.reserve().await;
288 if slot >= deadline {
289 break;
290 }
291 tokio::time::sleep_until(slot.into()).await;
292 }
293 let req_started = Instant::now();
294 match send_one_chunked_request(&client, &cfg).await {
295 Ok(SendResult { status, sample }) => {
296 successful.fetch_add(1, Ordering::Relaxed);
297 bytes_sent.fetch_add(cfg.body_len() as u64, Ordering::Relaxed);
298 let elapsed_ms = req_started.elapsed().as_millis() as u64;
299 latencies.lock().await.push(elapsed_ms);
300 *status_counts.lock().await.entry(status).or_insert(0) += 1;
301 if let Some(s) = sample {
302 let mut g = error_samples.lock().await;
303 if g.len() < MAX_ERROR_SAMPLES {
304 g.push(s);
305 }
306 }
307 }
308 Err(e) => {
309 failed.fetch_add(1, Ordering::Relaxed);
310 if e.downcast_ref::<reqwest::Error>()
311 .is_some_and(reqwest::Error::is_timeout)
312 {
313 timed_out.fetch_add(1, Ordering::Relaxed);
314 }
315 }
316 }
317 total_requests.fetch_add(1, Ordering::Relaxed);
318 }
319 }));
320 }
321
322 for w in workers {
323 let _ = w.await;
324 }
325
326 let elapsed = started_at.elapsed();
327 let load_window = cfg.duration;
328 let total = total_requests.load(Ordering::Relaxed);
329 let sent = bytes_sent.load(Ordering::Relaxed);
330 let mut samples: Vec<u64> = {
331 let mut g = latencies.lock().await;
332 std::mem::take(&mut *g)
333 };
334 let final_status_counts: HashMap<u16, u64> = {
335 let mut g = status_counts.lock().await;
336 std::mem::take(&mut *g)
337 };
338 let final_error_samples: Vec<ErrorSample> = {
339 let mut g = error_samples.lock().await;
340 std::mem::take(&mut *g)
341 };
342 samples.sort_unstable();
343 let avg = if samples.is_empty() {
344 0.0
345 } else {
346 samples.iter().copied().sum::<u64>() as f64 / samples.len() as f64
347 };
348 let p = |q: f64| -> u64 {
349 if samples.is_empty() {
350 return 0;
351 }
352 let idx = ((samples.len() as f64 - 1.0) * q).round() as usize;
353 samples[idx]
354 };
355
356 Ok(ChunkedBenchResult {
357 total_requests: total,
358 successful: successful.load(Ordering::Relaxed),
359 failed: failed.load(Ordering::Relaxed),
360 bytes_sent: sent,
361 elapsed,
362 req_per_sec: per_second(total as f64, elapsed),
363 load_window,
364 start_rate_per_sec: per_second(total as f64, load_window),
365 upload_mbps: per_second(sent as f64 * 8.0, elapsed) / 1_000_000.0,
366 avg_latency_ms: avg,
367 p50_ms: p(0.50),
368 p95_ms: p(0.95),
369 p99_ms: p(0.99),
370 latencies_ms: samples,
371 status_counts: final_status_counts,
372 timed_out: timed_out.load(Ordering::Relaxed),
373 request_timeout: cfg.request_timeout,
374 error_samples: final_error_samples,
375 })
376}
377
378pub(crate) fn per_second(amount: f64, window: Duration) -> f64 {
379 if window.is_zero() {
380 0.0
381 } else {
382 amount / window.as_secs_f64()
383 }
384}
385
386struct SendResult {
390 status: u16,
391 sample: Option<ErrorSample>,
392}
393
394async fn send_one_chunked_request(
395 client: &reqwest::Client,
396 cfg: &ChunkedBenchConfig,
397) -> anyhow::Result<SendResult> {
398 let chunk_size = cfg.chunk_size_bytes;
399 let total = cfg.body_len();
400 let interval_ms = cfg.chunk_interval_ms;
401 let prebuilt = cfg.body.clone();
402
403 let body_stream = stream! {
408 let mut sent: usize = 0;
409 let filler = vec![b'X'; chunk_size];
410 while sent < total {
411 if interval_ms > 0 && sent > 0 {
413 tokio::time::sleep(Duration::from_millis(interval_ms)).await;
414 }
415 let next = std::cmp::min(chunk_size, total - sent);
416 let chunk = match &prebuilt {
417 Some(b) => b[sent..sent + next].to_vec(),
418 None => filler[..next].to_vec(),
419 };
420 sent += next;
421 yield Ok::<_, std::io::Error>(chunk);
422 }
423 };
424
425 let body = reqwest::Body::wrap_stream(body_stream.boxed());
426
427 let mut req = client.request(cfg.method.clone(), &cfg.target_url).body(body);
428 if let Some(timeout) = cfg.request_timeout {
429 req = req.timeout(timeout);
430 }
431 for (k, v) in &cfg.headers {
432 req = req.header(k, v);
433 }
434 let resp = req.send().await?;
435 let status = resp.status().as_u16();
436
437 let sample = if !(200..300).contains(&status) {
442 let server_header = resp
443 .headers()
444 .get(reqwest::header::SERVER)
445 .and_then(|v| v.to_str().ok())
446 .map(str::to_owned);
447 let bytes = resp.bytes().await.unwrap_or_default();
448 let take = std::cmp::min(bytes.len(), ERROR_BODY_EXCERPT_BYTES);
449 let body_excerpt = String::from_utf8_lossy(&bytes[..take]).trim().to_owned();
450 Some(ErrorSample {
451 status,
452 server_header,
453 body_excerpt,
454 })
455 } else {
456 None
457 };
458
459 Ok(SendResult { status, sample })
460}
461
462#[cfg(test)]
463mod tests {
464 use super::*;
465
466 #[test]
467 fn error_sample_struct_holds_diagnostic_fields() {
468 let s = ErrorSample {
471 status: 503,
472 server_header: Some("nginx/1.21.0".into()),
473 body_excerpt: "upstream timed out".into(),
474 };
475 assert_eq!(s.status, 503);
476 assert_eq!(s.server_header.as_deref(), Some("nginx/1.21.0"));
477 assert_eq!(s.body_excerpt, "upstream timed out");
478 }
479
480 #[tokio::test]
481 async fn rejects_zero_concurrency() {
482 let cfg = ChunkedBenchConfig {
483 target_url: "http://127.0.0.1:1".into(),
484 method: reqwest::Method::POST,
485 concurrency: 0,
486 duration: Duration::from_millis(10),
487 chunk_size_bytes: 1024,
488 total_size_bytes: 4096,
489 chunk_interval_ms: 0,
490 headers: HashMap::new(),
491 skip_tls_verify: false,
492 rps: None,
493 no_keep_alive: false,
494 body: None,
495 request_timeout: None,
496 };
497 assert!(run(cfg).await.is_err());
498 }
499
500 #[tokio::test]
501 async fn rejects_zero_chunk_size() {
502 let cfg = ChunkedBenchConfig {
503 target_url: "http://127.0.0.1:1".into(),
504 method: reqwest::Method::POST,
505 concurrency: 1,
506 duration: Duration::from_millis(10),
507 chunk_size_bytes: 0,
508 total_size_bytes: 4096,
509 chunk_interval_ms: 0,
510 headers: HashMap::new(),
511 skip_tls_verify: false,
512 rps: None,
513 no_keep_alive: false,
514 body: None,
515 request_timeout: None,
516 };
517 assert!(run(cfg).await.is_err());
518 }
519
520 #[tokio::test]
521 async fn pacer_spaces_slots_evenly() {
522 let pacer = Pacer::new(10);
523 let a = pacer.reserve().await;
524 let b = pacer.reserve().await;
525 let c = pacer.reserve().await;
526 assert_eq!(b - a, Duration::from_millis(100));
527 assert_eq!(c - b, Duration::from_millis(100));
528 }
529
530 #[tokio::test]
531 async fn rps_caps_request_starts_across_workers() {
532 let mut server = mockito::Server::new_async().await;
533 let _m = server
534 .mock("POST", "/upload")
535 .with_status(200)
536 .expect_at_least(1)
537 .create_async()
538 .await;
539 let cfg = ChunkedBenchConfig {
540 target_url: format!("{}/upload", server.url()),
541 method: reqwest::Method::POST,
542 concurrency: 8,
543 duration: Duration::from_millis(1000),
544 chunk_size_bytes: 64,
545 total_size_bytes: 256,
546 chunk_interval_ms: 0,
547 headers: HashMap::new(),
548 skip_tls_verify: false,
549 rps: Some(5),
550 no_keep_alive: true,
551 body: None,
552 request_timeout: None,
553 };
554 let r = run(cfg).await.unwrap();
555 assert!((4..=6).contains(&r.total_requests), "got {} requests", r.total_requests);
558 assert_eq!(r.failed, 0);
559 assert_eq!(r.load_window, Duration::from_millis(1000));
561 assert!((r.start_rate_per_sec - r.total_requests as f64).abs() < 1e-9);
562 assert!(r.start_rate_per_sec <= 5.0 + 1e-9);
563 }
564
565 fn stalled_upload_cfg(url: String, request_timeout: Option<Duration>) -> ChunkedBenchConfig {
566 ChunkedBenchConfig {
567 target_url: url,
568 method: reqwest::Method::POST,
569 concurrency: 1,
570 duration: Duration::from_millis(200),
571 chunk_size_bytes: 64,
572 total_size_bytes: 256,
573 chunk_interval_ms: 0,
574 headers: HashMap::new(),
575 skip_tls_verify: false,
576 rps: None,
577 no_keep_alive: true,
578 body: None,
579 request_timeout,
580 }
581 }
582
583 #[tokio::test]
584 async fn request_timeout_frees_a_worker_stuck_on_a_silent_server() {
585 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
588 let url = format!("http://{}/upload", listener.local_addr().unwrap());
589 tokio::spawn(async move {
590 loop {
591 let (mut sock, _) = listener.accept().await.unwrap();
592 tokio::spawn(async move {
593 let _ = tokio::io::copy(&mut sock, &mut tokio::io::sink()).await;
594 });
595 }
596 });
597
598 let started = Instant::now();
599 let r = run(stalled_upload_cfg(url, Some(Duration::from_millis(300)))).await.unwrap();
600 assert!(started.elapsed() < Duration::from_secs(5), "took {:?}", started.elapsed());
601 assert_eq!(r.total_requests, 1);
602 assert_eq!(r.failed, 1);
603 assert_eq!(r.timed_out, 1);
604 }
605
606 #[test]
607 fn auto_request_timeout_covers_the_minimum_upload_time() {
608 let mut cfg = stalled_upload_cfg(String::new(), None);
610 cfg.total_size_bytes = 1_048_576;
611 cfg.chunk_size_bytes = 4096;
612 cfg.chunk_interval_ms = 50;
613 assert_eq!(cfg.min_upload_time(), Duration::from_millis(12_750));
614 assert_eq!(cfg.auto_request_timeout(), Duration::from_millis(85_500));
615 cfg.chunk_interval_ms = 0;
616 assert_eq!(cfg.auto_request_timeout(), Duration::from_secs(60));
617 }
618
619 #[tokio::test]
620 async fn rejects_zero_rps() {
621 let cfg = ChunkedBenchConfig {
622 target_url: "http://127.0.0.1:1".into(),
623 method: reqwest::Method::POST,
624 concurrency: 1,
625 duration: Duration::from_millis(10),
626 chunk_size_bytes: 1024,
627 total_size_bytes: 4096,
628 chunk_interval_ms: 0,
629 headers: HashMap::new(),
630 skip_tls_verify: false,
631 rps: Some(0),
632 no_keep_alive: false,
633 body: None,
634 request_timeout: None,
635 };
636 assert!(run(cfg).await.is_err());
637 }
638
639 #[test]
640 fn json_body_is_valid_and_exact_size() {
641 let seed = serde_json::json!({"name": "widget", "count": 3});
642 for total in [64usize, 4096, 1_048_576] {
643 let body = build_json_body(Some(&seed), total);
644 assert_eq!(body.len(), total);
645 let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
646 assert_eq!(v["name"], "widget");
647 assert_eq!(v["count"], 3);
648 }
649 }
650
651 #[test]
652 fn json_body_without_seed_or_non_object_seed() {
653 let body = build_json_body(None, 100);
654 assert_eq!(body.len(), 100);
655 assert!(serde_json::from_slice::<serde_json::Value>(&body).unwrap().is_object());
656
657 let body = build_json_body(Some(&serde_json::json!([1, 2])), 100);
658 assert_eq!(body.len(), 100);
659 let v: serde_json::Value = serde_json::from_slice(&body).unwrap();
660 assert_eq!(v["data"], serde_json::json!([1, 2]));
661 }
662
663 #[test]
664 fn json_body_smaller_than_minimum_stays_valid() {
665 let body = build_json_body(None, 1);
666 assert_eq!(body, br#"{"_padding":""}"#);
667 }
668
669 #[test]
670 fn detects_json_content_types() {
671 let h = |v: &str| HashMap::from([("content-type".to_string(), v.to_string())]);
672 assert!(is_json_content_type(&h("application/json")));
673 assert!(is_json_content_type(&h("Application/JSON; charset=utf-8")));
674 assert!(is_json_content_type(&h("application/vnd.api+json")));
675 assert!(!is_json_content_type(&h("application/octet-stream")));
676 assert!(!is_json_content_type(&HashMap::new()));
677 let upper = HashMap::from([("Content-Type".to_string(), "application/json".to_string())]);
678 assert!(is_json_content_type(&upper));
679 }
680
681 #[tokio::test]
682 async fn streams_prebuilt_json_body_intact() {
683 let mut server = mockito::Server::new_async().await;
684 let m = server
685 .mock("POST", "/items")
686 .match_body(mockito::Matcher::PartialJson(serde_json::json!({"name": "widget"})))
687 .with_status(201)
688 .expect_at_least(1)
689 .create_async()
690 .await;
691 let body = build_json_body(Some(&serde_json::json!({"name": "widget"})), 10_000);
692 let cfg = ChunkedBenchConfig {
693 target_url: format!("{}/items", server.url()),
694 method: reqwest::Method::POST,
695 concurrency: 1,
696 duration: Duration::from_millis(300),
697 chunk_size_bytes: 777,
698 total_size_bytes: 10_000,
699 chunk_interval_ms: 0,
700 headers: HashMap::from([("Content-Type".into(), "application/json".into())]),
701 skip_tls_verify: false,
702 rps: None,
703 no_keep_alive: false,
704 body: Some(Arc::new(body)),
705 request_timeout: None,
706 };
707 let r = run(cfg).await.unwrap();
708 assert_eq!(r.failed, 0);
709 assert!(r.status_counts.contains_key(&201), "statuses: {:?}", r.status_counts);
710 assert_eq!(r.bytes_sent, r.successful * 10_000);
711 m.assert_async().await;
712 }
713}