QuicPacketBuilder changes for RangeChainHead

Summary: See title. In this diff, I'm adding `insert` functions that allow us to pass in a `RangeChainHead` rather than a `BufQueue`.

Reviewed By: mjoras

Differential Revision: D57992385

fbshipit-source-id: 4c22727881059cdc8b66c1af3129aee99563b8dd

quic/codec/BUCK

@@ -116,6 +116,7 @@ mvfst_cpp_library(
     ],
     deps = [
         "//folly:random",
+        "//folly/io:iobuf",
     ],
     exported_deps = [
         ":packet_number",

quic/codec/QuicPacketBuilder.cpp

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <folly/io/IOBuf.h>
 #include <quic/codec/QuicPacketBuilder.h>
 #include <algorithm>
 
@@ -210,6 +211,28 @@ void RegularQuicPacketBuilder::insert(const BufQueue& buf, size_t limit) {
   insert(std::move(streamData));
 }
 
+void RegularQuicPacketBuilder::insert(
+    const ChainedByteRangeHead& buf,
+    size_t limit) {
+  limit = std::min(limit, buf.chainLength());
+  std::unique_ptr<folly::IOBuf> streamData = folly::IOBuf::wrapBuffer(
+      buf.getHead()->getRange().begin(),
+      std::min(limit, buf.getHead()->length()));
+  limit -= std::min(limit, buf.getHead()->length());
+
+  auto* current = buf.getHead()->getNext();
+  while (limit > 0) {
+    size_t amountToChopOff = std::min(limit, current->length());
+    auto tempBuf =
+        folly::IOBuf::wrapBuffer(current->getRange().begin(), amountToChopOff);
+    streamData->appendToChain(std::move(tempBuf));
+    limit -= amountToChopOff;
+    current = current->getNext();
+  }
+
+  insert(std::move(streamData));
+}
+
 void RegularQuicPacketBuilder::appendFrame(QuicWriteFrame frame) {
   packet_.empty = false;
   packet_.frames.push_back(std::move(frame));
@@ -662,6 +685,13 @@ void InplaceQuicPacketBuilder::insert(const BufQueue& buf, size_t limit) {
   bufWriter_.insert(buf.front(), limit);
 }
 
+void InplaceQuicPacketBuilder::insert(
+    const ChainedByteRangeHead& buf,
+    size_t limit) {
+  remainingBytes_ -= limit;
+  bufWriter_.insert(&buf, limit);
+}
+
 void InplaceQuicPacketBuilder::appendFrame(QuicWriteFrame frame) {
   packet_.empty = false;
   packet_.frames.push_back(std::move(frame));

quic/codec/QuicPacketBuilder.h

@@ -75,6 +75,7 @@ class PacketBuilderInterface {
   appendBytes(BufWriter& writer, PacketNum value, uint8_t byteNumber) = 0;
   virtual void insert(std::unique_ptr<folly::IOBuf> buf) = 0;
   virtual void insert(std::unique_ptr<folly::IOBuf> buf, size_t limit) = 0;
+  virtual void insert(const ChainedByteRangeHead& buf, size_t limit) = 0;
   virtual void insert(const BufQueue& buf, size_t limit) = 0;
   virtual void push(const uint8_t* data, size_t len) = 0;
 
@@ -145,6 +146,7 @@ class InplaceQuicPacketBuilder final : public PacketBuilderInterface {
       override;
   void insert(std::unique_ptr<folly::IOBuf> buf) override;
   void insert(std::unique_ptr<folly::IOBuf> buf, size_t limit) override;
+  void insert(const ChainedByteRangeHead& buf, size_t limit) override;
   void insert(const BufQueue& buf, size_t limit) override;
   void push(const uint8_t* data, size_t len) override;
 
@@ -224,6 +226,7 @@ class RegularQuicPacketBuilder final : public PacketBuilderInterface {
   void insert(std::unique_ptr<folly::IOBuf> buf) override;
   void insert(std::unique_ptr<folly::IOBuf> buf, size_t limit) override;
   void insert(const BufQueue& buf, size_t limit) override;
+  void insert(const ChainedByteRangeHead& buf, size_t limit) override;
 
   void push(const uint8_t* data, size_t len) override;
 
@@ -521,6 +524,10 @@ class PacketBuilderWrapper : public PacketBuilderInterface {
     builder.insert(buf, limit);
   }
 
+  void insert(const ChainedByteRangeHead& buf, size_t limit) override {
+    builder.insert(buf, limit);
+  }
+
   void appendFrame(QuicWriteFrame frame) override {
     builder.appendFrame(std::move(frame));
   }

quic/codec/test/Mocks.h

@@ -44,11 +44,18 @@ class MockQuicPacketBuilder : public PacketBuilderInterface {
   void insert(std::unique_ptr<folly::IOBuf> buf, size_t limit) override {
     _insert(buf, limit);
   }
+
+  void insert(const ChainedByteRangeHead& ChainedByteRangeHead, size_t limit)
+      override {
+    _insertRch(ChainedByteRangeHead, limit);
+  }
+
   MOCK_METHOD(void, appendFrame, (QuicWriteFrame));
   MOCK_METHOD(void, appendPaddingFrame, ());
   MOCK_METHOD(void, markNonEmpty, ());
   MOCK_METHOD(void, _insert, (std::unique_ptr<folly::IOBuf>&));
   MOCK_METHOD(void, _insert, (std::unique_ptr<folly::IOBuf>&, size_t));
+  MOCK_METHOD(void, _insertRch, (const ChainedByteRangeHead&, size_t));
   MOCK_METHOD(void, insert, (const BufQueue&, size_t));
   MOCK_METHOD(void, push, (const uint8_t*, size_t));
   MOCK_METHOD(void, write, (const QuicInteger&));