Invalidations: Move invalidation locks to Context, extend InvalidateGuestCodeRange with optional callback

External/FEXCore/Source/Interface/Context/Context.h

@@ -126,6 +126,8 @@ namespace FEXCore::Context {
     Event PauseWait;
     bool Running{};
 
+    std::shared_mutex CodeInvalidationMutex; 
+    
     FEXCore::CPUIDEmu CPUID;
     FEXCore::HLE::SyscallHandler *SyscallHandler{};
     std::unique_ptr<FEXCore::ThunkHandler> ThunkHandler;

External/FEXCore/Source/Interface/Core/Core.cpp

@@ -967,8 +967,8 @@ namespace FEXCore::Context {
   uintptr_t Context::CompileBlock(FEXCore::Core::CpuStateFrame *Frame, uint64_t GuestRIP) {
     auto Thread = Frame->Thread;
 
-    // Needs to be held for SMC interactions around concurrent compile and invalidation hazards
-    auto InvalidationLk = Thread->CTX->SyscallHandler->CompileCodeLock(GuestRIP);
+    // Invalidate might take a unique lock on this, to guarantee that during invalidation no code gets compiled
+    std::shared_lock lk(CodeInvalidationMutex);
 
     // Is the code in the cache?
     // The backends only check L1 and L2, not L3
@@ -1134,7 +1134,8 @@ namespace FEXCore::Context {
   }
 
   void InvalidateGuestCodeRange(FEXCore::Context::Context *CTX, uint64_t Start, uint64_t Length) {
-    std::lock_guard<std::mutex> lk(CTX->ThreadCreationMutex);
+    std::lock_guard lk(CTX->ThreadCreationMutex);
+    
     for (auto &Thread : CTX->Threads) {
       if (Thread->RunningEvents.Running.load()) {
         InvalidateGuestThreadCodeRange(Thread, Start, Length);
@@ -1142,6 +1143,13 @@ namespace FEXCore::Context {
     }
   }
 
+  void InvalidateGuestCodeRange(FEXCore::Context::Context *CTX, uint64_t Start, uint64_t Length, std::function<void(uint64_t start, uint64_t Length)> CallAfter) {
+    std::unique_lock CodeInvalidationLock(CTX->CodeInvalidationMutex);
+
+    InvalidateGuestCodeRange(CTX, Start, Length);
+    CallAfter(Start, Length);
+  }
+
   void Context::MarkMemoryShared() {
     if (!IsMemoryShared) {
       IsMemoryShared = true;

External/FEXCore/include/FEXCore/Core/Context.h

@@ -250,6 +250,7 @@ namespace FEXCore::Context {
   FEX_DEFAULT_VISIBILITY void FinalizeAOTIRCache(FEXCore::Context::Context *CTX);
   FEX_DEFAULT_VISIBILITY void WriteFilesWithCode(FEXCore::Context::Context *CTX, std::function<void(const std::string& fileid, const std::string& filename)> Writer);
   FEX_DEFAULT_VISIBILITY void InvalidateGuestCodeRange(FEXCore::Context::Context *CTX, uint64_t Start, uint64_t Length);
+  FEX_DEFAULT_VISIBILITY void InvalidateGuestCodeRange(FEXCore::Context::Context *CTX, uint64_t Start, uint64_t Length, std::function<void(uint64_t start, uint64_t Length)> callback);
   FEX_DEFAULT_VISIBILITY void MarkMemoryShared(FEXCore::Context::Context *CTX);
 
   FEX_DEFAULT_VISIBILITY void ConfigureAOTGen(FEXCore::Core::InternalThreadState *Thread, std::set<uint64_t> *ExternalBranches, uint64_t SectionMaxAddress);

External/FEXCore/include/FEXCore/HLE/SyscallHandler.h

@@ -79,7 +79,6 @@ namespace FEXCore::HLE {
     virtual FEXCore::CodeLoader *GetCodeLoader() const { return nullptr; }
     virtual void MarkGuestExecutableRange(uint64_t Start, uint64_t Length) { }
     virtual AOTIRCacheEntryLookupResult LookupAOTIRCacheEntry(uint64_t GuestAddr) = 0;
-    virtual std::shared_lock<std::shared_mutex> CompileCodeLock(uint64_t Start) = 0;
 
   protected:
     SyscallOSABI OSABI;

Source/Tests/IRLoader.cpp

@@ -138,11 +138,6 @@ class DummySyscallHandler: public FEXCore::HLE::SyscallHandler {
   FEXCore::HLE::AOTIRCacheEntryLookupResult LookupAOTIRCacheEntry(uint64_t GuestAddr) override {
     return {0, 0, FHU::ScopedSignalMaskWithSharedLock {StubMutex}};
   }
-
-  std::shared_mutex StubMutex2;
-  std::shared_lock<std::shared_mutex> CompileCodeLock(uint64_t Start) override {
-    return std::shared_lock(StubMutex2);
-  }
 };
 
 int main(int argc, char **argv, char **const envp)

Source/Tests/LinuxSyscalls/Syscalls.h

@@ -187,7 +187,6 @@ public:
   void MarkGuestExecutableRange(uint64_t Start, uint64_t Length) override;
   // AOTIRCacheEntryLookupResult also includes a shared lock guard, so the pointed AOTIRCacheEntry return can be safely used
   FEXCore::HLE::AOTIRCacheEntryLookupResult LookupAOTIRCacheEntry(uint64_t GuestAddr) final override;
-  std::shared_lock<std::shared_mutex> CompileCodeLock(uint64_t Start) override;
 
   ///// FORK tracking /////
   void LockBeforeFork();
@@ -295,9 +294,6 @@ private:
     using VMAEntry = SyscallHandler::VMAEntry;
     // Held while reading/writing this struct
     std::shared_mutex Mutex;
-
-    // Held unique {invalidate, mprotect change} to guarantee mt invalidation correctness
-    std::shared_mutex InvalidationMutex;
     
     // Memory ranges indexed by page aligned starting address
     std::map<uint64_t, VMAEntry> VMAs;

Source/Tests/LinuxSyscalls/SyscallsSMCTracking.cpp

@@ -63,10 +63,6 @@ bool SyscallHandler::HandleSegfault(FEXCore::Core::InternalThreadState *Thread,
 
     auto FaultBase = FEXCore::AlignDown(FaultAddress, FHU::FEX_PAGE_SIZE);
 
-    // take lock, no code compilation happens between invalidation and prot change
-    // no need for signal mask, we have an outer one
-    std::unique_lock lk2(VMATracking->InvalidationMutex);
-
     if (Entry->second.Flags.Shared) {
       LOGMAN_THROW_A_FMT(Entry->second.Resource, "VMA tracking error");
 
@@ -79,27 +75,28 @@ bool SyscallHandler::HandleSegfault(FEXCore::Core::InternalThreadState *Thread,
       do {
         if (VMA->Offset <= Offset && (VMA->Offset + VMA->Length) > Offset) {
           auto FaultBaseMirrored = Offset - VMA->Offset + VMA->Base;
-          FEXCore::Context::InvalidateGuestCodeRange(CTX, FaultBaseMirrored, FHU::FEX_PAGE_SIZE);
+
           if (VMA->Prot.Writable) {
-            auto rv = mprotect((void *)FaultBaseMirrored, FHU::FEX_PAGE_SIZE, PROT_READ | PROT_WRITE);
-            LogMan::Throw::AFmt(rv == 0, "mprotect({}, {}) failed", FaultBaseMirrored, FHU::FEX_PAGE_SIZE);
+            FEXCore::Context::InvalidateGuestCodeRange(CTX, FaultBaseMirrored, FHU::FEX_PAGE_SIZE, [](uintptr_t Start, uintptr_t Length) {
+              auto rv = mprotect((void *)Start, Length, PROT_READ | PROT_WRITE);
+              LogMan::Throw::AFmt(rv == 0, "mprotect({}, {}) failed", Start, Length);
+            });
+          } else {
+            FEXCore::Context::InvalidateGuestCodeRange(CTX, FaultBaseMirrored, FHU::FEX_PAGE_SIZE);
           }
         }
       } while ((VMA = VMA->ResourceNextVMA));
     } else {
-      FEXCore::Context::InvalidateGuestCodeRange(CTX, FaultBase, FHU::FEX_PAGE_SIZE);
-      auto rv = mprotect((void *)FaultBase, FHU::FEX_PAGE_SIZE, PROT_READ | PROT_WRITE);
-      LogMan::Throw::AFmt(rv == 0, "mprotect({}, {}) failed", FaultBase, FHU::FEX_PAGE_SIZE);
+      FEXCore::Context::InvalidateGuestCodeRange(CTX, FaultBase, FHU::FEX_PAGE_SIZE, [](uintptr_t Start, uintptr_t Length) {
+        auto rv = mprotect((void *)Start, Length, PROT_READ | PROT_WRITE);
+        LogMan::Throw::AFmt(rv == 0, "mprotect({}, {}) failed", Start, Length);
+      });
     }
 
     return true;
   }
 }
 
-std::shared_lock<std::shared_mutex> SyscallHandler::CompileCodeLock(uint64_t Start) {
-  return std::shared_lock(VMATracking.InvalidationMutex);
-}
-
 void SyscallHandler::MarkGuestExecutableRange(uint64_t Start, uint64_t Length) {
   const auto Base = Start & FHU::FEX_PAGE_MASK;
   const auto Top = FEXCore::AlignUp(Start + Length, FHU::FEX_PAGE_SIZE);