ThreadCachedInts use after free fix

Summary: ThreadCachedInts use after free fix

Reviewed By: djwatson

Differential Revision: D9445106

fbshipit-source-id: 989969e54178b5f908f6ca94cde4088ffb2da490

folly/synchronization/detail/ThreadCachedInts.h

@@ -44,9 +44,10 @@ class ThreadCachedInts {
    public:
     ThreadCachedInts* ints_;
     constexpr Integer(ThreadCachedInts* ints) noexcept
-        : ints_(ints), inc_{}, dec_{} {}
+        : ints_(ints), inc_{}, dec_{}, cache_(ints->int_cache_) {}
     std::atomic<int64_t> inc_[2];
     std::atomic<int64_t> dec_[2];
+    Integer*& cache_; // reference to the cached ptr
     ~Integer() noexcept {
       // Increment counts must be set before decrement counts
       ints_->orphan_inc_[0].fetch_add(
@@ -60,6 +61,8 @@ class ThreadCachedInts {
           dec_[1].load(std::memory_order_relaxed), std::memory_order_relaxed);
       ints_->waiting_.store(0, std::memory_order_release);
       ints_->waiting_.futexWake();
+      // reset the cache_ on destructor so we can handle the delete/recreate
+      cache_ = nullptr;
     }
   };
   folly::ThreadLocalPtr<Integer, Tag> cs_;

folly/test/ThreadLocalTest.cpp

@@ -42,6 +42,7 @@
 #include <folly/portability/GTest.h>
 #include <folly/portability/Unistd.h>
 #include <folly/synchronization/Baton.h>
+#include <folly/synchronization/detail/ThreadCachedInts.h>
 #include <folly/system/ThreadId.h>
 
 using namespace folly;
@@ -409,6 +410,39 @@ TEST(ThreadLocal, Movable2) {
   EXPECT_EQ(4, tls.size());
 }
 
+namespace {
+class ThreadCachedIntWidget {
+ public:
+  ThreadCachedIntWidget() {}
+
+  ~ThreadCachedIntWidget() {
+    if (ints_) {
+      ints_->increment(0);
+    }
+  }
+
+  void set(detail::ThreadCachedInts<void>* ints) {
+    ints_ = ints;
+  }
+
+ private:
+  detail::ThreadCachedInts<void>* ints_{nullptr};
+};
+} // namespace
+
+TEST(ThreadLocal, TCICreateOnThreadExit) {
+  detail::ThreadCachedInts<void> ints;
+  ThreadLocal<ThreadCachedIntWidget> w;
+
+  std::thread([&] {
+    // make sure the ints object is created
+    ints.increment(1);
+    // now the widget
+    w->set(&ints);
+  })
+      .join();
+}
+
 namespace {
 
 constexpr size_t kFillObjectSize = 300;