Add bench to masque proxy fragment handler

Measures performance of interleaved fragmented packets

5 files changed, 305 insertions, 3 deletions

diff --git a/Cargo.lock b/Cargo.lock
index c6b6a79e26..2285a260a2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -101,6 +101,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "anes"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
+
+[[package]]
 name = "anstream"
 version = "0.6.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -572,6 +578,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "cast"
+version = "0.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
+
+[[package]]
 name = "cbindgen"
 version = "0.28.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -662,6 +674,33 @@ dependencies = [
 ]
 
 [[package]]
+name = "ciborium"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
+dependencies = [
+ "ciborium-io",
+ "ciborium-ll",
+ "serde",
+]
+
+[[package]]
+name = "ciborium-io"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
+
+[[package]]
+name = "ciborium-ll"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
+dependencies = [
+ "ciborium-io",
+ "half",
+]
+
+[[package]]
 name = "cipher"
 version = "0.4.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -850,6 +889,39 @@ dependencies = [
 ]
 
 [[package]]
+name = "criterion"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e1c047a62b0cc3e145fa84415a3191f628e980b194c2755aa12300a4e6cbd928"
+dependencies = [
+ "anes",
+ "cast",
+ "ciborium",
+ "clap",
+ "criterion-plot",
+ "itertools 0.13.0",
+ "num-traits",
+ "oorandom",
+ "plotters",
+ "rayon",
+ "regex",
+ "serde",
+ "serde_json",
+ "tinytemplate",
+ "walkdir",
+]
+
+[[package]]
+name = "criterion-plot"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b1bcc0dc7dfae599d84ad0b1a55f80cde8af3725da8313b528da95ef783e338"
+dependencies = [
+ "cast",
+ "itertools 0.13.0",
+]
+
+[[package]]
 name = "crossbeam-channel"
 version = "0.5.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -859,12 +931,37 @@ dependencies = [
 ]
 
 [[package]]
+name = "crossbeam-deque"
+version = "0.8.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
 name = "crossbeam-utils"
 version = "0.8.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345"
 
 [[package]]
+name = "crunchy"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
+
+[[package]]
 name = "crypto-bigint"
 version = "0.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -1709,6 +1806,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "half"
+version = "2.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "459196ed295495a68f7d7fe1d84f6c4b7ff0e21fe3017b2f283c6fac3ad803c9"
+dependencies = [
+ "cfg-if",
+ "crunchy",
+]
+
+[[package]]
 name = "hashbrown"
 version = "0.12.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -2419,6 +2526,15 @@ dependencies = [
 
 [[package]]
 name = "itertools"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
+dependencies = [
+ "either",
+]
+
+[[package]]
+name = "itertools"
 version = "0.14.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
