IOBuf: add a method to signal the underlying buffer as externally shared
[folly.git] / folly / io / test / IOBufTest.cpp
index 4ee43cf2cfb39b1e1966b19517e913e37bafbc52..c012debdf16037416806b0cd2fb69c67fe614386 100644 (file)
@@ -1,5 +1,5 @@
 /*
- * Copyright 2015 Facebook, Inc.
+ * Copyright 2016 Facebook, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,7 +17,8 @@
 #include <folly/io/IOBuf.h>
 #include <folly/io/TypedIOBuf.h>
 
-#include <gflags/gflags.h>
+#include <cstddef>
+
 #include <boost/random.hpp>
 #include <gtest/gtest.h>
 
@@ -185,6 +186,14 @@ TEST(IOBuf, WrapBuffer) {
   EXPECT_EQ(size3, iobuf3.length());
   EXPECT_EQ(buf3.get(), iobuf3.buffer());
   EXPECT_EQ(size3, iobuf3.capacity());
+
+  const uint32_t size4 = 2345;
+  unique_ptr<uint8_t[]> buf4(new uint8_t[size4]);
+  IOBuf iobuf4 = IOBuf::wrapBufferAsValue(buf4.get(), size4);
+  EXPECT_EQ(buf4.get(), iobuf4.data());
+  EXPECT_EQ(size4, iobuf4.length());
+  EXPECT_EQ(buf4.get(), iobuf4.buffer());
+  EXPECT_EQ(size4, iobuf4.capacity());
 }
 
 TEST(IOBuf, CreateCombined) {
@@ -626,7 +635,6 @@ TEST(IOBuf, Reserve) {
     EXPECT_EQ(0, iob->headroom());
     EXPECT_EQ(100, iob->length());
     const void* p1 = iob->buffer();
-    const uint8_t* d1 = iob->data();
     iob->reserve(100, 2512);  // allocation sizes are multiples of 256
     EXPECT_LE(100, iob->headroom());
     if (folly::usingJEMalloc()) {
@@ -792,11 +800,7 @@ TEST(IOBuf, takeOwnershipUniquePtr) {
 }
 
 TEST(IOBuf, Alignment) {
-  // max_align_t doesn't exist in gcc 4.6.2
-  struct MaxAlign {
-    char c;
-  } __attribute__((__aligned__));
-  size_t alignment = alignof(MaxAlign);
+  size_t alignment = alignof(std::max_align_t);
 
   std::vector<size_t> sizes {0, 1, 64, 256, 1024, 1 << 10};
   for (size_t size : sizes) {
@@ -1152,9 +1156,173 @@ TEST(IOBuf, CopyConstructorAndAssignmentOperator) {
   EXPECT_FALSE(buf->isShared());
 }
 
-int main(int argc, char** argv) {
-  testing::InitGoogleTest(&argc, argv);
-  gflags::ParseCommandLineFlags(&argc, &argv, true);
+TEST(IOBuf, CloneAsValue) {
+  auto buf = IOBuf::create(4096);
+  append(buf, "hello world");
+  {
+    auto buf2 = IOBuf::create(4096);
+    append(buf2, " goodbye");
+    buf->prependChain(std::move(buf2));
+    EXPECT_FALSE(buf->isShared());
+  }
+
+  {
+    auto copy = buf->cloneOneAsValue();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_TRUE(copy.isShared());
+    EXPECT_EQ((void*)buf->data(), (void*)copy.data());
+    EXPECT_TRUE(buf->isChained());
+    EXPECT_FALSE(copy.isChained());
+
+    auto copy2 = buf->cloneAsValue();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_TRUE(copy.isShared());
+    EXPECT_TRUE(copy2.isShared());
+    EXPECT_TRUE(buf->isChained());
+    EXPECT_TRUE(copy2.isChained());
+
+    copy.unshareOne();
+    EXPECT_TRUE(buf->isShared());
+    EXPECT_FALSE(copy.isShared());
+    EXPECT_NE((void*)buf->data(), (void*)copy.data());
+    EXPECT_TRUE(copy2.isShared());
+
+    auto p = reinterpret_cast<const char*>(copy.data());
+    EXPECT_EQ("hello world", std::string(p, copy.length()));
+
+    copy2.coalesce();
+    EXPECT_FALSE(buf->isShared());
+    EXPECT_FALSE(copy.isShared());
+    EXPECT_FALSE(copy2.isShared());
+    EXPECT_FALSE(copy2.isChained());
+
+    auto p2 = reinterpret_cast<const char*>(copy2.data());
+    EXPECT_EQ("hello world goodbye", std::string(p2, copy2.length()));
+  }
+
+  EXPECT_FALSE(buf->isShared());
+}
+
+namespace {
+// Use with string literals only
+std::unique_ptr<IOBuf> wrap(const char* str) {
+  return IOBuf::wrapBuffer(str, strlen(str));
+}
+
+std::unique_ptr<IOBuf> copy(const char* str) {
+  // At least 1KiB of tailroom, to ensure an external buffer
+  return IOBuf::copyBuffer(str, strlen(str), 0, 1024);
+}
+
+std::string toString(const folly::IOBuf& buf) {
+  std::string result;
+  result.reserve(buf.computeChainDataLength());
+  for (auto& b : buf) {
+    result.append(reinterpret_cast<const char*>(b.data()), b.size());
+  }
+  return result;
+}
+
+char* writableStr(folly::IOBuf& buf) {
+  return reinterpret_cast<char*>(buf.writableData());
+}
+
+}  // namespace
+
+TEST(IOBuf, ExternallyShared) {
+  struct Item {
+    Item(const char* src, size_t len) : size(len) {
+      CHECK_LE(len, sizeof(buffer));
+      memcpy(buffer, src, len);
+    }
+    uint32_t refcount{0};
+    uint8_t size;
+    char buffer[256];
+  };
+
+  auto hello = "hello";
+  struct Item it(hello, strlen(hello));
+
+  {
+    auto freeFn = [](void* /* unused */, void* userData) {
+      auto it = static_cast<struct Item*>(userData);
+      it->refcount--;
+    };
+    it.refcount++;
+    auto buf1 = IOBuf::takeOwnership(it.buffer, it.size, freeFn, &it);
+    EXPECT_TRUE(buf1->isManagedOne());
+    EXPECT_FALSE(buf1->isSharedOne());
+
+    buf1->markExternallyShared();
+    EXPECT_TRUE(buf1->isSharedOne());
+
+    {
+      auto buf2 = buf1->clone();
+      EXPECT_TRUE(buf2->isManagedOne());
+      EXPECT_TRUE(buf2->isSharedOne());
+      EXPECT_EQ(buf1->data(), buf2->data());
+      EXPECT_EQ(it.refcount, 1);
+    }
+    EXPECT_EQ(it.refcount, 1);
+  }
+  EXPECT_EQ(it.refcount, 0);
+}
 
