From 392bbacb9c656152349a999d84e4c7f1bca8188e Mon Sep 17 00:00:00 2001
From: Lars Eggert <lars@eggert.org>
Date: Thu, 19 Sep 2024 15:55:54 +0300
Subject: [PATCH] fix: Mark all packets TX'ed before PTO as lost

We'd previously only mark 1 one or two packets as lost when a PTO fired.
That meant that we potentially didn't RTX all data that we could have
(i.e., that was in lost packets that we didn't mark lost).

This also changes the probing code to suppress redundant keep-alives,
i.e., PINGs that we sent for other reasons, which could double as
keep-alives but did not.

Broken out of #1998
---
 neqo-transport/src/connection/mod.rs | 23 +++++++++++-----------
 neqo-transport/src/recovery/mod.rs   | 29 +++++++++++-----------------
 2 files changed, 22 insertions(+), 30 deletions(-)

diff --git a/neqo-transport/src/connection/mod.rs b/neqo-transport/src/connection/mod.rs
index d8d9db2422..d4c5ad027c 100644
--- a/neqo-transport/src/connection/mod.rs
+++ b/neqo-transport/src/connection/mod.rs
@@ -2133,23 +2133,20 @@ impl Connection {
             return true;
         }
 
+        let pto = path.borrow().rtt().pto(PacketNumberSpace::ApplicationData);
         let probe = if untracked && builder.packet_empty() || force_probe {
             // If we received an untracked packet and we aren't probing already
             // or the PTO timer fired: probe.
             true
+        } else if !builder.packet_empty() {
+            // The packet only contains an ACK.  Check whether we want to
+            // force an ACK with a PING so we can stop tracking packets.
+            self.loss_recovery.should_probe(pto, now)
+        } else if self.streams.need_keep_alive() {
+            // We need to keep the connection alive, including sending a PING again.
+            self.idle_timeout.send_keep_alive(now, pto, tokens)
         } else {
-            let pto = path.borrow().rtt().pto(PacketNumberSpace::ApplicationData);
-            if !builder.packet_empty() {
-                // The packet only contains an ACK.  Check whether we want to
-                // force an ACK with a PING so we can stop tracking packets.
-                self.loss_recovery.should_probe(pto, now)
-            } else if self.streams.need_keep_alive() {
-                // We need to keep the connection alive, including sending
-                // a PING again.
-                self.idle_timeout.send_keep_alive(now, pto, tokens)
-            } else {
-                false
-            }
+            false
         };
         if probe {
             // Nothing ack-eliciting and we need to probe; send PING.
@@ -2157,6 +2154,8 @@ impl Connection {
             builder.encode_varint(crate::frame::FRAME_TYPE_PING);
             let stats = &mut self.stats.borrow_mut().frame_tx;
             stats.ping += 1;
+            // Any PING we send here can also double as a keep-alive.
+            self.idle_timeout.send_keep_alive(now, pto, tokens);
         }
         probe
     }
diff --git a/neqo-transport/src/recovery/mod.rs b/neqo-transport/src/recovery/mod.rs
index f2c3e8e298..e11fd651d0 100644
--- a/neqo-transport/src/recovery/mod.rs
+++ b/neqo-transport/src/recovery/mod.rs
@@ -165,18 +165,15 @@ impl LossRecoverySpace {
         self.in_flight_outstanding > 0
     }
 
-    pub fn pto_packets(&mut self, count: usize) -> impl Iterator<Item = &SentPacket> {
-        self.sent_packets
-            .iter_mut()
-            .filter_map(|sent| {
-                if sent.pto() {
-                    qtrace!("PTO: marking packet {} lost ", sent.pn());
-                    Some(&*sent)
-                } else {
-                    None
-                }
-            })
-            .take(count)
+    pub fn pto_packets(&mut self) -> impl Iterator<Item = &SentPacket> {
+        self.sent_packets.iter_mut().filter_map(|sent| {
+            if sent.pto() {
+                qtrace!("PTO: marking packet {} lost ", sent.pn());
+                Some(&*sent)
+            } else {
+                None
+            }
+        })
     }
 
     pub fn pto_base_time(&self) -> Option<Instant> {
@@ -823,7 +820,7 @@ impl LossRecovery {
     }
 
     /// This checks whether the PTO timer has fired and fires it if needed.
-    /// When it has, mark a few packets as "lost" for the purposes of having frames
+    /// When it has, mark packets as "lost" for the purposes of having frames
     /// regenerated in subsequent packets.  The packets aren't truly lost, so
     /// we have to clone the `SentPacket` instance.
     fn maybe_fire_pto(&mut self, rtt: &RttEstimate, now: Instant, lost: &mut Vec<SentPacket>) {
@@ -836,11 +833,7 @@ impl LossRecovery {
                 if t <= now {
                     qdebug!([self], "PTO timer fired for {}", pn_space);
                     let space = self.spaces.get_mut(*pn_space).unwrap();
-                    lost.extend(
-                        space
-                            .pto_packets(PtoState::pto_packet_count(*pn_space))
-                            .cloned(),
-                    );
+                    lost.extend(space.pto_packets().cloned());
 
                     pto_space = pto_space.or(Some(*pn_space));
                 }