Add write 100k micro-bench

Cargo.toml

@@ -67,6 +67,10 @@ rand = "0.8.4"
 name = "buffer"
 harness = false
 
+[[bench]]
+name = "write"
+harness = false
+
 [[example]]
 name = "client"
 required-features = ["handshake"]

benches/write.rs

@@ -0,0 +1,67 @@
+//! Benchmarks for write performance.
+use criterion::{BatchSize, Criterion};
+use std::{
+    hint,
+    io::{self, Read, Write},
+    time::{Duration, Instant},
+};
+use tungstenite::{Message, WebSocket};
+
+const MOCK_WRITE_LEN: usize = 8 * 1024 * 1024;
+
+/// `Write` impl that simulates fast writes and slow flushes.
+///
+/// Buffers up to 8 MiB fast on `write`. Each `flush` takes ~100ns.
+struct MockSlowFlushWrite(Vec<u8>);
+
+impl Read for MockSlowFlushWrite {
+    fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
+        Err(io::Error::new(io::ErrorKind::WouldBlock, "reads not supported"))
+    }
+}
+impl Write for MockSlowFlushWrite {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        if self.0.len() + buf.len() > MOCK_WRITE_LEN {
+            self.flush()?;
+        }
+        self.0.extend(buf);
+        Ok(buf.len())
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        if !self.0.is_empty() {
+            // simulate 100ns io
+            let a = Instant::now();
+            while a.elapsed() < Duration::from_nanos(100) {
+                hint::spin_loop();
+            }
+            self.0.clear();
+        }
+        Ok(())
+    }
+}
+
+fn benchmark(c: &mut Criterion) {
+    // Writes 100k small json text messages then calls `write_pending`
+    c.bench_function("write 100k small texts then flush", |b| {
+        let mut ws = WebSocket::from_raw_socket(
+            MockSlowFlushWrite(Vec::with_capacity(MOCK_WRITE_LEN)),
+            tungstenite::protocol::Role::Server,
+            None,
+        );
+
+        b.iter_batched(
+            || (0..100_000).map(|i| Message::Text(format!("{{\"id\":{i}}}"))),
+            |batch| {
+                for msg in batch {
+                    ws.write_message(msg).unwrap();
+                }
+                ws.write_pending().unwrap();
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+criterion::criterion_group!(write_benches, benchmark);
+criterion::criterion_main!(write_benches);