update moveToFbString()'s handling of flags_

Summary:
Improve moveToFbString()'s code which determines if we actually have a
buffer that was allocated with malloc().

The old code simply checked if flags_ == kFlagExt.  This check is rather
fragile, and relies on very specific behavior from the existing
construction methods.  It also unnecessarily forced reallocation in some
cases.

This updates the code to have more specific checks for the flags and
fields that actually matter.  In particular, even if we have an external
buffer, if sharedInfo->freeFn is set then it is not managed by malloc().
The old check happened to work because takeOwnership() was the only
function that set set freeFn, and it also set kFlagFreeSharedInfo.

This also improves the code so that it won't unnecessarily reallocate
the buffer if kFlagMaybeShared is set but the buffer isn't really
shared.  The code will also no longer unnecessarily reallocates the
buffer just because kFlagFreeSharedInfo is set.

Test Plan:
Updated the moveToFbString() test to also test with buffers created
using takeOwnership() and wrapBuffer().

Reviewed By: davejwatson@fb.com

FB internal diff: D1072304

@override-unit-failures

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index 5bab8716c9ce0722865dc4c6396d6990af5ef056..130e7d665493a0b4e6969cb392ae911c6f6b27bc 100644 (file)
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -648,9 +648,10 @@ void IOBuf::initExtBuffer(uint8_t* buf, size_t mallocSize,
 }
 
 fbstring IOBuf::moveToFbString() {
-  // Externally allocated buffers (malloc) are just fine, everything else needs
+  // malloc-allocated buffers are just fine, everything else needs
   // to be turned into one.
-  if (flags_ != kFlagExt ||  // not malloc()-ed
+  if ((flags_ & (kFlagExt | kFlagUserOwned)) != kFlagExt ||  // not malloc()-ed
+      ext_.sharedInfo->freeFn != nullptr || // not malloc()-ed
       headroom() != 0 ||     // malloc()-ed block doesn't start at beginning
       tailroom() == 0 ||     // no room for NUL terminator
       isShared() ||          // shared
@@ -666,6 +667,10 @@ fbstring IOBuf::moveToFbString() {
                length(),  capacity(),
                AcquireMallocatedString());
 
+  if (flags_ & kFlagFreeSharedInfo) {
+    delete ext_.sharedInfo;
+  }
+
   // Reset to internal buffer.
   flags_ = 0;
   clear();

diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index a0ca939fa401242cb5314ff522ca3f148041db9c..493c5fad7dbc8ca57943106bd333c9401af2dfb1 100644 (file)
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -494,10 +494,12 @@ TEST(IOBuf, Chaining) {
   EXPECT_EQ(3, iob2->computeChainDataLength());
 }
 
-void reserveTestFreeFn(void* buffer, void* ptr) {
+void testFreeFn(void* buffer, void* ptr) {
   uint32_t* freeCount = static_cast<uint32_t*>(ptr);;
   delete[] static_cast<uint8_t*>(buffer);
-  ++(*freeCount);
+  if (freeCount) {
+    ++(*freeCount);
+  }
 };
 
 TEST(IOBuf, Reserve) {
@@ -575,7 +577,7 @@ TEST(IOBuf, Reserve) {
   {
     uint32_t freeCount{0};
     uint8_t* buf = new uint8_t[100];
-    auto iob = IOBuf::takeOwnership(buf, 100, reserveTestFreeFn, &freeCount);
+    auto iob = IOBuf::takeOwnership(buf, 100, testFreeFn, &freeCount);
     iob->reserve(0, 2000);
     EXPECT_EQ(0, iob->headroom());
     EXPECT_LE(2000, iob->tailroom());
@@ -732,13 +734,19 @@ TEST(TypedIOBuf, Simple) {
     EXPECT_EQ(i, typed.data()[i]);
   }
 }
+enum BufType {
+  CREATE,
+  TAKE_OWNERSHIP_MALLOC,
+  TAKE_OWNERSHIP_CUSTOM,
+  USER_OWNED,
+};
 
 // chain element size, number of elements in chain, shared
 class MoveToFbStringTest
-  : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool>> {
+  : public ::testing::TestWithParam<std::tr1::tuple<int, int, bool, BufType>> {
  protected:
   void SetUp() {
-    std::tr1::tie(elementSize_, elementCount_, shared_) = GetParam();
+    std::tr1::tie(elementSize_, elementCount_, shared_, type_) = GetParam();
     buf_ = makeBuf();
     for (int i = 0; i < elementCount_ - 1; ++i) {
       buf_->prependChain(makeBuf());
@@ -753,9 +761,36 @@ class MoveToFbStringTest
   }
 
   std::unique_ptr<IOBuf> makeBuf() {
-    auto buf = IOBuf::create(elementSize_);
-    memset(buf->writableTail(), 'x', elementSize_);
-    buf->append(elementSize_);
+    unique_ptr<IOBuf> buf;
+    switch (type_) {
+      case CREATE:
+        buf = IOBuf::create(elementSize_);
+        buf->append(elementSize_);
+        break;
+      case TAKE_OWNERSHIP_MALLOC: {
+        void* data = malloc(elementSize_);
+        if (!data) {
+          throw std::bad_alloc();
+        }
+        buf = IOBuf::takeOwnership(data, elementSize_);
+        break;
+      }
+      case TAKE_OWNERSHIP_CUSTOM: {
+        uint8_t* data = new uint8_t[elementSize_];
+        buf = IOBuf::takeOwnership(data, elementSize_, testFreeFn);
+        break;
+      }
+      case USER_OWNED: {
+        unique_ptr<uint8_t[]> data(new uint8_t[elementSize_]);
+        buf = IOBuf::wrapBuffer(data.get(), elementSize_);
+        ownedBuffers_.emplace_back(std::move(data));
+        break;
+      }
+      default:
+        throw std::invalid_argument("unexpected buffer type parameter");
+        break;
+    }
+    memset(buf->writableData(), 'x', elementSize_);
     return buf;
   }
 
@@ -772,8 +807,10 @@ class MoveToFbStringTest
   int elementSize_;
   int elementCount_;
   bool shared_;
+  BufType type_;
   std::unique_ptr<IOBuf> buf_;
   std::unique_ptr<IOBuf> buf2_;
+  std::vector<std::unique_ptr<uint8_t[]>> ownedBuffers_;
 };
 
 TEST_P(MoveToFbStringTest, Simple) {
@@ -789,7 +826,9 @@ INSTANTIATE_TEST_CASE_P(
     ::testing::Combine(
         ::testing::Values(0, 1, 24, 256, 1 << 10, 1 << 20),  // element size
         ::testing::Values(1, 2, 10),                         // element count
-        ::testing::Bool()));                                 // shared
+        ::testing::Bool(),                                   // shared
+        ::testing::Values(CREATE, TAKE_OWNERSHIP_MALLOC,
+                          TAKE_OWNERSHIP_CUSTOM, USER_OWNED)));
 
 TEST(IOBuf, getIov) {
   uint32_t fillSeed = 0xdeadbeef;