X-Git-Url: http://plrg.eecs.uci.edu/git/?p=folly.git;a=blobdiff_plain;f=folly%2Fio%2Ftest%2FIOBufTest.cpp;h=493c5fad7dbc8ca57943106bd333c9401af2dfb1;hp=3bb86cf17927e17e10d84c95715486c01abffb5f;hb=d9298481e96658f28b9df5083fef1a70a399eba8;hpb=7fb1e3d7114111532bde7f15e17b1e5df89616cf

diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index 3bb86cf1..493c5fad 100644
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -28,6 +28,7 @@
 #include "folly/Range.h"
 
 using folly::fbstring;
+using folly::fbvector;
 using folly::IOBuf;
 using folly::TypedIOBuf;
 using folly::StringPiece;
@@ -493,6 +494,14 @@ TEST(IOBuf, Chaining) {
   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);
@@ -554,6 +563,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, testFreeFn, &freeCount);
+    iob->reserve(0, 2000);
+    EXPECT_EQ(0, iob->headroom());
+    EXPECT_LE(2000, iob->tailroom());
+    EXPECT_EQ(1, freeCount);
+  }
 }
 
 TEST(IOBuf, copyBuffer) {
@@ -705,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());
@@ -726,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;
   }
 
@@ -745,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) {
@@ -762,7 +826,54 @@ 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;
+  boost::mt19937 gen(fillSeed);
+
+  size_t len = 4096;
+  size_t count = 32;
+  auto buf = IOBuf::create(len + 1);
+  buf->append(rand() % len + 1);
+  fillBuf(buf.get(), gen);
+
+  for (size_t i = 0; i < count - 1; i++) {
+    auto buf2 = IOBuf::create(len + 1);
+    buf2->append(rand() % len + 1);
+    fillBuf(buf2.get(), gen);
+    buf->prependChain(std::move(buf2));
+  }
+  EXPECT_EQ(count, buf->countChainElements());
+
+  auto iov = buf->getIov();
+  EXPECT_EQ(count, iov.size());
+
+  IOBuf const* p = buf.get();
+  for (size_t i = 0; i < count; i++, p = p->next()) {
+    EXPECT_EQ(p->data(), iov[i].iov_base);
+    EXPECT_EQ(p->length(), iov[i].iov_len);
+  }
+
+  // an empty buf should be skipped in the iov.
+  buf->next()->clear();
+  iov = buf->getIov();
+  EXPECT_EQ(count - 1, iov.size());
+  EXPECT_EQ(buf->next()->next()->data(), iov[1].iov_base);
+
+  // same for the first one being empty
+  buf->clear();
+  iov = buf->getIov();
+  EXPECT_EQ(count - 2, iov.size());
+  EXPECT_EQ(buf->next()->next()->data(), iov[0].iov_base);
+
+  // and the last one
+  buf->prev()->clear();
+  iov = buf->getIov();
+  EXPECT_EQ(count - 3, iov.size());
+}
 
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);