Refactor AckScheduler to use it in both PacketScheduler and PacketRebuilder

Summary: The logic for deciding which ACK type to write was duplicated in QuicPacketScheduler and QuicPacketRebuilder. This refactors the logic out into a separate QuicAckScheduler so it can be tested for correction and reused in both places.

Reviewed By: sharmafb

Differential Revision: D69933311

fbshipit-source-id: e4f45688a5d258dd2a57f9f7844407f3efad5f49

cmake/mvfst-config.cmake.in

@@ -28,6 +28,7 @@ if(NOT TARGET mvfst::mvfst_transport)
 endif()
 
 set(mvfst_LIBRARIES
+  mvfst::mvfst_ack_scheduler
   mvfst::mvfst_constants
   mvfst::mvfst_exception
   mvfst::mvfst_transport

quic/api/BUCK

@@ -130,6 +130,20 @@ mvfst_cpp_library(
     ],
 )
 
+mvfst_cpp_library(
+    name = "ack_scheduler",
+    srcs = [
+        "QuicAckScheduler.cpp",
+    ],
+    headers = [
+        "QuicAckScheduler.h",
+    ],
+    exported_deps = [
+        "//quic:constants",
+        "//quic/state:quic_state_machine",
+    ],
+)
+
 mvfst_cpp_library(
     name = "transport_helpers",
     srcs = [
@@ -152,6 +166,7 @@ mvfst_cpp_library(
         "//quic/state:simple_frame_functions",
     ],
     exported_deps = [
+        ":ack_scheduler",
         ":quic_batch_writer",
         "//folly:expected",
         "//folly/lang:assume",

quic/api/CMakeLists.txt

@@ -41,6 +41,38 @@ target_link_libraries(
   mvfst_state_machine
 )
 
+add_library(
+  mvfst_ack_scheduler
+  QuicAckScheduler.cpp
+)
+
+set_property(TARGET mvfst_ack_scheduler PROPERTY VERSION ${PACKAGE_VERSION})
+
+target_include_directories(
+  mvfst_ack_scheduler PUBLIC
+  $<BUILD_INTERFACE:${QUIC_FBCODE_ROOT}>
+  $<INSTALL_INTERFACE:include/>
+)
+
+target_compile_options(
+  mvfst_ack_scheduler
+  PRIVATE
+  ${_QUIC_COMMON_COMPILE_OPTIONS}
+)
+
+add_dependencies(
+  mvfst_ack_scheduler
+  mvfst_constants
+  mvfst_state_machine
+)
+
+target_link_libraries(
+  mvfst_ack_scheduler PUBLIC
+  Folly::folly
+  mvfst_constants
+  mvfst_state_machine
+)
+
 add_library(
   mvfst_transport
   IoBufQuicBatch.cpp
@@ -67,6 +99,7 @@ target_compile_options(
 
 add_dependencies(
   mvfst_transport
+  mvfst_ack_scheduler
   mvfst_async_udp_socket
   mvfst_batch_writer
   mvfst_buf_accessor
@@ -98,6 +131,7 @@ add_dependencies(
 target_link_libraries(
   mvfst_transport PUBLIC
   Folly::folly
+  mvfst_ack_scheduler
   mvfst_batch_writer
   mvfst_buf_accessor
   mvfst_bufutil
@@ -150,4 +184,10 @@ install(
   DESTINATION ${CMAKE_INSTALL_LIBDIR}
 )
 
+install(
+  TARGETS mvfst_ack_scheduler
+  EXPORT mvfst-exports
+  DESTINATION ${CMAKE_INSTALL_LIBDIR}
+)
+
 add_subdirectory(test)

quic/api/QuicAckScheduler.cpp

@@ -0,0 +1,133 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <quic/api/QuicAckScheduler.h>
+
+namespace quic {
+
+bool hasAcksToSchedule(const AckState& ackState) {
+  Optional<PacketNum> largestAckSend = largestAckToSend(ackState);
+  if (!largestAckSend) {
+    return false;
+  }
+  if (!ackState.largestAckScheduled) {
+    // Never scheduled an ack, we need to send
+    return true;
+  }
+  return *largestAckSend > *(ackState.largestAckScheduled);
+}
+
+Optional<PacketNum> largestAckToSend(const AckState& ackState) {
+  if (ackState.acks.empty()) {
+    return none;
+  }
+  return ackState.acks.back().end;
+}
+
+AckScheduler::AckScheduler(
+    const QuicConnectionStateBase& conn,
+    const AckState& ackState)
+    : conn_(conn), ackState_(ackState) {}
+
+Optional<PacketNum> AckScheduler::writeNextAcks(
+    PacketBuilderInterface& builder) {
+  // Use default ack delay for long headers. Usually long headers are sent
+  // before crypto negotiation, so the peer might not know about the ack delay
+  // exponent yet, so we use the default.
+  uint8_t ackDelayExponentToUse =
+      builder.getPacketHeader().getHeaderForm() == HeaderForm::Long
+      ? kDefaultAckDelayExponent
+      : conn_.transportSettings.ackDelayExponent;
+  auto largestAckedPacketNum = *largestAckToSend(ackState_);
+  auto ackingTime = Clock::now();
+  DCHECK(ackState_.largestRecvdPacketTime.hasValue())
+      << "Missing received time for the largest acked packet";
+  // assuming that we're going to ack the largest received with highest pri
+  auto receivedTime = *ackState_.largestRecvdPacketTime;
+  std::chrono::microseconds ackDelay =
+      (ackingTime > receivedTime
+           ? std::chrono::duration_cast<std::chrono::microseconds>(
+                 ackingTime - receivedTime)
+           : 0us);
+
+  WriteAckFrameMetaData meta = {
+      ackState_, /* ackState*/
+      ackDelay, /* ackDelay */
+      static_cast<uint8_t>(ackDelayExponentToUse), /* ackDelayExponent */
+      conn_.connectionTime, /* connect timestamp */
+  };
+
+  Optional<WriteAckFrameResult> ackWriteResult;
+
+  bool isAckReceiveTimestampsSupported =
+      conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer &&
+      conn_.maybePeerAckReceiveTimestampsConfig;
+
+  uint64_t peerRequestedTimestampsCount =
+      conn_.maybePeerAckReceiveTimestampsConfig.has_value()
+      ? conn_.maybePeerAckReceiveTimestampsConfig.value()
+            .maxReceiveTimestampsPerAck
+      : 0;
+
+  uint64_t extendedAckSupportedAndEnabled =
+      conn_.peerAdvertisedExtendedAckFeatures &
+      conn_.transportSettings.enableExtendedAckFeatures;
+  // Disable the ECN fields if we are not reading them
+  if (!conn_.transportSettings.readEcnOnIngress) {
+    extendedAckSupportedAndEnabled &= ~static_cast<ExtendedAckFeatureMaskType>(
+        ExtendedAckFeatureMask::ECN_COUNTS);
+  }
+  // Disable the receive timestamps fields if we have not regoatiated receive
+  // timestamps support
+  if (!isAckReceiveTimestampsSupported || (peerRequestedTimestampsCount == 0)) {
+    extendedAckSupportedAndEnabled &= ~static_cast<ExtendedAckFeatureMaskType>(
+        ExtendedAckFeatureMask::RECEIVE_TIMESTAMPS);
+  }
+
+  if (extendedAckSupportedAndEnabled > 0) {
+    // The peer supports extended ACKs and we have them enabled.
+    ackWriteResult = writeAckFrame(
+        meta,
+        builder,
+        FrameType::ACK_EXTENDED,
+        conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer
+            .value_or(AckReceiveTimestampsConfig()),
+        peerRequestedTimestampsCount,
+        extendedAckSupportedAndEnabled);
+  } else if (
+      conn_.transportSettings.readEcnOnIngress &&
+      (meta.ackState.ecnECT0CountReceived ||
+       meta.ackState.ecnECT1CountReceived ||
+       meta.ackState.ecnCECountReceived)) {
+    // We have to report ECN counts, but we can't use the extended ACK
+    // frame. In this case, we give ACK_ECN precedence over
+    // ACK_RECEIVE_TIMESTAMPS.
+    ackWriteResult = writeAckFrame(meta, builder, FrameType::ACK_ECN);
+  } else if (
+      isAckReceiveTimestampsSupported && (peerRequestedTimestampsCount > 0)) {
+    // Use ACK_RECEIVE_TIMESTAMPS if its enabled on both endpoints AND the
+    // peer requests at least 1 timestamp
+    ackWriteResult = writeAckFrame(
+        meta,
+        builder,
+        FrameType::ACK_RECEIVE_TIMESTAMPS,
+        conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer
+            .value(),
+        peerRequestedTimestampsCount);
+  } else {
+    ackWriteResult = writeAckFrame(meta, builder, FrameType::ACK);
+  }
+  if (!ackWriteResult) {
+    return none;
+  }
+  return largestAckedPacketNum;
+}
+
+bool AckScheduler::hasPendingAcks() const {
+  return hasAcksToSchedule(ackState_);
+}
+} // namespace quic

quic/api/QuicAckScheduler.h

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <quic/QuicConstants.h>
+#include <quic/state/StateData.h>
+
+namespace quic {
+
+class AckScheduler {
+ public:
+  AckScheduler(const QuicConnectionStateBase& conn, const AckState& ackState);
+
+  Optional<PacketNum> writeNextAcks(PacketBuilderInterface& builder);
+
+  [[nodiscard]] bool hasPendingAcks() const;
+
+ private:
+  const QuicConnectionStateBase& conn_;
+  const AckState& ackState_;
+};
+
+/**
+ * Returns whether or not the Ack scheduler has acks to schedule. This does not
+ * tell you when the ACKs can be written.
+ */
+bool hasAcksToSchedule(const AckState& ackState);
+
+/**
+ * Returns the largest packet received which needs to be acked.
+ */
+Optional<PacketNum> largestAckToSend(const AckState& ackState);
+
+} // namespace quic

quic/api/QuicPacketScheduler.cpp

@@ -113,25 +113,6 @@ class MiddleStartingIterationWrapper {
 
 namespace quic {
 
-bool hasAcksToSchedule(const AckState& ackState) {
-  Optional<PacketNum> largestAckSend = largestAckToSend(ackState);
-  if (!largestAckSend) {
-    return false;
-  }
-  if (!ackState.largestAckScheduled) {
-    // Never scheduled an ack, we need to send
-    return true;
-  }
-  return *largestAckSend > *(ackState.largestAckScheduled);
-}
-
-Optional<PacketNum> largestAckToSend(const AckState& ackState) {
-  if (ackState.acks.empty()) {
-    return none;
-  }
-  return ackState.acks.back().end;
-}
-
 // Schedulers
 
 FrameScheduler::Builder::Builder(
@@ -569,108 +550,6 @@ bool StreamFrameScheduler::writeStreamFrame(
   return true;
 }
 
-AckScheduler::AckScheduler(
-    const QuicConnectionStateBase& conn,
-    const AckState& ackState)
-    : conn_(conn), ackState_(ackState) {}
-
-Optional<PacketNum> AckScheduler::writeNextAcks(
-    PacketBuilderInterface& builder) {
-  // Use default ack delay for long headers. Usually long headers are sent
-  // before crypto negotiation, so the peer might not know about the ack delay
-  // exponent yet, so we use the default.
-  uint8_t ackDelayExponentToUse =
-      builder.getPacketHeader().getHeaderForm() == HeaderForm::Long
-      ? kDefaultAckDelayExponent
-      : conn_.transportSettings.ackDelayExponent;
-  auto largestAckedPacketNum = *largestAckToSend(ackState_);
-  auto ackingTime = Clock::now();
-  DCHECK(ackState_.largestRecvdPacketTime.hasValue())
-      << "Missing received time for the largest acked packet";
-  // assuming that we're going to ack the largest received with highest pri
-  auto receivedTime = *ackState_.largestRecvdPacketTime;
-  std::chrono::microseconds ackDelay =
-      (ackingTime > receivedTime
-           ? std::chrono::duration_cast<std::chrono::microseconds>(
-                 ackingTime - receivedTime)
-           : 0us);
-
-  WriteAckFrameMetaData meta = {
-      ackState_, /* ackState*/
-      ackDelay, /* ackDelay */
-      static_cast<uint8_t>(ackDelayExponentToUse), /* ackDelayExponent */
-      conn_.connectionTime, /* connect timestamp */
-  };
-
-  Optional<WriteAckFrameResult> ackWriteResult;
-
-  bool isAckReceiveTimestampsSupported =
-      conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer &&
-      conn_.maybePeerAckReceiveTimestampsConfig;
-
-  uint64_t peerRequestedTimestampsCount =
-      conn_.maybePeerAckReceiveTimestampsConfig.has_value()
-      ? conn_.maybePeerAckReceiveTimestampsConfig.value()
-            .maxReceiveTimestampsPerAck
-      : 0;
-
-  uint64_t extendedAckSupportedAndEnabled =
-      conn_.peerAdvertisedExtendedAckFeatures &
-      conn_.transportSettings.enableExtendedAckFeatures;
-  // Disable the ECN fields if we are not reading them
-  if (!conn_.transportSettings.readEcnOnIngress) {
-    extendedAckSupportedAndEnabled &= ~static_cast<ExtendedAckFeatureMaskType>(
-        ExtendedAckFeatureMask::ECN_COUNTS);
-  }
-  // Disable the receive timestamps fields if we have not regoatiated receive
-  // timestamps support
-  if (!isAckReceiveTimestampsSupported || (peerRequestedTimestampsCount == 0)) {
-    extendedAckSupportedAndEnabled &= ~static_cast<ExtendedAckFeatureMaskType>(
-        ExtendedAckFeatureMask::RECEIVE_TIMESTAMPS);
-  }
-
-  if (extendedAckSupportedAndEnabled > 0) {
-    // The peer supports extended ACKs and we have them enabled.
-    ackWriteResult = writeAckFrame(
-        meta,
-        builder,
-        FrameType::ACK_EXTENDED,
-        conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer
-            .value_or(AckReceiveTimestampsConfig()),
-        peerRequestedTimestampsCount,
-        extendedAckSupportedAndEnabled);
-  } else if (
-      conn_.transportSettings.readEcnOnIngress &&
-      (meta.ackState.ecnECT0CountReceived ||
-       meta.ackState.ecnECT1CountReceived ||
-       meta.ackState.ecnCECountReceived)) {
-    // We have to report ECN counts, but we can't use the extended ACK frame. In
-    // this case, we give ACK_ECN precedence over ACK_RECEIVE_TIMESTAMPS.
-    ackWriteResult = writeAckFrame(meta, builder, FrameType::ACK_ECN);
-  } else if (
-      isAckReceiveTimestampsSupported && (peerRequestedTimestampsCount > 0)) {
-    // Use ACK_RECEIVE_TIMESTAMPS if its enabled on both endpoints AND the peer
-    // requests at least 1 timestamp
-    ackWriteResult = writeAckFrame(
-        meta,
-        builder,
-        FrameType::ACK_RECEIVE_TIMESTAMPS,
-        conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer
-            .value(),
-        peerRequestedTimestampsCount);
-  } else {
-    ackWriteResult = writeAckFrame(meta, builder, FrameType::ACK);
-  }
-  if (!ackWriteResult) {
-    return none;
-  }
-  return largestAckedPacketNum;
-}
-
-bool AckScheduler::hasPendingAcks() const {
-  return hasAcksToSchedule(ackState_);
-}
-
 RstStreamScheduler::RstStreamScheduler(const QuicConnectionStateBase& conn)
     : conn_(conn) {}
 

quic/api/QuicPacketScheduler.h

@@ -10,6 +10,7 @@
 #include <boost/iterator/iterator_facade.hpp>
 #include <quic/QuicConstants.h>
 #include <quic/QuicException.h>
+#include <quic/api/QuicAckScheduler.h>
 #include <quic/codec/QuicPacketBuilder.h>
 #include <quic/codec/QuicPacketRebuilder.h>
 #include <quic/codec/QuicWriteCodec.h>
@@ -126,30 +127,6 @@ class StreamFrameScheduler {
   bool nextStreamDsr_{false};
 };
 
-class AckScheduler {
- public:
-  AckScheduler(const QuicConnectionStateBase& conn, const AckState& ackState);
-
-  Optional<PacketNum> writeNextAcks(PacketBuilderInterface& builder);
-
-  bool hasPendingAcks() const;
-
- private:
-  const QuicConnectionStateBase& conn_;
-  const AckState& ackState_;
-};
-
-/**
- * Returns whether or not the Ack scheduler has acks to schedule. This does not
- * tell you when the ACKs can be written.
- */
-bool hasAcksToSchedule(const AckState& ackState);
-
-/**
- * Returns the largest packet received which needs to be acked.
- */
-Optional<PacketNum> largestAckToSend(const AckState& ackState);
-
 class RstStreamScheduler {
  public:
   explicit RstStreamScheduler(const QuicConnectionStateBase& conn);

quic/api/test/QuicPacketSchedulerTest.cpp

@@ -2717,4 +2717,392 @@ TEST_F(QuicPacketSchedulerTest, FixedShortHeaderPadding) {
   EXPECT_TRUE(frames[1].asWriteStreamFrame());
 }
 
+// This test class sets up a connection with all the fields that can be included
+// in an ACK. The fixtures for this class confirm that the scheduler writes the
+// correct frame type and fields enabled by the connection state.
+class QuicAckSchedulerTest : public Test {
+ protected:
+  QuicAckSchedulerTest()
+      : conn_(createConn(10, 100000, 100000)),
+        ackState_(getAckState(*conn_, PacketNumberSpace::AppData)),
+        builder_(setupMockPacketBuilder()) {}
+
+  void SetUp() override {
+    // One ack block
+    ackState_.acks.insert(1, 10);
+
+    // One receive timestamps
+    WriteAckFrameState::ReceivedPacket rpi;
+    rpi.pktNum = 10;
+    rpi.timings.receiveTimePoint = Clock::now();
+    ackState_.recvdPacketInfos.emplace_back(rpi);
+    ackState_.largestRecvdPacketNum = 10;
+    ackState_.largestRecvdPacketTime = rpi.timings.receiveTimePoint;
+
+    // Non-zero ECN values
+    ackState_.ecnECT0CountReceived = 1;
+    ackState_.ecnECT1CountReceived = 2;
+    ackState_.ecnCECountReceived = 3;
+
+    auto connId = getTestConnectionId();
+    mockPacketHeader_ = std::make_unique<PacketHeader>(ShortHeader(
+        ProtectionType::KeyPhaseZero,
+        connId,
+        getNextPacketNum(*conn_, PacketNumberSpace::AppData)));
+
+    EXPECT_CALL(*builder_, getPacketHeader())
+        .WillRepeatedly(ReturnRef(*mockPacketHeader_));
+  }
+
+  std::unique_ptr<QuicClientConnectionState> conn_;
+  AckState ackState_;
+  std::unique_ptr<MockQuicPacketBuilder> builder_;
+  std::unique_ptr<PacketHeader> mockPacketHeader_;
+};
+
+TEST_F(QuicAckSchedulerTest, DefaultAckFrame) {
+  // Default config writes ACK frame.
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 0);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 0);
+  EXPECT_EQ(ackFrame->ecnCECount, 0);
+}
+
+TEST_F(QuicAckSchedulerTest, WriteAckEcnWhenReadingEcnOnEgress) {
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_ECN);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
+TEST_F(QuicAckSchedulerTest, WriteAckReceiveTimestampsWhenEnabled) {
+  conn_->transportSettings.readEcnOnIngress = false;
+
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer =
+      AckReceiveTimestampsConfig();
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_RECEIVE_TIMESTAMPS);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  ASSERT_EQ(ackFrame->recvdPacketsTimestampRanges.size(), 1);
+  EXPECT_EQ(ackFrame->recvdPacketsTimestampRanges[0].timestamp_delta_count, 1);
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 0);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 0);
+  EXPECT_EQ(ackFrame->ecnCECount, 0);
+}
+
+TEST_F(QuicAckSchedulerTest, AckEcnTakesPrecedenceOverReceiveTimestamps) {
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer =
+      AckReceiveTimestampsConfig();
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_ECN);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
+TEST_F(QuicAckSchedulerTest, AckExtendedNotSentIfNotSupported) {
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  conn_->transportSettings.enableExtendedAckFeatures =
+      3; // ECN + ReceiveTimestamps
+  conn_->peerAdvertisedExtendedAckFeatures = 0;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_ECN);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
+TEST_F(QuicAckSchedulerTest, AckExtendedNotSentIfNotEnabled) {
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  conn_->transportSettings.enableExtendedAckFeatures = 0;
+  conn_->peerAdvertisedExtendedAckFeatures = 3; // ECN + ReceiveTimestamps;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_ECN);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
+TEST_F(
+    QuicAckSchedulerTest,
+    AckExtendedNotSentIfReceiveTimestampFeatureNotSupported) {
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  conn_->transportSettings.enableExtendedAckFeatures =
+      3; // We support ECN + ReceiveTimestamps
+  conn_->peerAdvertisedExtendedAckFeatures =
+      2; // Peer supports ReceiveTimestamps but not ECN in extended ack
+  // Peer sent ART config
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  // We don't have an ART config (i.e. we can't sent ART)
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer = none;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_ECN);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
+TEST_F(QuicAckSchedulerTest, AckExtendedNotSentIfECNFeatureNotSupported) {
+  conn_->transportSettings.enableExtendedAckFeatures =
+      3; // We support ECN + ReceiveTimestamps
+  conn_->peerAdvertisedExtendedAckFeatures =
+      1; // Peer supports ECN but not ReceiveTimestamps in extended ack
+
+  // ART support negotiated
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer =
+      AckReceiveTimestampsConfig();
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_RECEIVE_TIMESTAMPS);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  ASSERT_EQ(ackFrame->recvdPacketsTimestampRanges.size(), 1);
+  EXPECT_EQ(ackFrame->recvdPacketsTimestampRanges[0].timestamp_delta_count, 1);
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 0);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 0);
+  EXPECT_EQ(ackFrame->ecnCECount, 0);
+}
+
+TEST_F(QuicAckSchedulerTest, AckExtendedWithAllFeatures) {
+  conn_->transportSettings.enableExtendedAckFeatures =
+      3; // We support ECN + ReceiveTimestamps
+  conn_->peerAdvertisedExtendedAckFeatures =
+      3; // Peer supports ECN + ReceiveTimestamps
+
+  // ART support negotiated
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer =
+      AckReceiveTimestampsConfig();
+
+  // We can read ECN
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_EXTENDED);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  ASSERT_EQ(ackFrame->recvdPacketsTimestampRanges.size(), 1);
+  EXPECT_EQ(ackFrame->recvdPacketsTimestampRanges[0].timestamp_delta_count, 1);
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
+TEST_F(QuicAckSchedulerTest, AckExtendedTakesPrecedenceOverECN) {
+  conn_->transportSettings.enableExtendedAckFeatures =
+      3; // We support ECN + ReceiveTimestamps
+  conn_->peerAdvertisedExtendedAckFeatures =
+      2; // Peer supports extended ack with only ReceiveTimestamps
+
+  // ART support negotiated
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer =
+      AckReceiveTimestampsConfig();
+
+  // We can read ECN
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_EXTENDED);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  ASSERT_EQ(ackFrame->recvdPacketsTimestampRanges.size(), 1);
+  EXPECT_EQ(ackFrame->recvdPacketsTimestampRanges[0].timestamp_delta_count, 1);
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 0);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 0);
+  EXPECT_EQ(ackFrame->ecnCECount, 0);
+}
+
+TEST_F(QuicAckSchedulerTest, AckExtendedTakesPrecedenceOverReceiveTimestamps) {
+  conn_->transportSettings.enableExtendedAckFeatures =
+      3; // We support ECN + ReceiveTimestamps
+  conn_->peerAdvertisedExtendedAckFeatures =
+      1; // Peer supports extended ack with only ECN
+
+  // ART support negotiated
+  conn_->maybePeerAckReceiveTimestampsConfig = AckReceiveTimestampsConfig();
+  conn_->transportSettings.maybeAckReceiveTimestampsConfigSentToPeer =
+      AckReceiveTimestampsConfig();
+
+  // We can read ECN
+  conn_->transportSettings.readEcnOnIngress = true;
+
+  AckScheduler ackScheduler(*conn_, ackState_);
+  ASSERT_TRUE(ackScheduler.hasPendingAcks());
+
+  ASSERT_TRUE(ackScheduler.writeNextAcks(*builder_) != none);
+  ASSERT_EQ(builder_->frames_.size(), 1);
+
+  auto ackFrame = builder_->frames_[0].asWriteAckFrame();
+  ASSERT_TRUE(ackFrame != nullptr);
+
+  EXPECT_EQ(ackFrame->frameType, FrameType::ACK_EXTENDED);
+
+  EXPECT_EQ(ackFrame->ackBlocks.size(), 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].start, 1);
+  EXPECT_EQ(ackFrame->ackBlocks[0].end, 10);
+
+  EXPECT_TRUE(ackFrame->recvdPacketsTimestampRanges.empty());
+
+  EXPECT_EQ(ackFrame->ecnECT0Count, 1);
+  EXPECT_EQ(ackFrame->ecnECT1Count, 2);
+  EXPECT_EQ(ackFrame->ecnCECount, 3);
+}
+
 } // namespace quic::test

