Fix mem leak due to observer not being freed

Summary: Fix mem leak due to observer not being freed

Reviewed By: yfeldblum

Differential Revision: D16234386

fbshipit-source-id: 54ea9e904cf6cde0e24edf1511c2e3b46a58376f

folly/io/IOBuf.cpp

@@ -130,6 +130,22 @@ IOBuf::SharedInfo::SharedInfo(FreeFunction fn, void* arg, bool hfs)
   refcount.store(1, std::memory_order_relaxed);
 }
 
+void IOBuf::SharedInfo::invokeAndDeleteEachObserver(
+    SharedInfoObserverEntryBase* observerListHead,
+    ObserverCb cb) noexcept {
+  if (observerListHead && cb) {
+    // break the chain
+    observerListHead->prev->next = nullptr;
+    auto* entry = observerListHead;
+    while (entry) {
+      auto* tmp = entry->next;
+      cb(*entry);
+      delete entry;
+      entry = tmp;
+    }
+  }
+}
+
 void IOBuf::SharedInfo::releaseStorage(SharedInfo* info) noexcept {
   if (info->useHeapFullStorage) {
     auto* storageAddr =
@@ -962,17 +978,8 @@ void IOBuf::freeExtBuffer() noexcept {
     free(buf_);
   }
 
-  if (observerListHead) {
-    // break the chain
-    observerListHead->prev->next = nullptr;
-    auto* entry = observerListHead;
-    while (entry) {
-      auto* tmp = entry->next;
-      entry->afterFreeExtBuffer();
-      delete entry;
-      entry = tmp;
-    }
-  }
+  SharedInfo::invokeAndDeleteEachObserver(
+      observerListHead, [](auto& entry) { entry.afterFreeExtBuffer(); });
 }
 
 void IOBuf::allocExtBuffer(
@@ -1017,9 +1024,10 @@ void IOBuf::initExtBuffer(
 }
 
 fbstring IOBuf::moveToFbString() {
-  // we need to save useHeapFullStorage since
+  // we need to save useHeapFullStorage and the observerListHead since
   // sharedInfo() may not be valid after fbstring str
   bool useHeapFullStorage = false;
+  SharedInfoObserverEntryBase* observerListHead = nullptr;
   // malloc-allocated buffers are just fine, everything else needs
   // to be turned into one.
   if (!sharedInfo() || // user owned, not ours to give up
@@ -1032,9 +1040,19 @@ fbstring IOBuf::moveToFbString() {
     // to reallocate; we need 1 byte for NUL terminator.
     coalesceAndReallocate(0, computeChainDataLength(), this, 1);
   } else {
-    // if we do not call coalesceAndReallocate
-    // we might need to call SharedInfo::releaseStorage()
-    useHeapFullStorage = sharedInfo() && sharedInfo()->useHeapFullStorage;
+    auto* info = sharedInfo();
+    if (info) {
+      // if we do not call coalesceAndReallocate
+      // we might need to call SharedInfo::releaseStorage()
+      // and/or SharedInfo::invokeAndDeleteEachObserver()
+      useHeapFullStorage = info->useHeapFullStorage;
+      // save the observerListHead
+      // the coalesceAndReallocate path will call
+      // decrementRefcount and freeExtBuffer if needed
+      // so the observer lis notification is needed here
+      observerListHead = info->observerListHead;
+      info->observerListHead = nullptr;
+    }
   }
 
   // Ensure NUL terminated
@@ -1045,6 +1063,9 @@ fbstring IOBuf::moveToFbString() {
       capacity(),
       AcquireMallocatedString());
 
+  SharedInfo::invokeAndDeleteEachObserver(
+      observerListHead, [](auto& entry) { entry.afterReleaseExtBuffer(); });
+
   if (flags() & kFlagFreeSharedInfo) {
     delete sharedInfo();
   } else {

folly/io/IOBuf.h

@@ -29,6 +29,7 @@
 
 #include <folly/FBString.h>
 #include <folly/FBVector.h>
+#include <folly/Function.h>
 #include <folly/Portability.h>
 #include <folly/Range.h>
 #include <folly/detail/Iterators.h>
@@ -1419,6 +1420,7 @@ class IOBuf {
     virtual ~SharedInfoObserverEntryBase() = default;
 
     virtual void afterFreeExtBuffer() const noexcept = 0;
+    virtual void afterReleaseExtBuffer() const noexcept = 0;
   };
 
   template <typename Observer>
@@ -1432,6 +1434,10 @@ class IOBuf {
     void afterFreeExtBuffer() const noexcept final {
       observer.afterFreeExtBuffer();
     }
+
+    void afterReleaseExtBuffer() const noexcept final {
+      observer.afterReleaseExtBuffer();
+    }
   };
 
   struct SharedInfo {
@@ -1440,6 +1446,11 @@ class IOBuf {
 
     static void releaseStorage(SharedInfo* info) noexcept;
 
+    using ObserverCb = folly::FunctionRef<void(SharedInfoObserverEntryBase&)>;
+    static void invokeAndDeleteEachObserver(
+        SharedInfoObserverEntryBase* observerListHead,
+        ObserverCb cb) noexcept;
+
     // A pointer to a function to call to free the buffer when the refcount
     // hits 0.  If this is null, free() will be used instead.
     FreeFunction freeFn;

folly/io/test/IOBufTest.cpp

@@ -1609,49 +1609,71 @@ TEST(IOBuf, FreeFn) {
   class IOBufFreeObserver {
    public:
     using Func = std::function<void()>;
-    explicit IOBufFreeObserver(Func&& func) : func_(std::move(func)) {}
+    explicit IOBufFreeObserver(Func&& freeFunc, Func&& releaseFunc)
+        : freeFunc_(std::move(freeFunc)),
+          releaseFunc_(std::move(releaseFunc)) {}
     void afterFreeExtBuffer() const noexcept {
-      func_();
+      freeFunc_();
+    }
+    void afterReleaseExtBuffer() const noexcept {
+      releaseFunc_();
     }
 
    private:
-    Func func_;
+    Func freeFunc_;
+    Func releaseFunc_;
   };
 
+  int freeVal = 0;
+  int releaseVal = 0;
+  IOBufFreeObserver observer(
+      [&freeVal]() { freeVal += 1; }, [&releaseVal]() { releaseVal += 1; });
+
   // 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_TRUE(iobuf->appendSharedInfoObserver(observer));
   }
-  EXPECT_EQ(val, 1);
+  EXPECT_EQ(freeVal, 1);
+  EXPECT_EQ(releaseVal, 0);
 
-  val = 0;
-  // multiple observers
+  freeVal = 0;
+  releaseVal = 0;
+  // reserve
   {
     unique_ptr<IOBuf> iobuf(IOBuf::create(64));
+    EXPECT_TRUE(iobuf->appendSharedInfoObserver(observer));
+    iobuf->reserve(0, iobuf->capacity() + 1024);
+    EXPECT_EQ(freeVal, 1);
+    EXPECT_EQ(releaseVal, 0);
+  }
 
-    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
-      EXPECT_EQ(val, 0);
-      val += 1;
-    })));
+  freeVal = 0;
+  releaseVal = 0;
+  // one observer - call moveToFbString
+  {
+    unique_ptr<IOBuf> iobuf(IOBuf::create(64 * 1024));
 
-    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
-      EXPECT_EQ(val, 1);
-      val += 20;
-    })));
+    EXPECT_TRUE(iobuf->appendSharedInfoObserver(observer));
+    auto str = iobuf->moveToFbString().toStdString();
+  }
+  EXPECT_EQ(freeVal, 0);
+  EXPECT_EQ(releaseVal, 1);
 
-    EXPECT_TRUE(iobuf->appendSharedInfoObserver(IOBufFreeObserver([&val]() {
-      EXPECT_EQ(val, 21);
-      val += 300;
-    })));
+  freeVal = 0;
+  releaseVal = 0;
+  // multiple observers
+  {
+    unique_ptr<IOBuf> iobuf(IOBuf::create(64));
+
+    for (size_t i = 0; i < 3; i++) {
+      EXPECT_TRUE(iobuf->appendSharedInfoObserver(observer));
+    }
   }
-  EXPECT_EQ(val, 321);
+  EXPECT_EQ(freeVal, 3);
+  EXPECT_EQ(releaseVal, 0);
 }