Add the ability to add multiple IOBuf free callbacks

Summary: Add the ability to add multiple IOBuf free callbacks

Reviewed By: yfeldblum

Differential Revision: D15423850

fbshipit-source-id: ad0352b345835b0472ec090b5ff07e30e14532d8

folly/io/IOBuf.cpp

@@ -951,6 +951,18 @@ void IOBuf::freeExtBuffer() noexcept {
   SharedInfo* info = sharedInfo();
   DCHECK(info);
 
+  if (info->observerListHead) {
+    // break the chain
+    info->observerListHead->prev->next = nullptr;
+    auto* entry = info->observerListHead;
+    while (entry) {
+      auto* tmp = entry->next;
+      entry->beforeFreeExtBuffer();
+      delete entry;
+      entry = tmp;
+    }
+  }
+
   if (info->freeFn) {
     info->freeFn(buf_, info->userData);
   } else {

folly/io/IOBuf.h

@@ -34,6 +34,7 @@
 #include <folly/detail/Iterators.h>
 #include <folly/lang/Ordering.h>
 #include <folly/portability/SysUio.h>
+#include <folly/synchronization/MicroSpinLock.h>
 
 // Ignore shadowing warnings within this file, so includers can use -Wshadow.
 FOLLY_PUSH_WARNING
@@ -932,6 +933,29 @@ class IOBuf {
     return info ? info->freeFn : nullptr;
   }
 
+  template <typename Observer>
+  bool appendSharedInfoObserver(Observer&& observer) {
+    SharedInfo* info = sharedInfo();
+    if (!info) {
+      return false;
+    }
+
+    auto* entry =
+        new SharedInfoObserverEntry<Observer>(std::forward<Observer>(observer));
+    std::lock_guard<MicroSpinLock> guard(info->observerListLock);
+    if (!info->observerListHead) {
+      info->observerListHead = entry;
+    } else {
+      // prepend
+      entry->next = info->observerListHead;
+      entry->prev = info->observerListHead->prev;
+      info->observerListHead->prev->next = entry;
+      info->observerListHead->prev = entry;
+    }
+
+    return true;
+  }
+
   /**
    * Return true if all IOBufs in this chain are managed by the usual
    * refcounting mechanism (and so the lifetime of the underlying memory
@@ -1387,6 +1411,28 @@ class IOBuf {
     kFlagMask = kFlagFreeSharedInfo | kFlagMaybeShared
   };
 
+  struct SharedInfoObserverEntryBase {
+    SharedInfoObserverEntryBase* prev{this};
+    SharedInfoObserverEntryBase* next{this};
+
+    virtual ~SharedInfoObserverEntryBase() = default;
+
+    virtual void beforeFreeExtBuffer() const noexcept = 0;
+  };
+
+  template <typename Observer>
+  struct SharedInfoObserverEntry : SharedInfoObserverEntryBase {
+    std::decay_t<Observer> observer;
+
+    explicit SharedInfoObserverEntry(Observer&& obs) noexcept(
+        noexcept(Observer(std::forward<Observer>(obs))))
+        : observer(std::forward<Observer>(obs)) {}
+
+    void beforeFreeExtBuffer() const noexcept final {
+      observer.beforeFreeExtBuffer();
+    }
+  };
+
   struct SharedInfo {
     SharedInfo();
     SharedInfo(FreeFunction fn, void* arg, bool hfs = false);
@@ -1397,9 +1443,11 @@ class IOBuf {
     // hits 0.  If this is null, free() will be used instead.
     FreeFunction freeFn;
     void* userData;
+    SharedInfoObserverEntryBase* observerListHead{nullptr};
     std::atomic<uint32_t> refcount;
     bool externallyShared{false};
     bool useHeapFullStorage{false};
+    MicroSpinLock observerListLock{0};
   };
   // Helper structs for use by operator new and delete
   struct HeapPrefix;

folly/io/test/IOBufTest.cpp

@@ -1604,3 +1604,54 @@ TEST(IOBuf, fillIov2) {
   EXPECT_EQ(0, res.numIovecs);
   EXPECT_EQ(0, res.totalLength);
 }
+
+TEST(IOBuf, FreeFn) {
+  class IOBufFreeObserver {
+   public:
+    using Func = std::function<void()>;
+    explicit IOBufFreeObserver(Func&& func) : func_(std::move(func)) {}
+    void beforeFreeExtBuffer() const noexcept {
+      func_();
+    }
+
+   private:
+    Func func_;
+  };
+
+  // no observers
+  { unique_ptr<IOBuf> iobuf(IOBuf::create(64)); }
+
+  int val = 0;
+  // one observer
+  {
+    unique_ptr<IOBuf> iobuf(IOBuf::create(64));
+
+    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
+      EXPECT_EQ(val, 0);
+      val += 1;
+    })));
+  }
+  EXPECT_EQ(val, 1);
+
+  val = 0;
+  // multiple observers
+  {
+    unique_ptr<IOBuf> iobuf(IOBuf::create(64));
+
+    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
+      EXPECT_EQ(val, 0);
+      val += 1;
+    })));
+
+    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
+      EXPECT_EQ(val, 1);
+      val += 20;
+    })));
+
+    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
+      EXPECT_EQ(val, 21);
+      val += 300;
+    })));
+  }
+  EXPECT_EQ(val, 321);
+}