Merge pull request #5994 from cjacek/thread-term

WOW64: Use correct thread's CPU area in BTCpuThreadTerm

Source/Windows/ARM64EC/Module.cpp

@@ -1018,7 +1018,7 @@ NTSTATUS ThreadTerm(HANDLE Thread, LONG ExitCode) {
     return STATUS_ACCESS_DENIED;
   }
 
-  auto ThreadDup = FEX::Windows::DupHandle(Thread, THREAD_QUERY_INFORMATION | THREAD_SUSPEND_RESUME);
+  auto ThreadDup = FEX::Windows::DupHandle(Thread, THREAD_QUERY_INFORMATION | THREAD_SUSPEND_RESUME | THREAD_GET_CONTEXT);
 
   THREAD_BASIC_INFORMATION Info;
   if (auto Err = NtQueryInformationThread(*ThreadDup, ThreadBasicInformation, &Info, sizeof(Info), nullptr); Err) {
@@ -1032,6 +1032,7 @@ NTSTATUS ThreadTerm(HANDLE Thread, LONG ExitCode) {
     // If we are suspending a thread that isn't ourselves, try to suspend it first so we know internal JIT locks aren't being held.
     NtSuspendThread(*ThreadDup, NULL);
     // This will wait for the thread to be suspended
+    TmpContext.ContextFlags = CONTEXT_CONTROL;
     NtGetContextThread(*ThreadDup, &TmpContext);
   }
 

Source/Windows/WOW64/Module.cpp

@@ -102,11 +102,11 @@ struct TLS {
   uint64_t& CachedCallRetSp() const {
     return reinterpret_cast<uint64_t&>(TEB->TlsSlots[FEXCore::ToUnderlying(Slot::CACHED_CALLRET_SP)]);
   }
-};
 
-CHPE_V2_CPU_AREA_INFO*& GetCpuArea() {
-  return ((__TEB*)NtCurrentTeb())->ChpeV2CpuAreaInfo;
-}
+  CHPE_V2_CPU_AREA_INFO*& CpuArea() const {
+    return reinterpret_cast<__TEB*>(TEB)->ChpeV2CpuAreaInfo;
+  }
+};
 
 WOW64_CONTEXT* GetWow64Context() {
   WOW64_CONTEXT* WowContext;
@@ -118,13 +118,9 @@ struct FrontendThreadData {
   bool InLockedRWXRead {};
   CHPE_V2_CPU_AREA_INFO CpuArea {};
 
-  FrontendThreadData(FEXCore::Core::InternalThreadState* Thread) {
-    CpuArea.SuspendDoorbell = reinterpret_cast<ULONG*>(&Thread->CurrentFrame->SuspendDoorbell);
-    GetCpuArea() = &CpuArea;
-  }
-
-  ~FrontendThreadData() {
-    GetCpuArea() = nullptr;
+  FrontendThreadData(const TLS& TLS) {
+    CpuArea.SuspendDoorbell = reinterpret_cast<ULONG*>(&TLS.ThreadState()->CurrentFrame->SuspendDoorbell);
+    TLS.CpuArea() = &CpuArea;
   }
 };
 
@@ -355,12 +351,13 @@ bool HandleUnalignedAccess(CONTEXT* Context) {
 }
 
 void LockJITContext() {
-  GetCpuArea()->InSimulation = -1;
-  LoadStateFromWowContext(GetTLS().ThreadState());
+  TLS TLS = GetTLS();
+  TLS.CpuArea()->InSimulation = -1;
+  LoadStateFromWowContext(TLS.ThreadState());
 }
 
 void UnlockJITContext() {
-  CHPE_V2_CPU_AREA_INFO* CpuArea = GetCpuArea();
+  CHPE_V2_CPU_AREA_INFO* CpuArea = GetTLS().CpuArea();
   CpuArea->InSimulation = 0;
   if (*CpuArea->SuspendDoorbell) {
     CONTEXT ResumeContext;
@@ -611,7 +608,7 @@ void BTCpuThreadInit() {
   const auto TLS = GetTLS();
   TLS.ThreadState() = Thread;
 
-  Thread->FrontendPtr = new FrontendThreadData(Thread);
+  Thread->FrontendPtr = new FrontendThreadData(TLS);
 
   auto ThreadTID = GetCurrentThreadId();
   Threads.emplace(ThreadTID, Thread);
@@ -625,18 +622,23 @@ void BTCpuThreadTerm(HANDLE Thread, LONG ExitCode) {
     return;
   }
 
-  auto ThreadDup = FEX::Windows::DupHandle(Thread, THREAD_QUERY_INFORMATION | THREAD_SUSPEND_RESUME);
+  auto ThreadDup = FEX::Windows::DupHandle(Thread, THREAD_QUERY_INFORMATION | THREAD_SUSPEND_RESUME | THREAD_GET_CONTEXT);
 
   THREAD_BASIC_INFORMATION Info;
   if (auto Err = NtQueryInformationThread(*ThreadDup, ThreadBasicInformation, &Info, sizeof(Info), nullptr); Err) {
     return;
   }
 
+  LogMan::Msg::IFmt("Terminating {}", Info.ClientId.UniqueThread);
   const auto ThreadTID = reinterpret_cast<uint64_t>(Info.ClientId.UniqueThread);
   bool Self = ThreadTID == GetCurrentThreadId();
   if (!Self) {
+    CONTEXT TmpContext;
     // If we are suspending a thread that isn't ourselves, try to suspend it first so we know internal JIT locks aren't being held.
     RtlWow64SuspendThread(*ThreadDup, NULL);
+    // This will wait for the thread to be suspended
+    TmpContext.ContextFlags = CONTEXT_CONTROL;
+    NtGetContextThread(*ThreadDup, &TmpContext);
   }
 
   auto [Err, TLS] = GetThreadTLS(*ThreadDup);
@@ -659,6 +661,7 @@ void BTCpuThreadTerm(HANDLE Thread, LONG ExitCode) {
   }
   auto ThreadState = TLS.ThreadState();
 
+  TLS.CpuArea() = nullptr;
   delete GetFrontendThreadData(ThreadState);
 
   // GDT and LDT are mirrored, only free one.