LookupCache: Move CodePages to GuestToHostMap

Prevents invalidations being missed under the following circumstances:
Thread A JITs block A into the global codebuffer, adding the guest to host
mapping to its CodePages, thread A is then killed.
Thread B then performs SMC on block A. An exception will be triggered but
as CodePages was stored per-thread, and thread A is now killed when all
threads are iterated over by the frontend to perform invalidations it
will be missed.

The accumulator is introduced to handle the case where multiple threads
have the same code entry in their local caches but share the same codebuffer.
Consider a thread C in the above example that also has block A in its cache,
without an accumulator, when invalidating thread B the entrypoint of A is erased
from the shared guest to host map. So when C is invalidated, the local cache entry
for A is not removed since it was removed from CodePages when invalidating B.

FEXCore/Source/Interface/Context/Context.h

@@ -166,7 +166,8 @@ public:
 
   void OnCodeBufferAllocated(CPU::CodeBuffer&) override;
   void ClearCodeCache(FEXCore::Core::InternalThreadState* Thread, bool NewCodeBuffer = true) override;
-  void InvalidateGuestCodeRange(FEXCore::Core::InternalThreadState* Thread, uint64_t Start, uint64_t Length) override;
+  void InvalidateGuestCodeRange(FEXCore::Core::InternalThreadState* Thread, InvalidatedEntryAccumulator& Accumulator, uint64_t Start,
+                                uint64_t Length) override;
   FEXCore::ForkableSharedMutex& GetCodeInvalidationMutex() override {
     return CodeInvalidationMutex;
   }

FEXCore/Source/Interface/Core/Core.cpp

@@ -909,26 +909,32 @@ uintptr_t ContextImpl::CompileSingleStep(FEXCore::Core::CpuStateFrame* Frame, ui
   return (uintptr_t)CodePtr;
 }
 
