Don't arm PTO for ApplicationData until handshake confirmed.

Protected by FLAGS_quic_reloadable_flag_quic_fix_arm_pto_for_application_data.

PiperOrigin-RevId: 332100576
Change-Id: I74ee1bea3f99eaa62d38d59d3c36ed5e53f2b03a

quic/core/quic_connection_test.cc

diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 940a988..60226f2 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -3110,6 +3110,8 @@
   connection_.SetFromConfig(config);
 
   connection_.SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+  connection_.OnHandshakeComplete();
+
   EXPECT_CALL(visitor_, OnSuccessfulVersionNegotiation(_));
   EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(10);
 
@@ -9563,8 +9565,9 @@
   connection_.SetFromConfig(config);
   if (GetQuicReloadableFlag(quic_default_enable_5rto_blackhole_detection2)) {
     EXPECT_CALL(visitor_, GetHandshakeState())
-        .WillRepeatedly(Return(HANDSHAKE_COMPLETE));
+        .WillRepeatedly(Return(HANDSHAKE_CONFIRMED));
   }
+  connection_.OnHandshakeComplete();
   EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
 
   // Send stream data.
@@ -9614,8 +9617,9 @@
   connection_.SetFromConfig(config);
   if (GetQuicReloadableFlag(quic_default_enable_5rto_blackhole_detection2)) {
     EXPECT_CALL(visitor_, GetHandshakeState())
-        .WillRepeatedly(Return(HANDSHAKE_COMPLETE));
+        .WillRepeatedly(Return(HANDSHAKE_CONFIRMED));
   }
+  connection_.OnHandshakeComplete();
   EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
 
   // Send stream data.
@@ -9664,8 +9668,9 @@
   connection_.SetFromConfig(config);
   if (GetQuicReloadableFlag(quic_default_enable_5rto_blackhole_detection2)) {
     EXPECT_CALL(visitor_, GetHandshakeState())
-        .WillRepeatedly(Return(HANDSHAKE_COMPLETE));
+        .WillRepeatedly(Return(HANDSHAKE_CONFIRMED));
   }
+  connection_.OnHandshakeComplete();
   EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
 
   // Send stream data.
@@ -10603,6 +10608,7 @@
   }
   EXPECT_CALL(*send_algorithm_, SetFromConfig(_, _));
   connection_.SetFromConfig(config);
+  connection_.OnHandshakeComplete();
   EXPECT_FALSE(connection_.GetRetransmissionAlarm()->IsSet());
 
   EXPECT_EQ(MESSAGE_STATUS_SUCCESS, SendMessage("message"));

quic/core/quic_sent_packet_manager.cc

diff --git a/quic/core/quic_sent_packet_manager.cc b/quic/core/quic_sent_packet_manager.cc
index 74a91ea..c73fb9a 100644
--- a/quic/core/quic_sent_packet_manager.cc
+++ b/quic/core/quic_sent_packet_manager.cc
@@ -559,6 +559,11 @@
 
 bool QuicSentPacketManager::ShouldArmPtoForApplicationData() const {
   DCHECK(supports_multiple_packet_number_spaces());
+  if (GetQuicReloadableFlag(quic_fix_arm_pto_for_application_data)) {
+    QUIC_RELOADABLE_FLAG_COUNT(quic_fix_arm_pto_for_application_data);
+    // Do not arm PTO for application data until handshake gets confirmed.
+    return handshake_finished_;
+  }
   // Application data must be ignored before handshake completes (1-RTT key
   // is available). Not arming PTO for application data to prioritize the
   // completion of handshake. On the server side, handshake_finished_

quic/core/quic_sent_packet_manager_test.cc

diff --git a/quic/core/quic_sent_packet_manager_test.cc b/quic/core/quic_sent_packet_manager_test.cc
index 7448ee5..c5fb920 100644
--- a/quic/core/quic_sent_packet_manager_test.cc
+++ b/quic/core/quic_sent_packet_manager_test.cc
@@ -3327,18 +3327,32 @@
 
   // Send packet 7 in handshake.
   clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  const QuicTime packet7_sent_time = clock_.Now();
   SendDataPacket(7, ENCRYPTION_HANDSHAKE);
-  // Verify PTO timeout is now based on packet 6.
-  expected_pto_delay =
-      srtt + pto_rttvar_multiplier * rtt_stats->mean_deviation() +
-      QuicTime::Delta::FromMilliseconds(kDefaultDelayedAckTimeMs);
-  EXPECT_EQ(packet6_sent_time + expected_pto_delay * 2,
-            manager_.GetRetransmissionTime());
+
+  if (GetQuicReloadableFlag(quic_fix_arm_pto_for_application_data)) {
+    expected_pto_delay =
+        srtt + pto_rttvar_multiplier * rtt_stats->mean_deviation();
+    // Verify PTO timeout is now based on packet 7.
+    EXPECT_EQ(packet7_sent_time + expected_pto_delay * 2,
+              manager_.GetRetransmissionTime());
+
+  } else {
+    expected_pto_delay =
+        srtt + pto_rttvar_multiplier * rtt_stats->mean_deviation() +
+        QuicTime::Delta::FromMilliseconds(kDefaultDelayedAckTimeMs);
+    // Verify PTO timeout is now based on packet 6.
+    EXPECT_EQ(packet6_sent_time + expected_pto_delay * 2,
+              manager_.GetRetransmissionTime());
+  }
 
   // Neuter handshake key.
   manager_.SetHandshakeConfirmed();
   // Forward progress has been made, verify PTO counter gets reset. PTO timeout
   // is armed by left edge.
+  expected_pto_delay =
+      srtt + pto_rttvar_multiplier * rtt_stats->mean_deviation() +
+      QuicTime::Delta::FromMilliseconds(kDefaultDelayedAckTimeMs);
   EXPECT_EQ(packet4_sent_time + expected_pto_delay,
             manager_.GetRetransmissionTime());
 }