In QUIC, do not re-arm PTO timer on sending application data before handshake gets confirmed.

This change also includes 3 client side only changes:
1) Fix the encryption level of sending PINGs before handshake gets confirmed.
2) Do not postpone PTO alarm on receiving undecryptable packets.
3) Enqueue undecryptable HANDSHAKE packet in 0-RTT connections.

Protected by FLAGS_quic_reloadable_flag_quic_donot_rearm_pto_on_application_data_during_handshake.

PiperOrigin-RevId: 381566145
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index d525242..cd4e6be 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2065,6 +2065,28 @@
   }
 }
 
+bool QuicConnection::ShouldSetRetransmissionAlarmOnPacketSent(
+    bool in_flight, EncryptionLevel level) const {
+  if (!retransmission_alarm_->IsSet()) {
+    return true;
+  }
+  if (!in_flight) {
+    return false;
+  }
+
+  if (!SupportsMultiplePacketNumberSpaces()) {
+    return true;
+  }
+  // Before handshake gets confirmed, do not re-arm PTO timer on application
+  // data. Think about this scenario: on the client side, the CHLO gets
+  // acknowledged and the SHLO is not received yet. The PTO alarm is set when
+  // the CHLO acknowledge is received (and there is no in flight INITIAL
+  // packet). Re-arming PTO alarm on 0-RTT packet would keep postponing the PTO
+  // alarm.
+  return IsHandshakeConfirmed() || level == ENCRYPTION_INITIAL ||
+         level == ENCRYPTION_HANDSHAKE;
+}
+
 bool QuicConnection::OnNewConnectionIdFrameInner(
     const QuicNewConnectionIdFrame& frame) {
   QUICHE_DCHECK(support_multiple_connection_ids_);
@@ -2779,6 +2801,15 @@
   }
   if (version().KnowsWhichDecrypterToUse() &&
       decryption_level <= encryption_level_) {
+    if (decryption_level == ENCRYPTION_HANDSHAKE &&
+        encryption_level_ == ENCRYPTION_ZERO_RTT) {
+      // This is cient side only since only clients have
+      // encryption_level_ == ENCRYPTION_ZERO_RTT.
+      QUICHE_DCHECK_EQ(Perspective::IS_CLIENT, perspective_);
+      // Make sure we enqueue undecryptable HANDSHAKE packets when the
+      // encryption level is 0-RTT and we have not install HANDSHAKE key yet.
+      return true;
+    }
     // On versions that know which decrypter to use, we install keys in order
     // so we will not get newer keys for lower encryption levels.
     return false;
@@ -3794,8 +3825,15 @@
       return true;
     }
   }
-
-  if (in_flight || !retransmission_alarm_->IsSet()) {
+  if (GetQuicReloadableFlag(
+          quic_donot_rearm_pto_on_application_data_during_handshake)) {
+    QUIC_RELOADABLE_FLAG_COUNT(
+        quic_donot_rearm_pto_on_application_data_during_handshake);
+    if (ShouldSetRetransmissionAlarmOnPacketSent(in_flight,
+                                                 packet->encryption_level)) {
+      SetRetransmissionAlarm();
+    }
+  } else if (in_flight || !retransmission_alarm_->IsSet()) {
     SetRetransmissionAlarm();
   }
   SetPingAlarm();
@@ -4286,11 +4324,23 @@
                      sent_packet_manager_.pending_timer_transmission_count());
     EncryptionLevel level = encryption_level_;
     PacketNumberSpace packet_number_space = NUM_PACKET_NUMBER_SPACES;
-    if (SupportsMultiplePacketNumberSpaces() &&
-        sent_packet_manager_
-            .GetEarliestPacketSentTimeForPto(&packet_number_space)
-            .IsInitialized()) {
-      level = QuicUtils::GetEncryptionLevel(packet_number_space);
+    if (SupportsMultiplePacketNumberSpaces()) {
+      if (sent_packet_manager_
+              .GetEarliestPacketSentTimeForPto(&packet_number_space)
+              .IsInitialized()) {
+        level = QuicUtils::GetEncryptionLevel(packet_number_space);
+      }
+      if (level == ENCRYPTION_ZERO_RTT) {
+        QUICHE_DCHECK_EQ(Perspective::IS_CLIENT, perspective_);
+        // Do not send 0-RTT PING since it will only elicit (unprocessable)
+        // 1-RTT acknowledgement from server, which does not speed up recovery.
+        if (framer_.HasEncrypterOfEncryptionLevel(ENCRYPTION_INITIAL)) {
+          level = ENCRYPTION_INITIAL;
+        } else if (framer_.HasEncrypterOfEncryptionLevel(
+                       ENCRYPTION_HANDSHAKE)) {
+          level = ENCRYPTION_HANDSHAKE;
+        }
+      }
     }
     SendPingAtLevel(level);
   }
