fix: Add confirmed parameter to PTO calculation

neqo-transport/src/connection/mod.rs

@@ -629,13 +629,17 @@ impl Connection {
             .unwrap()
     }
 
+    fn confirmed(&self) -> bool {
+        self.state == State::Confirmed
+    }
+
     /// Get the simplest PTO calculation for all those cases where we need
     /// a value of this approximate order.  Don't use this for loss recovery,
     /// only use it where a more precise value is not important.
     fn pto(&self) -> Duration {
         self.paths.primary().map_or_else(
-            || RttEstimate::default().pto(PacketNumberSpace::ApplicationData),
-            |p| p.borrow().rtt().pto(PacketNumberSpace::ApplicationData),
+            || RttEstimate::default().pto(self.confirmed()),
+            |p| p.borrow().rtt().pto(self.confirmed()),
         )
     }
 
@@ -1058,7 +1062,7 @@ impl Connection {
         if let Some(p) = self.paths.primary() {
             let path = p.borrow();
             let rtt = path.rtt();
-            let pto = rtt.pto(PacketNumberSpace::ApplicationData);
+            let pto = rtt.pto(self.confirmed());
 
             let idle_time = self.idle_timeout.expiry(now, pto);
             qtrace!([self], "Idle/keepalive timer {:?}", idle_time);
@@ -1525,7 +1529,7 @@ impl Connection {
         let mut dcid = None;
 
         qtrace!([self], "{} input {}", path.borrow(), hex(&**d));
-        let pto = path.borrow().rtt().pto(PacketNumberSpace::ApplicationData);
+        let pto = path.borrow().rtt().pto(self.confirmed());
 
         // Handle each packet in the datagram.
         while !slc.is_empty() {
@@ -2141,7 +2145,7 @@ impl Connection {
             // or the PTO timer fired: probe.
             true
         } else {
-            let pto = path.borrow().rtt().pto(PacketNumberSpace::ApplicationData);
+            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.

neqo-transport/src/path.rs

@@ -30,7 +30,6 @@ use crate::{
     rtt::{RttEstimate, RttSource},
     sender::PacketSender,
     stats::FrameStats,
-    tracking::PacketNumberSpace,
     Stats,
 };
 
@@ -1020,7 +1019,7 @@ impl Path {
     pub fn on_packets_lost(
         &mut self,
         prev_largest_acked_sent: Option<Instant>,
-        space: PacketNumberSpace,
+        confirmed: bool,
         lost_packets: &[SentPacket],
         stats: &mut Stats,
         now: Instant,
@@ -1030,7 +1029,7 @@ impl Path {
         let cwnd_reduced = self.sender.on_packets_lost(
             self.rtt.first_sample_time(),
             prev_largest_acked_sent,
-            self.rtt.pto(space), // Important: the base PTO, not adjusted.
+            self.rtt.pto(confirmed), // Important: the base PTO, not adjusted.
             lost_packets,
             stats,
             now,

neqo-transport/src/recovery/mod.rs

@@ -568,6 +568,10 @@ impl LossRecovery {
         }
     }
 
+    const fn confirmed(&self) -> bool {
+        self.confirmed_time.is_some()
+    }
+
     /// Returns (acked packets, lost packets)
     #[allow(clippy::too_many_arguments)]
     pub fn on_ack_received<R>(
@@ -627,7 +631,7 @@ impl LossRecovery {
         // as we rely on the count of in-flight packets to determine whether to send
         // another probe.  Removing them too soon would result in not sending on PTO.
         let loss_delay = primary_path.borrow().rtt().loss_delay();
-        let cleanup_delay = self.pto_period(primary_path.borrow().rtt(), pn_space);
+        let cleanup_delay = self.pto_period(primary_path.borrow().rtt());
         let mut lost = Vec::new();
         self.spaces.get_mut(pn_space).unwrap().detect_lost_packets(
             now,
@@ -642,7 +646,7 @@ impl LossRecovery {
         // backoff, so that we can determine persistent congestion.
         primary_path.borrow_mut().on_packets_lost(
             prev_largest_acked,
-            pn_space,
+            self.confirmed(),
             &lost,
             &mut self.stats.borrow_mut(),
             now,
@@ -679,7 +683,7 @@ impl LossRecovery {
         dropped
     }
 
-    fn confirmed(&mut self, rtt: &RttEstimate, now: Instant) {
+    fn confirm(&mut self, rtt: &RttEstimate, now: Instant) {
         debug_assert!(self.confirmed_time.is_none());
         self.confirmed_time = Some(now);
         // Up until now, the ApplicationData space has been ignored for PTO.
@@ -716,7 +720,7 @@ impl LossRecovery {
         self.pto_state = None;
 
         if space == PacketNumberSpace::Handshake {
-            self.confirmed(path.rtt(), now);
+            self.confirm(path.rtt(), now);
         }
     }
 
@@ -757,41 +761,40 @@ impl LossRecovery {
     fn pto_period_inner(
         rtt: &RttEstimate,
         pto_state: Option<&PtoState>,
-        pn_space: PacketNumberSpace,
+        confirmed: bool,
         fast_pto: u8,
     ) -> Duration {
         // This is a complicated (but safe) way of calculating:
         //   base_pto * F * 2^pto_count
         // where F = fast_pto / FAST_PTO_SCALE (== 1 by default)
         let pto_count = pto_state.map_or(0, |p| u32::try_from(p.count).unwrap_or(0));
-        rtt.pto(pn_space)
+        rtt.pto(confirmed)
             .checked_mul(u32::from(fast_pto) << min(pto_count, u32::BITS - u8::BITS))
             .map_or(Duration::from_secs(3600), |p| p / u32::from(FAST_PTO_SCALE))
     }
 
     /// Get the current PTO period for the given packet number space.
     /// Unlike calling `RttEstimate::pto` directly, this includes exponential backoff.
-    fn pto_period(&self, rtt: &RttEstimate, pn_space: PacketNumberSpace) -> Duration {
-        Self::pto_period_inner(rtt, self.pto_state.as_ref(), pn_space, self.fast_pto)
+    fn pto_period(&self, rtt: &RttEstimate) -> Duration {
+        Self::pto_period_inner(
+            rtt,
+            self.pto_state.as_ref(),
+            self.confirmed(),
+            self.fast_pto,
+        )
     }
 
     // Calculate PTO time for the given space.
     fn pto_time(&self, rtt: &RttEstimate, pn_space: PacketNumberSpace) -> Option<Instant> {
-        if self.confirmed_time.is_none() && pn_space == PacketNumberSpace::ApplicationData {
-            None
-        } else {
-            self.spaces.get(pn_space).and_then(|space| {
-                space
-                    .pto_base_time()
-                    .map(|t| t + self.pto_period(rtt, pn_space))
-            })
-        }
+        self.spaces
+            .get(pn_space)
+            .and_then(|space| space.pto_base_time().map(|t| t + self.pto_period(rtt)))
     }
 
     /// Find the earliest PTO time for all active packet number spaces.
     /// Ignore Application if either Initial or Handshake have an active PTO.
     fn earliest_pto(&self, rtt: &RttEstimate) -> Option<Instant> {
-        if self.confirmed_time.is_some() {
+        if self.confirmed() {
             self.pto_time(rtt, PacketNumberSpace::ApplicationData)
         } else {
             self.pto_time(rtt, PacketNumberSpace::Initial)
@@ -859,21 +862,22 @@ impl LossRecovery {
         qtrace!([self], "timeout {:?}", now);
 
         let loss_delay = primary_path.borrow().rtt().loss_delay();
+        let confirmed = self.confirmed();
 
         let mut lost_packets = Vec::new();
         for space in self.spaces.iter_mut() {
             let first = lost_packets.len(); // The first packet lost in this space.
             let pto = Self::pto_period_inner(
                 primary_path.borrow().rtt(),
                 self.pto_state.as_ref(),
-                space.space,
+                confirmed,
                 self.fast_pto,
             );
             space.detect_lost_packets(now, loss_delay, pto, &mut lost_packets);
 
             primary_path.borrow_mut().on_packets_lost(
                 space.largest_acked_sent_time,
-                space.space,
+                confirmed,
                 &lost_packets[first..],
                 &mut self.stats.borrow_mut(),
                 now,
@@ -950,7 +954,6 @@ mod tests {
         ecn::EcnCount,
         packet::{PacketNumber, PacketType},
         path::{Path, PathRef},
-        rtt::RttEstimate,
         stats::{Stats, StatsCell},
     };
 
@@ -961,8 +964,8 @@ mod tests {
 
     const ON_SENT_SIZE: usize = 100;
     /// An initial RTT for using with `setup_lr`.
-    const TEST_RTT: Duration = ms(80);
-    const TEST_RTTVAR: Duration = ms(40);
+    const TEST_RTT: Duration = ms(7000);
+    const TEST_RTTVAR: Duration = ms(3500);
 
     struct Fixture {
         lr: LossRecovery,
@@ -1033,6 +1036,7 @@ mod tests {
                 ConnectionIdEntry::new(0, ConnectionId::from(&[1, 2, 3]), [0; 16]),
             );
             path.set_primary(true);
+            path.rtt_mut().set_initial(TEST_RTT);
             Self {
                 lr: LossRecovery::new(StatsCell::default(), FAST_PTO_SCALE),
                 path: Rc::new(RefCell::new(path)),
@@ -1510,13 +1514,13 @@ mod tests {
             ON_SENT_SIZE,
         ));
 
-        assert_eq!(lr.pto_time(PacketNumberSpace::ApplicationData), None);
+        assert!(lr.pto_time(PacketNumberSpace::ApplicationData).is_some());
         lr.discard(PacketNumberSpace::Initial, pn_time(1));
-        assert_eq!(lr.pto_time(PacketNumberSpace::ApplicationData), None);
+        assert!(lr.pto_time(PacketNumberSpace::ApplicationData).is_some());
 
         // Expiring state after the PTO on the ApplicationData space has
         // expired should result in setting a PTO state.
-        let default_pto = RttEstimate::default().pto(PacketNumberSpace::ApplicationData);
+        let default_pto = lr.path.borrow().rtt().pto(true);
         let expected_pto = pn_time(2) + default_pto;
         lr.discard(PacketNumberSpace::Handshake, expected_pto);
         let profile = lr.send_profile(now());
@@ -1548,7 +1552,7 @@ mod tests {
             ON_SENT_SIZE,
         ));
 
-        let handshake_pto = RttEstimate::default().pto(PacketNumberSpace::Handshake);
+        let handshake_pto = lr.path.borrow().rtt().pto(false);
         let expected_pto = now() + handshake_pto;
         assert_eq!(lr.pto_time(PacketNumberSpace::Initial), Some(expected_pto));
         let profile = lr.send_profile(now());

neqo-transport/src/rtt.rs

@@ -19,7 +19,6 @@ use crate::{
     qlog::{self, QlogMetric},
     recovery::RecoveryToken,
     stats::FrameStats,
-    tracking::PacketNumberSpace,
 };
 
 /// The smallest time that the system timer (via `sleep()`, `nanosleep()`,
@@ -163,9 +162,9 @@ impl RttEstimate {
         self.smoothed_rtt
     }
 
-    pub fn pto(&self, pn_space: PacketNumberSpace) -> Duration {
+    pub fn pto(&self, confirmed: bool) -> Duration {
         let mut t = self.estimate() + max(4 * self.rttvar, GRANULARITY);
-        if pn_space == PacketNumberSpace::ApplicationData {
+        if confirmed {
             t += self.ack_delay.max();
         }
         t