1#![forbid(unsafe_code)]
17
18use simoom_rt::{ConnHandle, Connection, Ctx, Message, Runtime, Socket};
19use std::cell::Cell;
20use std::net::SocketAddr;
21use std::sync::mpsc;
22use std::time::{Duration, Instant};
23
24const REQUEST: &[u8] = b"GET /plaintext HTTP/1.1\r\nHost: localhost\r\nUser-Agent: simoom-loadgen\r\nAccept: */*\r\n\r\n";
25
26thread_local! {
27 static DONE: Cell<u64> = const { Cell::new(0) };
28 static ERRORS: Cell<u64> = const { Cell::new(0) };
29}
30
31fn response_len(data: &[u8]) -> Option<usize> {
33 let head = data.windows(4).position(|w| w == b"\r\n\r\n")? + 4;
34 let mut clen = 0usize;
35 for line in data[..head].split(|&b| b == b'\n') {
36 if line.len() > 15 && line[..15].eq_ignore_ascii_case(b"content-length:") {
37 let v = std::str::from_utf8(&line[15..]).ok()?.trim();
38 clen = v.parse().ok()?;
39 }
40 }
41 if data.len() >= head + clen {
42 Some(head + clen)
43 } else {
44 None
45 }
46}
47
48struct Client {
49 pipeline: usize,
50}
51
52impl Connection for Client {
53 fn on_message(&mut self, cx: &mut Ctx<'_>, _msg: Message) {
54 for _ in 0..self.pipeline {
55 cx.send_static(REQUEST);
56 }
57 }
58
59 fn on_data(&mut self, cx: &mut Ctx<'_>, data: &[u8]) -> usize {
60 let mut pos = 0;
61 let mut n = 0u64;
62 while let Some(len) = response_len(&data[pos..]) {
63 pos += len;
64 n += 1;
65 cx.send_static(REQUEST);
66 }
67 DONE.with(|c| c.set(c.get() + n));
68 pos
69 }
70
71 fn on_close(&mut self, _reason: simoom_rt::CloseReason) {
72 ERRORS.with(|c| c.set(c.get() + 1));
73 }
74}
75
76fn main() {
77 let mut addr: SocketAddr = "127.0.0.1:8080".parse().unwrap();
78 let (mut conns, mut threads, mut secs, mut pipeline) = (64usize, 4usize, 10u64, 1usize);
79 let mut cpus: Option<Vec<usize>> = None;
80 let v: Vec<String> = std::env::args().skip(1).collect();
81 let mut i = 0;
82 while i < v.len() {
83 let val = v.get(i + 1).cloned().unwrap_or_default();
84 match v[i].as_str() {
85 "--addr" => addr = val.parse().expect("--addr"),
86 "--conns" => conns = val.parse().expect("--conns"),
87 "--threads" => threads = val.parse().expect("--threads"),
88 "--secs" => secs = val.parse().expect("--secs"),
89 "--pipeline" => pipeline = val.parse().expect("--pipeline"),
90 "--cpus" => cpus = Some(val.split(',').map(|s| s.trim().parse().expect("--cpus")).collect()),
91 other => panic!("unknown argument {other}"),
92 }
93 i += 2;
94 }
95 let rt = Runtime::builder()
96 .configure(|c| {
97 c.threads = threads;
98 c.thread_name = "loadgen".into();
99 if let Some(cpus) = cpus {
100 c.pin_threads = true;
101 c.cpus = Some(cpus);
102 }
103 })
104 .start()
105 .expect("start");
106 let (htx, hrx) = mpsc::channel::<ConnHandle>();
107 for k in 0..conns {
108 let sock = Socket::tcp_connect(addr).expect("connect");
109 let htx = htx.clone();
110 rt.loops()[k % threads].run(move |lcx| {
111 let h = lcx.adopt(sock, Box::new(Client { pipeline })).expect("adopt");
112 htx.send(h).unwrap();
113 });
114 }
115 drop(htx);
116 let handles: Vec<ConnHandle> = hrx.iter().take(conns).collect();
117 for h in &handles {
118 h.send(Box::new(()));
119 }
120 let count = |rt: &simoom_rt::RuntimeHandle| -> (u64, u64) {
121 let (tx, rx) = mpsc::channel();
122 for l in rt.loops() {
123 let tx = tx.clone();
124 l.run(move |_| tx.send((DONE.with(|c| c.get()), ERRORS.with(|c| c.get()))).unwrap());
125 }
126 drop(tx);
127 rx.iter().fold((0, 0), |a, b| (a.0 + b.0, a.1 + b.1))
128 };
129 std::thread::sleep(Duration::from_secs(1));
130 let (r0, _) = count(&rt);
131 let s0 = rt.stats().total().driver;
132 let t0 = Instant::now();
133 std::thread::sleep(Duration::from_secs(secs));
134 let (r1, errors) = count(&rt);
135 let el = t0.elapsed().as_secs_f64();
136 let s1 = rt.stats().total().driver;
137 let rps = (r1 - r0) as f64 / el;
138 let lat_us = (conns * pipeline) as f64 / rps * 1e6;
139 println!(
140 "{conns} conns x {pipeline} pipelined, {threads} threads: {rps:.0} req/s, mean latency {lat_us:.1} us (Little), closed connections {errors}, client syscalls/req {:.2}",
141 (s1.syscalls() - s0.syscalls()) as f64 / (r1 - r0).max(1) as f64
142 );
143 rt.stop(Duration::ZERO);
144 rt.join();
145}