@@ -4455,7 +4505,11 @@
   undecryptable_packets_.emplace_back(packet, decryption_level,
                                       last_received_packet_info_);
   if (perspective_ == Perspective::IS_CLIENT) {
-    SetRetransmissionAlarm();
+    if (!retransmission_alarm_->IsSet() ||
+        GetRetransmissionDeadline() < retransmission_alarm_->deadline()) {
+      // Re-arm PTO only if we can make it sooner to speed up recovery.
+      SetRetransmissionAlarm();
+    }
   }
 }
 
@@ -4527,7 +4581,14 @@
     undecryptable_packets_.clear();
   }
   if (perspective_ == Perspective::IS_CLIENT) {
-    SetRetransmissionAlarm();
+    if (!retransmission_alarm_->IsSet() || undecryptable_packets_.empty() ||
+        GetRetransmissionDeadline() < retransmission_alarm_->deadline()) {
+      // 1) If there is still undecryptable packet, only re-arm PTO to make it
+      // sooner to speed up recovery.
+      // 2) If all undecryptable packets get processed, re-arm (which may
+      // postpone) PTO since no immediate recovery is needed.
+      SetRetransmissionAlarm();
+    }
   }
 }
 
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 91e3358..371731b 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1873,6 +1873,11 @@
   // when a new client connection ID is received.
   void OnClientConnectionIdAvailable();
 
+  // Returns true if connection needs to set retransmission alarm after a packet
+  // gets sent.
+  bool ShouldSetRetransmissionAlarmOnPacketSent(bool in_flight,
+                                                EncryptionLevel level) const;
+
   QuicFramer framer_;
 
   // Contents received in the current packet, especially used to identify
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 95b9c48..f17345f 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -15074,6 +15074,166 @@
   }
 }
 
