In quic, coalesce packet of higher encryption level with initial or handshake retransmission. protected by gfe2_reloadable_flag_quic_coalesced_packet_of_higher_space.

The goal of this change is to mitigate inflated RTT sample.

PiperOrigin-RevId: 321556627
Change-Id: I7fc8675bb0f05a48b6e81ed4aa9c04b42fb0690a
diff --git a/quic/core/quic_coalesced_packet.cc b/quic/core/quic_coalesced_packet.cc
index 3b22180..b92cf75 100644
--- a/quic/core/quic_coalesced_packet.cc
+++ b/quic/core/quic_coalesced_packet.cc
@@ -69,6 +69,7 @@
     QUIC_CODE_COUNT(QUIC_SUCCESSFULLY_COALESCED_MULTIPLE_PACKETS);
   }
   length_ += packet.encrypted_length;
+  transmission_types_[packet.encryption_level] = packet.transmission_type;
   if (packet.encryption_level == ENCRYPTION_INITIAL) {
     // Save a copy of ENCRYPTION_INITIAL packet (excluding encrypted buffer, as
     // the packet will be re-serialized later).
@@ -90,6 +91,9 @@
   for (auto& packet : encrypted_buffers_) {
     packet.clear();
   }
+  for (size_t i = ENCRYPTION_INITIAL; i < NUM_ENCRYPTION_LEVELS; ++i) {
+    transmission_types_[i] = NOT_RETRANSMISSION;
+  }
   initial_packet_ = nullptr;
 }
 
@@ -118,6 +122,16 @@
          (level == ENCRYPTION_INITIAL && initial_packet_ != nullptr);
 }
 
