Some cleanups in Native AOT suspension area (#69885)

* some renames

* strongly typed transition frames

* remove remaining pieces of LoopHijack

* comments and some refactoring

* IsManaged

* removed m_CodeManagerList

* m_ModuleListLock

* better comment for a SafeHandle hack

* Renamed SetupHackPInvokeTunnel  (to not have "hack in the name)

* made IsInForbidBlockingRegion a debug-only assert

* removed ForbidBlockingHolder

* remove forbidBlocking tracking. It only affects 4 methods.

* Fix Unix build

* typo

* record managed code range in RuntimeInstance.

Author: VSadov

src/coreclr/nativeaot/Runtime/AsmOffsets.h

@@ -42,7 +42,7 @@ ASM_OFFSET(   14,    18, MethodTable, m_VTable)
 ASM_OFFSET(    0,     0, Thread, m_rgbAllocContextBuffer)
 ASM_OFFSET(   28,    38, Thread, m_ThreadStateFlags)
 ASM_OFFSET(   2c,    40, Thread, m_pTransitionFrame)
-ASM_OFFSET(   30,    48, Thread, m_pHackPInvokeTunnel)
+ASM_OFFSET(   30,    48, Thread, m_pDeferredTransitionFrame)
 ASM_OFFSET(   40,    68, Thread, m_ppvHijackedReturnAddressLocation)
 ASM_OFFSET(   44,    70, Thread, m_pvHijackedReturnAddress)
 #ifdef HOST_64BIT

src/coreclr/nativeaot/Runtime/EHHelpers.cpp

@@ -390,7 +390,7 @@ int32_t __stdcall RhpHardwareExceptionHandler(uintptr_t faultCode, uintptr_t fau
 {
     uintptr_t faultingIP = palContext->GetIp();
 
-    ICodeManager * pCodeManager = GetRuntimeInstance()->FindCodeManagerByAddress((PTR_VOID)faultingIP);
+    ICodeManager * pCodeManager = GetRuntimeInstance()->GetCodeManagerForAddress((PTR_VOID)faultingIP);
     bool translateToManagedException = false;
     if (pCodeManager != NULL)
     {
@@ -471,7 +471,7 @@ int32_t __stdcall RhpVectoredExceptionHandler(PEXCEPTION_POINTERS pExPtrs)
 
     uintptr_t faultingIP = pExPtrs->ContextRecord->GetIp();
 
-    ICodeManager * pCodeManager = GetRuntimeInstance()->FindCodeManagerByAddress((PTR_VOID)faultingIP);
+    ICodeManager * pCodeManager = GetRuntimeInstance()->GetCodeManagerForAddress((PTR_VOID)faultingIP);
     bool translateToManagedException = false;
     if (pCodeManager != NULL)
     {

src/coreclr/nativeaot/Runtime/GCHelpers.cpp

@@ -31,7 +31,7 @@ EXTERN_C NATIVEAOT_API void __cdecl RhpCollect(uint32_t uGeneration, uint32_t uM
 
     Thread * pCurThread = ThreadStore::GetCurrentThread();
 
-    pCurThread->SetupHackPInvokeTunnel();
+    pCurThread->DeferTransitionFrame();
     pCurThread->DisablePreemptiveMode();
 
     ASSERT(!pCurThread->IsDoNotTriggerGcSet());
@@ -46,7 +46,7 @@ EXTERN_C NATIVEAOT_API int64_t __cdecl RhpGetGcTotalMemory()
 
     Thread * pCurThread = ThreadStore::GetCurrentThread();
 
-    pCurThread->SetupHackPInvokeTunnel();
+    pCurThread->DeferTransitionFrame();
     pCurThread->DisablePreemptiveMode();
 
     int64_t ret = GCHeapUtilities::GetGCHeap()->GetTotalBytesInUse();
@@ -61,7 +61,7 @@ EXTERN_C NATIVEAOT_API int32_t __cdecl RhpStartNoGCRegion(int64_t totalSize, UIn
     Thread *pCurThread = ThreadStore::GetCurrentThread();
     ASSERT(!pCurThread->IsCurrentThreadInCooperativeMode());
 
-    pCurThread->SetupHackPInvokeTunnel();
+    pCurThread->DeferTransitionFrame();
     pCurThread->DisablePreemptiveMode();
 
     int result = GCHeapUtilities::GetGCHeap()->StartNoGCRegion(totalSize, hasLohSize, lohSize, disallowFullBlockingGC);
@@ -316,7 +316,7 @@ EXTERN_C NATIVEAOT_API void RhAllocateNewArray(MethodTable* pArrayEEType, uint32
 {
     Thread* pThread = ThreadStore::GetCurrentThread();
 
-    pThread->SetupHackPInvokeTunnel();
+    pThread->DeferTransitionFrame();
     pThread->DisablePreemptiveMode();
 
     ASSERT(!pThread->IsDoNotTriggerGcSet());
@@ -330,7 +330,7 @@ EXTERN_C NATIVEAOT_API void RhAllocateNewObject(MethodTable* pEEType, uint32_t f
 {
     Thread* pThread = ThreadStore::GetCurrentThread();
 
-    pThread->SetupHackPInvokeTunnel();
+    pThread->DeferTransitionFrame();
     pThread->DisablePreemptiveMode();
 
     ASSERT(!pThread->IsDoNotTriggerGcSet());

src/coreclr/nativeaot/Runtime/ICodeManager.h

@@ -1,17 +1,15 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
-#pragma once
 
-#define ICODEMANAGER_INCLUDED
+#pragma once
+#include <rhbinder.h>
 
 // TODO: Debugger/DAC support (look for TODO: JIT)
 
 struct REGDISPLAY;
 
 #define GC_CALL_INTERIOR            0x1
 #define GC_CALL_PINNED              0x2
-#define GC_CALL_CHECK_APP_DOMAIN    0x4
-#define GC_CALL_STATIC              0x8
 
 typedef void (*GCEnumCallback)(
     void *              hCallback,      // callback data
@@ -47,6 +45,27 @@ enum GCRefKind : unsigned char
 };
 
 #ifdef TARGET_ARM64
+// Verify that we can use bitwise shifts to convert from GCRefKind to PInvokeTransitionFrameFlags and back
+C_ASSERT(PTFF_X0_IS_GCREF == ((uint64_t)GCRK_Object << 32));
+C_ASSERT(PTFF_X0_IS_BYREF == ((uint64_t)GCRK_Byref << 32));
+C_ASSERT(PTFF_X1_IS_GCREF == ((uint64_t)GCRK_Scalar_Obj << 32));
+C_ASSERT(PTFF_X1_IS_BYREF == ((uint64_t)GCRK_Scalar_Byref << 32));
+
+inline uint64_t ReturnKindToTransitionFrameFlags(GCRefKind returnKind)
+{
+    if (returnKind == GCRK_Scalar)
+        return 0;
+
+    return PTFF_SAVE_X0 | PTFF_SAVE_X1 | ((uint64_t)returnKind << 32);
+}
+
+inline GCRefKind TransitionFrameFlagsToReturnKind(uint64_t transFrameFlags)
+{
+    GCRefKind returnKind = (GCRefKind)((transFrameFlags & (PTFF_X0_IS_GCREF | PTFF_X0_IS_BYREF | PTFF_X1_IS_GCREF | PTFF_X1_IS_BYREF)) >> 32);
+    ASSERT((returnKind == GCRK_Scalar) || ((transFrameFlags & PTFF_SAVE_X0) && (transFrameFlags & PTFF_SAVE_X1)));
+    return returnKind;
+}
+
 // Extract individual GCRefKind components from a composite return kind
 inline GCRefKind ExtractReg0ReturnKind(GCRefKind returnKind)
 {
@@ -135,7 +154,7 @@ class ICodeManager
 
     virtual bool UnwindStackFrame(MethodInfo *    pMethodInfo,
                                   REGDISPLAY *    pRegisterSet,                     // in/out
-                                  PTR_VOID *      ppPreviousTransitionFrame) = 0;   // out
+                                  PInvokeTransitionFrame**      ppPreviousTransitionFrame) = 0;   // out
 
     virtual uintptr_t GetConservativeUpperBoundForOutgoingArgs(MethodInfo *   pMethodInfo,
                                                                 REGDISPLAY *   pRegisterSet) = 0;

src/coreclr/nativeaot/Runtime/MiscHelpers.cpp

@@ -109,7 +109,7 @@ COOP_PINVOKE_HELPER(uint32_t, RhGetLoadedOSModules, (Array * pResultArray))
 
 COOP_PINVOKE_HELPER(HANDLE, RhGetOSModuleFromPointer, (PTR_VOID pPointerVal))
 {
-    ICodeManager * pCodeManager = GetRuntimeInstance()->FindCodeManagerByAddress(pPointerVal);
+    ICodeManager * pCodeManager = GetRuntimeInstance()->GetCodeManagerForAddress(pPointerVal);
 
     if (pCodeManager != NULL)
         return (HANDLE)pCodeManager->GetOsModuleHandle();
@@ -393,7 +393,7 @@ EXTERN_C NATIVEAOT_API int32_t __cdecl RhpGetCurrentThreadStackTrace(void* pOutp
 {
     // This must be called via p/invoke rather than RuntimeImport to make the stack crawlable.
 
-    ThreadStore::GetCurrentThread()->SetupHackPInvokeTunnel();
+    ThreadStore::GetCurrentThread()->DeferTransitionFrame();
 
     return RhpCalculateStackTraceWorker(pOutputBuffer, outputBufferLength, pAddressInCurrentFrame);
 }

src/coreclr/nativeaot/Runtime/RuntimeInstance.cpp

@@ -56,7 +56,7 @@ COOP_PINVOKE_HELPER(uint8_t *, RhFindMethodStartAddress, (void * codeAddr))
 
 PTR_UInt8 RuntimeInstance::FindMethodStartAddress(PTR_VOID ControlPC)
 {
-    ICodeManager * pCodeManager = FindCodeManagerByAddress(ControlPC);
+    ICodeManager * pCodeManager = GetCodeManagerForAddress(ControlPC);
     MethodInfo methodInfo;
     if (pCodeManager != NULL && pCodeManager->FindMethodInfo(ControlPC, &methodInfo))
     {
@@ -66,20 +66,23 @@ PTR_UInt8 RuntimeInstance::FindMethodStartAddress(PTR_VOID ControlPC)
     return NULL;
 }
 
-ICodeManager * RuntimeInstance::FindCodeManagerByAddress(PTR_VOID pvAddress)
+// WARNING: This method is called by suspension while one thread is interrupted
+//          in a random location, possibly holding random locks.
+//          It is unsafe to use blocking APIs or allocate in this method.
+//          Please ensure that all methods called by this one also have this warning.
+bool RuntimeInstance::IsManaged(PTR_VOID pvAddress)
 {
-    ReaderWriterLock::ReadHolder read(&m_ModuleListLock);
+    return (dac_cast<TADDR>(pvAddress) - dac_cast<TADDR>(m_pvManagedCodeStartRange) < m_cbManagedCodeRange);
+}
 
-    // TODO: ICodeManager support in DAC
-#ifndef DACCESS_COMPILE
-    for (CodeManagerEntry * pEntry = m_CodeManagerList.GetHead(); pEntry != NULL; pEntry = pEntry->m_pNext)
+ICodeManager * RuntimeInstance::GetCodeManagerForAddress(PTR_VOID pvAddress)
+{
+    if (!IsManaged(pvAddress))
     {
-        if (dac_cast<TADDR>(pvAddress) - dac_cast<TADDR>(pEntry->m_pvStartRange) < pEntry->m_cbRange)
-            return pEntry->m_pCodeManager;
+        return NULL;
     }
-#endif
 
-    return NULL;
+    return m_CodeManager;
 }
 
 #ifndef DACCESS_COMPILE
@@ -90,7 +93,7 @@ ICodeManager * RuntimeInstance::FindCodeManagerByAddress(PTR_VOID pvAddress)
 ICodeManager * RuntimeInstance::FindCodeManagerForClasslibFunction(PTR_VOID address)
 {
     // Try looking up the code manager assuming the address is for code first. This is expected to be most common.
-    ICodeManager * pCodeManager = FindCodeManagerByAddress(address);
+    ICodeManager * pCodeManager = GetCodeManagerForAddress(address);
     if (pCodeManager != NULL)
         return pCodeManager;
 
@@ -119,7 +122,7 @@ void * RuntimeInstance::GetClasslibFunctionFromCodeAddress(PTR_VOID address, Cla
 
 PTR_UInt8 RuntimeInstance::GetTargetOfUnboxingAndInstantiatingStub(PTR_VOID ControlPC)
 {
-    ICodeManager * pCodeManager = FindCodeManagerByAddress(ControlPC);
+    ICodeManager * pCodeManager = GetCodeManagerForAddress(ControlPC);
     if (pCodeManager != NULL)
     {
         PTR_UInt8 pData = (PTR_UInt8)pCodeManager->GetAssociatedData(ControlPC);
@@ -137,6 +140,10 @@ PTR_UInt8 RuntimeInstance::GetTargetOfUnboxingAndInstantiatingStub(PTR_VOID Cont
 
 GPTR_IMPL_INIT(RuntimeInstance, g_pTheRuntimeInstance, NULL);
 
+// WARNING: This method is called by suspension while one thread is interrupted
+//          in a random location, possibly holding random locks.
+//          It is unsafe to use blocking APIs or allocate in this method.
+//          Please ensure that all methods called by this one also have this warning.
 PTR_RuntimeInstance GetRuntimeInstance()
 {
     return g_pTheRuntimeInstance;
@@ -150,28 +157,21 @@ void RuntimeInstance::EnumAllStaticGCRefs(void * pfnCallback, void * pvCallbackD
     }
 }
 
-void RuntimeInstance::SetLoopHijackFlags(uint32_t flag)
-{
-    for (TypeManagerList::Iterator iter = m_TypeManagerList.Begin(); iter != m_TypeManagerList.End(); iter++)
-    {
-        iter->m_pTypeManager->SetLoopHijackFlag(flag);
-    }
-}
-
 RuntimeInstance::OsModuleList* RuntimeInstance::GetOsModuleList()
 {
     return dac_cast<DPTR(OsModuleList)>(dac_cast<TADDR>(this) + offsetof(RuntimeInstance, m_OsModuleList));
 }
 
 ReaderWriterLock& RuntimeInstance::GetTypeManagerLock()
 {
-    return m_ModuleListLock;
+    return m_TypeManagerLock;
 }
 
 #ifndef DACCESS_COMPILE
 
 RuntimeInstance::RuntimeInstance() :
     m_pThreadStore(NULL),
+    m_CodeManager(NULL),
     m_conservativeStackReportingEnabled(false),
     m_pUnboxingStubsRegion(NULL)
 {
@@ -196,56 +196,19 @@ void RuntimeInstance::EnableConservativeStackReporting()
     m_conservativeStackReportingEnabled = true;
 }
 
-bool RuntimeInstance::RegisterCodeManager(ICodeManager * pCodeManager, PTR_VOID pvStartRange, uint32_t cbRange)
+void RuntimeInstance::RegisterCodeManager(ICodeManager * pCodeManager, PTR_VOID pvStartRange, uint32_t cbRange)
 {
-    CodeManagerEntry * pEntry = new (nothrow) CodeManagerEntry();
-    if (NULL == pEntry)
-        return false;
+    _ASSERTE(m_CodeManager == NULL);
+    _ASSERTE(pCodeManager != NULL);
 
-    pEntry->m_pvStartRange = pvStartRange;
-    pEntry->m_cbRange = cbRange;
-    pEntry->m_pCodeManager = pCodeManager;
-
-    {
-        ReaderWriterLock::WriteHolder write(&m_ModuleListLock);
-
-        m_CodeManagerList.PushHead(pEntry);
-    }
-
-    return true;
-}
-
-void RuntimeInstance::UnregisterCodeManager(ICodeManager * pCodeManager)
-{
-    CodeManagerEntry * pEntry = NULL;
-
-    {
-        ReaderWriterLock::WriteHolder write(&m_ModuleListLock);
-
-        for (CodeManagerList::Iterator i = m_CodeManagerList.Begin(), end = m_CodeManagerList.End(); i != end; i++)
-        {
-            if (i->m_pCodeManager == pCodeManager)
-            {
-                pEntry = *i;
-
-                m_CodeManagerList.Remove(i);
-                break;
-            }
-        }
-    }
-
-    ASSERT(pEntry != NULL);
-    delete pEntry;
-}
-
-extern "C" bool __stdcall RegisterCodeManager(ICodeManager * pCodeManager, PTR_VOID pvStartRange, uint32_t cbRange)
-{
-    return GetRuntimeInstance()->RegisterCodeManager(pCodeManager, pvStartRange, cbRange);
+    m_CodeManager = pCodeManager;
+    m_pvManagedCodeStartRange = pvStartRange;
+    m_cbManagedCodeRange = cbRange;
 }
 
-extern "C" void __stdcall UnregisterCodeManager(ICodeManager * pCodeManager)
+extern "C" void __stdcall RegisterCodeManager(ICodeManager * pCodeManager, PTR_VOID pvStartRange, uint32_t cbRange)
 {
-    return GetRuntimeInstance()->UnregisterCodeManager(pCodeManager);
+    GetRuntimeInstance()->RegisterCodeManager(pCodeManager, pvStartRange, cbRange);
 }
 
 bool RuntimeInstance::RegisterUnboxingStubs(PTR_VOID pvStartRange, uint32_t cbRange)
@@ -297,7 +260,7 @@ bool RuntimeInstance::RegisterTypeManager(TypeManager * pTypeManager)
     pEntry->m_pTypeManager = pTypeManager;
 
     {
-        ReaderWriterLock::WriteHolder write(&m_ModuleListLock);
+        ReaderWriterLock::WriteHolder write(&m_TypeManagerLock);
 
         m_TypeManagerList.PushHead(pEntry);
     }

src/coreclr/nativeaot/Runtime/RuntimeInstance.h

@@ -25,7 +25,7 @@ class RuntimeInstance
 
     PTR_ThreadStore             m_pThreadStore;
     HANDLE                      m_hPalInstance; // this is the HANDLE passed into DllMain
-    ReaderWriterLock            m_ModuleListLock;
+    ReaderWriterLock            m_TypeManagerLock;
 
 public:
     struct OsModuleEntry;
@@ -40,19 +40,11 @@ class RuntimeInstance
 private:
     OsModuleList                m_OsModuleList;
 
-    struct CodeManagerEntry;
-    typedef DPTR(CodeManagerEntry) PTR_CodeManagerEntry;
+    ICodeManager*               m_CodeManager;
 
-    struct CodeManagerEntry
-    {
-        PTR_CodeManagerEntry    m_pNext;
-        PTR_VOID                m_pvStartRange;
-        uint32_t                  m_cbRange;
-        ICodeManager *          m_pCodeManager;
-    };
-
-    typedef SList<CodeManagerEntry> CodeManagerList;
-    CodeManagerList             m_CodeManagerList;
+    // we support only one code manager for now, so we just record the range.
+    void*                       m_pvManagedCodeStartRange;
+    uint32_t                    m_cbManagedCodeRange;
 
 public:
     struct TypeManagerEntry
@@ -99,10 +91,9 @@ class RuntimeInstance
     void EnableConservativeStackReporting();
     bool IsConservativeStackReportingEnabled() { return m_conservativeStackReportingEnabled; }
 
-    bool RegisterCodeManager(ICodeManager * pCodeManager, PTR_VOID pvStartRange, uint32_t cbRange);
-    void UnregisterCodeManager(ICodeManager * pCodeManager);
+    void RegisterCodeManager(ICodeManager * pCodeManager, PTR_VOID pvStartRange, uint32_t cbRange);
 
-    ICodeManager * FindCodeManagerByAddress(PTR_VOID ControlPC);
+    ICodeManager * GetCodeManagerForAddress(PTR_VOID ControlPC);
     PTR_VOID GetClasslibFunctionFromCodeAddress(PTR_VOID address, ClasslibFunctionId functionId);
 
     bool RegisterTypeManager(TypeManager * pTypeManager);
@@ -113,14 +104,15 @@ class RuntimeInstance
     bool RegisterUnboxingStubs(PTR_VOID pvStartRange, uint32_t cbRange);
     bool IsUnboxingStub(uint8_t* pCode);
 
+    bool IsManaged(PTR_VOID pvAddress);
+
     static bool Initialize(HANDLE hPalInstance);
     void Destroy();
 
     void EnumAllStaticGCRefs(void * pfnCallback, void * pvCallbackData);
 
     bool ShouldHijackCallsiteForGcStress(uintptr_t CallsiteIP);
     bool ShouldHijackLoopForGcStress(uintptr_t CallsiteIP);
-    void SetLoopHijackFlags(uint32_t flag);
 };
 typedef DPTR(RuntimeInstance) PTR_RuntimeInstance;