From a9c78aa8cbed43a075fe6abdca65d8124d618a3c Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Thu, 5 May 2016 16:41:33 -0700
Subject: [PATCH] Allow optional atfork hook for singleton destruction

Summary: Grouping all folly::ThreadLocal atfork hooks in one. This allows proper ordering between folly::ThreadLocal hooks and hooks added by other libraries (we always want folly::ThreadLocal hooks to be registered first).

Reviewed By: ericniebler

Differential Revision: D3262666

fb-gh-sync-id: 1aa928b9ddd484580f4a4a7b320e9e64164bb04a
fbshipit-source-id: 1aa928b9ddd484580f4a4a7b320e9e64164bb04a
---
 folly/detail/ThreadLocalDetail.cpp | 76 ++++++++++++++++++++++++++++++
 folly/detail/ThreadLocalDetail.h   | 19 ++++----
 2 files changed, 84 insertions(+), 11 deletions(-)

diff --git a/folly/detail/ThreadLocalDetail.cpp b/folly/detail/ThreadLocalDetail.cpp
index 4a8f2b66..7dcb22d4 100644
--- a/folly/detail/ThreadLocalDetail.cpp
+++ b/folly/detail/ThreadLocalDetail.cpp
@@ -15,6 +15,9 @@
  */
 #include <folly/ThreadLocal.h>
 
+#include <list>
+#include <mutex>
+
 namespace folly { namespace threadlocal_detail {
 
 StaticMetaBase::StaticMetaBase(ThreadEntry* (*threadEntry)())
@@ -222,6 +225,79 @@ void StaticMetaBase::reserve(EntryID* id) {
   free(reallocated);
 }
 
+namespace {
+
+struct AtForkTask {
+  folly::Function<void()> prepare;
+  folly::Function<void()> parent;
+  folly::Function<void()> child;
+};
+
+class AtForkList {
+ public:
+  static AtForkList& instance() {
+    static auto instance = new AtForkList();
+    return *instance;
+  }
+
+  static void prepare() noexcept {
+    instance().tasksLock.lock();
+    auto& tasks = instance().tasks;
+    for (auto task = tasks.rbegin(); task != tasks.rend(); ++task) {
+      task->prepare();
+    }
+  }
+
+  static void parent() noexcept {
+    auto& tasks = instance().tasks;
+    for (auto& task : tasks) {
+      task.parent();
+    }
+    instance().tasksLock.unlock();
+  }
+
+  static void child() noexcept {
+    auto& tasks = instance().tasks;
+    for (auto& task : tasks) {
+      task.child();
+    }
+    instance().tasksLock.unlock();
+  }
+
+  std::mutex tasksLock;
+  std::list<AtForkTask> tasks;
+
+ private:
+  AtForkList() {
+#if FOLLY_HAVE_PTHREAD_ATFORK
+    int ret = pthread_atfork(
+        &AtForkList::prepare, &AtForkList::parent, &AtForkList::child);
+    checkPosixError(ret, "pthread_atfork failed");
+#elif !__ANDROID__ && !defined(_MSC_VER)
+// pthread_atfork is not part of the Android NDK at least as of n9d. If
+// something is trying to call native fork() directly at all with Android's
+// process management model, this is probably the least of the problems.
+//
+// But otherwise, this is a problem.
+#warning pthread_atfork unavailable
+#endif
+  }
+};
+}
+
+void StaticMetaBase::initAtFork() {
+  AtForkList::instance();
+}
+
+void StaticMetaBase::registerAtFork(
+    folly::Function<void()> prepare,
+    folly::Function<void()> parent,
+    folly::Function<void()> child) {
+  std::lock_guard<std::mutex> lg(AtForkList::instance().tasksLock);
+  AtForkList::instance().tasks.push_back(
+      {std::move(prepare), std::move(parent), std::move(child)});
+}
+
 FOLLY_STATIC_CTOR_PRIORITY_MAX
 PthreadKeyUnregister PthreadKeyUnregister::instance_;
 }}
diff --git a/folly/detail/ThreadLocalDetail.h b/folly/detail/ThreadLocalDetail.h
index 65ecfc67..29e06b1f 100644
--- a/folly/detail/ThreadLocalDetail.h
+++ b/folly/detail/ThreadLocalDetail.h
@@ -29,6 +29,7 @@
 
 #include <folly/Exception.h>
 #include <folly/Foreach.h>
+#include <folly/Function.h>
 #include <folly/Malloc.h>
 #include <folly/MicroSpinLock.h>
 #include <folly/Portability.h>
@@ -272,6 +273,12 @@ struct StaticMetaBase {
 
   ElementWrapper& get(EntryID* ent);
 
+  static void initAtFork();
+  static void registerAtFork(
+      folly::Function<void()> prepare,
+      folly::Function<void()> parent,
+      folly::Function<void()> child);
+
   uint32_t nextId_;
   std::vector<uint32_t> freeIds_;
   std::mutex lock_;
@@ -290,20 +297,10 @@ struct StaticMetaBase {
 template <class Tag>
 struct StaticMeta : StaticMetaBase {
   StaticMeta() : StaticMetaBase(&StaticMeta::getThreadEntrySlow) {
-#if FOLLY_HAVE_PTHREAD_ATFORK
-    int ret = pthread_atfork(
+    registerAtFork(
         /*prepare*/ &StaticMeta::preFork,
         /*parent*/ &StaticMeta::onForkParent,
         /*child*/ &StaticMeta::onForkChild);
-    checkPosixError(ret, "pthread_atfork failed");
-#elif !__ANDROID__ && !defined(_MSC_VER)
-    // pthread_atfork is not part of the Android NDK at least as of n9d. If
-    // something is trying to call native fork() directly at all with Android's
-    // process management model, this is probably the least of the problems.
-    //
-    // But otherwise, this is a problem.
-    #warning pthread_atfork unavailable
-#endif
   }
 
   static StaticMeta<Tag>& instance() {
-- 
2.34.1