+TransmissionType QuicCoalescedPacket::TransmissionTypeOfPacket(
+    EncryptionLevel level) const {
+  if (!ContainsPacketOfEncryptionLevel(level)) {
+    QUIC_BUG << "Coalesced packet does not contain packet of encryption level: "
+             << EncryptionLevelToString(level);
+    return NOT_RETRANSMISSION;
+  }
+  return transmission_types_[level];
+}
+
 std::string QuicCoalescedPacket::ToString(size_t serialized_length) const {
   // Total length and padding size.
   std::string info = quiche::QuicheStrCat(
diff --git a/quic/core/quic_coalesced_packet.h b/quic/core/quic_coalesced_packet.h
index 447ab95..6cf712d 100644
--- a/quic/core/quic_coalesced_packet.h
+++ b/quic/core/quic_coalesced_packet.h
@@ -39,6 +39,10 @@
   // Returns true if this coalesced packet contains packet of |level|.
   bool ContainsPacketOfEncryptionLevel(EncryptionLevel level) const;
 
+  // Returns transmission type of packet of |level|. This should only be called
+  // when this coalesced packet contains packet of |level|.
+  TransmissionType TransmissionTypeOfPacket(EncryptionLevel level) const;
+
   const SerializedPacket* initial_packet() const {
     return initial_packet_.get();
   }
@@ -64,6 +68,8 @@
   // Copies of packets' encrypted buffers according to different encryption
   // levels.
   std::string encrypted_buffers_[NUM_ENCRYPTION_LEVELS];
+  // Recorded transmission type according to different encryption levels.
+  TransmissionType transmission_types_[NUM_ENCRYPTION_LEVELS];
 
   // A copy of ENCRYPTION_INITIAL packet if this coalesced packet contains one.
   // Null otherwise. Please note, the encrypted_buffer field is not copied. The
diff --git a/quic/core/quic_coalesced_packet_test.cc b/quic/core/quic_coalesced_packet_test.cc
index 0a7a51e..81e6b27 100644
--- a/quic/core/quic_coalesced_packet_test.cc
+++ b/quic/core/quic_coalesced_packet_test.cc
@@ -24,12 +24,15 @@
   QuicSocketAddress peer_address(QuicIpAddress::Loopback4(), 2);
   SerializedPacket packet1(QuicPacketNumber(1), PACKET_4BYTE_PACKET_NUMBER,
                            buffer, 500, false, false);
+  packet1.transmission_type = PTO_RETRANSMISSION;
   QuicAckFrame ack_frame(InitAckFrame(1));
   packet1.nonretransmittable_frames.push_back(QuicFrame(&ack_frame));
   packet1.retransmittable_frames.push_back(
       QuicFrame(QuicStreamFrame(1, true, 0, 100)));
   ASSERT_TRUE(coalesced.MaybeCoalescePacket(packet1, self_address, peer_address,
                                             &allocator, 1500));
+  EXPECT_EQ(PTO_RETRANSMISSION,
+            coalesced.TransmissionTypeOfPacket(ENCRYPTION_INITIAL));
   EXPECT_EQ(1500u, coalesced.max_packet_length());
   EXPECT_EQ(500u, coalesced.length());
   EXPECT_EQ(
@@ -46,10 +49,13 @@
                            buffer, 500, false, false);
   packet3.nonretransmittable_frames.push_back(QuicFrame(QuicPaddingFrame(100)));
   packet3.encryption_level = ENCRYPTION_ZERO_RTT;
+  packet3.transmission_type = LOSS_RETRANSMISSION;
   ASSERT_TRUE(coalesced.MaybeCoalescePacket(packet3, self_address, peer_address,
                                             &allocator, 1500));
   EXPECT_EQ(1500u, coalesced.max_packet_length());
   EXPECT_EQ(1000u, coalesced.length());
+  EXPECT_EQ(LOSS_RETRANSMISSION,
+            coalesced.TransmissionTypeOfPacket(ENCRYPTION_ZERO_RTT));
   EXPECT_EQ(
       "total_length: 1500 padding_size: 500 packets: {ENCRYPTION_INITIAL, "
       "ENCRYPTION_ZERO_RTT}",
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index e60d22b..63a79db 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2431,6 +2431,12 @@
     return false;
   }
 
+  if (fill_coalesced_packet_) {
+    // Try to coalesce packet, only allow to write when creator is on soft max
+    // packet length.
+    return packet_creator_.HasSoftMaxPacketLength();
+  }
+
   if (sent_packet_manager_.pending_timer_transmission_count() > 0) {
     // Force sending the retransmissions for HANDSHAKE, TLP, RTO, PROBING cases.
     return true;
@@ -3647,6 +3653,10 @@
     }
     connection_->packet_creator_.Flush();
     if (connection_->version().CanSendCoalescedPackets()) {
+      if (connection_->packet_creator().coalesced_packet_of_higher_space()) {
+        QUIC_RELOADABLE_FLAG_COUNT(quic_coalesced_packet_of_higher_space);
+        connection_->MaybeCoalescePacketOfHigherSpace();
+      }
       connection_->FlushCoalescedPacket();
     }
     connection_->FlushPackets();
@@ -4339,7 +4349,37 @@
   return false;
 }
 
+void QuicConnection::MaybeCoalescePacketOfHigherSpace() {
+  if (!packet_creator_.HasSoftMaxPacketLength()) {
+    return;
+  }
+  // INITIAL or HANDSHAKE retransmission could cause peer to derive new
+  // keys, such that the buffered undecryptable packets may be processed.
+  // This endpoint would derive an inflated RTT sample (which includes the PTO
+  // timeout) when receiving ACKs of those undecryptable packets. To mitigate
+  // this, tries to coalesce a packet of higher encryption level.
+  for (EncryptionLevel retransmission_level :
+       {ENCRYPTION_INITIAL, ENCRYPTION_HANDSHAKE}) {
+    // Coalesce HANDSHAKE with INITIAL retransmission, and coalesce 1-RTT with
+    // HANDSHAKE retransmission.
+    const EncryptionLevel coalesced_level =
+        retransmission_level == ENCRYPTION_INITIAL ? ENCRYPTION_HANDSHAKE
+                                                   : ENCRYPTION_FORWARD_SECURE;
+    if (coalesced_packet_.ContainsPacketOfEncryptionLevel(
+            retransmission_level) &&
+        coalesced_packet_.TransmissionTypeOfPacket(retransmission_level) !=
+            NOT_RETRANSMISSION &&
+        framer_.HasEncrypterOfEncryptionLevel(coalesced_level) &&
+        !coalesced_packet_.ContainsPacketOfEncryptionLevel(coalesced_level)) {
+      fill_coalesced_packet_ = true;
+      sent_packet_manager_.RetransmitDataOfSpaceIfAny(
+          QuicUtils::GetPacketNumberSpace(coalesced_level));
+    }
+  }
+}
+
 bool QuicConnection::FlushCoalescedPacket() {
+  fill_coalesced_packet_ = false;
   ScopedCoalescedPacketClearer clearer(&coalesced_packet_);
   if (!version().CanSendCoalescedPackets()) {
     QUIC_BUG_IF(coalesced_packet_.length() > 0);
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 23f2d12..36ef402 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1248,6 +1248,9 @@
   // Called to update ACK timeout when an retransmittable frame has been parsed.
   void MaybeUpdateAckTimeout();
 
+  // Tries to fill coalesced packet with data of higher packet space.
+  void MaybeCoalescePacketOfHigherSpace();
+
   // Serialize and send coalesced_packet. Returns false if serialization fails
   // or the write causes errors, otherwise, returns true.
   bool FlushCoalescedPacket();
@@ -1680,6 +1683,9 @@
   bool legacy_version_encapsulation_in_progress_ = false;
   // SNI to send when using Legacy Version Encapsulation.
   std::string legacy_version_encapsulation_sni_;
+  // True if next packet is intended to consume remaining space in the
+  // coalescer.
+  bool fill_coalesced_packet_ = false;
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 62aa83f..de7e450 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -10399,7 +10399,12 @@
                            _, _));
   EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
   connection_.GetRetransmissionAlarm()->Fire();
-  EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    // Verify 1-RTT packet gets coalesced with handshake retransmission.
+    EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
+  } else {
+    EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+  }
 
   // Send application data.
   connection_.SendApplicationDataAtLevel(ENCRYPTION_FORWARD_SECURE, 5, "data",
@@ -10410,15 +10415,24 @@
 
   // Retransmit handshake data again.
   clock_.AdvanceTime(retransmission_time - clock_.Now());
+  QuicPacketNumber handshake_retransmission =
+      GetQuicReloadableFlag(quic_default_on_pto) ? QuicPacketNumber(5)
+                                                 : QuicPacketNumber(7);
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    handshake_retransmission += 1;
+    EXPECT_CALL(*send_algorithm_,
+                OnPacketSent(_, _, handshake_retransmission + 1, _, _));
+  }
   EXPECT_CALL(*send_algorithm_,
-              OnPacketSent(_, _,
-                           GetQuicReloadableFlag(quic_default_on_pto)
-                               ? QuicPacketNumber(5)
-                               : QuicPacketNumber(7),
-                           _, _));
+              OnPacketSent(_, _, handshake_retransmission, _, _));
   EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
   connection_.GetRetransmissionAlarm()->Fire();
