fix: Mark all packets TX'ed before PTO as lost

neqo-transport/src/connection/mod.rs

@@ -2216,30 +2216,29 @@ impl Connection {
             return true;
         }
 
+        let pto = path.borrow().rtt().pto(self.confirmed());
         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(self.confirmed());
-            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.
             debug_assert_ne!(builder.remaining(), 0);
             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
     }

neqo-transport/src/recovery/mod.rs

@@ -177,18 +177,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
+            }
+        })
     }
 
     #[must_use]
@@ -861,7 +858,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>) {
@@ -874,11 +871,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));
                 }