gfe-relnote: In QUIC, make PTO per packet number space. This only affects QUIC versions with TLS handshake. Protected by disabled gfe2_reloadable_flag_quic_enable_version_t* flags.

PiperOrigin-RevId: 291947359
Change-Id: I18bff16ecbacf2f12e1a1442ccf52a8ebdd292fe
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 035f457..7c20394 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -383,6 +383,8 @@
         framer_.framer()->InstallDecrypter(
             ENCRYPTION_INITIAL, std::make_unique<TaggingDecrypter>());
         framer_.framer()->InstallDecrypter(
+            ENCRYPTION_HANDSHAKE, std::make_unique<TaggingDecrypter>());
+        framer_.framer()->InstallDecrypter(
             ENCRYPTION_ZERO_RTT, std::make_unique<TaggingDecrypter>());
         framer_.framer()->InstallDecrypter(
             ENCRYPTION_FORWARD_SECURE, std::make_unique<TaggingDecrypter>());
@@ -763,18 +765,24 @@
 
   QuicConsumedData SendCryptoDataWithString(quiche::QuicheStringPiece data,
                                             QuicStreamOffset offset) {
+    return SendCryptoDataWithString(data, offset, ENCRYPTION_INITIAL);
+  }
+
+  QuicConsumedData SendCryptoDataWithString(quiche::QuicheStringPiece data,
+                                            QuicStreamOffset offset,
+                                            EncryptionLevel encryption_level) {
     if (!QuicVersionUsesCryptoFrames(transport_version())) {
       return SendStreamDataWithString(
           QuicUtils::GetCryptoStreamId(transport_version()), data, offset,
           NO_FIN);
     }
-    producer_.SaveCryptoData(ENCRYPTION_INITIAL, offset, data);
+    producer_.SaveCryptoData(encryption_level, offset, data);
     size_t bytes_written;
     if (notifier_) {
       bytes_written =
-          notifier_->WriteCryptoData(ENCRYPTION_INITIAL, data.length(), offset);
+          notifier_->WriteCryptoData(encryption_level, data.length(), offset);
     } else {
-      bytes_written = QuicConnection::SendCryptoData(ENCRYPTION_INITIAL,
+      bytes_written = QuicConnection::SendCryptoData(encryption_level,
                                                      data.length(), offset);
     }
     return QuicConsumedData(bytes_written, /*fin_consumed*/ false);
@@ -872,7 +880,7 @@
     if (QuicConnectionPeer::GetSentPacketManager(this)->pto_enabled()) {
       // PTO mode is default enabled for T099. And TLP/RTO related tests are
       // stale.
-      DCHECK_EQ(ParsedQuicVersion(PROTOCOL_TLS1_3, QUIC_VERSION_99), version());
+      DCHECK_EQ(PROTOCOL_TLS1_3, version().handshake_protocol);
       return true;
     }
     return false;
@@ -9736,6 +9744,57 @@
               IsError(IETF_QUIC_PROTOCOL_VIOLATION));
 }
 
+TEST_P(QuicConnectionTest, MultiplePacketNumberSpacePto) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  use_tagging_decrypter();
+  // Send handshake packet.
+  connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+  connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
+  EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+
+  // Send application data.
+  connection_.SendApplicationDataAtLevel(ENCRYPTION_FORWARD_SECURE, 5, "data",
+                                         0, NO_FIN);
+  EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
+  QuicTime retransmission_time =
+      connection_.GetRetransmissionAlarm()->deadline();
+  EXPECT_NE(QuicTime::Zero(), retransmission_time);
+
+  // Retransmit handshake data.
+  clock_.AdvanceTime(retransmission_time - clock_.Now());
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, QuicPacketNumber(3), _, _));
+  connection_.GetRetransmissionAlarm()->Fire();
+  EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+
+  // Send application data.
+  connection_.SendApplicationDataAtLevel(ENCRYPTION_FORWARD_SECURE, 5, "data",
+                                         4, NO_FIN);
+  EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
+  retransmission_time = connection_.GetRetransmissionAlarm()->deadline();
+  EXPECT_NE(QuicTime::Zero(), retransmission_time);
+
+  // Retransmit handshake data again.
+  clock_.AdvanceTime(retransmission_time - clock_.Now());
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, QuicPacketNumber(5), _, _));
+  connection_.GetRetransmissionAlarm()->Fire();
+  EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+
+  // Discard handshake key.
+  connection_.OnHandshakeComplete();
+  retransmission_time = connection_.GetRetransmissionAlarm()->deadline();
+  EXPECT_NE(QuicTime::Zero(), retransmission_time);
+
+  // Retransmit application data.
+  clock_.AdvanceTime(retransmission_time - clock_.Now());
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, QuicPacketNumber(6), _, _));
+  connection_.GetRetransmissionAlarm()->Fire();
+  EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_sent_packet_manager.cc b/quic/core/quic_sent_packet_manager.cc
index 767566c..9246499 100644
--- a/quic/core/quic_sent_packet_manager.cc
+++ b/quic/core/quic_sent_packet_manager.cc
@@ -493,6 +493,31 @@
   return true;
 }
 
