1 //! Benchmarks for write performance.
2 use criterion::{BatchSize, Criterion};
3 use std::{
4     hint,
5     io::{self, Read, Write},
6     time::{Duration, Instant},
7 };
8 use tungstenite::{Message, WebSocket};
9 
10 const MOCK_WRITE_LEN: usize = 8 * 1024 * 1024;
11 
12 /// `Write` impl that simulates slowish writes and slow flushes.
13 ///
14 /// Each `write` can buffer up to 8 MiB before flushing but takes an additional **~80ns**
15 /// to simulate stuff going on in the underlying stream.
16 /// Each `flush` takes **~8µs** to simulate flush io.
17 struct MockWrite(Vec<u8>);
18 
19 impl Read for MockWrite {
read(&mut self, _: &mut [u8]) -> io::Result<usize>20     fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
21         Err(io::Error::new(io::ErrorKind::WouldBlock, "reads not supported"))
22     }
23 }
24 impl Write for MockWrite {
write(&mut self, buf: &[u8]) -> io::Result<usize>25     fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
26         if self.0.len() + buf.len() > MOCK_WRITE_LEN {
27             self.flush()?;
28         }
29         // simulate io
30         spin(Duration::from_nanos(80));
31         self.0.extend(buf);
32         Ok(buf.len())
33     }
34 
flush(&mut self) -> io::Result<()>35     fn flush(&mut self) -> io::Result<()> {
36         if !self.0.is_empty() {
37             // simulate io
38             spin(Duration::from_micros(8));
39             self.0.clear();
40         }
41         Ok(())
42     }
43 }
44 
spin(duration: Duration)45 fn spin(duration: Duration) {
46     let a = Instant::now();
47     while a.elapsed() < duration {
48         hint::spin_loop();
49     }
50 }
51 
benchmark(c: &mut Criterion)52 fn benchmark(c: &mut Criterion) {
53     // Writes 100k small json text messages then flushes
54     c.bench_function("write 100k small texts then flush", |b| {
55         let mut ws = WebSocket::from_raw_socket(
56             MockWrite(Vec::with_capacity(MOCK_WRITE_LEN)),
57             tungstenite::protocol::Role::Server,
58             None,
59         );
60 
61         b.iter_batched(
62             || (0..100_000).map(|i| Message::Text(format!("{{\"id\":{i}}}"))),
63             |batch| {
64                 for msg in batch {
65                     ws.write(msg).unwrap();
66                 }
67                 ws.flush().unwrap();
68             },
69             BatchSize::SmallInput,
70         )
71     });
72 }
73 
74 criterion::criterion_group!(write_benches, benchmark);
75 criterion::criterion_main!(write_benches);
76