Make JIT mutexes safe for fork

Summary:
Currently we can have various JIT mutexes locked when a fork happens. This was a mostly a non-issue in the GIL build because Python code couldn't fork while the JIT was running. But w/ background compile it can. It would have been problematic if C code forked though.

It's also presumably a lurking issue in the free-threaded builds.

This makes it so we have proper safety for our locks around forking.

Reviewed By: alexmalyshev

Differential Revision: D115804360

fbshipit-source-id: adc835ed8787f64f36d6060dc3a242405a23e8ea

Author: DinoV

cinderx/Common/hugepages.cpp

@@ -16,6 +16,7 @@
 #include <cstdlib>
 #include <cstring>
 #include <memory>
+#include <new>
 #include <vector>
 
 #ifndef WIN32
@@ -119,4 +120,19 @@ void HugePageArena::afterForkChild() {
 #endif
 }
 
+void HugePageArena::atForkPrepare() {
+  mutex_.lock();
+}
+
+void HugePageArena::atForkParent() {
+  mutex_.unlock();
+}
+
+void HugePageArena::atForkChild() {
+  // Reuse the storage to get a fresh, unlocked mutex.  The inherited one is
+  // still locked by atForkPrepare() and destroying a locked mutex is
+  // undefined, so its lifetime is ended without running its destructor.
+  new (&mutex_) std::mutex{};
+}
+
 } // namespace cinderx

cinderx/Common/hugepages.h

