IOBuf: add a method to signal the underlying buffer as externally shared

Summary:
There are use cases where 1). the underlying buffer is externally managed (e.g. by a slab allocator) and 2). we need to do bookkeeping when the wrapped IOBuf gets destroyed (e.g. reference counting). This diff adds a another method to mark the underlying buffer as shared with the external memory management mechanism.

The `takeOwnership` doesn't meet the criteria since it assumes the ownership of the buffer, while in this case we need to signal it as externally managed so that hopefully callers won't try to modify the underlying buffer.

Reviewed By: simpkins

Differential Revision: D2662954

fbshipit-source-id: e908c3ebeeefe9a5d332c75070f377fb1dad5acb

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index 763522b74abc1b2134c9d20966a522f69f0641af..bd2bab0412d4293c833fcdf1edc02721f74c8e0e 100644 (file)
--- a/folly/io/IOBuf.cpp
+++ b/folly/io/IOBuf.cpp
@@ -581,6 +581,14 @@ void IOBuf::unshareChained() {
   coalesceSlow();
 }
 
+void IOBuf::markExternallyShared() {
+  IOBuf* current = this;
+  do {
+    current->markExternallySharedOne();
+    current = current->next_;
+  } while (current != this);
+}
+
 void IOBuf::makeManagedChained() {
   assert(isChained());
 

diff --git a/folly/io/IOBuf.h b/folly/io/IOBuf.h
index 1fb26b974145e75500171863f60e4e4bb01f23b3..1727b6e46bcb516054a39a04a4a4cc573932eba6 100644 (file)
--- a/folly/io/IOBuf.h
+++ b/folly/io/IOBuf.h
@@ -923,6 +923,10 @@ class IOBuf {
       return true;
     }
 
+    if (UNLIKELY(sharedInfo()->externallyShared)) {
+      return true;
+    }
+
     if (LIKELY(!(flags() & kFlagMaybeShared))) {
       return false;
     }
@@ -982,6 +986,30 @@ class IOBuf {
     }
   }
 
+  /**
+   * Mark the underlying buffers in this chain as shared with external memory
+   * management mechanism. This will make isShared() always returns true.
+   *
+   * This function is not thread-safe, and only safe to call immediately after
+   * creating an IOBuf, before it has been shared with other threads.
+   */
+  void markExternallyShared();
+
+  /**
+   * Mark the underlying buffer that this IOBuf refers to as shared with
+   * external memory management mechanism. This will make isSharedOne() always
+   * returns true.
+   *
+   * This function is not thread-safe, and only safe to call immediately after
+   * creating an IOBuf, before it has been shared with other threads.
+   */
+  void markExternallySharedOne() {
+    SharedInfo* info = sharedInfo();
+    if (info) {
+      info->externallyShared = true;
+    }
+  }
+
   /**
    * Ensure that the memory that IOBufs in this chain refer to will continue to
    * be allocated for as long as the IOBufs of the chain (or any clone()s
@@ -1223,6 +1251,7 @@ class IOBuf {
     FreeFunction freeFn;
     void* userData;
     std::atomic<uint32_t> refcount;
+    bool externallyShared{false};
   };
   // Helper structs for use by operator new and delete
   struct HeapPrefix;

diff --git a/folly/io/test/IOBufTest.cpp b/folly/io/test/IOBufTest.cpp
index f2308b7c1b73019272108f3958eff8f4f4219178..c012debdf16037416806b0cd2fb69c67fe614386 100644 (file)
--- a/folly/io/test/IOBufTest.cpp
+++ b/folly/io/test/IOBufTest.cpp
@@ -1229,6 +1229,45 @@ char* writableStr(folly::IOBuf& buf) {
 
 }  // 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);
+}
+
 TEST(IOBuf, Managed) {
   auto hello = "hello";
   auto buf1UP = wrap(hello);