From 864dbdc3cf6384cdc96869816f2e93521724c0b4 Mon Sep 17 00:00:00 2001
From: toidiu <apoorv@toidiu.com>
Date: Thu, 20 Apr 2023 20:51:35 -0700
Subject: [PATCH] fix(s2n-quic-transport): reset PTO backoff once per space
 discard (#1717)

---
 quic/s2n-quic-transport/src/space/mod.rs |  4 +-
 quic/s2n-quic/src/tests.rs               | 84 ++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/quic/s2n-quic-transport/src/space/mod.rs b/quic/s2n-quic-transport/src/space/mod.rs
index 569e9150a8..397cf3894c 100644
--- a/quic/s2n-quic-transport/src/space/mod.rs
+++ b/quic/s2n-quic-transport/src/space/mod.rs
@@ -115,8 +115,10 @@ macro_rules! packet_space_api {
                 //# detection timers MUST be reset, because discarding keys indicates
                 //# forward progress and the loss detection timer might have been set for
                 //# a now discarded packet number space.
-                path.reset_pto_backoff();
                 if let Some(mut space) = self.$field.take() {
+                    // reset the PTO backoff value as part of resetting the PTO timer.
+                    path.reset_pto_backoff();
+
                     space.on_discard(path, path_id, publisher);
                 }
 
diff --git a/quic/s2n-quic/src/tests.rs b/quic/s2n-quic/src/tests.rs
index 36a5ba0a6a..cdb939f188 100644
--- a/quic/s2n-quic/src/tests.rs
+++ b/quic/s2n-quic/src/tests.rs
@@ -361,6 +361,15 @@ event_recorder!(
         }
     }
 );
+event_recorder!(
+    PtoRecorder,
+    RecoveryMetrics,
+    on_recovery_metrics,
+    u32,
+    |event: &RecoveryMetrics, storage: &mut Vec<u32>| {
+        storage.push(event.pto_count);
+    }
+);
 
 #[test]
 fn packet_sent_event_test() {
@@ -607,3 +616,78 @@ fn client_path_handle_update() {
     // there should only be a single update to the path handle
     assert_eq!(events_handle.len(), 2);
 }
+
+#[test]
+fn increasing_pto_count_under_loss() {
+    let delay_time = Duration::from_millis(10);
+
+    let model = Model::default();
+    model.set_delay(delay_time);
+    let subscriber = PtoRecorder::new();
+    let pto_events = subscriber.events();
+
+    test(model.clone(), |handle| {
+        spawn(async move {
+            // allow for 1 RTT worth of data and then drop all packet after
+            // the client gets an initial ACK from the server
+            delay(delay_time * 2).await;
+            model.set_drop_rate(1.0);
+        });
+
+        let mut server = Server::builder()
+            .with_io(handle.builder().build()?)?
+            .with_tls((certificates::CERT_PEM, certificates::KEY_PEM))?
+            .start()?;
+
+        let addr = server.local_addr()?;
+        spawn(async move {
+            if let Some(conn) = server.accept().await {
+                delay(Duration::from_secs(10)).await;
+                let _ = conn;
+            }
+        });
+
+        let client = Client::builder()
+            .with_io(handle.builder().build().unwrap())?
+            .with_tls(certificates::CERT_PEM)?
+            .with_event(subscriber)?
+            .start()?;
+
+        primary::spawn(async move {
+            let connect = Connect::new(addr).with_server_name("localhost");
+            let conn = client.connect(connect).await.unwrap();
+
+            delay(Duration::from_secs(10)).await;
+            let _ = conn;
+        });
+
+        Ok(addr)
+    })
+    .unwrap();
+
+    let mut pto_events = pto_events.lock().unwrap();
+
+    // assert that sufficient recovery events were captured
+    let pto_len = pto_events.len();
+    assert!(pto_len > 10);
+    // the last recovery event is fired after we discard the handshake space so ignore it
+    pto_events.truncate(pto_len - 1);
+
+    let pto_count: u32 = *pto_events
+        .iter()
+        .reduce(|prev, new| {
+            // assert that the value is monotonically increasing
+            assert!(new >= prev, "prev_value {}, new_value {}", prev, new);
+            new
+        })
+        .unwrap();
+
+    // assert that the final pto_count increased to some large value over the
+    // duration of the test
+    assert!(
+        pto_count > 5,
+        "delay: {:?}. pto_count: {}",
+        delay_time,
+        pto_count
+    );
+}