From: Adam Simpkins <simpkins@fb.com>
Date: Mon, 25 Nov 2013 00:30:22 +0000 (-0800)
Subject: update moveToFbString()'s handling of flags_
X-Git-Tag: v0.22.0~780
X-Git-Url: http://plrg.eecs.uci.edu/git/?p=folly.git;a=commitdiff_plain;h=d9298481e96658f28b9df5083fef1a70a399eba8

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 5bab8716..130e7d66 100644
--- 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 a0ca939f..493c5fad 100644
--- 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;