perf(transport/recovery): optimize SentPackets::take_ranges

neqo-transport/Cargo.toml

@@ -58,3 +58,8 @@ required-features = ["bench"]
 name = "range_tracker"
 harness = false
 required-features = ["bench"]
+
+[[bench]]
+name = "sent_packets"
+harness = false
+required-features = ["bench"]

neqo-transport/benches/sent_packets.rs

@@ -0,0 +1,52 @@
+// 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::time::Instant;
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use neqo_common::IpTosEcn;
+use neqo_transport::{
+    packet::{PacketNumber, PacketType},
+    recovery::sent::{SentPacket, SentPackets},
+};
+
+fn sent_packets() -> SentPackets {
+    let mut pkts = SentPackets::default();
+    let now = Instant::now();
+    // Simulate high bandwidth-delay-product connection.
+    for i in 0..2_000u64 {
+        pkts.track(SentPacket::new(
+            PacketType::Short,
+            PacketNumber::from(i),
+            IpTosEcn::default(),
+            now,
+            true,
+            Vec::new(),
+            100,
+        ))
+    }
+    pkts
+}
+
+/// Confirm that taking a small number of ranges from the front of
+/// a large span of sent packets is performant.
+/// This is the most common pattern when sending at a higher rate.
+/// New acknowledgments will include low-numbered packets,
+/// while the acknowledgment rate will ensure that most of the
+/// outstanding packets remain in flight.
+fn take_ranges(c: &mut Criterion) {
+    c.bench_function("SentPackets::take_ranges", |b| {
+        b.iter_batched_ref(
+            || sent_packets(),
+            // Take the first 90 packets, minus some gaps.
+            |pkts| pkts.take_ranges([70..=89, 40..=59, 10..=29]),
+            criterion::BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group!(benches, take_ranges);
+criterion_main!(benches);

neqo-transport/src/lib.rs

@@ -23,14 +23,17 @@ pub mod frame;
 #[cfg(not(fuzzing))]
 mod frame;
 mod pace;
-#[cfg(fuzzing)]
+#[cfg(any(fuzzing, feature = "bench"))]
 pub mod packet;
-#[cfg(not(fuzzing))]
+#[cfg(not(any(fuzzing, feature = "bench")))]
 mod packet;
 mod path;
 mod pmtud;
 mod qlog;
 mod quic_datagrams;
+#[cfg(feature = "bench")]
+pub mod recovery;
+#[cfg(not(feature = "bench"))]
 mod recovery;
 #[cfg(feature = "bench")]
 pub mod recv_stream;

neqo-transport/src/recovery/mod.rs

@@ -6,6 +6,9 @@
 
 // Tracking of sent packets and detecting their loss.
 
+#[cfg(feature = "bench")]
+pub mod sent;
+#[cfg(not(feature = "bench"))]
 mod sent;
 mod token;
 

neqo-transport/src/recovery/sent.rs

@@ -186,7 +186,7 @@
         self.packets.values_mut()
     }
 
-    /// Take values from a specified ranges of packet numbers.
+    /// Take values from specified ranges of packet numbers.
     /// The values returned will be reversed, so that the most recent packet appears first.
     /// This is because ACK frames arrive with ranges starting from the largest acknowledged
     /// and we want to match that.
@@ -196,19 +196,60 @@
         R::IntoIter: ExactSizeIterator,
     {
         let mut result = Vec::new();
-        // Remove all packets. We will add them back as we don't need them.
+
+        // Start with all packets. We will add unacknowledged packets back.
+        //  [---------------------------packets----------------------------]
         let mut packets = std::mem::take(&mut self.packets);
+
+        let mut previous_range_start: Option<PacketNumber> = None;
+
         for range in acked_ranges {
-            // For each acked range, split off the acknowledged part,
-            // then split off the part that hasn't been acknowledged.
-            // This order works better when processing ranges that
-            // have already been processed, which is common.
-            let mut acked = packets.split_off(range.start());
-            let keep = acked.split_off(&(*range.end() + 1));
-            self.packets.extend(keep);
-            result.extend(acked.into_values().rev());
+            // Split off at the end of the acked range.
+            //
+            //  [---------packets--------][----------after_acked_range---------]
+            let after_acked_range = packets.split_off(&(*range.end() + 1));
+
+            // Split off at the start of the acked range.
+            //
+            //  [-packets-][-acked_range-][----------after_acked_range---------]
+            let acked_range = packets.split_off(range.start());
+
+            // According to RFC 9000 19.3.1 ACK ranges are in descending order:
+            //
+            // > Each ACK Range consists of alternating Gap and ACK Range Length
+            // > values in **descending packet number order**.
+            //
+            // <https://www.rfc-editor.org/rfc/rfc9000.html#section-19.3.1>
+            if let Some(prev_start) = previous_range_start {
+                debug_assert!(
+                    *range.end() < prev_start,
+                    "ACK ranges not in descending order: current range end {}, previous range start {prev_start}",
+                    *range.end(),
+                );
+            }
+            previous_range_start = Some(*range.start());
+
+            // Thus none of the following ACK ranges will acknowledge packets in
+            // `after_acked_range`. Let's put those back early.
+            //
+            //  [-packets-][-acked_range-][------------self.packets------------]
+            if self.packets.is_empty() {
+                // Don't re-insert un-acked packets into empty collection, but
+                // instead replace the empty one entirely.
+                self.packets = after_acked_range;
+            } else {
+                // Need to extend existing one. Not the first iteration, thus
+                // `after_acked_range` should be small.
+                self.packets.extend(after_acked_range);
+            }
+
+            // Take the acked packets.
+            result.extend(acked_range.into_values().rev());
         }
+
+        // Put remaining non-acked packets back.
         self.packets.extend(packets);
+
         result
     }