-  return RUN_ALL_TESTS();
+TEST(IOBuf, Managed) {
+  auto hello = "hello";
+  auto buf1UP = wrap(hello);
+  auto buf1 = buf1UP.get();
+  EXPECT_FALSE(buf1->isManagedOne());
+  auto buf2UP = copy("world");
+  auto buf2 = buf2UP.get();
+  EXPECT_TRUE(buf2->isManagedOne());
+  auto buf3UP = wrap(hello);
+  auto buf3 = buf3UP.get();
+  auto buf4UP = buf2->clone();
+  auto buf4 = buf4UP.get();
+
+  // buf1 and buf3 share the same memory (but are unmanaged)
+  EXPECT_FALSE(buf1->isManagedOne());
+  EXPECT_FALSE(buf3->isManagedOne());
+  EXPECT_TRUE(buf1->isSharedOne());
+  EXPECT_TRUE(buf3->isSharedOne());
+  EXPECT_EQ(buf1->data(), buf3->data());
+
+  // buf2 and buf4 share the same memory (but are managed)
+  EXPECT_TRUE(buf2->isManagedOne());
+  EXPECT_TRUE(buf4->isManagedOne());
+  EXPECT_TRUE(buf2->isSharedOne());
+  EXPECT_TRUE(buf4->isSharedOne());
+  EXPECT_EQ(buf2->data(), buf4->data());
+
+  buf1->prependChain(std::move(buf2UP));
+  buf1->prependChain(std::move(buf3UP));
+  buf1->prependChain(std::move(buf4UP));
+
+  EXPECT_EQ("helloworldhelloworld", toString(*buf1));
+  EXPECT_FALSE(buf1->isManaged());
+
+  buf1->makeManaged();
+  EXPECT_TRUE(buf1->isManaged());
+
+  // buf1 and buf3 are now unshared (because they were unmanaged)
+  EXPECT_TRUE(buf1->isManagedOne());
+  EXPECT_TRUE(buf3->isManagedOne());
+  EXPECT_FALSE(buf1->isSharedOne());
+  EXPECT_FALSE(buf3->isSharedOne());
+  EXPECT_NE(buf1->data(), buf3->data());
+
+  // buf2 and buf4 are still shared
+  EXPECT_TRUE(buf2->isManagedOne());
+  EXPECT_TRUE(buf4->isManagedOne());
+  EXPECT_TRUE(buf2->isSharedOne());
+  EXPECT_TRUE(buf4->isSharedOne());
+  EXPECT_EQ(buf2->data(), buf4->data());
+
+  // And verify that the truth is what we expect: modify a byte in buf1 and
+  // buf2, see that the change from buf1 is *not* reflected in buf3, but the
+  // change from buf2 is reflected in buf4.
+  writableStr(*buf1)[0] = 'j';
+  writableStr(*buf2)[0] = 'x';
+  EXPECT_EQ("jelloxorldhelloxorld", toString(*buf1));
 }