From fe381c8b35cb0eafcd65b9b692a557eb78085d8c Mon Sep 17 00:00:00 2001
From: Tudor Bosman <tudorb@fb.com>
Date: Thu, 30 Apr 2015 17:17:13 -0700
Subject: [PATCH] IOBuf::reserve would return less tailroom than requested
 under certain circumstances

Test Plan: test added

Reviewed By: lxiong@fb.com

Subscribers: lxiong, net-systems@, folly-diffs@, yfeldblum, chalfant, pamelavagata, kma

FB internal diff: D2036967

Tasks: 6925950

Signature: t1:2036967:1430431606:3e115f7ed76b207572db26d352bebefe7a3d306d
---
 folly/io/IOBuf.cpp          |  6 ++++--
 folly/io/test/IOBufTest.cpp | 27 +++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index ff60ac83..07c784b4 100644
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -698,7 +698,7 @@ void IOBuf::reserveSlow(uint64_t minHeadroom, uint64_t minTailroom) {
     return;
   }
 
-  size_t newAllocatedCapacity = goodExtBufferSize(newCapacity);
+  size_t newAllocatedCapacity = 0;
   uint8_t* newBuffer = nullptr;
   uint64_t newHeadroom = 0;
   uint64_t oldHeadroom = headroom();
@@ -708,8 +708,9 @@ void IOBuf::reserveSlow(uint64_t minHeadroom, uint64_t minTailroom) {
   SharedInfo* info = sharedInfo();
   if (info && (info->freeFn == nullptr) && length_ != 0 &&
       oldHeadroom >= minHeadroom) {
+    size_t headSlack = oldHeadroom - minHeadroom;
+    newAllocatedCapacity = goodExtBufferSize(newCapacity + headSlack);
     if (usingJEMalloc()) {
-      size_t headSlack = oldHeadroom - minHeadroom;
       // We assume that tailroom is more useful and more important than
       // headroom (not least because realloc / xallocx allow us to grow the
       // buffer at the tail, but not at the head)  So, if we have more headroom
@@ -743,6 +744,7 @@ void IOBuf::reserveSlow(uint64_t minHeadroom, uint64_t minTailroom) {
   // None of the previous reallocation strategies worked (or we're using
   // an internal buffer).  malloc/copy/free.
   if (newBuffer == nullptr) {
+    newAllocatedCapacity = goodExtBufferSize(newCapacity);
     void* p = malloc(newAllocatedCapacity);
     if (UNLIKELY(p == nullptr)) {
       throw std::bad_alloc();
diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index df644a31..0cc3ffd0 100644
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -1081,6 +1081,33 @@ TEST(IOBuf, HashAndEqual) {
   EXPECT_EQ(hash(e), hash(f));
 }
 
+// reserveSlow() had a bug when reallocating the buffer in place. It would
+// preserve old headroom if it's not too much (heuristically) but wouldn't
+// adjust the requested amount of memory to account for that; the end result
+// would be that reserve() would return with less tailroom than requested.
+TEST(IOBuf, ReserveWithHeadroom) {
+  // This is assuming jemalloc, where we know that 4096 and 8192 bytes are
+  // valid (and consecutive) allocation sizes. We're hoping that our
+  // 4096-byte buffer can be expanded in place to 8192 (in practice, this
+  // usually happens).
+  const char data[] = "Lorem ipsum dolor sit amet, consectetur adipiscing elit";
+  constexpr size_t reservedSize = 24;  // sizeof(SharedInfo)
+  // chosen carefully so that the buffer is exactly 4096 bytes
+  IOBuf buf(IOBuf::CREATE, 4096 - reservedSize);
+  buf.advance(10);
+  memcpy(buf.writableData(), data, sizeof(data));
+  buf.append(sizeof(data));
+  EXPECT_EQ(sizeof(data), buf.length());
+
+  // Grow the buffer (hopefully in place); this would incorrectly reserve
+  // the 10 bytes of headroom, giving us 10 bytes less than requested.
+  size_t tailroom = 8192 - reservedSize - sizeof(data);
+  buf.reserve(0, tailroom);
+  EXPECT_LE(tailroom, buf.tailroom());
+  EXPECT_EQ(sizeof(data), buf.length());
+  EXPECT_EQ(0, memcmp(data, buf.data(), sizeof(data)));
+}
+
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   gflags::ParseCommandLineFlags(&argc, &argv, true);
-- 
2.34.1