From 7acf192db1befc5a54c8a77e2450176fe587fa1f Mon Sep 17 00:00:00 2001
From: Dave Watson <davejwatson@fb.com>
Date: Mon, 4 Dec 2017 12:57:07 -0800
Subject: [PATCH] add AtFork::unregisterHandler

Summary: Adds an AdFork::unregisterHandler command, such that we can register and unregister fork handlers.

Reviewed By: yfeldblum

Differential Revision: D6439796

fbshipit-source-id: b710152f8cf98371c330b484cd5fe42de44e74ac
---
 folly/detail/AtFork.cpp          | 15 +++++++++-
 folly/detail/AtFork.h            |  3 +-
 folly/detail/ThreadLocalDetail.h |  1 +
 folly/test/AtForkTest.cpp        | 48 ++++++++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 folly/test/AtForkTest.cpp

diff --git a/folly/detail/AtFork.cpp b/folly/detail/AtFork.cpp
index e33b5792..48c42fd3 100644
--- a/folly/detail/AtFork.cpp
+++ b/folly/detail/AtFork.cpp
@@ -27,6 +27,7 @@ namespace detail {
 namespace {
 
 struct AtForkTask {
+  void* object;
   folly::Function<void()> prepare;
   folly::Function<void()> parent;
   folly::Function<void()> child;
@@ -89,12 +90,24 @@ void AtFork::init() {
 }
 
 void AtFork::registerHandler(
+    void* object,
     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)});
+      {object, std::move(prepare), std::move(parent), std::move(child)});
+}
+
+void AtFork::unregisterHandler(void* object) {
+  auto& list = AtForkList::instance();
+  std::lock_guard<std::mutex> lg(list.tasksLock);
+  for (auto it = list.tasks.begin(); it != list.tasks.end(); ++it) {
+    if (it->object == object) {
+      list.tasks.erase(it);
+      return;
+    }
+  }
 }
 
 } // namespace detail
diff --git a/folly/detail/AtFork.h b/folly/detail/AtFork.h
index c714dd84..814115ad 100644
--- a/folly/detail/AtFork.h
+++ b/folly/detail/AtFork.h
@@ -24,10 +24,11 @@ namespace detail {
 struct AtFork {
   static void init();
   static void registerHandler(
+      void* object,
       folly::Function<void()> prepare,
       folly::Function<void()> parent,
       folly::Function<void()> child);
-  static void unregisterHandler();
+  static void unregisterHandler(void* object);
 };
 
 } // namespace detail
diff --git a/folly/detail/ThreadLocalDetail.h b/folly/detail/ThreadLocalDetail.h
index ecf4e67d..4a88e1bf 100644
--- a/folly/detail/ThreadLocalDetail.h
+++ b/folly/detail/ThreadLocalDetail.h
@@ -317,6 +317,7 @@ struct StaticMeta : StaticMetaBase {
             &StaticMeta::getThreadEntrySlow,
             std::is_same<AccessMode, AccessModeStrict>::value) {
     detail::AtFork::registerHandler(
+        this,
         /*prepare*/ &StaticMeta::preFork,
         /*parent*/ &StaticMeta::onForkParent,
         /*child*/ &StaticMeta::onForkChild);
diff --git a/folly/test/AtForkTest.cpp b/folly/test/AtForkTest.cpp
new file mode 100644
index 00000000..11c8a946
--- /dev/null
+++ b/folly/test/AtForkTest.cpp
@@ -0,0 +1,48 @@
+/*
+ * Copyright 2017-present Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <folly/detail/AtFork.h>
+
+#include <folly/portability/GTest.h>
+#include <glog/logging.h>
+
+TEST(ThreadLocal, AtFork) {
+  int foo;
+  bool forked = false;
+  folly::detail::AtFork::registerHandler(
+      &foo, [&] { forked = true; }, [] {}, [] {});
+  auto pid = fork();
+  if (pid) {
+    int status;
+    auto pid2 = wait(&status);
+    EXPECT_EQ(status, 0);
+    EXPECT_EQ(pid, pid2);
+  } else {
+    exit(0);
+  }
+  EXPECT_TRUE(forked);
+  forked = false;
+  folly::detail::AtFork::unregisterHandler(&foo);
+  pid = fork();
+  if (pid) {
+    int status;
+    auto pid2 = wait(&status);
+    EXPECT_EQ(status, 0);
+    EXPECT_EQ(pid, pid2);
+  } else {
+    exit(0);
+  }
+  EXPECT_FALSE(forked);
+}
-- 
2.34.1