Preempt stream data with handshake packet before handshake gets confirmed.

Protected by FLAGS_quic_reloadable_flag_quic_preempt_stream_data_with_handshake_packet.

PiperOrigin-RevId: 360493291
Change-Id: I5d827932b223621eff3b7384f91896cfe03a22d6
diff --git a/quic/core/http/end_to_end_test.cc b/quic/core/http/end_to_end_test.cc
index 7e9a679..1b300e9 100644
--- a/quic/core/http/end_to_end_test.cc
+++ b/quic/core/http/end_to_end_test.cc
@@ -2385,6 +2385,50 @@
   SendSynchronousBarRequestAndCheckResponse();
 }
 
+// Regression test for b/180737158.
+TEST_P(
+    EndToEndTest,
+    HalfRttResponseBlocksShloRetransmissionWithoutTokenBasedAddressValidation) {
+  // Turn off token based address validation to make the server get constrained
+  // by amplification factor during handshake.
+  // TODO(fayang): Keep this test while deprecating
+  // quic_enable_token_based_address_validation. For example, consider always
+  // rejecting the received address token.
+  SetQuicReloadableFlag(quic_enable_token_based_address_validation, false);
+  ASSERT_TRUE(Initialize());
+  if (!version_.SupportsAntiAmplificationLimit()) {
+    return;
+  }
+  // Perform a full 1-RTT handshake to get the new session ticket such that the
+  // next connection will perform a 0-RTT handshake.
+  EXPECT_TRUE(client_->client()->WaitForHandshakeConfirmed());
+  client_->Disconnect();
+
+  server_thread_->Pause();
+  // Drop the 1st server packet which is the coalesced INITIAL + HANDSHAKE +
+  // 1RTT.
+  PacketDroppingTestWriter* writer = new PacketDroppingTestWriter();
+  writer->set_fake_drop_first_n_packets(1);
+  QuicDispatcherPeer::UseWriter(
+      QuicServerPeer::GetDispatcher(server_thread_->server()), writer);
+  server_thread_->Resume();
+
+  // Large response (100KB) for 0-RTT request.
+  std::string large_body(102400, 'a');
+  AddToCache("/large_response", 200, large_body);
+  if (GetQuicReloadableFlag(quic_preempt_stream_data_with_handshake_packet)) {
+    SendSynchronousRequestAndCheckResponse(client_.get(), "/large_response",
+                                           large_body);
+  } else {
+    // Server consistently gets constrained by amplification factor, hence PTO
+    // never gets armed. The CHLO retransmission would trigger the
+    // retransmission of SHLO, however, the ENCRYPTION_HANDSHAKE packet NEVER
+    // gets retransmitted since half RTT data consumes the remaining space in
+    // the coalescer.
+    EXPECT_EQ("", client_->SendSynchronousRequest("/large_response"));
+  }
+}
+
 TEST_P(EndToEndTest, MaxStreamsUberTest) {
   // Connect with lower fake packet loss than we'd like to test.  Until
   // b/10126687 is fixed, losing handshake packets is pretty brutal.
diff --git a/quic/core/quic_coalesced_packet.cc b/quic/core/quic_coalesced_packet.cc
index 2ae9b41..28f3ece 100644
--- a/quic/core/quic_coalesced_packet.cc
+++ b/quic/core/quic_coalesced_packet.cc
@@ -152,6 +152,16 @@
   return transmission_types_[level];
 }
 
