Add an API that avoids unnecessary cloning of small IOBufs

Summary:
`IOBufQueue` has packing logic, that would copy data instead of chaining `IOBufs`. However, that API unnecessary requires exclusive ownership of the chain.

This diff adds an API that takes reference to IOBuf, copies as much as necessary and then clones/chains the rest.

Reviewed By: yfeldblum

Differential Revision: D14715581

fbshipit-source-id: e047d1c69ca19eec95a7adbca6df740440d23db3

folly/io/Cursor.h

@@ -1173,7 +1173,7 @@ class QueueAppender : public detail::Writable<QueueAppender> {
   }
 
   void insert(const folly::IOBuf& buf) {
-    insert(buf.clone());
+    queueCache_.queue()->append(buf, true);
   }
 
  private:

folly/io/IOBufQueue.cpp

@@ -153,6 +153,42 @@ void IOBufQueue::append(unique_ptr<IOBuf>&& buf, bool pack) {
   appendToChain(head_, std::move(buf), pack);
 }
 
+void IOBufQueue::append(const folly::IOBuf& buf, bool pack) {
+  if (!head_ || !pack) {
+    append(buf.clone(), pack);
+    return;
+  }
+
+  auto guard = updateGuard();
+  if (options_.cacheChainLength) {
+    chainLength_ += buf.computeChainDataLength();
+  }
+
+  size_t copyRemaining = MAX_PACK_COPY;
+  std::size_t n;
+  const folly::IOBuf* src = &buf;
+  folly::IOBuf* tail = head_->prev();
+  while ((n = src->length()) <= copyRemaining && n <= tail->tailroom()) {
+    if (n > 0) {
+      memcpy(tail->writableTail(), src->data(), n);
+      tail->append(n);
+      copyRemaining -= n;
+    }
+    src = src->next();
+
+    // Consumed full input.
+    if (src == &buf) {
+      return;
+    }
+  }
+
+  // Clone the rest.
+  do {
+    head_->prependChain(src->cloneOne());
+    src = src->next();
+  } while (src != &buf);
+}
+
 void IOBufQueue::append(IOBufQueue& other, bool pack) {
   if (!other.head_) {
     return;

folly/io/IOBufQueue.h

@@ -292,6 +292,7 @@ class IOBufQueue {
    * the chain topology unchanged.
    */
   void append(std::unique_ptr<folly::IOBuf>&& buf, bool pack = false);
+  void append(const folly::IOBuf& buf, bool pack = false);
 
   /**
    * Add a queue to the end of this queue. The queue takes ownership of

folly/io/test/IOBufQueueTest.cpp

@@ -57,6 +57,12 @@ void checkConsistency(const IOBufQueue& queue) {
   }
 }
 
+std::string queueToString(const IOBufQueue& queue) {
+  std::string out;
+  queue.appendToString(out);
+  return out;
+}
+
 } // namespace
 
 TEST(IOBufQueue, Simple) {
@@ -120,6 +126,50 @@ TEST(IOBufQueue, AppendStringPiece) {
   EXPECT_EQ(0, memcmp(chain->data(), chain2->data(), s.length()));
 }
 
+TEST(IOBufQueue, AppendIOBufRef) {
+  IOBufQueue queue(clOptions);
+  queue.append(*stringToIOBuf("abc", 3), true);
+  EXPECT_EQ(3, queue.chainLength());
+  EXPECT_EQ(1, queue.front()->countChainElements());
+  EXPECT_EQ("abc", queueToString(queue));
+  // Make sure we have enough space to copy over next data.
+  queue.preallocate(10, 10, 10);
+  EXPECT_LE(10, queue.tailroom());
+  auto numElements = queue.front()->countChainElements();
+  queue.append(*stringToIOBuf("FooBar", 6), true);
+  EXPECT_EQ(9, queue.chainLength());
+  // Make sure that we performed copy and not append chain.
+  EXPECT_EQ(numElements, queue.front()->countChainElements());
+  EXPECT_EQ("abcFooBar", queueToString(queue));
+}
+
+TEST(IOBufQueue, AppendIOBufRefChain) {
+  IOBufQueue queue(clOptions);
+  queue.append(*stringToIOBuf("abc", 3), true);
+  queue.preallocate(10, 10, 10);
+  auto numElements = queue.front()->countChainElements();
+  auto chain = stringToIOBuf("Hello", 5);
+  chain->prependChain(stringToIOBuf("World", 5));
+  queue.append(*chain, true);
+  // Make sure that we performed a copy and not append chain.
+  EXPECT_EQ(numElements, queue.front()->countChainElements());
+  EXPECT_EQ("abcHelloWorld", queueToString(queue));
+}
+
+TEST(IOBufQueue, AppendIOBufRefChainPartial) {
+  IOBufQueue queue(clOptions);
+  queue.append(*stringToIOBuf("abc", 3), true);
+  queue.preallocate(10, 10, 10);
+  auto numElements = queue.front()->countChainElements();
+  auto chain = stringToIOBuf("Hello", 5);
+  chain->prependChain(stringToIOBuf("World!", 6));
+  chain->prependChain(stringToIOBuf("Test", 4));
+  queue.append(*chain, true);
+  // Make sure that we performed a copy of first IOBuf and cloned the rest.
+  EXPECT_EQ(numElements + 2, queue.front()->countChainElements());
+  EXPECT_EQ("abcHelloWorld!Test", queueToString(queue));
+}
+
 TEST(IOBufQueue, Split) {
   IOBufQueue queue(clOptions);
   queue.append(stringToIOBuf(SCL("Hello")));