From 81fa7fd7ac3de80e645dae5243987f05049627b5 Mon Sep 17 00:00:00 2001
From: Philip Pronin <philipp@fb.com>
Date: Fri, 14 Nov 2014 14:00:47 -0800
Subject: [PATCH] fix potential memory leak in ThreadLocal

Summary:
See this LSan abort: https://phabricator.fb.com/P17233565.

Destructor of the object stored in `folly::ThreadLocal` itself may be using
`folly::ThreadLocal` with the same tag, with the current implementation these
objects may escape cleanup happening on thread exit.

Test Plan:
_build/opt/folly/test/thread_local_test --gtest_filter=ThreadLocalPtr.CreateOnThreadExit

Reviewed By: lucian@fb.com

Subscribers: njormrod, folly-diffs@

FB internal diff: D1682698

Tasks: 5596043

Signature: t1:1682698:1416006810:100aaa5c17cecceeea568165d552d9d7907f38d0
---
 folly/detail/ThreadLocalDetail.h | 29 ++++++++++++++++++++---------
 folly/test/ThreadLocalTest.cpp   | 17 +++++++++++++++++
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/folly/detail/ThreadLocalDetail.h b/folly/detail/ThreadLocalDetail.h
index efaadc57..0c05802c 100644
--- a/folly/detail/ThreadLocalDetail.h
+++ b/folly/detail/ThreadLocalDetail.h
@@ -73,13 +73,15 @@ class CustomDeleter : public DeleterBase {
  * This must be POD, as we memset() it to 0 and memcpy() it around.
  */
 struct ElementWrapper {
-  void dispose(TLPDestructionMode mode) {
-    if (ptr != nullptr) {
-      DCHECK(deleter != nullptr);
-      deleter->dispose(ptr, mode);
-
-      cleanup();
+  bool dispose(TLPDestructionMode mode) {
+    if (ptr == nullptr) {
+      return false;
     }
+
+    DCHECK(deleter != nullptr);
+    deleter->dispose(ptr, mode);
+    cleanup();
+    return true;
   }
 
   void* release() {
@@ -242,7 +244,7 @@ struct StaticMeta {
   }
 
   static void onThreadExit(void* ptr) {
-    auto & meta = instance();
+    auto& meta = instance();
 #if !__APPLE__
     ThreadEntry* threadEntry = getThreadEntry();
 
@@ -257,8 +259,17 @@ struct StaticMeta {
       // No need to hold the lock any longer; the ThreadEntry is private to this
       // thread now that it's been removed from meta.
     }
-    FOR_EACH_RANGE(i, 0, threadEntry->elementsCapacity) {
-      threadEntry->elements[i].dispose(TLPDestructionMode::THIS_THREAD);
+    // NOTE: User-provided deleter / object dtor itself may be using ThreadLocal
+    // with the same Tag, so dispose() calls below may (re)create some of the
+    // elements or even increase elementsCapacity, thus multiple cleanup rounds
+    // may be required.
+    for (bool shouldRun = true; shouldRun; ) {
+      shouldRun = false;
+      FOR_EACH_RANGE(i, 0, threadEntry->elementsCapacity) {
+        if (threadEntry->elements[i].dispose(TLPDestructionMode::THIS_THREAD)) {
+          shouldRun = true;
+        }
+      }
     }
     free(threadEntry->elements);
     threadEntry->elements = nullptr;
diff --git a/folly/test/ThreadLocalTest.cpp b/folly/test/ThreadLocalTest.cpp
index 4493c774..5c92d40c 100644
--- a/folly/test/ThreadLocalTest.cpp
+++ b/folly/test/ThreadLocalTest.cpp
@@ -101,6 +101,23 @@ TEST(ThreadLocalPtr, TestRelease) {
   EXPECT_EQ(10, Widget::totalVal_);
 }
 
+TEST(ThreadLocalPtr, CreateOnThreadExit) {
+  Widget::totalVal_ = 0;
+  ThreadLocal<Widget> w;
+  ThreadLocalPtr<int> tl;
+
+  std::thread([&] {
+      tl.reset(new int(1), [&] (int* ptr, TLPDestructionMode mode) {
+        delete ptr;
+        // This test ensures Widgets allocated here are not leaked.
+        ++w.get()->val_;
+        ThreadLocal<Widget> wl;
+        ++wl.get()->val_;
+      });
+    }).join();
+  EXPECT_EQ(2, Widget::totalVal_);
+}
+
 // Test deleting the ThreadLocalPtr object
 TEST(ThreadLocalPtr, CustomDeleter2) {
   Widget::totalVal_ = 0;
-- 
2.34.1