+size_t QuicCoalescedPacket::NumberOfPackets() const {
+  size_t num_of_packets = 0;
+  for (int8_t i = ENCRYPTION_INITIAL; i < NUM_ENCRYPTION_LEVELS; ++i) {
+    if (ContainsPacketOfEncryptionLevel(static_cast<EncryptionLevel>(i))) {
+      ++num_of_packets;
+    }
+  }
+  return num_of_packets;
+}
+
 std::string QuicCoalescedPacket::ToString(size_t serialized_length) const {
   // Total length and padding size.
   std::string info = absl::StrCat(
diff --git a/quic/core/quic_coalesced_packet.h b/quic/core/quic_coalesced_packet.h
index e3f674e..e8bc2c3 100644
--- a/quic/core/quic_coalesced_packet.h
+++ b/quic/core/quic_coalesced_packet.h
@@ -46,6 +46,9 @@
   // when this coalesced packet contains packet of |level|.
   TransmissionType TransmissionTypeOfPacket(EncryptionLevel level) const;
 
+  // Returns number of packets contained in this coalesced packet.
+  size_t NumberOfPackets() const;
+
   const SerializedPacket* initial_packet() const {
     return initial_packet_.get();
   }
diff --git a/quic/core/quic_coalesced_packet_test.cc b/quic/core/quic_coalesced_packet_test.cc
index 1d629ae..34f6da3 100644
--- a/quic/core/quic_coalesced_packet_test.cc
+++ b/quic/core/quic_coalesced_packet_test.cc
@@ -19,6 +19,7 @@
             coalesced.ToString(0));
   SimpleBufferAllocator allocator;
   EXPECT_EQ(0u, coalesced.length());
+  EXPECT_EQ(0u, coalesced.NumberOfPackets());
   char buffer[1000];
   QuicSocketAddress self_address(QuicIpAddress::Loopback4(), 1);
   QuicSocketAddress peer_address(QuicIpAddress::Loopback4(), 2);
@@ -35,6 +36,7 @@
             coalesced.TransmissionTypeOfPacket(ENCRYPTION_INITIAL));
   EXPECT_EQ(1500u, coalesced.max_packet_length());
   EXPECT_EQ(500u, coalesced.length());
+  EXPECT_EQ(1u, coalesced.NumberOfPackets());
   EXPECT_EQ(
       "total_length: 1500 padding_size: 1000 packets: {ENCRYPTION_INITIAL}",
       coalesced.ToString(1500));
@@ -54,6 +56,7 @@
                                             &allocator, 1500));
   EXPECT_EQ(1500u, coalesced.max_packet_length());
   EXPECT_EQ(1000u, coalesced.length());
+  EXPECT_EQ(2u, coalesced.NumberOfPackets());
   EXPECT_EQ(LOSS_RETRANSMISSION,
             coalesced.TransmissionTypeOfPacket(ENCRYPTION_ZERO_RTT));
   EXPECT_EQ(
@@ -77,6 +80,7 @@
                                              peer_address, &allocator, 1500));
   EXPECT_EQ(1500u, coalesced.max_packet_length());
   EXPECT_EQ(1000u, coalesced.length());
+  EXPECT_EQ(2u, coalesced.NumberOfPackets());
 
   // Max packet number length changed.
   SerializedPacket packet6(QuicPacketNumber(6), PACKET_4BYTE_PACKET_NUMBER,
@@ -87,6 +91,7 @@
                   "Max packet length changes in the middle of the write path");
   EXPECT_EQ(1500u, coalesced.max_packet_length());
   EXPECT_EQ(1000u, coalesced.length());
+  EXPECT_EQ(2u, coalesced.NumberOfPackets());
 }
 
 TEST(QuicCoalescedPacketTest, CopyEncryptedBuffers) {
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 471331b..bbdb64c 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2315,6 +2315,21 @@
       QuicUtils::IsCryptoStreamId(transport_version(), id)) {
     MaybeActivateLegacyVersionEncapsulation();
   }
+  if (GetQuicReloadableFlag(quic_preempt_stream_data_with_handshake_packet) &&
+      perspective_ == Perspective::IS_SERVER &&
+      version().CanSendCoalescedPackets() && !IsHandshakeConfirmed()) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_preempt_stream_data_with_handshake_packet,
+                                 1, 2);
+    if (coalesced_packet_.ContainsPacketOfEncryptionLevel(ENCRYPTION_INITIAL) &&
+        coalesced_packet_.NumberOfPackets() == 1u) {
+      // Handshake is not confirmed yet, if there is only an initial packet in
+      // the coalescer, try to bundle an ENCRYPTION_HANDSHAKE packet before
+      // sending stream data.
+      QUIC_RELOADABLE_FLAG_COUNT_N(
+          quic_preempt_stream_data_with_handshake_packet, 2, 2);
+      sent_packet_manager_.RetransmitDataOfSpaceIfAny(HANDSHAKE_DATA);
+    }
+  }
   QuicConsumedData consumed_data(0, false);
   {
     // Opportunistically bundle an ack with every outgoing packet.
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index 5b4cbbc..19a466f 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -47,6 +47,7 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_fix_willing_and_able_to_write2, true)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_h3_datagram, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_pass_path_response_to_validator, false)
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_preempt_stream_data_with_handshake_packet, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_require_handshake_confirmation, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_send_path_response, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_send_timestamps, false)