-  EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    // Verify 1-RTT packet gets coalesced with handshake retransmission.
+    EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
+  } else {
+    EXPECT_EQ(0x02020202u, writer_->final_bytes_of_last_packet());
+  }
 
   // Discard handshake key.
   connection_.OnHandshakeComplete();
@@ -10427,12 +10441,14 @@
 
   // Retransmit application data.
   clock_.AdvanceTime(retransmission_time - clock_.Now());
+  QuicPacketNumber application_retransmission =
+      GetQuicReloadableFlag(quic_default_on_pto) ? QuicPacketNumber(6)
+                                                 : QuicPacketNumber(9);
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    application_retransmission += 2;
+  }
   EXPECT_CALL(*send_algorithm_,
-              OnPacketSent(_, _,
-                           GetQuicReloadableFlag(quic_default_on_pto)
-                               ? QuicPacketNumber(6)
-                               : QuicPacketNumber(9),
-                           _, _));
+              OnPacketSent(_, _, application_retransmission, _, _));
   connection_.GetRetransmissionAlarm()->Fire();
   EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
 }
@@ -11217,7 +11233,11 @@
   connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
                            std::make_unique<TaggingEncrypter>(0x02));
   connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
-  EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+  } else {
+    EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+  }
   connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
   // Verify PTO time does not change.
   EXPECT_EQ(expected_pto_time,
@@ -11356,6 +11376,89 @@
   }
 }
 