-static void InvalidateGuestThreadCodeRange(FEXCore::Core::InternalThreadState* Thread, uint64_t Start, uint64_t Length) {
+static void InvalidateGuestThreadCodeRange(FEXCore::Core::InternalThreadState* Thread, InvalidatedEntryAccumulator& Accumulator,
+                                           uint64_t Start, uint64_t Length) {
   // Ensures now-modified mappings aren't cached as being in their previous non-executable state.
   // Accessing FrontendDecoder is safe as the thread's code invalidation mutex must be locked here.
   Thread->FrontendDecoder->ResetExecutableRangeCache();
 
   auto lk = Thread->LookupCache->AcquireLock();
+  auto& CodePages = Thread->LookupCache->Shared->CodePages;
 
-  auto lower = Thread->LookupCache->CodePages.lower_bound(Start >> 12);
-  auto upper = Thread->LookupCache->CodePages.upper_bound((Start + Length - 1) >> 12);
+  auto lower = CodePages.lower_bound(Start >> 12);
+  auto upper = CodePages.upper_bound((Start + Length - 1) >> 12);
 
   for (auto it = lower; it != upper; it++) {
-    for (auto Address : it->second) {
-      ContextImpl::ThreadRemoveCodeEntry(Thread, Address);
+    Accumulator.emplace_back(std::move(it->second));
+  }
+
+  for (const auto& PageEntries : Accumulator) {
+    for (const auto& Entry : PageEntries) {
+      ContextImpl::ThreadRemoveCodeEntry(Thread, Entry);
     }
-    it->second.clear();
   }
 }
 
-void ContextImpl::InvalidateGuestCodeRange(FEXCore::Core::InternalThreadState* Thread, uint64_t Start, uint64_t Length) {
-  InvalidateGuestThreadCodeRange(Thread, Start, Length);
+void ContextImpl::InvalidateGuestCodeRange(FEXCore::Core::InternalThreadState* Thread, InvalidatedEntryAccumulator& Accumulator,
+                                           uint64_t Start, uint64_t Length) {
+  InvalidateGuestThreadCodeRange(Thread, Accumulator, Start, Length);
 }
 
 void ContextImpl::MarkMemoryShared(FEXCore::Core::InternalThreadState* Thread) {
@@ -1033,7 +1039,8 @@ void ContextImpl::RemoveCustomIREntrypoint(uintptr_t Entrypoint) {
 
   std::scoped_lock lk(CustomIRMutex);
 
-  InvalidateGuestCodeRange(nullptr, Entrypoint, 1);
+  InvalidatedEntryAccumulator Accumulator;
+  InvalidateGuestCodeRange(nullptr, Accumulator, Entrypoint, 1);
   CustomIRHandlers.erase(Entrypoint);
 
   HasCustomIRHandlers = !CustomIRHandlers.empty();

FEXCore/Source/Interface/Core/LookupCache.h

@@ -56,6 +56,8 @@ struct GuestToHostMap {
 
   fextl::robin_map<uint64_t, uint64_t> BlockList;
 
+  fextl::map<uint64_t, fextl::vector<uint64_t>> CodePages;
+
   GuestToHostMap();
 
   // Adds to Guest -> Host code mapping
@@ -93,6 +95,18 @@ struct GuestToHostMap {
     BlockLinks->insert({{GuestDestination, HostLink}, delinker});
   }
 
+  bool AddBlockExecutableRange(const fextl::set<uint64_t>& Addresses, uint64_t Start, uint64_t Length, const LockToken&) {
+    bool rv = false;
+
+    for (auto CurrentPage = Start >> 12, EndPage = (Start + Length - 1) >> 12; CurrentPage <= EndPage; CurrentPage++) {
+      auto& CodePage = CodePages[CurrentPage];
+      rv |= CodePage.empty();
+      CodePage.insert(CodePage.end(), Addresses.begin(), Addresses.end());
+    }
+
+    return rv;
+  }
+
   void ClearCache(const LockToken&);
 };
 
@@ -155,22 +169,11 @@ public:
 
   GuestToHostMap* Shared = nullptr;
 
-  fextl::map<uint64_t, fextl::vector<uint64_t>> CodePages;
-
   // Appends a list of Block {Address} to CodePages [Start, Start + Length)
   // Returns true if new pages are marked as containing code
   bool AddBlockExecutableRange(const fextl::set<uint64_t>& Addresses, uint64_t Start, uint64_t Length) {
     auto lk = Shared->AcquireLock();
-
-    bool rv = false;
-
-    for (auto CurrentPage = Start >> 12, EndPage = (Start + Length - 1) >> 12; CurrentPage <= EndPage; CurrentPage++) {
-      auto& CodePage = CodePages[CurrentPage];
-      rv |= CodePage.empty();
-      CodePage.insert(CodePage.end(), Addresses.begin(), Addresses.end());
-    }
-
-    return rv;
+    return Shared->AddBlockExecutableRange(Addresses, Start, Length, lk);
   }
 
   // Adds to Guest -> Host code mapping

FEXCore/include/FEXCore/Core/Context.h

@@ -64,6 +64,9 @@ public:
 
 using CodeRangeInvalidationFn = std::function<void(uint64_t start, uint64_t Length)>;
 
+// Nested vector of guest block entrypoints
+using InvalidatedEntryAccumulator = fextl::vector<fextl::vector<uint64_t>>;
+
 using CustomIREntrypointHandler = std::function<void(uintptr_t Entrypoint, IR::IREmitter*)>;
 
 using ExitHandler = std::function<void(Core::InternalThreadState* Thread)>;
@@ -172,7 +175,8 @@ public:
   FEX_DEFAULT_VISIBILITY virtual void WriteFilesWithCode(AOTIRCodeFileWriterFn Writer) = 0;
 
   FEX_DEFAULT_VISIBILITY virtual void ClearCodeCache(FEXCore::Core::InternalThreadState* Thread, bool NewCodeBuffer = true) = 0;
-  FEX_DEFAULT_VISIBILITY virtual void InvalidateGuestCodeRange(FEXCore::Core::InternalThreadState* Thread, uint64_t Start, uint64_t Length) = 0;
+  FEX_DEFAULT_VISIBILITY virtual void InvalidateGuestCodeRange(
+    FEXCore::Core::InternalThreadState* Thread, InvalidatedEntryAccumulator& Accumulator, uint64_t Start, uint64_t Length) = 0;
   FEX_DEFAULT_VISIBILITY virtual FEXCore::ForkableSharedMutex& GetCodeInvalidationMutex() = 0;
 
   FEX_DEFAULT_VISIBILITY virtual void MarkMemoryShared(FEXCore::Core::InternalThreadState* Thread) = 0;

Source/Tools/CodeSizeValidation/Main.cpp

@@ -52,7 +52,8 @@ public:
 
     {
       auto CodeInvalidationlk = FEXCore::GuardSignalDeferringSection(CTX->GetCodeInvalidationMutex(), Thread);
-      CTX->InvalidateGuestCodeRange(Thread, reinterpret_cast<uint64_t>(CodeStart), MAX_CODE_SIZE);
+      FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
+      CTX->InvalidateGuestCodeRange(Thread, Accumulator, reinterpret_cast<uint64_t>(CodeStart), MAX_CODE_SIZE);
     }
 
     ClearStats();

Source/Tools/LinuxEmulation/LinuxSyscalls/ThreadManager.h

@@ -184,9 +184,10 @@ public:
     // Thread object isn't valid very early in frontend's initialization.
     // To be more optimal the frontend should provide this code with a valid Thread object earlier.
     auto CodeInvalidationlk = GuardSignalDeferringSectionWithFallback(CTX->GetCodeInvalidationMutex(), CallingThread);
+    FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
 
     for (auto& Thread : Threads) {
-      CTX->InvalidateGuestCodeRange(Thread->Thread, Start, Length);
+      CTX->InvalidateGuestCodeRange(Thread->Thread, Accumulator, Start, Length);
     }
   }
 
@@ -198,9 +199,10 @@ public:
     // Thread object isn't valid very early in frontend's initialization.
     // To be more optimal the frontend should provide this code with a valid Thread object earlier.
     auto CodeInvalidationlk = GuardSignalDeferringSectionWithFallback(CTX->GetCodeInvalidationMutex(), CallingThread);
+    FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
 
     for (auto& Thread : Threads) {
-      CTX->InvalidateGuestCodeRange(Thread->Thread, Start, Length);
+      CTX->InvalidateGuestCodeRange(Thread->Thread, Accumulator, Start, Length);
     }
 
     // Callback while holding the locks.

Source/Windows/Common/InvalidationTracker.cpp

@@ -53,8 +53,9 @@ void InvalidationTracker::HandleMemoryProtectionNotification(uint64_t Address, u
   if (NeedsInvalidate) {
     // IntervalsLock cannot be held during invalidation
     std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
+    FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
     for (auto Thread : Threads) {
-      CTX.InvalidateGuestCodeRange(Thread.second, AlignedBase, AlignedSize);
+      CTX.InvalidateGuestCodeRange(Thread.second, Accumulator, AlignedBase, AlignedSize);
     }
   }
 }
@@ -94,8 +95,9 @@ InvalidationTracker::InvalidateContainingSectionResult InvalidationTracker::Inva
   }
   {
     std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
+    FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
     for (auto Thread : Threads) {
-      CTX.InvalidateGuestCodeRange(Thread.second, SectionBase, SectionSize);
+      CTX.InvalidateGuestCodeRange(Thread.second, Accumulator, SectionBase, SectionSize);
     }
   }
 
@@ -119,8 +121,9 @@ void InvalidationTracker::InvalidateAlignedInterval(uint64_t Address, uint64_t S
 
   {
     std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
+    FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
     for (auto Thread : Threads) {
-      CTX.InvalidateGuestCodeRange(Thread.second, AlignedBase, AlignedSize);
+      CTX.InvalidateGuestCodeRange(Thread.second, Accumulator, AlignedBase, AlignedSize);
     }
   }
 
@@ -170,8 +173,9 @@ bool InvalidationTracker::HandleRWXAccessViolation(uint64_t FaultAddress) {
   if (NeedsInvalidate) {
     // IntervalsLock cannot be held during invalidation
     std::scoped_lock Lock(CTX.GetCodeInvalidationMutex());
+    FEXCore::Context::InvalidatedEntryAccumulator Accumulator;
     for (auto Thread : Threads) {
-      CTX.InvalidateGuestCodeRange(Thread.second, FaultAddress & FEXCore::Utils::FEX_PAGE_MASK, FEXCore::Utils::FEX_PAGE_SIZE);
+      CTX.InvalidateGuestCodeRange(Thread.second, Accumulator, FaultAddress & FEXCore::Utils::FEX_PAGE_MASK, FEXCore::Utils::FEX_PAGE_SIZE);
     }
     return true;
   }