From 03d240dc01f3821797183d8eae6237ef33c2edfc Mon Sep 17 00:00:00 2001
From: Nick Terrell <terrelln@fb.com>
Date: Mon, 13 Mar 2017 12:50:55 -0700
Subject: [PATCH] Fill buffer before allocating more

Summary:
Modify `QueueAppender::pushAtMost(const uint8_t*, size_t)` to fill the end of
the current buffer before it allocates another.

Reviewed By: yfeldblum

Differential Revision: D4687832

fbshipit-source-id: 7c21f4359da9f0cf26d9577abddfd640b0c877e9
---
 folly/io/Cursor.h                 | 10 +++++-
 folly/io/test/IOBufCursorTest.cpp | 55 +++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/folly/io/Cursor.h b/folly/io/Cursor.h
index c16c6a45..32d65605 100644
--- a/folly/io/Cursor.h
+++ b/folly/io/Cursor.h
@@ -975,7 +975,15 @@ class QueueAppender : public detail::Writable<QueueAppender> {
 
   using detail::Writable<QueueAppender>::pushAtMost;
   size_t pushAtMost(const uint8_t* buf, size_t len) {
-    size_t remaining = len;
+    // Fill the current buffer
+    const size_t copyLength = std::min(len, length());
+    if (copyLength != 0) {
+      memcpy(writableData(), buf, copyLength);
+      append(copyLength);
+      buf += copyLength;
+    }
+    // Allocate more buffers as necessary
+    size_t remaining = len - copyLength;
     while (remaining != 0) {
       auto p = queue_->preallocate(std::min(remaining, growth_),
                                    growth_,
diff --git a/folly/io/test/IOBufCursorTest.cpp b/folly/io/test/IOBufCursorTest.cpp
index 2dc1766c..89e12052 100644
--- a/folly/io/test/IOBufCursorTest.cpp
+++ b/folly/io/test/IOBufCursorTest.cpp
@@ -20,6 +20,8 @@
 #include <folly/Range.h>
 #include <folly/io/Cursor.h>
 #include <folly/portability/GTest.h>
+#include <numeric>
+#include <vector>
 
 using folly::ByteRange;
 using folly::format;
@@ -530,6 +532,59 @@ TEST(IOBuf, QueueAppender) {
   EXPECT_THROW({cursor.readBE<uint32_t>();}, std::out_of_range);
 }
 
+TEST(IOBuf, QueueAppenderPushAtMostFillBuffer) {
+  folly::IOBufQueue queue;
+  // There should be a goodMallocSize between 125 and 1000
+  QueueAppender appender{&queue, 125};
+  std::vector<uint8_t> data;
+  data.resize(1000);
+  std::iota(data.begin(), data.end(), 0);
+  // Add 100 byte
+  appender.pushAtMost(data.data(), 100);
+  // Add 900 bytes
+  appender.pushAtMost(data.data() + 100, data.size() - 100);
+  const auto buf = queue.front();
+  // Should fill the current buffer before adding another
+  EXPECT_LE(2, buf->countChainElements());
+  EXPECT_EQ(0, buf->tailroom());
+  EXPECT_LE(125, buf->length());
+  EXPECT_EQ(1000, buf->computeChainDataLength());
+  const StringPiece sp{(const char*)data.data(), data.size()};
+  EXPECT_EQ(sp, toString(*buf));
+}
+
+TEST(IOBuf, QueueAppenderInsertOwn) {
+  auto buf = IOBuf::create(10);
+  folly::IOBufQueue queue;
+  QueueAppender appender{&queue, 128};
+  appender.insert(std::move(buf));
+
+  std::vector<uint8_t> data;
+  data.resize(256);
+  std::iota(data.begin(), data.end(), 0);
+  appender.pushAtMost(folly::range(data));
+  // Buffer is owned, so we should write to it
+  EXPECT_LE(2, queue.front()->countChainElements());
+  EXPECT_EQ(0, queue.front()->tailroom());
+  const StringPiece sp{(const char*)data.data(), data.size()};
+  EXPECT_EQ(sp, toString(*queue.front()));
+}
+
+TEST(IOBuf, QueueAppenderInsertClone) {
+  IOBuf buf{IOBuf::CREATE, 100};
+  folly::IOBufQueue queue;
+  QueueAppender appender{&queue, 100};
+  // Buffer is shared, so we create a new buffer to write to
+  appender.insert(buf);
+  uint8_t x = 42;
+  appender.pushAtMost(&x, 1);
+  EXPECT_EQ(2, queue.front()->countChainElements());
+  EXPECT_EQ(0, queue.front()->length());
+  EXPECT_LT(0, queue.front()->tailroom());
+  EXPECT_EQ(1, queue.front()->next()->length());
+  EXPECT_EQ(x, queue.front()->next()->data()[0]);
+}
+
 TEST(IOBuf, CursorOperators) {
   // Test operators on a single-item chain
   {
-- 
2.34.1