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

diff --git a/folly/detail/AtFork.cpp b/folly/detail/AtFork.cpp
index e33b57922f39f6441ae08e0ce7ed2a24c07d8563..48c42fd31afc3852e7818c1040136a7a8926d487 100644 (file)
--- 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 c714dd84056c77564186214d56911f717a351239..814115ad24e7add7dbb00052b15699a498bcbb0c 100644 (file)
--- 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 ecf4e67d956006b9850fe67012dc36b61c1b0ca8..4a88e1bff26663d216912ed1ca4871d0713bfd85 100644 (file)
--- 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 (file)
index 0000000..11c8a94
--- /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);
+}