@@ -35,6 +35,10 @@ class HugePageArena {
   // Re-establish huge page backing for every chunk after a fork().
   void afterForkChild();
 
+  void atForkPrepare();
+  void atForkParent();
+  void atForkChild();
+
  private:
   struct Chunk {
     void* ptr;

cinderx/Common/slab_arena.cpp

@@ -4,6 +4,9 @@
 
 #include "cinderx/module_state.h"
 
+#include <algorithm>
+#include <new>
+
 #if defined(__linux__) && defined(__aarch64__)
 // On ARM64 we see huge dTLB misses on our inline caches so
 // we put them on huge pages
@@ -22,4 +25,51 @@ std::shared_ptr<HugePageArena> getSharedHugePageArena() {
   return nullptr;
 }
 
+SlabArenaForkRegistry& SlabArenaForkRegistry::get() {
+  // Deliberately leaked: SlabArenas are owned by the CinderX module state,
+  // which can outlive static destructors and would then unregister into a
+  // destroyed object.
+  static auto* registry = new SlabArenaForkRegistry;
+  return *registry;
+}
+
+void SlabArenaForkRegistry::add(std::mutex* mutex) {
+  std::lock_guard<std::mutex> guard{lock_};
+  mutexes_.push_back(mutex);
+}
+
+void SlabArenaForkRegistry::remove(std::mutex* mutex) {
+  auto it = std::find(mutexes_.begin(), mutexes_.end(), mutex);
+  if (it != mutexes_.end()) {
+    *it = std::move(mutexes_.back());
+    mutexes_.pop_back();
+  }
+}
+
+void SlabArenaForkRegistry::atForkPrepare() {
+  // Holding lock_ across the fork also stops the list itself from being
+  // mutated while it's being walked.
+  lock_.lock();
+  for (std::mutex* mutex : mutexes_) {
+    mutex->lock();
+  }
+}
+
+void SlabArenaForkRegistry::atForkParent() {
+  for (std::mutex* mutex : mutexes_) {
+    mutex->unlock();
+  }
+  lock_.unlock();
+}
+
+void SlabArenaForkRegistry::atForkChild() {
+  // Reuse each mutex's storage to get a fresh, unlocked one.  They're all
+  // still locked by atForkPrepare() and destroying a locked mutex is
+  // undefined, so their lifetimes end without running their destructors.
+  for (std::mutex* mutex : mutexes_) {
+    new (mutex) std::mutex{};
+  }
+  new (&lock_) std::mutex{};
+}
+
 } // namespace cinderx

cinderx/Common/slab_arena.h

@@ -92,6 +92,28 @@ class SlabArenaIterator {
 
 std::shared_ptr<HugePageArena> getSharedHugePageArena();
 
+// The mutexes of every live SlabArena, so that pthread_atfork() handlers can
+// quiesce them across a fork().
+//
+// SlabArena is a template with instances scattered across JIT state, so they
+// register themselves here instead of being enumerated by hand.  No SlabArena
+// ever locks another, so the handlers may take them in any order.
+class SlabArenaForkRegistry {
+ public:
+  static SlabArenaForkRegistry& get();
+
+  void add(std::mutex* mutex);
+  void remove(std::mutex* mutex);
+
+  void atForkPrepare();
+  void atForkParent();
+  void atForkChild();
+
+ private:
+  std::mutex lock_;
+  std::vector<std::mutex*> mutexes_;
+};
+
 // SlabArena is a simple arena allocator, using slabs that are multiples of the
 // system's page size. Allocated objects never move after creation, and all
 // objects will be kept alive until the SlabArena they came from is destroyed.
@@ -118,8 +140,20 @@ class SlabArena {
 
   SlabArena() {
     slabs_.emplace_back(SizeTrait::size(), getSharedHugePageArena());
+    // Registered last so a throwing constructor can't leave a dangling pointer
+    // behind, as the destructor won't run for a half-constructed arena.
+    SlabArenaForkRegistry::get().add(&mutex_);
   }
 
+  ~SlabArena() {
+    SlabArenaForkRegistry::get().remove(&mutex_);
+  }
+
+  SlabArena(const SlabArena&) = delete;
+  SlabArena(SlabArena&&) = delete;
+  SlabArena& operator=(const SlabArena&) = delete;
+  SlabArena& operator=(SlabArena&&) = delete;
+
   // Allocate a new instance of T using the given constructor arguments.
   template <typename... Args>
   T* allocate(Args&&... args) {

cinderx/Jit/code_allocator.cpp

@@ -6,6 +6,7 @@
 #include "cinderx/Jit/codegen/code_section.h"
 #include "cinderx/Jit/config.h"
 #include "cinderx/Jit/jit_rt.h"
+#include "cinderx/module_state.h"
 
 #ifdef WIN32
 #include <Windows.h>
@@ -22,6 +23,7 @@
 #endif
 
 #include <cstring>
+#include <new>
 
 namespace cinderx::jit {
 
@@ -211,6 +213,8 @@ ICodeAllocator* CodeAllocator::make() {
 }
 
 AllocateResult CodeAllocator::addCode(asmjit::CodeHolder* code) {
+  std::lock_guard lock{runtime_mutex_};
+
   void* addr = nullptr;
   asmjit::Error error = runtime_.add(&addr, code);
 
@@ -222,6 +226,8 @@ AllocateResult CodeAllocator::addCode(asmjit::CodeHolder* code) {
 }
 
 asmjit::Error CodeAllocator::releaseCode(void* code) {
+  std::lock_guard lock{runtime_mutex_};
+
   // Find the size of the allocated region.
   asmjit::JitAllocator* inner = runtime_.allocator();
   asmjit::JitAllocator::Span span;
@@ -247,9 +253,12 @@ asmjit::Error CodeAllocator::releaseCode(void* code) {
 }
 
 bool CodeAllocator::contains(const void* ptr) const {
+  // query() is internally thread-safe, but it takes asmjit's own lock, which
+  // can't be recovered in a forked child.  Going through runtime_mutex_ keeps
+  // that lock free whenever a fork can happen.
+  std::lock_guard lock{runtime_mutex_};
+
   asmjit::JitAllocator::Span unused;
-  // asmjit docs don't say that query() is thread-safe, but peeking at the
-  // implementation shows that it is.
   return runtime_.allocator()->query(unused, const_cast<void*>(ptr)) ==
       asmjit::kErrorOk;
 }
@@ -262,6 +271,21 @@ const asmjit::Environment& CodeAllocator::asmJitEnvironment() const {
   return runtime_.environment();
 }
 
+void CodeAllocator::atForkPrepare() {
+  runtime_mutex_.lock();
+}
+
+void CodeAllocator::atForkParent() {
+  runtime_mutex_.unlock();
+}
+
+void CodeAllocator::atForkChild() {
+  // Reuse the storage to get a fresh, unlocked mutex.  The inherited one is
+  // still locked by atForkPrepare() and destroying a locked mutex is
+  // undefined, so its lifetime is ended without running its destructor.
+  new (&runtime_mutex_) std::mutex{};
+}
+
 CodeAllocatorCinder::~CodeAllocatorCinder() {
   for (std::span<uint8_t> alloc : allocations_) {
 #ifndef WIN32
@@ -484,4 +508,50 @@ bool CodeAllocatorCinder::contains(const void* ptr) const {
   return false;
 }
 
+void CodeAllocatorCinder::atForkPrepare() {
+  CodeAllocator::atForkPrepare();
+  allocator_mutex_.lock();
+}
+
+void CodeAllocatorCinder::atForkParent() {
+  allocator_mutex_.unlock();
+  CodeAllocator::atForkParent();
+}
+
+void CodeAllocatorCinder::atForkChild() {
+  new (&allocator_mutex_) std::mutex{};
+  CodeAllocator::atForkChild();
+}
+
+void codeAllocatorAtForkPrepare() {
+  ModuleState* state = cinderx::getModuleState();
+  if (state != nullptr && state->code_allocator != nullptr) {
+    state->code_allocator->atForkPrepare();
+  }
+#if defined(__linux__)
+  // Innermost: ensureSpace() reaches this while holding allocator_mutex_.
+  cinder_jit_region_mutex_.lock();
+#endif
+}
+
+void codeAllocatorAtForkParent() {
+#if defined(__linux__)
+  cinder_jit_region_mutex_.unlock();
+#endif
+  ModuleState* state = cinderx::getModuleState();
+  if (state != nullptr && state->code_allocator != nullptr) {
+    state->code_allocator->atForkParent();
+  }
+}
+
+void codeAllocatorAtForkChild() {
+#if defined(__linux__)
+  new (&cinder_jit_region_mutex_) std::mutex{};
+#endif
+  ModuleState* state = cinderx::getModuleState();
+  if (state != nullptr && state->code_allocator != nullptr) {
+    state->code_allocator->atForkChild();
+  }
+}
+
 } // namespace cinderx::jit

cinderx/Jit/code_allocator.h

@@ -41,9 +41,20 @@ class CodeAllocator : public ICodeAllocator {
   size_t usedBytes() const override;
   const asmjit::Environment& asmJitEnvironment() const override;
 
+  void atForkPrepare() override;
+  void atForkParent() override;
+  void atForkChild() override;
+
  protected:
   asmjit::JitRuntime runtime_;
   std::atomic<size_t> used_bytes_{0};
+
+  // Serializes every operation that reaches asmjit's own JitAllocator lock.
+  // That lock is private to asmjit with no way to reset it in a forked child,
+  // so this makes it observably free at fork time instead: atForkPrepare()
+  // holds this, which means no thread can be inside asmjit when the fork
+  // happens.
+  mutable std::mutex runtime_mutex_;
 };
 
 // A code allocator which tries to allocate all code on huge pages.
@@ -71,6 +82,10 @@ class CodeAllocatorCinder : public CodeAllocator {
   asmjit::Error releaseCode(void* code) override;
   bool contains(const void* ptr) const override;
 
+  void atForkPrepare() override;
+  void atForkParent() override;
+  void atForkChild() override;
+
  private:
   // Add code with hot/cold section splitting. Called by addCode() when
   // multiple_code_sections is enabled. Caller must hold allocator_mutex_.
@@ -114,6 +129,12 @@ class CodeAllocatorCinder : public CodeAllocator {
   std::atomic<size_t> fragmented_allocs_{0};
 };
 
+// pthread_atfork() handlers covering both the process-global code-allocation
+// state and the current ICodeAllocator, if one exists.
+void codeAllocatorAtForkPrepare();
+void codeAllocatorAtForkParent();
+void codeAllocatorAtForkChild();
+
 void populateCodeSections(
     std::vector<std::pair<void*, std::size_t>>& output_vector,
     asmjit::CodeHolder& code,

cinderx/Jit/code_allocator_iface.h

@@ -34,6 +34,14 @@ class ICodeAllocator {
 
   // Get the asmjit environment used by this allocator.
   virtual const asmjit::Environment& asmJitEnvironment() const = 0;
+
+  // pthread_atfork() handlers, called via codeAllocatorAtFork*().  Compile
+  // threads allocate code with the GIL released, so a child forked at the
+  // wrong moment would otherwise inherit an allocator lock held by a thread
+  // that no longer exists.
+  virtual void atForkPrepare() {}
+  virtual void atForkParent() {}
+  virtual void atForkChild() {}
 };
 
 } // namespace cinderx::jit

cinderx/Jit/compilation_lock.cpp

@@ -2,11 +2,31 @@
 
 #include "cinderx/Jit/compilation_lock.h"
 
+#include <new>
+
 namespace cinderx::jit {
 
 std::recursive_mutex& jitCompilationMutex() {
   static std::recursive_mutex mutex;
   return mutex;
 }
 
+void jitCompilationAtForkPrepare() {
+  jitCompilationMutex().lock();
+}
+
+void jitCompilationAtForkParent() {
+  jitCompilationMutex().unlock();
+}
+
+void jitCompilationAtForkChild() {
+  // Reuse the storage to get a fresh, unlocked mutex.  Unlocking is not an
+  // option even though the forking thread is the owner: a recursive mutex
+  // identifies its owner by thread id, and the surviving thread gets a new one
+  // across the fork, so unlock() would fail with EPERM and leave the lock held
+  // forever.  Destroying it isn't an option either, as it is still locked by
+  // atForkPrepare(), so its lifetime ends without running its destructor.
+  new (&jitCompilationMutex()) std::recursive_mutex{};
+}
+
 } // namespace cinderx::jit

cinderx/Jit/compilation_lock.h

@@ -12,6 +12,13 @@ namespace cinderx::jit {
 // logic for GIL handling during threaded compiles.
 std::recursive_mutex& jitCompilationMutex();
 
+// pthread_atfork() handlers for the compilation lock.  Compile threads hold it
+// with the GIL released, so a child forked at the wrong moment would otherwise
+// inherit it locked by a thread that no longer exists.
+void jitCompilationAtForkPrepare();
+void jitCompilationAtForkParent();
+void jitCompilationAtForkChild();
+
 // Uses to track if the current thread holds the lock for assertion purposes.
 inline thread_local int jitCompilationLockDepth = 0;