feat: Use a BinaryHeap for the RxStreamOrderer

neqo-transport/Cargo.toml

@@ -76,3 +76,8 @@ required-features = ["bench"]
 name = "min_bandwidth"
 harness = false
 required-features = ["bench"]
+
+[[bench]]
+name = "rx_stream_orderer_comparison"
+harness = false
+required-features = ["bench"]

neqo-transport/benches/rx_stream_orderer.rs

@@ -12,10 +12,10 @@
 use std::hint::black_box;
 
 use criterion::{criterion_group, criterion_main, Criterion};
-use neqo_transport::recv_stream::RxStreamOrderer;
+use neqo_transport::rx_stream_orderer_heap::RxStreamOrdererView;
 
 fn rx_stream_orderer() {
-    let mut rx = RxStreamOrderer::new();
+    let mut rx = RxStreamOrdererView::new();
     let data: &[u8] = &[0; 1337];
 
     for i in 0..100_000 {

neqo-transport/benches/rx_stream_orderer_comparison.rs

@@ -0,0 +1,245 @@
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use std::hint::black_box;
+
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use neqo_transport::{
+    recv_stream::RxStreamOrdererBTreeMap as BTreeMapOrderer,
+    rx_stream_orderer_heap::RxStreamOrdererView as HeapOrderer,
+};
+
+const FRAME_SIZE: usize = 1337;
+
+fn benchmark_in_order_btreemap(c: &mut Criterion) {
+    c.bench_function("BTreeMap: in-order frames", |b| {
+        b.iter(|| {
+            let mut rx = BTreeMapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_in_order_heap(c: &mut Criterion) {
+    c.bench_function("BinaryHeap: in-order frames", |b| {
+        b.iter(|| {
+            let mut rx = HeapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_reverse_order_btreemap(c: &mut Criterion) {
+    c.bench_function("BTreeMap: reverse-order frames", |b| {
+        b.iter(|| {
+            let mut rx = BTreeMapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in (0..1000).rev() {
+                rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_reverse_order_heap(c: &mut Criterion) {
+    c.bench_function("BinaryHeap: reverse-order frames", |b| {
+        b.iter(|| {
+            let mut rx = HeapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in (0..1000).rev() {
+                rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_random_order_btreemap(c: &mut Criterion) {
+    // Create a deterministic "random" order by using a simple permutation.
+    let mut order: Vec<usize> = (0..1000).collect();
+    // Simple shuffle using XOR swap.
+    for i in 0..order.len() {
+        let j = (i * 7 + 13) % order.len();
+        order.swap(i, j);
+    }
+
+    c.bench_function("BTreeMap: random-order frames", |b| {
+        b.iter(|| {
+            let mut rx = BTreeMapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for &i in &order {
+                rx.inbound_frame(black_box(i as u64 * FRAME_SIZE as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_random_order_heap(c: &mut Criterion) {
+    let mut order: Vec<usize> = (0..1000).collect();
+    for i in 0..order.len() {
+        let j = (i * 7 + 13) % order.len();
+        order.swap(i, j);
+    }
+
+    c.bench_function("BinaryHeap: random-order frames", |b| {
+        b.iter(|| {
+            let mut rx = HeapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for &i in &order {
+                rx.inbound_frame(black_box(i as u64 * FRAME_SIZE as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_with_gaps_btreemap(c: &mut Criterion) {
+    c.bench_function("BTreeMap: frames with gaps", |b| {
+        b.iter(|| {
+            let mut rx = BTreeMapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            // Send every other frame.
+            for i in 0..1000 {
+                if i % 2 == 0 {
+                    rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+                }
+            }
+        });
+    });
+}
+
+fn benchmark_with_gaps_heap(c: &mut Criterion) {
+    c.bench_function("BinaryHeap: frames with gaps", |b| {
+        b.iter(|| {
+            let mut rx = HeapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                if i % 2 == 0 {
+                    rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+                }
+            }
+        });
+    });
+}
+
+fn benchmark_overlapping_btreemap(c: &mut Criterion) {
+    c.bench_function("BTreeMap: overlapping frames", |b| {
+        b.iter(|| {
+            let mut rx = BTreeMapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                // Each frame overlaps with the previous by half.
+                rx.inbound_frame(black_box(i * (FRAME_SIZE / 2) as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_overlapping_heap(c: &mut Criterion) {
+    c.bench_function("BinaryHeap: overlapping frames", |b| {
+        b.iter(|| {
+            let mut rx = HeapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                rx.inbound_frame(black_box(i * (FRAME_SIZE / 2) as u64), black_box(data));
+            }
+        });
+    });
+}
+
+fn benchmark_read_to_end_btreemap(c: &mut Criterion) {
+    c.bench_function("BTreeMap: read_to_end after in-order insert", |b| {
+        b.iter(|| {
+            let mut rx = BTreeMapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                rx.inbound_frame(i * FRAME_SIZE as u64, data);
+            }
+
+            let mut buf = Vec::new();
+            black_box(rx.read_to_end(&mut buf));
+        });
+    });
+}
+
+fn benchmark_read_to_end_heap(c: &mut Criterion) {
+    c.bench_function("BinaryHeap: read_to_end after in-order insert", |b| {
+        b.iter(|| {
+            let mut rx = HeapOrderer::new();
+            let data: &[u8] = &[0; FRAME_SIZE];
+
+            for i in 0..1000 {
+                rx.inbound_frame(i * FRAME_SIZE as u64, data);
+            }
+
+            let mut buf = Vec::new();
+            black_box(rx.read_to_end(&mut buf));
+        });
+    });
+}
+
+fn benchmark_varying_sizes(c: &mut Criterion) {
+    let mut group = c.benchmark_group("varying_frame_counts");
+
+    for count in [100, 500, 1000, 5000, 10000].iter() {
+        group.bench_with_input(BenchmarkId::new("BTreeMap", count), count, |b, &count| {
+            b.iter(|| {
+                let mut rx = BTreeMapOrderer::new();
+                let data: &[u8] = &[0; FRAME_SIZE];
+
+                for i in 0..count {
+                    rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+                }
+            });
+        });
+
+        group.bench_with_input(BenchmarkId::new("BinaryHeap", count), count, |b, &count| {
+            b.iter(|| {
+                let mut rx = HeapOrderer::new();
+                let data: &[u8] = &[0; FRAME_SIZE];
+
+                for i in 0..count {
+                    rx.inbound_frame(black_box(i * FRAME_SIZE as u64), black_box(data));
+                }
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(
+    benches,
+    benchmark_in_order_btreemap,
+    benchmark_in_order_heap,
+    benchmark_reverse_order_btreemap,
+    benchmark_reverse_order_heap,
+    benchmark_random_order_btreemap,
+    benchmark_random_order_heap,
+    benchmark_with_gaps_btreemap,
+    benchmark_with_gaps_heap,
+    benchmark_overlapping_btreemap,
+    benchmark_overlapping_heap,
+    benchmark_read_to_end_btreemap,
+    benchmark_read_to_end_heap,
+    benchmark_varying_sizes,
+);
+criterion_main!(benches);

neqo-transport/src/connection/mod.rs

@@ -2543,10 +2543,14 @@ impl Connection {
         max_datagrams: NonZeroUsize,
     ) -> Res<SendOptionBatch> {
         let packet_tos = path.borrow().tos();
-        let mut send_buffer = Vec::new();
+        let mtu = path.borrow().plpmtu();
+
+        // Pre-allocate for typical batch of 2 packets to avoid reallocations.
+        let initial_capacity = mtu.saturating_mul(2);
+        let mut send_buffer = Vec::with_capacity(initial_capacity);
+
         let mut max_datagram_size = None;
         let mut num_datagrams = 0;
-        let mtu = path.borrow().plpmtu();
         let address_family_max_mtu = path.borrow().pmtud().address_family_max_mtu();
 
         loop {
@@ -3821,7 +3825,7 @@ impl Connection {
         };
         let path = self.paths.primary().ok_or(Error::NotAvailable)?;
         let mtu = path.borrow().plpmtu();
-        let mut buffer = Vec::new();
+        let mut buffer = Vec::with_capacity(mtu);
         let encoder = Encoder::new_borrowed_vec(&mut buffer);
 
         let (_, builder, _) = Self::build_packet_header(

neqo-transport/src/lib.rs

@@ -36,16 +36,11 @@ mod quic_datagrams;
 pub mod recovery;
 #[cfg(not(feature = "bench"))]
 mod recovery;
-mod saved;
-// #[cfg(feature = "bench")]
 pub mod recv_stream;
-// #[cfg(not(feature = "bench"))]
-// mod recv_stream;
 mod rtt;
-// #[cfg(feature = "bench")]
+pub mod rx_stream_orderer_heap;
+mod saved;
 pub mod send_stream;
-// #[cfg(not(feature = "bench"))]
-// mod send_stream;
 mod sender;
 pub mod server;
 mod sni;