1use std::sync::atomic::{AtomicU64, Ordering};
22use std::sync::Arc;
23use std::time::{Duration, Instant};
24
25use serde::Serialize;
26use tokio::io::{AsyncReadExt, AsyncWriteExt};
27use tokio::net::TcpStream;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum BodyKind {
32 Json,
33 Xml,
34 UrlEncoded,
35 Multipart,
36}
37
38impl BodyKind {
39 pub fn content_type(self) -> &'static str {
41 match self {
42 BodyKind::Json => "application/json",
43 BodyKind::Xml => "application/xml",
44 BodyKind::UrlEncoded => "application/x-www-form-urlencoded",
45 BodyKind::Multipart => "multipart/form-data; boundary=mockforge-pipeline",
46 }
47 }
48}
49
50pub fn parse_body_kind(name: &str) -> Result<BodyKind, String> {
52 match name.to_ascii_lowercase().as_str() {
53 "json" | "application/json" => Ok(BodyKind::Json),
54 "xml" | "application/xml" => Ok(BodyKind::Xml),
55 "urlencoded" | "form" | "application/x-www-form-urlencoded" => Ok(BodyKind::UrlEncoded),
56 "multipart" | "multipart/form-data" => Ok(BodyKind::Multipart),
57 other => Err(format!(
58 "unsupported content type '{other}' (use json | xml | urlencoded | multipart)"
59 )),
60 }
61}
62
63pub fn parse_body_size(spec: &str) -> Result<u64, String> {
66 let s = spec.trim();
67 let lower = s.to_ascii_lowercase();
68 let suffix_mult: Option<(&str, u64)> = [
70 ("kib", 1u64 << 10),
71 ("mib", 1u64 << 20),
72 ("gib", 1u64 << 30),
73 ("kb", 1_000),
74 ("mb", 1_000_000),
75 ("gb", 1_000_000_000),
76 ]
77 .into_iter()
78 .find_map(|(suf, mult)| lower.strip_suffix(suf).map(|num| (num, mult)));
79 let (num_part, mult): (&str, u64) = suffix_mult.unwrap_or((s, 1));
80 let n: f64 = num_part
81 .trim()
82 .parse()
83 .map_err(|_| format!("invalid body size '{spec}' (examples: 500KB, 10MB, 2GB)"))?;
84 if n < 0.0 {
85 return Err(format!("body size must be non-negative, got '{spec}'"));
86 }
87 let bytes = (n * mult as f64) as u64;
88 if bytes == 0 {
89 return Err(format!("body size must be > 0, got '{spec}'"));
90 }
91 Ok(bytes)
92}
93
94struct StreamBody {
98 kind: BodyKind,
99 prefix: Vec<u8>,
100 suffix: Vec<u8>,
101 fill_total: u64,
102 written_fill: u64,
103 phase: u8,
104 done_prefix: bool,
105 done_fill: bool,
106 done_suffix: bool,
107}
108
109impl StreamBody {
110 const CHUNK: usize = 64 * 1024;
111
112 fn new(kind: BodyKind, size: u64) -> Result<Self, String> {
115 use BodyKind::*;
116 let (prefix, suffix): (Vec<u8>, Vec<u8>) = match kind {
117 Json => (b"{\"data\":\"".to_vec(), b"\"}".to_vec()),
118 Xml => (b"<root><data>".to_vec(), b"</data></root>".to_vec()),
119 UrlEncoded => (Vec::new(), Vec::new()),
120 Multipart => (
123 b"--mockforge-pipeline\r\n\
124 Content-Disposition: form-data; name=\"file\"; filename=\"synthetic.bin\"\r\n\
125 Content-Type: application/octet-stream\r\n\r\n"
126 .to_vec(),
127 b"\r\n--mockforge-pipeline--\r\n".to_vec(),
128 ),
129 };
130 let overhead = (prefix.len() + suffix.len()) as u64;
131 if size < overhead {
132 return Err(format!(
133 "body size {size} is smaller than the {kind:?} framing overhead ({overhead} bytes)"
134 ));
135 }
136 Ok(Self {
137 kind,
138 prefix,
139 suffix,
140 fill_total: size - overhead,
141 written_fill: 0,
142 phase: b'a',
143 done_prefix: false,
144 done_fill: false,
145 done_suffix: false,
146 })
147 }
148
149 #[cfg_attr(not(test), allow(dead_code))]
150 fn exhausted(&self) -> bool {
151 self.done_prefix && self.done_fill && self.done_suffix
152 }
153
154 fn next_chunk(&mut self) -> Option<Vec<u8>> {
156 if !self.done_prefix {
157 self.done_prefix = true;
158 return Some(std::mem::take(&mut self.prefix));
159 }
160 if !self.done_fill && self.fill_total > 0 {
161 let remaining = self.fill_total - self.written_fill;
162 let take = remaining.min(Self::CHUNK as u64) as usize;
163 let mut buf = vec![0u8; take];
164 for b in buf.iter_mut() {
165 *b = self.phase;
168 self.phase = if self.phase >= b'z' {
169 b'a'
170 } else {
171 self.phase + 1
172 };
173 }
174 self.written_fill += take as u64;
175 if self.written_fill == self.fill_total {
176 self.done_fill = true;
177 }
178 return Some(buf);
179 }
180 if !self.done_fill && self.fill_total == 0 {
181 self.done_fill = true;
182 }
183 if !self.done_suffix {
184 self.done_suffix = true;
185 return Some(std::mem::take(&mut self.suffix));
186 }
187 None
188 }
189}
190
191#[derive(Debug, Clone)]
193pub struct PipelineBenchConfig {
194 pub target_url: String,
196 pub method: String,
199 pub body_kind: BodyKind,
201 pub body_size: u64,
203 pub pipeline_depth: usize,
205 pub connections: usize,
207 pub duration: Duration,
209}
210
211#[derive(Debug, Clone, Default, Serialize)]
213pub struct PipelineBenchResult {
214 pub requests_sent: u64,
216 pub responses_received: u64,
218 pub status_counts: std::collections::BTreeMap<String, u64>,
220 pub bytes_sent: u64,
222 pub bytes_received: u64,
224 pub connection_closed_early: u64,
228 pub connection_errors: u64,
230 pub reconnects: u64,
232}
233
234fn parse_http_target(url: &str) -> Result<(String, u16, String), String> {
237 let rest = url
238 .strip_prefix("http://")
239 .ok_or_else(|| format!("bench-pipeline needs a plain http:// target, got '{url}'"))?;
240 let (authority, path) = match rest.find('/') {
241 Some(i) => (&rest[..i], &rest[i..]),
242 None => (rest, "/"),
243 };
244 let (host, port) = match authority.rsplit_once(':') {
245 Some((h, p)) => {
246 let port: u16 = p.parse().map_err(|_| format!("invalid port in '{url}'"))?;
247 (h.to_string(), port)
248 }
249 None => (authority.to_string(), 80),
250 };
251 if host.is_empty() {
252 return Err(format!("empty host in '{url}'"));
253 }
254 Ok((
255 host,
256 port,
257 if path.is_empty() {
258 "/".into()
259 } else {
260 path.to_string()
261 },
262 ))
263}
264
265async fn read_response(stream: &mut TcpStream, buf: &mut Vec<u8>) -> std::io::Result<(u16, u64)> {
268 let mut consumed: u64 = 0;
269 let head_end = loop {
271 if let Some(pos) = find_head_end(buf) {
272 break pos;
273 }
274 if buf.len() > 128 * 1024 {
275 return Err(std::io::Error::new(
276 std::io::ErrorKind::InvalidData,
277 "response head exceeds 128 KiB",
278 ));
279 }
280 let mut chunk = [0u8; 8192];
281 let n = stream.read(&mut chunk).await?;
282 if n == 0 {
283 return Err(std::io::Error::new(
284 std::io::ErrorKind::UnexpectedEof,
285 "connection closed during response head",
286 ));
287 }
288 consumed += n as u64;
289 buf.extend_from_slice(&chunk[..n]);
290 };
291 let head = String::from_utf8_lossy(&buf[..head_end]).to_string();
292 buf.drain(..head_end + 4);
293
294 let status = head
295 .lines()
296 .next()
297 .and_then(|l| l.split_whitespace().nth(1))
298 .and_then(|s| s.parse::<u16>().ok())
299 .ok_or_else(|| {
300 std::io::Error::new(std::io::ErrorKind::InvalidData, "malformed status line")
301 })?;
302
303 let lower = head.to_ascii_lowercase();
304 let content_length: Option<u64> = lower
305 .lines()
306 .find_map(|l| l.strip_prefix("content-length:"))
307 .and_then(|v| v.trim().parse().ok());
308 let chunked = lower.contains("transfer-Encoding:")
309 || lower.contains("transfer-encoding:")
310 && lower.split("transfer-encoding:").nth(1).is_some_and(|v| v.contains("chunked"));
311
312 if chunked {
314 loop {
315 let size_line_end = loop {
317 if let Some(p) = buf.windows(2).position(|w| w == b"\r\n") {
318 break p;
319 }
320 let mut chunk = [0u8; 1024];
321 let n = stream.read(&mut chunk).await?;
322 if n == 0 {
323 return Err(std::io::Error::new(
324 std::io::ErrorKind::UnexpectedEof,
325 "eof inside chunked body",
326 ));
327 }
328 consumed += n as u64;
329 buf.extend_from_slice(&chunk[..n]);
330 };
331 let size_str = String::from_utf8_lossy(&buf[..size_line_end]);
332 let size = u64::from_str_radix(size_str.split(';').next().unwrap_or("0").trim(), 16)
333 .unwrap_or(0);
334 buf.drain(..size_line_end + 2);
335 if size == 0 {
336 loop {
338 if let Some(p) = buf.windows(2).position(|w| w == b"\r\n") {
339 buf.drain(..p + 2);
340 if p == 0 {
341 break;
342 }
343 } else {
344 let mut chunk = [0u8; 1024];
345 let n = stream.read(&mut chunk).await?;
346 if n == 0 {
347 break;
348 }
349 consumed += n as u64;
350 buf.extend_from_slice(&chunk[..n]);
351 }
352 }
353 break;
354 }
355 let mut to_read = size + 2; while to_read > 0 {
357 if !buf.is_empty() {
358 let take = to_read.min(buf.len() as u64) as usize;
359 buf.drain(..take);
360 to_read -= take as u64;
361 } else {
362 let mut chunk = [0u8; 16384];
363 let cap = to_read.min(chunk.len() as u64) as usize;
364 let n = stream.read(&mut chunk[..cap]).await?;
365 if n == 0 {
366 return Err(std::io::Error::new(
367 std::io::ErrorKind::UnexpectedEof,
368 "eof inside chunk data",
369 ));
370 }
371 consumed += n as u64;
372 to_read -= n as u64;
373 }
374 }
375 }
376 } else if let Some(len) = content_length {
377 let mut to_read = len;
378 while to_read > 0 {
379 if !buf.is_empty() {
380 let take = to_read.min(buf.len() as u64) as usize;
381 buf.drain(..take);
382 to_read -= take as u64;
383 } else {
384 let mut chunk = [0u8; 16384];
385 let cap = to_read.min(chunk.len() as u64) as usize;
386 let n = stream.read(&mut chunk[..cap]).await?;
387 if n == 0 {
388 return Err(std::io::Error::new(
389 std::io::ErrorKind::UnexpectedEof,
390 "eof inside content-length body",
391 ));
392 }
393 consumed += n as u64;
394 to_read -= n as u64;
395 }
396 }
397 }
398 Ok((status, consumed))
402}
403
404fn find_head_end(buf: &[u8]) -> Option<usize> {
407 buf.windows(4).position(|w| w == b"\r\n\r\n")
408}
409
410impl StreamBody {
411 fn total_len(&self) -> u64 {
412 self.prefix.len() as u64 + self.fill_total + self.suffix.len() as u64
413 }
414}
415
416pub async fn run(cfg: PipelineBenchConfig) -> anyhow::Result<PipelineBenchResult> {
418 let (host, port, path) = parse_http_target(&cfg.target_url).map_err(anyhow::Error::msg)?;
419 if cfg.pipeline_depth == 0 {
420 anyhow::bail!("--pipeline-depth must be >= 1");
421 }
422 if cfg.connections == 0 {
423 anyhow::bail!("--connections must be >= 1");
424 }
425 StreamBody::new(cfg.body_kind, cfg.body_size).map_err(anyhow::Error::msg)?;
427
428 let addr = tokio::net::lookup_host((host.as_str(), port))
429 .await
430 .anyhow_err()?
431 .next()
432 .ok_or_else(|| anyhow::anyhow!("could not resolve {host}"))?;
433
434 let deadline = Instant::now() + cfg.duration;
435 let counters = Arc::new(Counters::default());
436
437 let mut handles = Vec::with_capacity(cfg.connections);
438 for _ in 0..cfg.connections {
439 let counters = counters.clone();
440 let host = host.clone();
441 let path = path.clone();
442 let method = cfg.method.clone();
443 handles.push(tokio::spawn(async move {
444 connection_loop(
445 addr,
446 &host,
447 port,
448 &path,
449 &method,
450 cfg.body_kind,
451 cfg.body_size,
452 cfg.pipeline_depth,
453 deadline,
454 &counters,
455 )
456 .await;
457 }));
458 }
459 for h in handles {
460 let _ = h.await;
461 }
462
463 let c = Arc::try_unwrap(counters).unwrap_or_else(|c| (*c).clone_snapshot());
464 Ok(PipelineBenchResult {
465 requests_sent: c.requests_sent.load(Ordering::Relaxed),
466 responses_received: c.responses_received.load(Ordering::Relaxed),
467 status_counts: c.status_snapshot(),
468 bytes_sent: c.bytes_sent.load(Ordering::Relaxed),
469 bytes_received: c.bytes_received.load(Ordering::Relaxed),
470 connection_closed_early: c.closed_early.load(Ordering::Relaxed),
471 connection_errors: c.connection_errors.load(Ordering::Relaxed),
472 reconnects: c.reconnects.load(Ordering::Relaxed),
473 })
474}
475
476trait AnyhowExt<T> {
477 fn anyhow_err(self) -> anyhow::Result<T>;
478}
479impl<T> AnyhowExt<T> for std::io::Result<T> {
480 fn anyhow_err(self) -> anyhow::Result<T> {
481 self.map_err(|e| anyhow::anyhow!(e))
482 }
483}
484
485#[derive(Default)]
486struct Counters {
487 requests_sent: AtomicU64,
488 responses_received: AtomicU64,
489 bytes_sent: AtomicU64,
490 bytes_received: AtomicU64,
491 closed_early: AtomicU64,
492 connection_errors: AtomicU64,
493 reconnects: AtomicU64,
494 statuses: std::sync::Mutex<std::collections::BTreeMap<String, u64>>,
495}
496
497impl Counters {
498 fn status_snapshot(&self) -> std::collections::BTreeMap<String, u64> {
499 self.statuses.lock().map(|m| m.clone()).unwrap_or_default()
500 }
501 fn clone_snapshot(&self) -> Self {
502 Self {
503 requests_sent: AtomicU64::new(self.requests_sent.load(Ordering::Relaxed)),
504 responses_received: AtomicU64::new(self.responses_received.load(Ordering::Relaxed)),
505 bytes_sent: AtomicU64::new(self.bytes_sent.load(Ordering::Relaxed)),
506 bytes_received: AtomicU64::new(self.bytes_received.load(Ordering::Relaxed)),
507 closed_early: AtomicU64::new(self.closed_early.load(Ordering::Relaxed)),
508 connection_errors: AtomicU64::new(self.connection_errors.load(Ordering::Relaxed)),
509 reconnects: AtomicU64::new(self.reconnects.load(Ordering::Relaxed)),
510 statuses: std::sync::Mutex::new(self.status_snapshot()),
511 }
512 }
513}
514
515#[allow(clippy::too_many_arguments)]
516async fn connection_loop(
517 addr: std::net::SocketAddr,
518 host: &str,
519 _port: u16,
520 path: &str,
521 method: &str,
522 kind: BodyKind,
523 body_size: u64,
524 depth: usize,
525 deadline: Instant,
526 counters: &Counters,
527) {
528 let mut stream: Option<TcpStream> = None;
529 while Instant::now() < deadline {
530 if stream.is_none() {
531 match TcpStream::connect(addr).await {
532 Ok(s) => stream = Some(s),
533 Err(_) => {
534 counters.connection_errors.fetch_add(1, Ordering::Relaxed);
535 tokio::time::sleep(Duration::from_millis(100)).await;
536 continue;
537 }
538 }
539 }
540 let s = stream.as_mut().expect("stream just set");
541
542 let mut batch_sent = 0u64;
544 let mut ok = true;
545 for _ in 0..depth {
546 let mut body = match StreamBody::new(kind, body_size) {
550 Ok(b) => b,
551 Err(_) => return, };
553 match write_batch_request(s, method, host, path, &mut body).await {
554 Ok(bytes) => {
555 batch_sent += bytes;
556 counters.requests_sent.fetch_add(1, Ordering::Relaxed);
557 }
558 Err(_) => {
559 counters.connection_errors.fetch_add(1, Ordering::Relaxed);
560 ok = false;
561 break;
562 }
563 }
564 }
565 if ok {
566 if let Err(e) = s.flush().await {
567 let _ = e;
568 counters.connection_errors.fetch_add(1, Ordering::Relaxed);
569 ok = false;
570 }
571 }
572 counters.bytes_sent.fetch_add(batch_sent, Ordering::Relaxed);
573 if !ok {
574 stream = None;
575 continue;
576 }
577
578 let mut recv_buf: Vec<u8> = Vec::with_capacity(16 * 1024);
580 for _ in 0..depth {
581 match read_response(s, &mut recv_buf).await {
582 Ok((status, consumed)) => {
583 counters.responses_received.fetch_add(1, Ordering::Relaxed);
584 counters.bytes_received.fetch_add(consumed, Ordering::Relaxed);
585 if let Ok(mut map) = counters.statuses.lock() {
586 *map.entry(status.to_string()).or_insert(0) += 1;
587 }
588 }
589 Err(e) => {
590 if e.kind() == std::io::ErrorKind::UnexpectedEof {
593 counters.closed_early.fetch_add(1, Ordering::Relaxed);
594 } else {
595 counters.connection_errors.fetch_add(1, Ordering::Relaxed);
596 }
597 break;
598 }
599 }
600 }
601
602 stream = None;
603 counters.reconnects.fetch_add(1, Ordering::Relaxed);
604 }
605}
606
607async fn write_batch_request(
609 stream: &mut TcpStream,
610 method: &str,
611 host: &str,
612 path: &str,
613 body: &mut StreamBody,
614) -> std::io::Result<u64> {
615 let ct = body.kind.content_type();
616 let head = format!(
617 "{method} {path} HTTP/1.1\r\nHost: {host}\r\nUser-Agent: mockforge-bench-pipeline\r\n\
618 Content-Type: {ct}\r\nContent-Length: {}\r\n\r\n",
619 body.total_len(),
620 );
621 let mut written = head.len() as u64;
622 stream.write_all(head.as_bytes()).await?;
623 while let Some(chunk) = body.next_chunk() {
624 stream.write_all(&chunk).await?;
625 written += chunk.len() as u64;
626 }
627 Ok(written)
628}
629
630pub fn render_report(res: &PipelineBenchResult) -> String {
632 let mut out = String::new();
633 out.push_str("\n=== bench-pipeline results ===\n");
634 out.push_str(&format!("requests sent : {}\n", res.requests_sent));
635 out.push_str(&format!("responses received : {}\n", res.responses_received));
636 if res.requests_sent != res.responses_received {
637 out.push_str(&format!(
638 " NOTE: {} request(s) never got a response — the target likely \
639 does not support pipelining (serialized or closed early).\n",
640 res.requests_sent - res.responses_received
641 ));
642 }
643 out.push_str(&format!("connection closed early: {}\n", res.connection_closed_early));
644 out.push_str(&format!("connection errors : {}\n", res.connection_errors));
645 out.push_str(&format!("reconnects : {}\n", res.reconnects));
646 out.push_str(&format!(
647 "bytes sent / received : {} / {}\n",
648 res.bytes_sent, res.bytes_received
649 ));
650 out.push_str("status histogram :\n");
651 for (code, n) in &res.status_counts {
652 out.push_str(&format!(" {code}: {n}\n"));
653 }
654 out
655}
656
657#[cfg(test)]
658mod tests {
659 use super::*;
660
661 #[test]
662 fn body_size_parses_suffixes() {
663 assert_eq!(parse_body_size("4096").unwrap(), 4096);
664 assert_eq!(parse_body_size("500KB").unwrap(), 500_000);
665 assert_eq!(parse_body_size("1MB").unwrap(), 1_000_000);
666 assert_eq!(parse_body_size("2GB").unwrap(), 2_000_000_000);
667 assert_eq!(parse_body_size("1KiB").unwrap(), 1024);
668 assert_eq!(parse_body_size("1.5KB").unwrap(), 1500);
669 assert!(parse_body_size("abc").is_err());
670 assert!(parse_body_size("0").is_err());
671 }
672
673 #[test]
674 fn body_kinds_parse() {
675 assert_eq!(parse_body_kind("json").unwrap(), BodyKind::Json);
676 assert_eq!(
677 parse_body_kind("application/x-www-form-urlencoded").unwrap(),
678 BodyKind::UrlEncoded
679 );
680 assert!(parse_body_kind("grpc").is_err());
681 }
682
683 #[test]
684 fn stream_body_is_exact_size_and_well_formed() {
685 for (kind, size) in [
686 (BodyKind::Json, 1000u64),
687 (BodyKind::Xml, 2048u64),
688 (BodyKind::UrlEncoded, 777u64),
689 (BodyKind::Multipart, 50_000u64),
690 ] {
691 let mut body = StreamBody::new(kind, size).unwrap();
692 let mut total = 0usize;
693 let mut assembled: Vec<u8> = Vec::new();
694 while let Some(chunk) = body.next_chunk() {
695 total += chunk.len();
696 if assembled.len() < 512 {
697 assembled.extend_from_slice(&chunk);
698 }
699 }
700 assert_eq!(total as u64, size, "{kind:?} must produce exactly {size}");
701 assert!(body.exhausted());
702 match kind {
703 BodyKind::Json => {
704 assert!(assembled.starts_with(b"{\""));
705 }
706 BodyKind::Xml => assert!(assembled.starts_with(b"<root>")),
707 BodyKind::UrlEncoded => {}
708 BodyKind::Multipart => assert!(assembled.starts_with(b"--mockforge-pipeline")),
709 }
710 }
711 }
712
713 #[test]
714 fn json_body_roundtrips_through_parser() {
715 let size = 4096u64;
716 let mut body = StreamBody::new(BodyKind::Json, size).unwrap();
717 let mut full: Vec<u8> = Vec::with_capacity(size as usize);
718 while let Some(chunk) = body.next_chunk() {
719 full.extend_from_slice(&chunk);
720 }
721 assert_eq!(full.len() as u64, size);
722 let v: serde_json::Value = serde_json::from_slice(&full).expect("well-formed JSON");
723 assert!(v.get("data").and_then(|d| d.as_str()).is_some());
724 }
725
726 #[test]
727 fn too_small_body_rejected() {
728 assert!(StreamBody::new(BodyKind::Json, 4).is_err());
729 }
730
731 #[test]
732 fn target_parses() {
733 let (h, p, path) = parse_http_target("http://localhost:3000/upload").unwrap();
734 assert_eq!((h.as_str(), p, path.as_str()), ("localhost", 3000, "/upload"));
735 let (h, p, path) = parse_http_target("http://example.com").unwrap();
736 assert_eq!((h.as_str(), p, path.as_str()), ("example.com", 80, "/"));
737 assert!(parse_http_target("https://x/").is_err());
738 }
739}