From 1dad9543b989ba743a052682d8de6d22f7f594d1 Mon Sep 17 00:00:00 2001
From: Jolene Tan <lene@meta.com>
Date: Thu, 24 Oct 2024 15:30:20 -0700
Subject: [PATCH] Immediately retransmit initial packets in
 writeCryptoAndAckDataToSocket

Summary: In `writeCryptoAndAckDataToSocket`, add an additional `writeProbingDataToSocket` call at the end that is limited to the number of CRYPTO frame-containing packets just written, gated by the new `TransportSetttings` field `immediatelyRetransmitInitialPackets`.

Reviewed By: mjoras

Differential Revision: D64485616

fbshipit-source-id: f0927a3796767700fd46673195e1cd4e1bbbcbeb
---
 quic/api/QuicTransportFunctions.cpp          | 23 +++++++++++++-
 quic/api/test/QuicTransportFunctionsTest.cpp | 32 ++++++++++++++++++++
 quic/state/TransportSettings.h               |  1 +
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/quic/api/QuicTransportFunctions.cpp b/quic/api/QuicTransportFunctions.cpp
index fd60bb7a6..c6344b8a3 100644
--- a/quic/api/QuicTransportFunctions.cpp
+++ b/quic/api/QuicTransportFunctions.cpp
@@ -1107,7 +1107,7 @@ WriteQuicDataResult writeCryptoAndAckDataToSocket(
       connection,
       srcConnId,
       dstConnId,
-      std::move(builder),
+      builder,
       LongHeader::typeToPacketNumberSpace(packetType),
       scheduler,
       congestionControlWritableBytes,
@@ -1121,6 +1121,27 @@ WriteQuicDataResult writeCryptoAndAckDataToSocket(
   packetsWritten += writeResult.packetsWritten;
   bytesWritten += writeResult.bytesWritten;
 
+  if (connection.transportSettings.immediatelyRetransmitInitialPackets &&
+      packetsWritten > 0 && packetsWritten < packetLimit) {
+    auto remainingLimit = packetLimit - packetsWritten;
+    auto cloneResult = writeProbingDataToSocket(
+        sock,
+        connection,
+        srcConnId,
+        dstConnId,
+        builder,
+        encryptionLevel,
+        LongHeader::typeToPacketNumberSpace(packetType),
+        scheduler,
+        packetsWritten < remainingLimit ? packetsWritten : remainingLimit,
+        cleartextCipher,
+        headerCipher,
+        version,
+        token);
+    probesWritten += cloneResult.probesWritten;
+    bytesWritten += cloneResult.bytesWritten;
+  }
+
   VLOG_IF(10, packetsWritten || probesWritten)
       << nodeToString(connection.nodeType)
       << " written crypto and acks data type=" << packetType
diff --git a/quic/api/test/QuicTransportFunctionsTest.cpp b/quic/api/test/QuicTransportFunctionsTest.cpp
index 28d7beba5..4b22a9198 100644
--- a/quic/api/test/QuicTransportFunctionsTest.cpp
+++ b/quic/api/test/QuicTransportFunctionsTest.cpp
@@ -3372,6 +3372,38 @@ TEST_F(QuicTransportFunctionsTest, NoCryptoProbeWriteIfNoProbeCredit) {
   EXPECT_EQ(0, res.probesWritten);
 }
 
+TEST_F(QuicTransportFunctionsTest, ImmediatelyRetransmitInitialPackets) {
+  auto conn = createConn();
+  conn->transportSettings.immediatelyRetransmitInitialPackets = true;
+  auto cryptoStream = &conn->cryptoState->initialStream;
+  auto buf = buildRandomInputData(1600);
+  writeDataToQuicStream(*cryptoStream, buf->clone());
+  EventBase evb;
+  std::shared_ptr<FollyQuicEventBase> qEvb =
+      std::make_shared<FollyQuicEventBase>(&evb);
+  auto socket =
+      std::make_unique<NiceMock<quic::test::MockAsyncUDPSocket>>(qEvb);
+  auto rawSocket = socket.get();
+  auto res = writeCryptoAndAckDataToSocket(
+      *rawSocket,
+      *conn,
+      *conn->clientConnectionId,
+      *conn->serverConnectionId,
+      LongHeader::Types::Initial,
+      *conn->initialWriteCipher,
+      *conn->initialHeaderCipher,
+      getVersion(*conn),
+      conn->transportSettings.writeConnectionDataPacketsLimit);
+  EXPECT_GE(res.bytesWritten, buf->computeChainDataLength());
+
+  EXPECT_EQ(2, res.packetsWritten);
+  EXPECT_EQ(2, res.probesWritten);
+  EXPECT_EQ(conn->udpSendPacketLen * 4, res.bytesWritten);
+  ASSERT_EQ(4, conn->outstandings.packets.size());
+  ASSERT_EQ(2, cryptoStream->retransmissionBuffer.size());
+  ASSERT_TRUE(cryptoStream->pendingWrites.empty());
+}
+
 TEST_F(QuicTransportFunctionsTest, ResetNumProbePackets) {
   auto conn = createConn();
   EventBase evb;
diff --git a/quic/state/TransportSettings.h b/quic/state/TransportSettings.h
index d3d383cfb..71f21f5f7 100644
--- a/quic/state/TransportSettings.h
+++ b/quic/state/TransportSettings.h
@@ -415,6 +415,7 @@ struct TransportSettings {
   // Raise read callbacks for all unidirectional streams first on data
   // reception.
   bool unidirectionalStreamsReadCallbacksFirst{false};
+  bool immediatelyRetransmitInitialPackets{false};
 };
 
 } // namespace quic