fix IOBuf::reserve() when operating on a user-supplied buffer

Summary:
The IOBuf::reserveSlow() code assumed that external buffers were always
allocated with malloc.  This would cause problems when if reserve() was
ever used on a buffer created with IOBuf::takeOwnership().

This changes the code to now check if a user-specified free function has
been supplied.  If so, then it does not try to use realloc()/rallocm(),
and it now frees the old buffer using the specified free function rather
than free().

User-supplied buffers also have a separately allocated SharedInfo
object, which must be freed when we no longer need it.

Test Plan:
Added additional unit tests to check calling reserve() on IOBufs created
with IOBuf::takeOwnership().

Reviewed By: [email protected]

FB internal diff: D1072292

folly/io/IOBuf.cpp

@@ -453,18 +453,7 @@ void IOBuf::decrementRefcount() {
   }
 
   // We were the last user.  Free the buffer
-  if (ext_.sharedInfo->freeFn != NULL) {
-    try {
-      ext_.sharedInfo->freeFn(ext_.buf, ext_.sharedInfo->userData);
-    } catch (...) {
-      // The user's free function should never throw.  Otherwise we might
-      // throw from the IOBuf destructor.  Other code paths like coalesce()
-      // also assume that decrementRefcount() cannot throw.
-      abort();
-    }
-  } else {
-    free(ext_.buf);
-  }
+  freeExtBuffer();
 
   // Free the SharedInfo if it was allocated separately.
   //
@@ -485,6 +474,11 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) {
   size_t newCapacity = (size_t)length_ + minHeadroom + minTailroom;
   DCHECK_LT(newCapacity, UINT32_MAX);
 
+  // reserveSlow() is dangerous if anyone else is sharing the buffer, as we may
+  // reallocate and free the original buffer.  It should only ever be called if
+  // we are the only user of the buffer.
+  DCHECK(!isSharedOne());
+
   // We'll need to reallocate the buffer.
   // There are a few options.
   // - If we have enough total room, move the data around in the buffer
@@ -511,11 +505,15 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) {
   uint32_t newHeadroom = 0;
   uint32_t oldHeadroom = headroom();
 
-  if ((flags_ & kFlagExt) && length_ != 0 && oldHeadroom >= minHeadroom) {
+  // If we have a buffer allocated with malloc and we just need more tailroom,
+  // try to use realloc()/rallocm() to grow the buffer in place.
+  if ((flags_ & (kFlagExt | kFlagUserOwned)) == kFlagExt &&
+      (ext_.sharedInfo->freeFn == nullptr) &&
+      length_ != 0 && oldHeadroom >= minHeadroom) {
     if (usingJEMalloc()) {
       size_t headSlack = oldHeadroom - minHeadroom;
       // We assume that tailroom is more useful and more important than
-      // tailroom (not least because realloc / rallocm allow us to grow the
+      // headroom (not least because realloc / rallocm allow us to grow the
       // buffer at the tail, but not at the head)  So, if we have more headroom
       // than we need, we consider that "wasted".  We arbitrarily define "too
       // much" headroom to be 25% of the capacity.
@@ -559,8 +557,8 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) {
     }
     newBuffer = static_cast<uint8_t*>(p);
     memcpy(newBuffer + minHeadroom, data_, length_);
-    if (flags_ & kFlagExt) {
-      free(ext_.buf);
+    if ((flags_ & (kFlagExt | kFlagUserOwned)) == kFlagExt) {
+      freeExtBuffer();
     }
     newHeadroom = minHeadroom;
   }
@@ -569,8 +567,11 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) {
   uint32_t cap;
   initExtBuffer(newBuffer, newAllocatedCapacity, &info, &cap);
 
-  flags_ = kFlagExt;
+  if (flags_ & kFlagFreeSharedInfo) {
+    delete ext_.sharedInfo;
+  }
 
+  flags_ = kFlagExt;
   ext_.capacity = cap;
   ext_.type = kExtAllocated;
   ext_.buf = newBuffer;
@@ -579,6 +580,23 @@ void IOBuf::reserveSlow(uint32_t minHeadroom, uint32_t minTailroom) {
   // length_ is unchanged
 }
 
+void IOBuf::freeExtBuffer() {
+  DCHECK((flags_ & (kFlagExt | kFlagUserOwned)) == kFlagExt);
+
+  if (ext_.sharedInfo->freeFn) {
+    try {
+      ext_.sharedInfo->freeFn(ext_.buf, ext_.sharedInfo->userData);
+    } catch (...) {
+      // The user's free function should never throw.  Otherwise we might
+      // throw from the IOBuf destructor.  Other code paths like coalesce()
+      // also assume that decrementRefcount() cannot throw.
+      abort();
+    }
+  } else {
+    free(ext_.buf);
+  }
+}
+
 void IOBuf::allocExtBuffer(uint32_t minCapacity,
                            uint8_t** bufReturn,
                            SharedInfo** infoReturn,

folly/io/IOBuf.h

@@ -1078,6 +1078,7 @@ class IOBuf {
       size_t newTailroom);
   void decrementRefcount();
   void reserveSlow(uint32_t minHeadroom, uint32_t minTailroom);
+  void freeExtBuffer();
 
   static size_t goodExtBufferSize(uint32_t minCapacity);
   static void initExtBuffer(uint8_t* buf, size_t mallocSize,

folly/io/test/IOBufTest.cpp

@@ -494,6 +494,12 @@ TEST(IOBuf, Chaining) {
   EXPECT_EQ(3, iob2->computeChainDataLength());
 }
 
+void reserveTestFreeFn(void* buffer, void* ptr) {
+  uint32_t* freeCount = static_cast<uint32_t*>(ptr);;
+  delete[] static_cast<uint8_t*>(buffer);
+  ++(*freeCount);
+};
+
 TEST(IOBuf, Reserve) {
   uint32_t fillSeed = 0x23456789;
   boost::mt19937 gen(fillSeed);
@@ -555,6 +561,26 @@ TEST(IOBuf, Reserve) {
     EXPECT_EQ(0, iob->headroom());
     EXPECT_LE(2000, iob->tailroom());
   }
+
+  // Test reserving from a user-allocated buffer.
+  {
+    uint8_t* buf = static_cast<uint8_t*>(malloc(100));
+    auto iob = IOBuf::takeOwnership(buf, 100);
+    iob->reserve(0, 2000);
+    EXPECT_EQ(0, iob->headroom());
+    EXPECT_LE(2000, iob->tailroom());
+  }
+
+  // Test reserving from a user-allocated with a custom free function.
+  {
+    uint32_t freeCount{0};
+    uint8_t* buf = new uint8_t[100];
+    auto iob = IOBuf::takeOwnership(buf, 100, reserveTestFreeFn, &freeCount);
+    iob->reserve(0, 2000);
+    EXPECT_EQ(0, iob->headroom());
+    EXPECT_LE(2000, iob->tailroom());
+    EXPECT_EQ(1, freeCount);
+  }
 }
 
 TEST(IOBuf, copyBuffer) {