Guard JIT allocator state in FT mode

Summary:
Add FT-only allocator mutexes around allocator-owned mutable state in `CodeAllocatorCinder` and `MultipleSectionCodeAllocator`.

ThreadSanitizer: data race fbcode/cinderx/Jit/code_allocator.cpp:189 in jit::CodeAllocatorCinder::addCode(asmjit::_abi_1_13::CodeHolder*)
==================
```

Reviewed By: alexmalyshev

Differential Revision: D96948112

fbshipit-source-id: d4ec434a75425927569cbb3287222830bbae9db5

Author: yoney

cinderx/Jit/code_allocator.cpp

@@ -137,6 +137,9 @@ CodeAllocatorCinder::~CodeAllocatorCinder() {
 
 AllocateResult CodeAllocatorCinder::addCode(asmjit::CodeHolder* code) {
   ThreadedCompileSerialize guard;
+#ifdef Py_GIL_DISABLED
+  std::lock_guard lock{allocator_mutex_};
+#endif
 
   PROPAGATE_ERROR(code->flatten());
   PROPAGATE_ERROR(code->resolveUnresolvedLinks());
@@ -153,12 +156,7 @@ AllocateResult CodeAllocatorCinder::addCode(asmjit::CodeHolder* code) {
       huge_allocs_++;
     }
     current_alloc_ = static_cast<uint8_t*>(res);
-    {
-#ifdef Py_GIL_DISABLED
-      std::lock_guard lock{allocations_mutex_};
-#endif
-      allocations_.emplace_back(res, alloc_size);
-    }
+    allocations_.emplace_back(res, alloc_size);
     current_alloc_free_ = alloc_size;
   }
 
@@ -200,7 +198,7 @@ asmjit::Error CodeAllocatorCinder::releaseCode([[maybe_unused]] void* code) {
 
 bool CodeAllocatorCinder::contains(const void* ptr) const {
 #ifdef Py_GIL_DISABLED
-  std::lock_guard lock{allocations_mutex_};
+  std::lock_guard lock{allocator_mutex_};
 #endif
   for (std::span<uint8_t> alloc : allocations_) {
     if (alloc.data() <= ptr && ptr < alloc.data() + alloc.size()) {
@@ -257,6 +255,9 @@ void MultipleSectionCodeAllocator::createSlabs() noexcept {
 
 AllocateResult MultipleSectionCodeAllocator::addCode(asmjit::CodeHolder* code) {
   ThreadedCompileSerialize guard;
+#ifdef Py_GIL_DISABLED
+  std::lock_guard lock{allocator_mutex_};
+#endif
 
   if (code_sections_.empty()) {
     createSlabs();
@@ -338,6 +339,9 @@ bool MultipleSectionCodeAllocator::contains(const void* ptr) const {
   // is already thread-safe.
   {
     ThreadedCompileSerialize guard;
+#ifdef Py_GIL_DISABLED
+    std::lock_guard lock{allocator_mutex_};
+#endif
     if (code_alloc_ <= ptr && ptr < code_alloc_ + total_allocation_size_) {
       return true;
     }

cinderx/Jit/code_allocator.h

@@ -69,11 +69,9 @@ class CodeAllocatorCinder : public CodeAllocator {
   bool contains(const void* ptr) const override;
 
  private:
-  // Protects allocations_ from concurrent access between addCode() and
-  // contains(). In free-threaded Python, one thread may be compiling a
-  // function while another checks isJitCompiled() via a function watcher.
 #ifdef Py_GIL_DISABLED
-  mutable std::mutex allocations_mutex_;
+  // Protects all allocator-owned state used by addCode()/contains().
+  mutable std::mutex allocator_mutex_;
 #endif
 
   // List of chunks allocated for use in deallocation
@@ -104,6 +102,11 @@ class MultipleSectionCodeAllocator : public CodeAllocator {
  private:
   void createSlabs() noexcept;
 
+#ifdef Py_GIL_DISABLED
+  // Protects all allocator-owned state used by addCode()/contains().
+  mutable std::mutex allocator_mutex_;
+#endif
+
   std::unordered_map<codegen::CodeSection, uint8_t*> code_sections_;
   std::unordered_map<codegen::CodeSection, size_t> code_section_free_sizes_;