Optimize frame encoding, mux session

README.md

@@ -32,3 +32,11 @@ See [examples][examples-url].
 ### WebAssembly
 
 The `websock-wasm-demo` crate contains a small browser demo that connects to an echo server. 
+
+## Benchmarking
+
+Criterion benchmarks are available for `websock-mux-proto`.
+
+```bash
+cargo bench -p websock-mux-proto
+```

devcert/generate.ps1

@@ -0,0 +1,76 @@
+Set-StrictMode -Version Latest
+$ErrorActionPreference = "Stop"
+
+$scriptDir = Split-Path -Parent $MyInvocation.MyCommand.Path
+Set-Location $scriptDir
+
+$certName = "localhost"
+$days = 10
+$conf = "openssl.conf"
+
+Write-Host "==> Generating self-signed certificate ($certName)"
+Write-Host "    validity: $days days"
+Write-Host "    config:   $conf"
+
+$openssl = if ($env:OPENSSL_BIN) { $env:OPENSSL_BIN } else { "openssl" }
+if (-not (Get-Command $openssl -ErrorAction SilentlyContinue)) {
+    throw "openssl command not found. Set OPENSSL_BIN to openssl.exe path or add openssl to PATH."
+}
+
+# Generate ECDSA P-256 private key
+& $openssl ecparam `
+    -genkey `
+    -name prime256v1 `
+    -out "$certName.key"
+
+if ($LASTEXITCODE -ne 0) {
+    throw "Failed to generate private key."
+}
+
+# Generate self-signed certificate
+& $openssl req `
+    -x509 `
+    -sha256 `
+    -nodes `
+    -days "$days" `
+    -key "$certName.key" `
+    -out "$certName.crt" `
+    -config "$conf" `
+    -extensions v3_req
+
+if ($LASTEXITCODE -ne 0) {
+    throw "Failed to generate certificate."
+}
+
+# Generate raw SHA-256 hash (DER -> hex, no colons)
+$derPath = "$certName.der"
+& $openssl x509 `
+    -in "$certName.crt" `
+    -outform der `
+    -out $derPath
+
+if ($LASTEXITCODE -ne 0) {
+    throw "Failed to export certificate in DER format."
+}
+
+$hex = (Get-FileHash -Path $derPath -Algorithm SHA256).Hash.ToLowerInvariant()
+Set-Content -Path "$certName.hex" -Value $hex -NoNewline
+Remove-Item $derPath -Force
+
+# Also print human-readable fingerprint
+& $openssl x509 `
+    -in "$certName.crt" `
+    -noout `
+    -fingerprint `
+    -sha256 `
+    > "$certName.fingerprint"
+
+if ($LASTEXITCODE -ne 0) {
+    throw "Failed to generate fingerprint."
+}
+
+Write-Host "==> Done"
+Write-Host "  - $certName.crt"
+Write-Host "  - $certName.key"
+Write-Host "  - $certName.hex      (for serverCertificateHashes)"
+Write-Host "  - $certName.fingerprint"

websock-mux-proto/Cargo.toml

@@ -17,3 +17,10 @@ websock-proto = { workspace = true }
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 tokio = { version = "1", features = ["io-util"] }
+
+[dev-dependencies]
+criterion = "0.8"
+
+[[bench]]
+name = "proto_bench"
+harness = false

websock-mux-proto/benches/proto_bench.rs

@@ -0,0 +1,122 @@
+use bytes::{Bytes, BytesMut};
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use std::io::Cursor;
+use std::hint::black_box;
+use websock_mux_proto::{Frame, StreamDir, StreamId, VarInt};
+
+fn bench_varint_encode(c: &mut Criterion) {
+    let values: [u64; 6] = [0, 63, 64, 16_383, 16_384, (1u64 << 62) - 1];
+    let mut group = c.benchmark_group("varint_encode");
+
+    for &v in &values {
+        group.bench_with_input(BenchmarkId::from_parameter(v), &v, |b, &v| {
+            b.iter(|| {
+                let mut buf = BytesMut::with_capacity(8);
+                VarInt::from_u64(v)
+                    .expect("valid varint value")
+                    .encode(&mut buf);
+                black_box(buf);
+            });
+        });
+    }
+
+    group.finish();
+}
+
+fn bench_varint_decode(c: &mut Criterion) {
+    let values: [u64; 6] = [0, 63, 64, 16_383, 16_384, (1u64 << 62) - 1];
+    let mut group = c.benchmark_group("varint_decode");
+
+    for &v in &values {
+        let mut encoded = BytesMut::with_capacity(8);
+        VarInt::from_u64(v)
+            .expect("valid varint value")
+            .encode(&mut encoded);
+        let encoded = encoded.freeze();
+
+        group.bench_with_input(BenchmarkId::from_parameter(v), &encoded, |b, encoded| {
+            b.iter(|| {
+                let mut cur = Cursor::new(encoded.clone());
+                black_box(VarInt::decode(&mut cur).expect("decode succeeds"));
+            });
+        });
+    }
+
+    group.finish();
+}
+
+fn bench_frame_encode(c: &mut Criterion) {
+    let id = StreamId::new(7, false, StreamDir::Bi).expect("stream id");
+    let small = Frame::Stream {
+        id,
+        data: Bytes::from_static(b"hello"),
+        fin: false,
+    };
+    let large = Frame::Stream {
+        id,
+        data: Bytes::from(vec![7u8; 64 * 1024]),
+        fin: true,
+    };
+
+    let mut group = c.benchmark_group("frame_encode");
+    group.throughput(Throughput::Bytes(small.encoded_len() as u64));
+    group.bench_function("stream_small", |b| {
+        b.iter(|| {
+            black_box(small.encode());
+        });
+    });
+
+    group.throughput(Throughput::Bytes(large.encoded_len() as u64));
+    group.bench_function("stream_large_64k", |b| {
+        b.iter(|| {
+            black_box(large.encode());
+        });
+    });
+    group.finish();
+}
+
+fn bench_frame_decode(c: &mut Criterion) {
+    let id = StreamId::new(7, false, StreamDir::Bi).expect("stream id");
+    let small = Frame::Stream {
+        id,
+        data: Bytes::from_static(b"hello"),
+        fin: false,
+    }
+    .encode()
+    .freeze();
+
+    let large = Frame::Stream {
+        id,
+        data: Bytes::from(vec![7u8; 64 * 1024]),
+        fin: true,
+    }
+    .encode()
+    .freeze();
+
+    let mut group = c.benchmark_group("frame_decode");
+    group.throughput(Throughput::Bytes(small.len() as u64));
+    group.bench_function("stream_small", |b| {
+        b.iter(|| {
+            let mut cur = Cursor::new(small.clone());
+            black_box(Frame::decode(&mut cur).expect("decode succeeds"));
+        });
+    });
+
+    group.throughput(Throughput::Bytes(large.len() as u64));
+    group.bench_function("stream_large_64k", |b| {
+        b.iter(|| {
+            let mut cur = Cursor::new(large.clone());
+            black_box(Frame::decode(&mut cur).expect("decode succeeds"));
+        });
+    });
+    group.finish();
+}
+
+criterion_group!(
+    benches,
+    bench_varint_encode,
+    bench_varint_decode,
+    bench_frame_encode,
+    bench_frame_decode,
+);
+criterion_main!(benches);

websock-mux-proto/src/stream.rs

@@ -70,8 +70,29 @@ pub enum Frame {
 }
 
 impl Frame {
+    /// Return the exact encoded frame length in bytes.
+    ///
+    /// This is primarily useful for pre-allocating encode buffers.
+    pub fn encoded_len(&self) -> usize {
+        match self {
+            Frame::OpenUni { id } | Frame::OpenBi { id } => 1 + VarInt(id.0).size(),
+            Frame::Stream { id, data, fin } => {
+                1 + VarInt(id.0).size()
+                    + VarInt(u64::from(*fin)).size()
+                    + VarInt(data.len() as u64).size()
+                    + data.len()
+            }
+            Frame::ResetStream { id, code } | Frame::StopSending { id, code } => {
+                1 + VarInt(id.0).size() + VarInt(*code).size()
+            }
+            Frame::ConnectionClose { code, reason } => {
+                1 + VarInt(*code).size() + VarInt(reason.len() as u64).size() + reason.len()
+            }
+        }
+    }
+
     pub fn encode(&self) -> BytesMut {
-        let mut buf = BytesMut::new();
+        let mut buf = BytesMut::with_capacity(self.encoded_len());
         match self {
             Frame::OpenUni { id } => {
                 VarInt(0).encode(&mut buf);
@@ -124,13 +145,8 @@ impl Frame {
                 if buf.remaining() < len {
                     return Err(FrameDecodeError::UnexpectedEnd);
                 }
-                let mut data = vec![0u8; len];
-                buf.copy_to_slice(&mut data);
-                Ok(Frame::Stream {
-                    id,
-                    data: Bytes::from(data),
-                    fin,
-                })
+                let data = buf.copy_to_bytes(len);
+                Ok(Frame::Stream { id, data, fin })
             }
             3 => Ok(Frame::ResetStream {
                 id: StreamId(VarInt::decode(buf)?.into_inner()),
@@ -146,9 +162,10 @@ impl Frame {
                 if buf.remaining() < len {
                     return Err(FrameDecodeError::UnexpectedEnd);
                 }
-                let mut data = vec![0u8; len];
-                buf.copy_to_slice(&mut data);
-                let reason = String::from_utf8(data).map_err(|_| FrameDecodeError::InvalidUtf8)?;
+                let data = buf.copy_to_bytes(len);
+                let reason = std::str::from_utf8(data.as_ref())
+                    .map_err(|_| FrameDecodeError::InvalidUtf8)?
+                    .to_owned();
                 Ok(Frame::ConnectionClose { code, reason })
             }
             _ => Err(FrameDecodeError::UnknownTag(tag)),

websock-mux-proto/src/varint.rs

@@ -22,6 +22,10 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
 pub struct VarInt(pub(crate) u64);
 
 impl VarInt {
+    const MAX_1BYTE: u64 = (1 << 6) - 1;
+    const MAX_2BYTE: u64 = (1 << 14) - 1;
+    const MAX_4BYTE: u64 = (1 << 30) - 1;
+
     /// The largest representable value.
     pub const MAX: Self = Self((1 << 62) - 1);
     /// The largest encoded value length.
@@ -34,7 +38,7 @@ impl VarInt {
 
     /// Succeeds if `x` < 2^62.
     pub fn from_u64(x: u64) -> Result<Self, VarIntBoundsExceeded> {
-        if x < 2u64.pow(62) {
+        if x <= Self::MAX.0 {
             Ok(Self(x))
         } else {
             Err(VarIntBoundsExceeded)
@@ -58,13 +62,13 @@ impl VarInt {
     /// Compute the number of bytes needed to encode this value.
     pub fn size(self) -> usize {
         let x = self.0;
-        if x < 2u64.pow(6) {
+        if x <= Self::MAX_1BYTE {
             1
-        } else if x < 2u64.pow(14) {
+        } else if x <= Self::MAX_2BYTE {
             2
-        } else if x < 2u64.pow(30) {
+        } else if x <= Self::MAX_4BYTE {
             4
-        } else if x < 2u64.pow(62) {
+        } else if x <= Self::MAX.0 {
             8
         } else {
             unreachable!("malformed VarInt");
@@ -137,32 +141,38 @@ impl VarInt {
         if !r.has_remaining() {
             return Err(VarIntUnexpectedEnd);
         }
-        let mut buf = [0; 8];
-        buf[0] = r.get_u8();
-        let tag = buf[0] >> 6;
-        buf[0] &= 0b0011_1111;
+        let first = r.get_u8();
+        let tag = first >> 6;
+        let body = (first & 0b0011_1111) as u64;
         let x = match tag {
-            0b00 => u64::from(buf[0]),
+            0b00 => body,
             0b01 => {
                 if r.remaining() < 1 {
                     return Err(VarIntUnexpectedEnd);
                 }
-                r.copy_to_slice(&mut buf[1..2]);
-                u64::from(u16::from_be_bytes(buf[..2].try_into().unwrap()))
+                (body << 8) | u64::from(r.get_u8())
             }
             0b10 => {
                 if r.remaining() < 3 {
                     return Err(VarIntUnexpectedEnd);
                 }
-                r.copy_to_slice(&mut buf[1..4]);
-                u64::from(u32::from_be_bytes(buf[..4].try_into().unwrap()))
+                (body << 24)
+                    | (u64::from(r.get_u8()) << 16)
+                    | (u64::from(r.get_u8()) << 8)
+                    | u64::from(r.get_u8())
             }
             0b11 => {
                 if r.remaining() < 7 {
                     return Err(VarIntUnexpectedEnd);
                 }
-                r.copy_to_slice(&mut buf[1..8]);
-                u64::from_be_bytes(buf)
+                (body << 56)
+                    | (u64::from(r.get_u8()) << 48)
+                    | (u64::from(r.get_u8()) << 40)
+                    | (u64::from(r.get_u8()) << 32)
+                    | (u64::from(r.get_u8()) << 24)
+                    | (u64::from(r.get_u8()) << 16)
+                    | (u64::from(r.get_u8()) << 8)
+                    | u64::from(r.get_u8())
             }
             _ => unreachable!(),
         };
@@ -197,13 +207,13 @@ impl VarInt {
 
     pub fn encode<B: BufMut>(&self, w: &mut B) {
         let x = self.0;
-        if x < 2u64.pow(6) {
+        if x <= Self::MAX_1BYTE {
             w.put_u8(x as u8);
-        } else if x < 2u64.pow(14) {
+        } else if x <= Self::MAX_2BYTE {
             w.put_u16((0b01 << 14) | x as u16);
-        } else if x < 2u64.pow(30) {
+        } else if x <= Self::MAX_4BYTE {
             w.put_u32((0b10 << 30) | x as u32);
-        } else if x < 2u64.pow(62) {
+        } else if x <= Self::MAX.0 {
             w.put_u64((0b11 << 62) | x);
         } else {
             unreachable!("malformed VarInt")