From 77eccc1925a21b1aba50dd4e0bd672b54be68a17 Mon Sep 17 00:00:00 2001
From: Alessandro Ghedini <alessandro@ghedini.me>
Date: Thu, 11 Apr 2024 14:48:59 +0100
Subject: [PATCH] recovery: ensure lost PMTUD probe bytes are removed from
 in-flight count

This ensure that when a PMTUD probe is lost we decrease the
`bytes_in_flight` counter appropriately.
---
 quiche/src/recovery/mod.rs | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/quiche/src/recovery/mod.rs b/quiche/src/recovery/mod.rs
index dd89ecb79e..088458b65a 100644
--- a/quiche/src/recovery/mod.rs
+++ b/quiche/src/recovery/mod.rs
@@ -958,7 +958,15 @@ impl Recovery {
                 self.lost[epoch].extend(unacked.frames.drain(..));
                 unacked.time_lost = Some(now);
 
-                if unacked.in_flight && !unacked.pmtud {
+                if unacked.pmtud {
+                    self.bytes_in_flight =
+                        self.bytes_in_flight.saturating_sub(unacked.size);
+
+                    // Do not track PMTUD probes losses.
+                    continue;
+                }
+
+                if unacked.in_flight {
                     lost_bytes += unacked.size;
 
                     // Frames have already been removed from the packet, so
@@ -976,11 +984,8 @@ impl Recovery {
                     );
                 }
 
-                if !unacked.pmtud {
-                    // Do not add PMTUD probes to loss statistics
-                    lost_packets += 1;
-                    self.lost_count += 1;
-                }
+                lost_packets += 1;
+                self.lost_count += 1;
             } else {
                 let loss_time = match self.loss_time[epoch] {
                     None => unacked.time_sent + loss_delay,
@@ -2438,7 +2443,7 @@ mod tests {
         assert_eq!(r.loss_probes[packet::Epoch::Application], 0);
 
         assert_eq!(r.sent[packet::Epoch::Application].len(), 2);
-        assert_eq!(r.bytes_in_flight, 1000);
+        assert_eq!(r.bytes_in_flight, 0);
         assert_eq!(r.congestion_window, 12000);
 
         assert_eq!(r.lost_count, 0);
@@ -2449,6 +2454,8 @@ mod tests {
         r.detect_lost_packets(packet::Epoch::Application, now, "");
 
         assert_eq!(r.sent[packet::Epoch::Application].len(), 0);
+        assert_eq!(r.bytes_in_flight, 0);
+        assert_eq!(r.lost_count, 0);
     }
 }