+// Regression test for b/161228202
+TEST_P(QuicConnectionTest, InflatedRttSample) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  // 30ms RTT.
+  const QuicTime::Delta kTestRTT = QuicTime::Delta::FromMilliseconds(30);
+  set_perspective(Perspective::IS_SERVER);
+  RttStats* rtt_stats = const_cast<RttStats*>(manager_->GetRttStats());
+  use_tagging_decrypter();
+  // Receives packet 1000 in initial data.
+  if (QuicVersionUsesCryptoFrames(connection_.transport_version())) {
+    EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
+  }
+  EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(AnyNumber());
+  ProcessCryptoPacketAtLevel(1000, ENCRYPTION_INITIAL);
+  EXPECT_TRUE(connection_.HasPendingAcks());
+
+  connection_.SetEncrypter(ENCRYPTION_INITIAL,
+                           std::make_unique<TaggingEncrypter>(0x01));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_INITIAL);
+  // Send INITIAL 1.
+  connection_.SendCryptoDataWithString(std::string(512, 'a'), 0,
+                                       ENCRYPTION_INITIAL);
+  ASSERT_TRUE(connection_.sent_packet_manager()
+                  .GetRetransmissionTime()
+                  .IsInitialized());
+  QuicTime::Delta pto_timeout =
+      connection_.sent_packet_manager().GetRetransmissionTime() - clock_.Now();
+  // Send Handshake 2.
+  connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    // Verify HANDSHAKE packet is coalesced with INITIAL retransmission.
+    EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+  } else {
+    EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+  }
+  connection_.SendCryptoDataWithString(std::string(1024, 'a'), 0,
+                                       ENCRYPTION_HANDSHAKE);
+
+  // INITIAL 1 gets lost and PTO fires.
+  clock_.AdvanceTime(pto_timeout);
+  connection_.GetRetransmissionAlarm()->Fire();
+
+  clock_.AdvanceTime(kTestRTT);
+  // Assume retransmitted INITIAL gets received.
+  QuicFrames frames;
+  QuicPacketNumber initial_retransmission =
+      GetQuicReloadableFlag(quic_default_on_pto) ? QuicPacketNumber(3)
+                                                 : QuicPacketNumber(4);
+  auto ack_frame =
+      InitAckFrame({{initial_retransmission, initial_retransmission + 1}});
+  frames.push_back(QuicFrame(&ack_frame));
+  EXPECT_CALL(*send_algorithm_, OnCongestionEvent(_, _, _, _, _))
+      .Times(AnyNumber());
+  ProcessFramesPacketAtLevel(1001, frames, ENCRYPTION_INITIAL);
+  EXPECT_EQ(kTestRTT, rtt_stats->latest_rtt());
+  // Because retransmitted INITIAL gets received so HANDSHAKE 2 gets processed.
+  frames.clear();
+  QuicAckFrame ack_frame2;
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    // HANDSHAKE 5 is also processed.
+    ack_frame2 = InitAckFrame(
+        {{QuicPacketNumber(2), QuicPacketNumber(3)},
+         {initial_retransmission + 1, initial_retransmission + 2}});
+  } else {
+    ack_frame2 = InitAckFrame({{QuicPacketNumber(2), QuicPacketNumber(3)}});
+  }
+  ack_frame2.ack_delay_time = QuicTime::Delta::Zero();
+  frames.push_back(QuicFrame(&ack_frame2));
+  ProcessFramesPacketAtLevel(1, frames, ENCRYPTION_HANDSHAKE);
+  if (GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space)) {
+    // Verify RTT inflation gets mitigated.
+    EXPECT_EQ(rtt_stats->latest_rtt(), kTestRTT);
+  } else {
+    // Verify this RTT sample gets inflated as it includes the PTO timeout and
+    // the actual RTT.
+    EXPECT_GE(rtt_stats->latest_rtt(), pto_timeout + kTestRTT);
+  }
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index 066763a..131ef5e 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -203,6 +203,9 @@
     // Please note: this would not guarantee to fit next packet if the size of
     // packet header increases (e.g., encryption level changes).
     QUIC_DLOG(INFO) << length << " is too small to fit packet header";
+    if (coalesced_packet_of_higher_space_) {
+      RemoveSoftMaxPacketLength();
+    }
     return;
   }
   QUIC_DVLOG(1) << "Setting soft max packet length to: " << length;
@@ -2018,5 +2021,9 @@
   return flusher_attached_;
 }
 
+bool QuicPacketCreator::HasSoftMaxPacketLength() const {
+  return latched_hard_max_packet_length_ != 0;
+}
+
 #undef ENDPOINT  // undef for jumbo builds
 }  // namespace quic
diff --git a/quic/core/quic_packet_creator.h b/quic/core/quic_packet_creator.h
index 80e35be..976f4f6 100644
--- a/quic/core/quic_packet_creator.h
+++ b/quic/core/quic_packet_creator.h
@@ -445,6 +445,9 @@
                                   char* buffer,
                                   size_t buffer_len);
 
+  // Returns true if max_packet_length_ is currently a soft value.
+  bool HasSoftMaxPacketLength() const;
+
   void set_disable_padding_override(bool should_disable_padding) {
     disable_padding_override_ = should_disable_padding;
   }
@@ -453,6 +456,10 @@
     return determine_serialized_packet_fate_early_;
   }
 
+  bool coalesced_packet_of_higher_space() const {
+    return coalesced_packet_of_higher_space_;
+  }
+
  private:
   friend class test::QuicPacketCreatorPeer;
 
@@ -648,6 +655,9 @@
 
   const bool determine_serialized_packet_fate_early_ =
       GetQuicReloadableFlag(quic_determine_serialized_packet_fate_early);
+
+  const bool coalesced_packet_of_higher_space_ =
+      GetQuicReloadableFlag(quic_coalesced_packet_of_higher_space);
 };
 
 }  // namespace quic