quic/codec/BUCK

@@ -141,6 +141,7 @@ mvfst_cpp_library(
     ],
     deps = [
         ":codec",
+        "//quic/api:ack_scheduler",
         "//quic/flowcontrol:flow_control",
         "//quic/state:simple_frame_functions",
         "//quic/state:state_functions",

quic/codec/CMakeLists.txt

@@ -159,6 +159,7 @@ target_compile_options(
 
 add_dependencies(
   mvfst_codec_pktrebuilder
+  mvfst_ack_scheduler
   mvfst_codec
   mvfst_codec_pktbuilder
   mvfst_flowcontrol
@@ -170,6 +171,7 @@ add_dependencies(
 target_link_libraries(
   mvfst_codec_pktrebuilder PUBLIC
   Folly::folly
+  mvfst_ack_scheduler
   mvfst_codec
   mvfst_codec_pktbuilder
   mvfst_flowcontrol

quic/codec/QuicPacketRebuilder.cpp

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <quic/api/QuicAckScheduler.h>
 #include <quic/codec/QuicPacketRebuilder.h>
 #include <quic/codec/QuicWriteCodec.h>
 #include <quic/flowcontrol/QuicFlowController.h>
@@ -216,96 +217,11 @@ Optional<ClonedPacketIdentifier> PacketRebuilder::rebuildFromPacket(
   // cloned packet.
   if (shouldRebuildWriteAckFrame) {
     auto& packetHeader = builder_.getPacketHeader();
-    uint64_t ackDelayExponent =
+    const AckState& ackState = getAckState(
-        (packetHeader.getHeaderForm() == HeaderForm::Long)
-        ? kDefaultAckDelayExponent
-        : conn_.transportSettings.ackDelayExponent;
-    const AckState& ackState_ = getAckState(
         conn_,
         protectionTypeToPacketNumberSpace(packetHeader.getProtectionType()));
-    auto ackingTime = Clock::now();
+    AckScheduler ackScheduler(conn_, ackState);
-    DCHECK(ackState_.largestRecvdPacketTime.hasValue())
+    ackScheduler.writeNextAcks(builder_);
-        << "Missing received time for the largest acked packet";
-    auto receivedTime = *ackState_.largestRecvdPacketTime;
-    std::chrono::microseconds ackDelay =
-        (ackingTime > receivedTime
-             ? std::chrono::duration_cast<std::chrono::microseconds>(
-                   ackingTime - receivedTime)
-             : 0us);
-
-    WriteAckFrameMetaData meta = {
-        ackState_, /* ackState*/
-        ackDelay, /* ackDelay */
-        static_cast<uint8_t>(ackDelayExponent), /* ackDelayExponent */
-        conn_.connectionTime, /* connect timestamp */
-    };
-
-    // TODO: This code needs refactoring. The logic below duplicated from
-    // PacketScheduler::writeNextAcks().
-
-    // Write the AckFrame ignoring the result. This is best-effort.
-    Optional<WriteAckFrameResult> ackWriteResult;
-
-    uint64_t peerRequestedTimestampsCount =
-        conn_.maybePeerAckReceiveTimestampsConfig.has_value()
-        ? conn_.maybePeerAckReceiveTimestampsConfig.value()
-              .maxReceiveTimestampsPerAck
-        : 0;
-
-    bool isAckReceiveTimestampsSupported =
-        conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer &&
-        conn_.maybePeerAckReceiveTimestampsConfig;
-
-    uint64_t extendedAckSupportedAndEnabled =
-        conn_.peerAdvertisedExtendedAckFeatures &
-        conn_.transportSettings.enableExtendedAckFeatures;
-    // Disable the ECN fields if we are not reading them
-    if (!conn_.transportSettings.readEcnOnIngress) {
-      extendedAckSupportedAndEnabled &=
-          ~static_cast<ExtendedAckFeatureMaskType>(
-              ExtendedAckFeatureMask::ECN_COUNTS);
-    }
-    // Disable the receive timestamps fields if we have not regoatiated receive
-    // timestamps support
-    if (!isAckReceiveTimestampsSupported ||
-        (peerRequestedTimestampsCount == 0)) {
-      extendedAckSupportedAndEnabled &=
-          ~static_cast<ExtendedAckFeatureMaskType>(
-              ExtendedAckFeatureMask::RECEIVE_TIMESTAMPS);
-    }
-
-    if (extendedAckSupportedAndEnabled > 0) {
-      // The peer supports extended ACKs and we have them enabled.
-      ackWriteResult = writeAckFrame(
-          meta,
-          builder_,
-          FrameType::ACK_EXTENDED,
-          conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer
-              .value_or(AckReceiveTimestampsConfig()),
-          peerRequestedTimestampsCount,
-          extendedAckSupportedAndEnabled);
-    } else if (
-        conn_.transportSettings.readEcnOnIngress &&
-        (meta.ackState.ecnECT0CountReceived ||
-         meta.ackState.ecnECT1CountReceived ||
-         meta.ackState.ecnCECountReceived)) {
-      // We have to report ECN counts, but we can't use the extended ACK frame.
-      // In this case, we give ACK_ECN precedence over ACK_RECEIVE_TIMESTAMPS.
-      ackWriteResult = writeAckFrame(meta, builder_, FrameType::ACK_ECN);
-    } else if (
-        isAckReceiveTimestampsSupported && (peerRequestedTimestampsCount > 0)) {
-      // Use ACK_RECEIVE_TIMESTAMPS if its enabled on both endpoints AND the
-      // peer requests at least 1 timestamp
-      ackWriteResult = writeAckFrame(
-          meta,
-          builder_,
-          FrameType::ACK_RECEIVE_TIMESTAMPS,
-          conn_.transportSettings.maybeAckReceiveTimestampsConfigSentToPeer
-              .value(),
-          peerRequestedTimestampsCount);
-    } else {
-      ackWriteResult = writeAckFrame(meta, builder_, FrameType::ACK);
-    }
   }
   // We shouldn't clone if:
   // (1) we only end up cloning only acks, ping, or paddings.