From 24e1c247327b88fbd3f63d024ba9b53530d64431 Mon Sep 17 00:00:00 2001
From: Andrii Grynenko <andrii@fb.com>
Date: Thu, 3 Sep 2020 01:28:04 -0700
Subject: [PATCH] Don't rely on pthread_atfork when possible

Summary:
Because of https://github.com/google/sanitizers/issues/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    | 19 +++++++++--
 folly/detail/AtFork.cpp | 70 +++++++++++++++++++++++++++++++++++++++++
 folly/detail/AtFork.h   |  3 ++
 3 files changed, 89 insertions(+), 3 deletions(-)

diff --git a/folly/Subprocess.cpp b/folly/Subprocess.cpp
index 8b3f4cf7657..cd9c22d4848 100644
--- a/folly/Subprocess.cpp
+++ b/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
diff --git a/folly/detail/AtFork.cpp b/folly/detail/AtFork.cpp
index 8afb90dcaf7..efa9704b71d 100644
--- a/folly/detail/AtFork.cpp
+++ b/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
diff --git a/folly/detail/AtFork.h b/folly/detail/AtFork.h
index b47efe6e6f1..b37f60638b3 100644
--- a/folly/detail/AtFork.h
+++ b/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