+QuicTime QuicSentPacketManager::GetEarliestPacketSentTimeForPto(
+    PacketNumberSpace* packet_number_space) const {
+  DCHECK(supports_multiple_packet_number_spaces());
+  QuicTime earliest_sent_time = QuicTime::Zero();
+  for (int8_t i = 0; i < NUM_PACKET_NUMBER_SPACES; ++i) {
+    const QuicTime sent_time = unacked_packets_.GetLastInFlightPacketSentTime(
+        static_cast<PacketNumberSpace>(i));
+    if (!earliest_sent_time.IsInitialized()) {
+      earliest_sent_time = sent_time;
+      *packet_number_space = static_cast<PacketNumberSpace>(i);
+      continue;
+    }
+    if (i == APPLICATION_DATA) {
+      // Not arming PTO for application data to prioritize the completion of
+      // handshake.
+      break;
+    }
+    if (sent_time.IsInitialized()) {
+      earliest_sent_time = std::min(earliest_sent_time, sent_time);
+      *packet_number_space = static_cast<PacketNumberSpace>(i);
+    }
+  }
+  return earliest_sent_time;
+}
+
 void QuicSentPacketManager::MarkForRetransmission(
     QuicPacketNumber packet_number,
     TransmissionType transmission_type) {
@@ -776,12 +801,25 @@
   if (pending_timer_transmission_count_ == 0) {
     return;
   }
+  PacketNumberSpace packet_number_space;
+  if (supports_multiple_packet_number_spaces()) {
+    // Find out the packet number space to send probe packets.
+    if (!GetEarliestPacketSentTimeForPto(&packet_number_space)
+             .IsInitialized()) {
+      QUIC_BUG << "earlist_sent_time not initialized when trying to send PTO "
+                  "retransmissions";
+      return;
+    }
+  }
   QuicPacketNumber packet_number = unacked_packets_.GetLeastUnacked();
   std::vector<QuicPacketNumber> probing_packets;
   for (QuicUnackedPacketMap::const_iterator it = unacked_packets_.begin();
        it != unacked_packets_.end(); ++it, ++packet_number) {
     if (it->state == OUTSTANDING &&
-        unacked_packets_.HasRetransmittableFrames(*it)) {
+        unacked_packets_.HasRetransmittableFrames(*it) &&
+        (!supports_multiple_packet_number_spaces() ||
+         unacked_packets_.GetPacketNumberSpace(it->encryption_level) ==
+             packet_number_space)) {
       DCHECK(it->in_flight);
       probing_packets.push_back(packet_number);
       if (probing_packets.size() == pending_timer_transmission_count_) {
@@ -948,9 +986,16 @@
       return std::max(tlp_time, rto_time);
     }
     case PTO_MODE: {
-      // Ensure PTO never gets set to a time in the past.
+      if (!supports_multiple_packet_number_spaces()) {
+        // Ensure PTO never gets set to a time in the past.
+        return std::max(clock_->ApproximateNow(),
+                        unacked_packets_.GetLastInFlightPacketSentTime() +
+                            GetProbeTimeoutDelay());
+      }
+
+      PacketNumberSpace packet_number_space;
       return std::max(clock_->ApproximateNow(),
-                      unacked_packets_.GetLastInFlightPacketSentTime() +
+                      GetEarliestPacketSentTimeForPto(&packet_number_space) +
                           GetProbeTimeoutDelay());
     }
   }
@@ -1267,6 +1312,7 @@
 }
 
 void QuicSentPacketManager::EnableMultiplePacketNumberSpacesSupport() {
+  EnableIetfPtoAndLossDetection();
   unacked_packets_.EnableMultiplePacketNumberSpacesSupport();
 }
 
diff --git a/quic/core/quic_sent_packet_manager.h b/quic/core/quic_sent_packet_manager.h
index ab53066..f70b525 100644
--- a/quic/core/quic_sent_packet_manager.h
+++ b/quic/core/quic_sent_packet_manager.h
@@ -511,6 +511,11 @@
   // timeout.
   bool ShouldAddMaxAckDelay() const;
 
+  // Gets the earliest in flight packet sent time to calculate PTO. Also
+  // updates |packet_number_space| if a PTO timer should be armed.
+  QuicTime GetEarliestPacketSentTimeForPto(
+      PacketNumberSpace* packet_number_space) const;
+
   // Newly serialized retransmittable packets are added to this map, which
   // contains owning pointers to any contained frames.  If a packet is
   // retransmitted, this map will contain entries for both the old and the new
diff --git a/quic/core/quic_sent_packet_manager_test.cc b/quic/core/quic_sent_packet_manager_test.cc
index 49d0b1d..247c03f 100644
--- a/quic/core/quic_sent_packet_manager_test.cc
+++ b/quic/core/quic_sent_packet_manager_test.cc
@@ -3294,6 +3294,119 @@
   EXPECT_EQ(10u, manager_.initial_congestion_window());
 }
 
