Fix thread-safely issue in IsolatedManagers

Summary:
There's a thread safety issue here:

Thread 1 swaps in it's map and swaps out the global one.
Thread 2 swaps in it's map and swaps out thread's 1.
Thread 1 completes and swaps the global one back in.
Thread 2 completes and swap's back in thread 1's - the value is stack allocated so now we have garbage.

Claude's solution is to use a thread-local to store the preload manager on the stack and then clear it out when done.

Reviewed By: jbower-fb

Differential Revision: D92114066

fbshipit-source-id: 59ce532c4b8c32fc1b5104dade9b1d909d1665b9

Author: DinoV

cinderx/Jit/hir/preload.cpp

@@ -119,6 +119,7 @@ static std::unique_ptr<NativeTarget> resolve_native_target(
 }
 
 PreloaderManager s_manager;
+thread_local PreloaderManager* tls_manager = nullptr;
 
 } // namespace
 
@@ -535,33 +536,31 @@ bool PreloaderManager::empty() const {
   return preloaders_.empty();
 }
 
+size_t PreloaderManager::size() const {
+  return preloaders_.size();
+}
+
 void PreloaderManager::clear() {
   preloaders_.clear();
 }
 
-void PreloaderManager::swap(PreloaderMap& replacement) {
-  // Should never be called from within the actual multi-threaded compile;
-  // it's not safe to mess with the global preloaders map in that context.
-  JIT_CHECK(
-      !getThreadedCompileContext().compileRunning(),
-      "cannot preload single func from within multi-threaded compile");
-  preloaders_.swap(replacement);
+bool PreloaderManager::isGlobalManager() const {
+  return tls_manager == nullptr;
 }
 
 PreloaderManager& preloaderManager() {
+  if (tls_manager != nullptr) {
+    return *tls_manager;
+  }
   return s_manager;
 }
 
-IsolatedPreloaders::IsolatedPreloaders() {
-  swap();
+IsolatedPreloaders::IsolatedPreloaders() : prev_manager_(tls_manager) {
+  tls_manager = &local_manager_;
 }
 
 IsolatedPreloaders::~IsolatedPreloaders() {
-  swap();
-}
-
-void IsolatedPreloaders::swap() {
-  s_manager.swap(orig_preloaders_);
+  tls_manager = prev_manager_;
 }
 
 } // namespace jit::hir

cinderx/Jit/hir/preload.h

@@ -269,13 +269,13 @@ class PreloaderManager {
   // Check if there are any preloaders registered.
   bool empty() const;
 
+  size_t size() const;
+
+  bool isGlobalManager() const;
+
   // Clear out all preloaders.
   void clear();
 
-  // Swap the inner map with the one passed as an argument.  Needed by
-  // IsolatedPreloaders.
-  void swap(PreloaderMap& replacement);
-
  private:
   PreloaderMap preloaders_;
 };
@@ -284,15 +284,16 @@ class PreloaderManager {
 PreloaderManager& preloaderManager();
 
 // RAII device for isolating preloaders state.
+// Uses thread-local storage to give each thread its own PreloaderManager
+// while isolation is active, avoiding race conditions between threads.
 class IsolatedPreloaders {
  public:
   IsolatedPreloaders();
   ~IsolatedPreloaders();
 
  private:
-  void swap();
-
-  PreloaderMap orig_preloaders_;
+  PreloaderManager local_manager_;
+  PreloaderManager* prev_manager_;
 };
 
 } // namespace jit::hir

cinderx/Jit/pyjit.cpp

@@ -3420,7 +3420,10 @@ void finalize() {
   jit_reg_units.clear();
   JIT_CHECK(
       hir::preloaderManager().empty(),
-      "JIT cannot be finalized while batch compilation is active");
+      "JIT cannot be finalized while batch compilation is active size:{} "
+      "is_global:{}",
+      hir::preloaderManager().size(),
+      hir::preloaderManager().isGlobalManager());
 
   for (auto func : jitCtx()->compiledFuncs()) {
     deoptFuncImpl(func);