Don't rely on pthread_atfork when possible

Summary:
Because of google/sanitizers#1116, any races detected during pthread_atfork deadlock the process.
For forks issued from folly we can use an instrumented fork version to minimize the amount of code running via pthread_atfork.

Reviewed By: yfeldblum

Differential Revision: D23281093

fbshipit-source-id: fd4d3bf06b7992ff314f631ed899854a8b3f6c4b

folly/Subprocess.cpp

@@ -39,6 +39,7 @@
 #include <folly/Exception.h>
 #include <folly/ScopeGuard.h>
 #include <folly/String.h>
+#include <folly/detail/AtFork.h>
 #include <folly/io/Cursor.h>
 #include <folly/lang/Assume.h>
 #include <folly/portability/Sockets.h>
@@ -444,9 +445,15 @@ void Subprocess::spawnInternal(
     if (options.detach_) {
       // If we are detaching we must use fork() instead of vfork() for the first
       // fork, since we aren't going to simply call exec() in the child.
-      pid = fork();
+      pid = detail::AtFork::forkInstrumented(fork);
     } else {
-      pid = vfork();
+      if (kIsSanitizeThread) {
+        // TSAN treats vfork as fork, so use the instrumented version
+        // instead
+        pid = detail::AtFork::forkInstrumented(fork);
+      } else {
+        pid = vfork();
+      }
     }
 #ifdef __linux__
   }
@@ -461,7 +468,13 @@ void Subprocess::spawnInternal(
         pid = syscall(SYS_clone, *options.cloneFlags_, 0, nullptr, nullptr);
       } else {
 #endif
-        pid = vfork();
+        if (kIsSanitizeThread) {
+          // TSAN treats vfork as fork, so use the instrumented version
+          // instead
+          pid = detail::AtFork::forkInstrumented(fork);
+        } else {
+          pid = vfork();
+        }
 #ifdef __linux__
       }
 #endif

folly/detail/AtFork.cpp

@@ -30,6 +30,51 @@ namespace detail {
 
 namespace {
 
+#if !FOLLY_MOBILE && !defined(__APPLE__) && !defined(_MSC_VER)
+#define FOLLY_ATFORK_USE_FOLLY_TLS 1
+#else
+#define FOLLY_ATFORK_USE_FOLLY_TLS 0
+#endif
+
+template <bool enabled = FOLLY_ATFORK_USE_FOLLY_TLS>
+struct SkipAtForkHandlers;
+
+template <>
+struct SkipAtForkHandlers<false> {
+  static bool available() {
+    return false;
+  }
+
+  static bool get() {
+    return false;
+  }
+
+  [[noreturn]] static void set(bool) {
+    std::terminate();
+  }
+};
+
+#if FOLLY_ATFORK_USE_FOLLY_TLS
+template <>
+struct SkipAtForkHandlers<true> {
+  static bool available() {
+    return true;
+  }
+
+  static bool get() {
+    return value_;
+  }
+
+  static void set(bool value) {
+    value_ = value;
+  }
+
+ private:
+  static FOLLY_TLS bool value_;
+};
+FOLLY_TLS bool SkipAtForkHandlers<true>::value_{false};
+#endif
+
 struct AtForkTask {
   void const* handle;
   folly::Function<bool()> prepare;
@@ -45,6 +90,9 @@ class AtForkList {
   }
 
   static void prepare() noexcept {
+    if (SkipAtForkHandlers<>::get()) {
+      return;
+    }
     instance().tasksLock.lock();
     while (true) {
       auto& tasks = instance().tasks;
@@ -64,6 +112,9 @@ class AtForkList {
   }
 
   static void parent() noexcept {
+    if (SkipAtForkHandlers<>::get()) {
+      return;
+    }
     auto& tasks = instance().tasks;
     for (auto& task : tasks) {
       task.parent();
@@ -72,6 +123,9 @@ class AtForkList {
   }
 
   static void child() noexcept {
+    if (SkipAtForkHandlers<>::get()) {
+      return;
+    }
     // if we fork a multithreaded process
     // some of the TSAN mutexes might be locked
     // so we just enable ignores for everything
@@ -138,5 +192,21 @@ void AtFork::unregisterHandler(void const* handle) {
   }
 }
 
+pid_t AtFork::forkInstrumented(fork_t forkFn) {
+  if (SkipAtForkHandlers<>::available()) {
+    AtForkList::prepare();
+    SkipAtForkHandlers<>::set(true);
+    auto ret = forkFn();
+    SkipAtForkHandlers<>::set(false);
+    if (ret) {
+      AtForkList::parent();
+    } else {
+      AtForkList::child();
+    }
+    return ret;
+  } else {
+    return forkFn();
+  }
+}
 } // namespace detail
 } // namespace folly

folly/detail/AtFork.h

@@ -30,6 +30,9 @@ struct AtFork {
       folly::Function<void()> parent,
       folly::Function<void()> child);
   static void unregisterHandler(void const* handle);
+
+  using fork_t = pid_t();
+  static pid_t forkInstrumented(fork_t forkFn);
 };
 
 } // namespace detail