+TEST_F(QuicSentPacketManagerTest, ClientMultiplePacketNumberSpacePtoTimeout) {
+  manager_.EnableMultiplePacketNumberSpacesSupport();
+  EnablePto(k1PTO);
+  EXPECT_CALL(*send_algorithm_, PacingRate(_))
+      .WillRepeatedly(Return(QuicBandwidth::Zero()));
+  EXPECT_CALL(*send_algorithm_, GetCongestionWindow())
+      .WillRepeatedly(Return(10 * kDefaultTCPMSS));
+  RttStats* rtt_stats = const_cast<RttStats*>(manager_.GetRttStats());
+  rtt_stats->UpdateRtt(QuicTime::Delta::FromMilliseconds(100),
+                       QuicTime::Delta::Zero(), QuicTime::Zero());
+  QuicTime::Delta srtt = rtt_stats->smoothed_rtt();
+
+  // Send packet 1.
+  SendDataPacket(1, ENCRYPTION_INITIAL);
+  // Verify PTO is correctly set.
+  QuicTime::Delta expected_pto_delay =
+      srtt + 4 * rtt_stats->mean_deviation() +
+      QuicTime::Delta::FromMilliseconds(kDefaultDelayedAckTimeMs);
+  EXPECT_EQ(clock_.Now() + expected_pto_delay,
+            manager_.GetRetransmissionTime());
+
+  // Discard initial key and send packet 2 in handshake.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  EXPECT_CALL(notifier_, IsFrameOutstanding(_)).WillRepeatedly(Return(false));
+  manager_.NeuterUnencryptedPackets();
+
+  EXPECT_CALL(notifier_, IsFrameOutstanding(_)).WillRepeatedly(Return(true));
+  SendDataPacket(2, ENCRYPTION_HANDSHAKE);
+  // Verify PTO is correctly set based on sent time of packet 2.
+  EXPECT_EQ(clock_.Now() + expected_pto_delay,
+            manager_.GetRetransmissionTime());
+  // Invoke PTO.
+  clock_.AdvanceTime(expected_pto_delay);
+  manager_.OnRetransmissionTimeout();
+  EXPECT_EQ(QuicTime::Delta::Zero(), manager_.TimeUntilSend(clock_.Now()));
+  EXPECT_EQ(1u, stats_.pto_count);
+
+  // Verify probe packet gets sent.
+  EXPECT_CALL(notifier_, RetransmitFrames(_, _))
+      .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
+        RetransmitDataPacket(3, type, ENCRYPTION_HANDSHAKE);
+      })));
+  manager_.MaybeSendProbePackets();
+  // Verify PTO period gets set to twice the current value.
+  const QuicTime packet3_sent_time = clock_.Now();
+  EXPECT_EQ(packet3_sent_time + expected_pto_delay * 2,
+            manager_.GetRetransmissionTime());
+
+  // Send packet 4 in application data.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  SendDataPacket(4, ENCRYPTION_FORWARD_SECURE);
+  const QuicTime packet4_sent_time = clock_.Now();
+  // Verify PTO timeout is still based on packet 3.
+  EXPECT_EQ(packet3_sent_time + expected_pto_delay * 2,
+            manager_.GetRetransmissionTime());
+
+  // Send packet 5 in handshake.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  SendDataPacket(5, ENCRYPTION_HANDSHAKE);
+  // Verify PTO timeout is now based on packet 5.
+  EXPECT_EQ(clock_.Now() + expected_pto_delay * 2,
+            manager_.GetRetransmissionTime());
+
+  // Neuter handshake key.
+  manager_.SetHandshakeConfirmed();
+  // Verify PTO timeout is now based on packet 4.
+  EXPECT_EQ(packet4_sent_time + expected_pto_delay * 2,
+            manager_.GetRetransmissionTime());
+}
+
+TEST_F(QuicSentPacketManagerTest, ServerMultiplePacketNumberSpacePtoTimeout) {
+  manager_.EnableMultiplePacketNumberSpacesSupport();
+  EnablePto(k1PTO);
+  EXPECT_CALL(*send_algorithm_, PacingRate(_))
+      .WillRepeatedly(Return(QuicBandwidth::Zero()));
+  EXPECT_CALL(*send_algorithm_, GetCongestionWindow())
+      .WillRepeatedly(Return(10 * kDefaultTCPMSS));
+  RttStats* rtt_stats = const_cast<RttStats*>(manager_.GetRttStats());
+  rtt_stats->UpdateRtt(QuicTime::Delta::FromMilliseconds(100),
+                       QuicTime::Delta::Zero(), QuicTime::Zero());
+  QuicTime::Delta srtt = rtt_stats->smoothed_rtt();
+
+  // Send packet 1.
+  SendDataPacket(1, ENCRYPTION_INITIAL);
+  const QuicTime packet1_sent_time = clock_.Now();
+  // Verify PTO is correctly set.
+  QuicTime::Delta expected_pto_delay =
+      srtt + 4 * rtt_stats->mean_deviation() +
+      QuicTime::Delta::FromMilliseconds(kDefaultDelayedAckTimeMs);
+  EXPECT_EQ(packet1_sent_time + expected_pto_delay,
+            manager_.GetRetransmissionTime());
+
+  // Send packet 2 in handshake.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  SendDataPacket(2, ENCRYPTION_HANDSHAKE);
+  // Verify PTO timeout is still based on packet 1.
+  EXPECT_EQ(packet1_sent_time + expected_pto_delay,
+            manager_.GetRetransmissionTime());
+
+  // Discard both initial and handshake keys.
+  EXPECT_CALL(notifier_, IsFrameOutstanding(_)).WillRepeatedly(Return(false));
+  manager_.NeuterUnencryptedPackets();
+  manager_.SetHandshakeConfirmed();
+  EXPECT_EQ(QuicTime::Zero(), manager_.GetRetransmissionTime());
+
+  // Send packet 3 in application data.
+  clock_.AdvanceTime(QuicTime::Delta::FromMilliseconds(10));
+  SendDataPacket(3, ENCRYPTION_FORWARD_SECURE);
+  // Verify PTO timeout is correctly set.
+  EXPECT_EQ(clock_.Now() + expected_pto_delay,
+            manager_.GetRetransmissionTime());
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_unacked_packet_map.cc b/quic/core/quic_unacked_packet_map.cc
index 50f5c85..ef08535 100644
--- a/quic/core/quic_unacked_packet_map.cc
+++ b/quic/core/quic_unacked_packet_map.cc
@@ -30,6 +30,9 @@
       bytes_in_flight_(0),
       packets_in_flight_(0),
       last_inflight_packet_sent_time_(QuicTime::Zero()),
+      last_inflight_packets_sent_time_{{QuicTime::Zero()},
+                                       {QuicTime::Zero()},
+                                       {QuicTime::Zero()}},
       last_crypto_packet_sent_time_(QuicTime::Zero()),
       session_notifier_(nullptr),
       supports_multiple_packet_number_spaces_(false) {}
@@ -66,14 +69,14 @@
 
   largest_sent_packet_ = packet_number;
   if (set_in_flight) {
+    const PacketNumberSpace packet_number_space =
+        GetPacketNumberSpace(info.encryption_level);
     bytes_in_flight_ += bytes_sent;
     ++packets_in_flight_;
     info.in_flight = true;
-    largest_sent_retransmittable_packets_[GetPacketNumberSpace(
-        info.encryption_level)] = packet_number;
-    // TODO(ianswett): Should this field be per packet number space or should
-    // GetInFlightPacketSentTime() use largest_sent_retransmittable_packets_?
+    largest_sent_retransmittable_packets_[packet_number_space] = packet_number;
     last_inflight_packet_sent_time_ = sent_time;
+    last_inflight_packets_sent_time_[packet_number_space] = sent_time;
   }
   unacked_packets_.push_back(info);
   // Swap the retransmittable frames to avoid allocations.
@@ -218,6 +221,9 @@
       DCHECK(!HasRetransmittableFrames(*it));
     }
   }
