update moveToFbString()'s handling of flags_

[folly.git] / folly / io / test / IOBufTest.cpp

diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index b17be70887072539d06b5fadbe6e8f258a719a2f..493c5fad7dbc8ca57943106bd333c9401af2dfb1 100644 (file)
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -494,6 +494,14 @@ TEST(IOBuf, Chaining) {
   EXPECT_EQ(3, iob2->computeChainDataLength());
 }
 
   EXPECT_EQ(3, iob2->computeChainDataLength());
 }
 

+void testFreeFn(void* buffer, void* ptr) {
+  uint32_t* freeCount = static_cast<uint32_t*>(ptr);;
+  delete[] static_cast<uint8_t*>(buffer);
+  if (freeCount) {
+    ++(*freeCount);
+  }
+};
+

 TEST(IOBuf, Reserve) {
   uint32_t fillSeed = 0x23456789;
   boost::mt19937 gen(fillSeed);
 TEST(IOBuf, Reserve) {
   uint32_t fillSeed = 0x23456789;
   boost::mt19937 gen(fillSeed);


@@ -555,6 +563,26 @@ TEST(IOBuf, Reserve) {
     EXPECT_EQ(0, iob->headroom());
     EXPECT_LE(2000, iob->tailroom());
   }
     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, testFreeFn, &freeCount);
+    iob->reserve(0, 2000);
+    EXPECT_EQ(0, iob->headroom());
+    EXPECT_LE(2000, iob->tailroom());
+    EXPECT_EQ(1, freeCount);
+  }

 }
 
 TEST(IOBuf, copyBuffer) {
 }
 
 TEST(IOBuf, copyBuffer) {


@@ -706,13 +734,19 @@ TEST(TypedIOBuf, Simple) {
     EXPECT_EQ(i, typed.data()[i]);
   }
 }
     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
 
 // 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() {
  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());
     buf_ = makeBuf();
     for (int i = 0; i < elementCount_ - 1; ++i) {
       buf_->prependChain(makeBuf());


@@ -727,9 +761,36 @@ class MoveToFbStringTest
   }
 
   std::unique_ptr<IOBuf> makeBuf() {
   }
 
   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;
   }
 
     return buf;
   }
 


@@ -746,8 +807,10 @@ class MoveToFbStringTest
   int elementSize_;
   int elementCount_;
   bool shared_;
   int elementSize_;
   int elementCount_;
   bool shared_;

+  BufType type_;

   std::unique_ptr<IOBuf> buf_;
   std::unique_ptr<IOBuf> buf2_;
   std::unique_ptr<IOBuf> buf_;
   std::unique_ptr<IOBuf> buf2_;

+  std::vector<std::unique_ptr<uint8_t[]>> ownedBuffers_;

 };
 
 TEST_P(MoveToFbStringTest, Simple) {
 };
 
 TEST_P(MoveToFbStringTest, Simple) {


@@ -763,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::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;
 
 TEST(IOBuf, getIov) {
   uint32_t fillSeed = 0xdeadbeef;