@@ -3141,6 +3257,7 @@ dependencies = [
  "anyhow",
  "bytes",
  "clap",
+ "criterion",
  "env_logger 0.11.7",
  "h3",
  "h3-datagram",
@@ -3152,6 +3269,7 @@ dependencies = [
  "rustls 0.23.18",
  "rustls-pemfile 2.1.3",
  "socket2 0.5.8",
+ "talpid-tunnel",
  "thiserror 2.0.9",
  "tokio",
  "typed-builder 0.21.0",
@@ -3785,6 +3903,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "945462a4b81e43c4e3ba96bd7b49d834c6f61198356aa858733bc4acf3cbe62e"
 
 [[package]]
+name = "oorandom"
+version = "11.1.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
+
+[[package]]
 name = "opaque-debug"
 version = "0.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -4097,6 +4221,34 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec"
 
 [[package]]
+name = "plotters"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747"
+dependencies = [
+ "num-traits",
+ "plotters-backend",
+ "plotters-svg",
+ "wasm-bindgen",
+ "web-sys",
+]
+
+[[package]]
+name = "plotters-backend"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a"
+
+[[package]]
+name = "plotters-svg"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670"
+dependencies = [
+ "plotters-backend",
+]
+
+[[package]]
 name = "pnet_base"
 version = "0.34.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -4601,6 +4753,26 @@ dependencies = [
 ]
 
 [[package]]
+name = "rayon"
+version = "1.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
+[[package]]
 name = "redox_syscall"
 version = "0.4.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -5813,6 +5985,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "tinytemplate"
+version = "1.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
+dependencies = [
+ "serde",
+ "serde_json",
+]
+
+[[package]]
 name = "tinyvec"
 version = "1.6.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/mullvad-masque-proxy/Cargo.toml b/mullvad-masque-proxy/Cargo.toml
index 6f54734d1d..98fc75f0e4 100644
--- a/mullvad-masque-proxy/Cargo.toml
+++ b/mullvad-masque-proxy/Cargo.toml
@@ -29,6 +29,14 @@ env_logger = { workspace = true }
 tokio = { workspace = true, features = ["fs", "macros", "io-util", "rt-multi-thread"] }
 clap = { workspace = true }
 rand = "0.8.5"
+criterion = { version = "0.7.0", features = ["html_reports"] }
+talpid-tunnel = { path = "../talpid-tunnel" }
+
 
 [lints]
 workspace = true
+
+[[bench]]
+name = "fragmentation"
+harness = false
+
diff --git a/mullvad-masque-proxy/benches/fragmentation.rs b/mullvad-masque-proxy/benches/fragmentation.rs
new file mode 100644
index 0000000000..6b4793dd5a
--- /dev/null
+++ b/mullvad-masque-proxy/benches/fragmentation.rs
@@ -0,0 +1,112 @@
+use bytes::Bytes;
+use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main};
+use mullvad_masque_proxy::{
+    FRAGMENT_HEADER_SIZE_FRAGMENTED,
+    fragment::{FRAGMENT_BUFFER_CAP, Fragments, fragment_packet},
+};
+use rand::{SeedableRng, rngs::StdRng, seq::SliceRandom};
+use talpid_tunnel::IPV4_HEADER_SIZE;
+
+const MAX_PAYLOAD_SIZE: u16 = 1280 - FRAGMENT_HEADER_SIZE_FRAGMENTED - IPV4_HEADER_SIZE;
+
+fn assemble_fragment_ordered(c: &mut Criterion) {
+    let mut group = c.benchmark_group("fragmentation_reconstruction");
+
+    for (n_packets, payload_len) in [(10, 30000u16), (100, 1500)] {
+        let mut fragment_buf = Vec::with_capacity(FRAGMENT_BUFFER_CAP);
+        for i in 0..n_packets {
+            let packet_id = i;
+            let mut payload = Bytes::from(vec![i as u8; payload_len as usize]);
+
+            fragment_buf.extend(
+                &mut fragment_packet(
+                    MAX_PAYLOAD_SIZE + FRAGMENT_HEADER_SIZE_FRAGMENTED,
+                    &mut payload,
+                    packet_id,
+                )
+                .unwrap(),
+            );
+        }
+        let n_fragments = fragment_buf.len();
+        assert!(
+            n_fragments <= FRAGMENT_BUFFER_CAP,
+            "Too many fragments generated"
+        );
+        group.throughput(criterion::Throughput::Bytes(
+            (n_packets * payload_len) as u64,
+        ));
+
+        group.bench_with_input(
+            BenchmarkId::new(
+                "assemble_fragment_ordered",
+                format!("{n_packets}pkts_{payload_len}B_{n_fragments}frags"),
+            ),
+            &fragment_buf,
+            |b, fragment_buf| {
+                b.iter_batched(
+                    || (fragment_buf.clone(), Fragments::default()),
+                    |(f, mut fragments)| {
+                        for frag in f {
+                            fragments.handle_incoming_packet(frag).unwrap();
+                        }
+                    },
+                    BatchSize::SmallInput,
+                )
+            },
+        );
+    }
+    group.finish();
+}
+
+fn assemble_fragment_random(c: &mut Criterion) {
+    let mut group = c.benchmark_group("fragmentation_reconstruction");
+
+    for (n_packets, payload_len) in [(10, 30000u16), (100, 1500)] {
+        let mut fragment_buf = Vec::with_capacity(FRAGMENT_BUFFER_CAP);
+        for i in 0..n_packets {
+            let packet_id = i;
+            let mut payload = Bytes::from(vec![i as u8; payload_len as usize]);
+
+            fragment_buf.extend(
+                &mut fragment_packet(
+                    MAX_PAYLOAD_SIZE + FRAGMENT_HEADER_SIZE_FRAGMENTED,
+                    &mut payload,
+                    packet_id,
+                )
+                .unwrap(),
+            );
+        }
+        let n_fragments = fragment_buf.len();
+        assert!(
+            n_fragments <= FRAGMENT_BUFFER_CAP,
+            "Too many fragments generated"
+        );
+        group.throughput(criterion::Throughput::Bytes(
+            (n_packets * payload_len) as u64,
+        ));
+        let mut rng = StdRng::seed_from_u64(42);
+        fragment_buf.shuffle(&mut rng);
+
+        group.bench_with_input(
+            BenchmarkId::new(
+                "assemble_fragment_random",
+                format!("{n_packets}pkts_{payload_len}B_{n_fragments}frags"),
+            ),
+            &fragment_buf,
+            |b, fragment_buf| {
+                b.iter_batched(
+                    || (fragment_buf.clone(), Fragments::default()),
+                    |(f, mut fragments)| {
+                        for frag in f {
+                            fragments.handle_incoming_packet(frag).unwrap();
+                        }
+                    },
+                    BatchSize::SmallInput,
+                )
+            },
+        );
+    }
+    group.finish();
+}
+criterion_group!(benches, assemble_fragment_ordered, assemble_fragment_random);
+criterion_main!(benches);
diff --git a/mullvad-masque-proxy/src/fragment.rs b/mullvad-masque-proxy/src/fragment.rs
index 5cb05d132c..9a47295fb3 100644
--- a/mullvad-masque-proxy/src/fragment.rs
+++ b/mullvad-masque-proxy/src/fragment.rs
@@ -10,7 +10,7 @@ const FRAGMENT_INDEX_START: u8 = 1;
 
 /// The maximum number of unassembled fragments that we buffer.
 // 255 is the theoretical maximum number of fragments for a single packet.