+  if (supports_multiple_packet_number_spaces_) {
+    last_inflight_packets_sent_time_[INITIAL_DATA] = QuicTime::Zero();
+  }
 }
 
 void QuicUnackedPacketMap::NeuterHandshakePackets() {
@@ -233,6 +239,9 @@
       NotifyFramesAcked(*it, QuicTime::Delta::Zero(), QuicTime::Zero());
     }
   }
+  if (supports_multiple_packet_number_spaces()) {
+    last_inflight_packets_sent_time_[HANDSHAKE_DATA] = QuicTime::Zero();
+  }
 }
 
 bool QuicUnackedPacketMap::HasInFlightPackets() const {
@@ -428,6 +437,15 @@
   return largest_acked_packets_[packet_number_space];
 }
 
+QuicTime QuicUnackedPacketMap::GetLastInFlightPacketSentTime(
+    PacketNumberSpace packet_number_space) const {
+  if (packet_number_space >= NUM_PACKET_NUMBER_SPACES) {
+    QUIC_BUG << "Invalid packet number space: " << packet_number_space;
+    return QuicTime::Zero();
+  }
+  return last_inflight_packets_sent_time_[packet_number_space];
+}
+
 QuicPacketNumber
 QuicUnackedPacketMap::GetLargestSentRetransmittableOfPacketNumberSpace(
     PacketNumberSpace packet_number_space) const {
diff --git a/quic/core/quic_unacked_packet_map.h b/quic/core/quic_unacked_packet_map.h
index 61261a5..1cad5eb 100644
--- a/quic/core/quic_unacked_packet_map.h
+++ b/quic/core/quic_unacked_packet_map.h
@@ -212,6 +212,10 @@
   QuicPacketNumber GetLargestSentPacketOfPacketNumberSpace(
       EncryptionLevel encryption_level) const;
 
+  // Returns last in flight packet sent time of |packet_number_space|.
+  QuicTime GetLastInFlightPacketSentTime(
+      PacketNumberSpace packet_number_space) const;
+
   void SetSessionNotifier(SessionNotifierInterface* session_notifier);
 
   void EnableMultiplePacketNumberSpacesSupport();
@@ -273,6 +277,8 @@
 
   // Time that the last inflight packet was sent.
   QuicTime last_inflight_packet_sent_time_;
+  // Time that the last in flight packet was sent per packet number space.
+  QuicTime last_inflight_packets_sent_time_[NUM_PACKET_NUMBER_SPACES];
 
   // Time that the last unacked crypto packet was sent.
   QuicTime last_crypto_packet_sent_time_;