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

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());
 

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;

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);