|
| 1 | +//! Benchmarks for end to end performance including real `Read` & `Write` impls. |
| 2 | +use bytes::Bytes; |
| 3 | +use criterion::{BatchSize, Criterion, Throughput}; |
| 4 | +use rand::{ |
| 5 | + distr::{Alphanumeric, SampleString}, |
| 6 | + rngs::SmallRng, |
| 7 | + SeedableRng, |
| 8 | +}; |
| 9 | +use std::net::TcpListener; |
| 10 | +use tungstenite::{accept_hdr_with_config, protocol::WebSocketConfig, Message}; |
| 11 | + |
| 12 | +/// Binary message meaning "stop". |
| 13 | +const B_STOP: Bytes = Bytes::from_static(b"stop"); |
| 14 | + |
| 15 | +fn benchmark(c: &mut Criterion) { |
| 16 | + /// Benchmark that starts a simple server and client then sends (writes+flush) a |
| 17 | + /// single text message client->server and reads a single response text message |
| 18 | + /// server->client. Both message will be of the given `msg_len` size. |
| 19 | + fn send_and_recv(msg_len: usize, b: &mut criterion::Bencher<'_>) { |
| 20 | + let socket = TcpListener::bind("127.0.0.1:0").unwrap(); |
| 21 | + let port = socket.local_addr().unwrap().port(); |
| 22 | + let conf = WebSocketConfig::default().max_message_size(None).max_frame_size(None); |
| 23 | + |
| 24 | + let server_thread = std::thread::spawn(move || { |
| 25 | + // single thread / single client server |
| 26 | + let (stream, _) = socket.accept().unwrap(); |
| 27 | + let mut websocket = |
| 28 | + accept_hdr_with_config(stream, |_: &_, res| Ok(res), Some(conf)).unwrap(); |
| 29 | + loop { |
| 30 | + let uppercase_txt = match websocket.read().unwrap() { |
| 31 | + Message::Text(msg) => msg.to_ascii_uppercase(), |
| 32 | + Message::Binary(msg) if msg == B_STOP => return, |
| 33 | + msg => panic!("Unexpected msg: {msg:?}"), |
| 34 | + }; |
| 35 | + websocket.send(Message::text(uppercase_txt)).unwrap(); |
| 36 | + } |
| 37 | + }); |
| 38 | + |
| 39 | + let (mut client, _) = tungstenite::client::connect_with_config( |
| 40 | + format!("ws://127.0.0.1:{port}"), |
| 41 | + Some(conf), |
| 42 | + 3, |
| 43 | + ) |
| 44 | + .unwrap(); |
| 45 | + let mut rng = SmallRng::seed_from_u64(123); |
| 46 | + |
| 47 | + b.iter_batched( |
| 48 | + || { |
| 49 | + let msg = Alphanumeric.sample_string(&mut rng, msg_len); |
| 50 | + let expected_response = msg.to_ascii_uppercase(); |
| 51 | + (msg, expected_response) |
| 52 | + }, |
| 53 | + |(txt, expected_response)| { |
| 54 | + client.send(Message::text(txt)).unwrap(); |
| 55 | + let response = client.read().unwrap(); |
| 56 | + match response { |
| 57 | + Message::Text(v) => assert_eq!(v, expected_response), |
| 58 | + msg => panic!("Unexpected response msg: {msg:?}"), |
| 59 | + }; |
| 60 | + }, |
| 61 | + BatchSize::PerIteration, |
| 62 | + ); |
| 63 | + |
| 64 | + // cleanup |
| 65 | + client.send(Message::binary(B_STOP)).unwrap(); |
| 66 | + server_thread.join().unwrap(); |
| 67 | + } |
| 68 | + |
| 69 | + // bench sending & receiving various sizes 512B to 1GiB. |
| 70 | + for len in (0..8).map(|n| 512 * 8_usize.pow(n)) { |
| 71 | + let mut group = c.benchmark_group("send+recv"); |
| 72 | + group |
| 73 | + .throughput(Throughput::Bytes(len as u64 * 2)) // *2 as we send and then recv it |
| 74 | + .bench_function(HumanLen(len).to_string(), |b| send_and_recv(len, b)); |
| 75 | + } |
| 76 | +} |
| 77 | + |
| 78 | +struct HumanLen(usize); |
| 79 | + |
| 80 | +impl std::fmt::Display for HumanLen { |
| 81 | + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
| 82 | + match self.0 { |
| 83 | + n if n < 1024 => write!(f, "{n} B"), |
| 84 | + n if n < 1024 * 1024 => write!(f, "{} KiB", n / 1024), |
| 85 | + n if n < 1024 * 1024 * 1024 => write!(f, "{} MiB", n / (1024 * 1024)), |
| 86 | + n => write!(f, "{} GiB", n / (1024 * 1024 * 1024)), |
| 87 | + } |
| 88 | + } |
| 89 | +} |
| 90 | + |
| 91 | +criterion::criterion_group!(read_benches, benchmark); |
| 92 | +criterion::criterion_main!(read_benches); |
0 commit comments