add AtFork::unregisterHandler
authorDave Watson <davejwatson@fb.com>
Mon, 4 Dec 2017 20:57:07 +0000 (12:57 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Mon, 4 Dec 2017 21:08:16 +0000 (13:08 -0800)
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
folly/detail/AtFork.h
folly/detail/ThreadLocalDetail.h
folly/test/AtForkTest.cpp [new file with mode: 0644]

index e33b579..48c42fd 100644 (file)
@@ -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
index c714dd8..814115a 100644 (file)
@@ -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
index ecf4e67..4a88e1b 100644 (file)
@@ -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
@@ -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);
+}