-const FRAGMENT_BUFFER_CAP: usize = 255;
+pub const FRAGMENT_BUFFER_CAP: usize = 255;
 
 pub struct Fragments {
     /// FIFO queue of fragment indices. Used to mitigate floods of unordered packet fragments.
diff --git a/mullvad-masque-proxy/src/lib.rs b/mullvad-masque-proxy/src/lib.rs
index 7d22dc11d7..37f8353acd 100644
--- a/mullvad-masque-proxy/src/lib.rs
+++ b/mullvad-masque-proxy/src/lib.rs
@@ -2,7 +2,7 @@ use h3::proto::varint::VarInt;
 use std::net::SocketAddr;
 
 pub mod client;
-mod fragment;
+pub mod fragment;
 pub mod server;
 mod stats;
 
@@ -19,7 +19,7 @@ const PACKET_BUFFER_SIZE: usize = (u16::MAX - UDP_HEADER_SIZE + 1) as usize;
 const MAX_INFLIGHT_PACKETS: usize = 100;
 
 /// Fragment headers size for fragmented packets
-const FRAGMENT_HEADER_SIZE_FRAGMENTED: u16 = 5;
+pub const FRAGMENT_HEADER_SIZE_FRAGMENTED: u16 = 5;
 
 /// UDP header overhead
 const UDP_HEADER_SIZE: u16 = 8;