From 3808dccd48ba2ba808b384882c7542473787fa4c Mon Sep 17 00:00:00 2001
From: Max Inden <mail@max-inden.de>
Date: Wed, 10 Jul 2024 09:31:47 +0200
Subject: [PATCH] Use filter fn

Follow-up on https://github.com/mozilla/neqo/pull/1903#discussion_r1617011637
---
 neqo-transport/src/cc/classic_cc.rs | 29 ++++++++++++++---------------
 neqo-transport/src/pmtud.rs         | 18 +++++++++++++++---
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/neqo-transport/src/cc/classic_cc.rs b/neqo-transport/src/cc/classic_cc.rs
index 68167ef9d8..ce062db801 100644
--- a/neqo-transport/src/cc/classic_cc.rs
+++ b/neqo-transport/src/cc/classic_cc.rs
@@ -308,15 +308,14 @@ impl<T: WindowAdjustment> CongestionControl for ClassicCongestionControl<T> {
             &[QlogMetric::BytesInFlight(self.bytes_in_flight)],
         );
 
+        let is_pmtud_probe = self.pmtud.is_pmtud_probe_filter();
+
         // Lost PMTUD probes do not elicit a congestion control reaction.
-        let probe_size = self.pmtud.probe_size();
-        let probe_sent = self.pmtud.probe_sent();
-        let mut lost_packets = lost_packets
+        let Some(last_lost_packet) = lost_packets
             .iter()
-            .filter(|pkt| !probe_sent || pkt.len() < probe_size)
-            .rev()
-            .peekable();
-        let Some(last_lost_packet) = lost_packets.peek() else {
+            .filter(|pkt| !is_pmtud_probe(pkt))
+            .last()
+        else {
             return false;
         };
 
@@ -325,7 +324,7 @@ impl<T: WindowAdjustment> CongestionControl for ClassicCongestionControl<T> {
             first_rtt_sample_time,
             prev_largest_acked_sent,
             pto,
-            lost_packets.rev(),
+            lost_packets.iter().filter(|pkt| !is_pmtud_probe(pkt)),
         );
         qdebug!(
             "on_packets_lost this={:p}, bytes_in_flight={}, cwnd={}, state={:?}",
@@ -469,16 +468,13 @@ impl<T: WindowAdjustment> ClassicCongestionControl<T> {
         }
     }
 
-    fn detect_persistent_congestion<'a, I>(
+    fn detect_persistent_congestion<'a>(
         &mut self,
         first_rtt_sample_time: Option<Instant>,
         prev_largest_acked_sent: Option<Instant>,
         pto: Duration,
-        lost_packets: I,
-    ) -> bool
-    where
-        I: Iterator<Item = &'a SentPacket>,
-    {
+        lost_packets: impl IntoIterator<Item = &'a SentPacket>,
+    ) -> bool {
         if first_rtt_sample_time.is_none() {
             return false;
         }
@@ -493,7 +489,10 @@ impl<T: WindowAdjustment> ClassicCongestionControl<T> {
         // Also, make sure to ignore any packets sent before we got an RTT estimate
         // as we might not have sent PTO packets soon enough after those.
         let cutoff = max(first_rtt_sample_time, prev_largest_acked_sent);
-        for p in lost_packets.skip_while(|p| Some(p.time_sent()) < cutoff) {
+        for p in lost_packets
+            .into_iter()
+            .skip_while(|p| Some(p.time_sent()) < cutoff)
+        {
             if p.pn() != last_pn + 1 {
                 // Not a contiguous range of lost packets, start over.
                 start = None;
diff --git a/neqo-transport/src/pmtud.rs b/neqo-transport/src/pmtud.rs
index cb452b4f2f..d44fb401a6 100644
--- a/neqo-transport/src/pmtud.rs
+++ b/neqo-transport/src/pmtud.rs
@@ -132,14 +132,26 @@ impl Pmtud {
         );
     }
 
+    /// Provides a [`Fn`] that returns true if the packet is a PMTUD probe.
+    ///
+    /// Allows filtering packets without holding a reference to [`Pmtud`]. When
+    /// in doubt, use [`Pmtud::is_pmtud_probe`].
+    #[must_use]
+    pub fn is_pmtud_probe_filter(&self) -> impl Fn(&SentPacket) -> bool {
+        let probe_state = Probe::Sent;
+        let probe_size = self.probe_size();
+
+        move |p: &SentPacket| -> bool { probe_state == Probe::Sent && p.len() == probe_size }
+    }
+
     /// Returns true if the packet is a PMTUD probe.
-    fn is_probe(&self, p: &SentPacket) -> bool {
-        self.probe_state == Probe::Sent && p.len() == self.probe_size()
+    fn is_pmtud_probe(&self, p: &SentPacket) -> bool {
+        self.is_pmtud_probe_filter()(p)
     }
 
     /// Count the PMTUD probes included in `pkts`.
     fn count_probes(&self, pkts: &[SentPacket]) -> usize {
-        pkts.iter().filter(|p| self.is_probe(p)).count()
+        pkts.iter().filter(|p| self.is_pmtud_probe(p)).count()
     }
 
     /// Checks whether a PMTUD probe has been acknowledged, and if so, updates the PMTUD state.