+TEST_P(QuicConnectionTest, SendingZeroRttPacketsDoesNotPostponePTO) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  // Send CHLO.
+  connection_.SendCryptoStreamData();
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  // Install 0-RTT keys.
+  connection_.SetEncrypter(ENCRYPTION_ZERO_RTT,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_ZERO_RTT);
+
+  // CHLO gets acknowledged after 10ms.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  QuicAckFrame frame1 = InitAckFrame(1);
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(_, _, _, _, _));
+  ProcessFramePacketAtLevel(1, QuicFrame(&frame1), ENCRYPTION_INITIAL);
+  // Verify PTO is still armed since address validation is not finished yet.
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  QuicTime pto_deadline = connection_.GetRetransmissionAlarm()->deadline();
+
+  // Send 0-RTT packet.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(5));
+  connection_.SetEncrypter(ENCRYPTION_ZERO_RTT,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_ZERO_RTT);
+  connection_.SendStreamDataWithString(2, "foo", 0, NO_FIN);
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  if (GetQuicReloadableFlag(
+          quic_donot_rearm_pto_on_application_data_during_handshake)) {
+    // PTO deadline should be unchanged.
+    EXPECT_EQ(pto_deadline, connection_.GetRetransmissionAlarm()->deadline());
+  } else {
+    // PTO gets re-armed.
+    EXPECT_NE(pto_deadline, connection_.GetRetransmissionAlarm()->deadline());
+  }
+}
+
+TEST_P(QuicConnectionTest, QueueingUndecryptablePacketsDoesntPostponePTO) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  EXPECT_CALL(*send_algorithm_, SetFromConfig(_, _));
+  QuicConfig config;
+  config.set_max_undecryptable_packets(3);
+  connection_.SetFromConfig(config);
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  connection_.RemoveDecrypter(ENCRYPTION_FORWARD_SECURE);
+  // Send CHLO.
+  connection_.SendCryptoStreamData();
+
+  // Send 0-RTT packet.
+  connection_.SetEncrypter(ENCRYPTION_ZERO_RTT,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_ZERO_RTT);
+  connection_.SendStreamDataWithString(2, "foo", 0, NO_FIN);
+
+  // CHLO gets acknowledged after 10ms.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  QuicAckFrame frame1 = InitAckFrame(1);
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(_, _, _, _, _));
+  ProcessFramePacketAtLevel(1, QuicFrame(&frame1), ENCRYPTION_INITIAL);
+  // Verify PTO is still armed since address validation is not finished yet.
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  QuicTime pto_deadline = connection_.GetRetransmissionAlarm()->deadline();
+
+  // Receive an undecryptable packets.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(5));
+  peer_framer_.SetEncrypter(ENCRYPTION_FORWARD_SECURE,
+                            std::make_unique<TaggingEncrypter>(0xFF));
+  ProcessDataPacketAtLevel(3, !kHasStopWaiting, ENCRYPTION_FORWARD_SECURE);
+  // Verify PTO deadline is sooner.
+  EXPECT_GT(pto_deadline, connection_.GetRetransmissionAlarm()->deadline());
+  pto_deadline = connection_.GetRetransmissionAlarm()->deadline();
+
+  // PTO fires.
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(1);
+  clock_.AdvanceTime(pto_deadline - clock_.ApproximateNow());
+  connection_.GetRetransmissionAlarm()->Fire();
+  // Verify PTO is still armed since address validation is not finished yet.
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  pto_deadline = connection_.GetRetransmissionAlarm()->deadline();
+
+  // Verify PTO deadline does not change.
+  ProcessDataPacketAtLevel(4, !kHasStopWaiting, ENCRYPTION_FORWARD_SECURE);
+  EXPECT_EQ(pto_deadline, connection_.GetRetransmissionAlarm()->deadline());
+}
+
+TEST_P(QuicConnectionTest, QueueUndecryptableHandshakePackets) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  EXPECT_CALL(*send_algorithm_, SetFromConfig(_, _));
+  QuicConfig config;
+  config.set_max_undecryptable_packets(3);
+  connection_.SetFromConfig(config);
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  connection_.RemoveDecrypter(ENCRYPTION_HANDSHAKE);
+  // Send CHLO.
+  connection_.SendCryptoStreamData();
+
+  // Send 0-RTT packet.
+  connection_.SetEncrypter(ENCRYPTION_ZERO_RTT,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_ZERO_RTT);
+  connection_.SendStreamDataWithString(2, "foo", 0, NO_FIN);
+  EXPECT_EQ(0u, QuicConnectionPeer::NumUndecryptablePackets(&connection_));
+
+  // Receive an undecryptable handshake packet.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(5));
+  peer_framer_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                            std::make_unique<TaggingEncrypter>(0xFF));
+  ProcessDataPacketAtLevel(3, !kHasStopWaiting, ENCRYPTION_HANDSHAKE);
+  // Verify this handshake packet gets queued.
+  EXPECT_EQ(1u, QuicConnectionPeer::NumUndecryptablePackets(&connection_));
+}
+
+TEST_P(QuicConnectionTest, PingNotSentAt0RTTLevelWhenInitialAvailable) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  // Send CHLO.
+  connection_.SendCryptoStreamData();
+  // Send 0-RTT packet.
+  connection_.SetEncrypter(ENCRYPTION_ZERO_RTT,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_ZERO_RTT);
+  connection_.SendStreamDataWithString(2, "foo", 0, NO_FIN);
+
+  // CHLO gets acknowledged after 10ms.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  QuicAckFrame frame1 = InitAckFrame(1);
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(_, _, _, _, _));
+  ProcessFramePacketAtLevel(1, QuicFrame(&frame1), ENCRYPTION_INITIAL);
+  // Verify PTO is still armed since address validation is not finished yet.
+  ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
+  QuicTime pto_deadline = connection_.GetRetransmissionAlarm()->deadline();
+
+  // PTO fires.
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(1);
+  clock_.AdvanceTime(pto_deadline - clock_.ApproximateNow());
+  connection_.GetRetransmissionAlarm()->Fire();
+  // Verify the PING gets sent in ENCRYPTION_INITIAL.
+  EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index 135b0ab..74aa4a1 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -61,6 +61,8 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_single_ack_in_packet2, false)
 // If true, do not count bytes sent/received on the alternative path into the bytes sent/received on the default path.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_count_bytes_on_alternative_path_seperately, true)
+// If true, do not re-arm PTO while sending application data during handshake.
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_donot_rearm_pto_on_application_data_during_handshake, true)
 // If true, do not send control frames before encryption is established.
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_encrypted_control_frames, false)
 // If true, do not send stream data when PTO fires.