From 523ca0edb14fd73137660cc7f460cac59170b941 Mon Sep 17 00:00:00 2001
From: Huapeng Zhou <hzhou@fb.com>
Date: Fri, 24 Jun 2016 10:45:37 -0700
Subject: [PATCH] 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          |  8 ++++++++
 folly/io/IOBuf.h            | 29 +++++++++++++++++++++++++++
 folly/io/test/IOBufTest.cpp | 39 +++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+)

diff --git a/folly/io/IOBuf.cpp b/folly/io/IOBuf.cpp
index 763522b7..bd2bab04 100644
--- 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 1fb26b97..1727b6e4 100644
--- 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 f2308b7c..c012debd 100644
--- 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);
-- 
2.34.1