[NativeAOT] Asynchronous thread suspension at GC safe points on Windows (#70316)

* simple case works

* FEATURE_SUSPEND_REDIRECTION

* AVX support

* refactor

* no need for asm

* ARM64

* stub `IsSafePoint` on Unix

* revert unnecessary changes to RhpGcPoll2

* no copy of redirect context

* Specialcase active frames in EnumGcRefs

* remove bogus assert.

Author: VSadov

src/coreclr/nativeaot/Runtime/CMakeLists.txt

@@ -245,6 +245,10 @@ add_definitions(-D_LIB)
 if(WIN32)
   add_definitions(-DFEATURE_ETW)
   add_definitions(-DFEATURE_EVENT_TRACE)
+  # redirection on ARM64 should work, but needs to be tested
+  if (CLR_CMAKE_TARGET_ARCH_AMD64)
+    add_definitions(-DFEATURE_SUSPEND_REDIRECTION)
+  endif()
 else()
   add_definitions(-DNO_UI_ASSERT)
   include(unix/configure.cmake)

src/coreclr/nativeaot/Runtime/ICodeManager.h

@@ -139,6 +139,8 @@ enum class AssociatedDataFlags : unsigned char
 class ICodeManager
 {
 public:
+    virtual bool IsSafePoint(PTR_VOID pvAddress) = 0;
+
     virtual bool FindMethodInfo(PTR_VOID        ControlPC,
                                 MethodInfo *    pMethodInfoOut) = 0;
 
@@ -150,7 +152,8 @@ class ICodeManager
     virtual void EnumGcRefs(MethodInfo *    pMethodInfo,
                             PTR_VOID        safePointAddress,
                             REGDISPLAY *    pRegisterSet,
-                            GCEnumContext * hCallback) = 0;
+                            GCEnumContext * hCallback,
+                            bool            isActiveStackFrame) = 0;
 
     virtual bool UnwindStackFrame(MethodInfo *    pMethodInfo,
                                   REGDISPLAY *    pRegisterSet,                     // in/out

src/coreclr/nativeaot/Runtime/PalRedhawk.h

@@ -405,6 +405,7 @@ typedef struct DECLSPEC_ALIGN(16) _CONTEXT {
     void SetArg1Reg(uintptr_t val) { X1 = val; }
     uintptr_t GetIp() { return Pc; }
     uintptr_t GetLr() { return Lr; }
+    uintptr_t GetSp() { return Sp; }
 } CONTEXT, *PCONTEXT;
 
 #elif defined(HOST_WASM)
@@ -590,6 +591,11 @@ REDHAWK_PALIMPORT bool REDHAWK_PALAPI PalInit();
 // Given the OS handle of a loaded module, compute the upper and lower virtual address bounds (inclusive).
 REDHAWK_PALIMPORT void REDHAWK_PALAPI PalGetModuleBounds(HANDLE hOsHandle, _Out_ uint8_t ** ppLowerBound, _Out_ uint8_t ** ppUpperBound);
 
+REDHAWK_PALIMPORT CONTEXT* PalAllocateCompleteOSContext(_Out_ uint8_t** contextBuffer);
+REDHAWK_PALIMPORT bool REDHAWK_PALAPI PalGetCompleteThreadContext(HANDLE hThread, _Out_ CONTEXT * pCtx);
+REDHAWK_PALIMPORT bool REDHAWK_PALAPI PalSetThreadContext(HANDLE hThread, _Out_ CONTEXT * pCtx);
+REDHAWK_PALIMPORT void REDHAWK_PALAPI PalRestoreContext(CONTEXT * pCtx);
+
 REDHAWK_PALIMPORT bool REDHAWK_PALAPI PalGetThreadContext(HANDLE hThread, _Out_ PAL_LIMITED_CONTEXT * pCtx);
 
 REDHAWK_PALIMPORT int32_t REDHAWK_PALAPI PalGetProcessCpuCount();

src/coreclr/nativeaot/Runtime/StackFrameIterator.cpp

@@ -92,7 +92,18 @@ StackFrameIterator::StackFrameIterator(Thread * pThreadToWalk, PInvokeTransition
 {
     STRESS_LOG0(LF_STACKWALK, LL_INFO10000, "----Init---- [ GC ]\n");
     ASSERT(!pThreadToWalk->DangerousCrossThreadIsHijacked());
-    InternalInit(pThreadToWalk, pInitialTransitionFrame, GcStackWalkFlags);
+
+#ifdef FEATURE_SUSPEND_REDIRECTION
+    if (pInitialTransitionFrame == REDIRECTED_THREAD_MARKER)
+    {
+        InternalInit(pThreadToWalk, pThreadToWalk->GetRedirectionContext(), GcStackWalkFlags | ActiveStackFrame);
+    }
+    else
+#endif
+    {
+        InternalInit(pThreadToWalk, pInitialTransitionFrame, GcStackWalkFlags);
+    }
+
     PrepareToYieldFrame();
 }
 
@@ -497,6 +508,112 @@ void StackFrameIterator::InternalInit(Thread * pThreadToWalk, PTR_PAL_LIMITED_CO
 #endif // TARGET_ARM
 }
 
+// Prepare to start a stack walk from the context listed in the supplied CONTEXT.
+// The supplied context can describe a location in either managed or unmanaged code.  In the
+// latter case the iterator is left in an invalid state when this function returns.
+void StackFrameIterator::InternalInit(Thread * pThreadToWalk, CONTEXT* pCtx, uint32_t dwFlags)
+{
+    ASSERT((dwFlags & MethodStateCalculated) == 0);
+
+    EnterInitialInvalidState(pThreadToWalk);
+
+    m_dwFlags = dwFlags;
+
+    // We need to walk the ExInfo chain in parallel with the stackwalk so that we know when we cross over
+    // exception throw points.  So we must find our initial point in the ExInfo chain here so that we can
+    // properly walk it in parallel.
+    ResetNextExInfoForSP(pCtx->GetSp());
+
+    // This codepath is used by the hijack stackwalk and we can get arbitrary ControlPCs from there.  If this
+    // context has a non-managed control PC, then we're done.
+    if (!m_pInstance->IsManaged(dac_cast<PTR_VOID>(pCtx->GetIp())))
+        return;
+
+    //
+    // control state
+    //
+    SetControlPC(dac_cast<PTR_VOID>(pCtx->GetIp()));
+    m_RegDisplay.SP   = pCtx->GetSp();
+    m_RegDisplay.IP   = pCtx->GetIp();
+
+#ifdef TARGET_ARM64
+
+    m_RegDisplay.pIP  = PTR_TO_MEMBER(CONTEXT, pCtx, Pc);
+
+    //
+    // preserved regs
+    //
+    m_RegDisplay.pX19 = PTR_TO_MEMBER(CONTEXT, pCtx, X19);
+    m_RegDisplay.pX20 = PTR_TO_MEMBER(CONTEXT, pCtx, X20);
+    m_RegDisplay.pX21 = PTR_TO_MEMBER(CONTEXT, pCtx, X21);
+    m_RegDisplay.pX22 = PTR_TO_MEMBER(CONTEXT, pCtx, X22);
+    m_RegDisplay.pX23 = PTR_TO_MEMBER(CONTEXT, pCtx, X23);
+    m_RegDisplay.pX24 = PTR_TO_MEMBER(CONTEXT, pCtx, X24);
+    m_RegDisplay.pX25 = PTR_TO_MEMBER(CONTEXT, pCtx, X25);
+    m_RegDisplay.pX26 = PTR_TO_MEMBER(CONTEXT, pCtx, X26);
+    m_RegDisplay.pX27 = PTR_TO_MEMBER(CONTEXT, pCtx, X27);
+    m_RegDisplay.pX28 = PTR_TO_MEMBER(CONTEXT, pCtx, X28);
+    m_RegDisplay.pFP = PTR_TO_MEMBER(CONTEXT, pCtx, Fp);
+    m_RegDisplay.pLR = PTR_TO_MEMBER(CONTEXT, pCtx, Lr);
+
+    //
+    // scratch regs
+    //
+    m_RegDisplay.pX0 = PTR_TO_MEMBER(CONTEXT, pCtx, X0);
+    m_RegDisplay.pX1 = PTR_TO_MEMBER(CONTEXT, pCtx, X1);
+    m_RegDisplay.pX2 = PTR_TO_MEMBER(CONTEXT, pCtx, X2);
+    m_RegDisplay.pX3 = PTR_TO_MEMBER(CONTEXT, pCtx, X3);
+    m_RegDisplay.pX4 = PTR_TO_MEMBER(CONTEXT, pCtx, X4);
+    m_RegDisplay.pX5 = PTR_TO_MEMBER(CONTEXT, pCtx, X5);
+    m_RegDisplay.pX6 = PTR_TO_MEMBER(CONTEXT, pCtx, X6);
+    m_RegDisplay.pX7 = PTR_TO_MEMBER(CONTEXT, pCtx, X7);
+    m_RegDisplay.pX8 = PTR_TO_MEMBER(CONTEXT, pCtx, X8);
+    m_RegDisplay.pX9 = PTR_TO_MEMBER(CONTEXT, pCtx, X9);
+    m_RegDisplay.pX10 = PTR_TO_MEMBER(CONTEXT, pCtx, X10);
+    m_RegDisplay.pX11 = PTR_TO_MEMBER(CONTEXT, pCtx, X11);
+    m_RegDisplay.pX12 = PTR_TO_MEMBER(CONTEXT, pCtx, X12);
+    m_RegDisplay.pX13 = PTR_TO_MEMBER(CONTEXT, pCtx, X13);
+    m_RegDisplay.pX14 = PTR_TO_MEMBER(CONTEXT, pCtx, X14);
+    m_RegDisplay.pX15 = PTR_TO_MEMBER(CONTEXT, pCtx, X15);
+    m_RegDisplay.pX16 = PTR_TO_MEMBER(CONTEXT, pCtx, X16);
+    m_RegDisplay.pX17 = PTR_TO_MEMBER(CONTEXT, pCtx, X17);
+    m_RegDisplay.pX18 = PTR_TO_MEMBER(CONTEXT, pCtx, X18);
+
+#elif defined(TARGET_X86) || defined(TARGET_AMD64)
+
+    m_RegDisplay.pIP  = (PTR_PCODE)PTR_TO_MEMBER(CONTEXT, pCtx, Rip);
+
+    //
+    // preserved regs
+    //
+    m_RegDisplay.pRbp = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rbp);
+    m_RegDisplay.pRsi = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rsi);
+    m_RegDisplay.pRdi = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rdi);
+    m_RegDisplay.pRbx = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rbx);
+#ifdef TARGET_AMD64     
+    m_RegDisplay.pR12 = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R12);
+    m_RegDisplay.pR13 = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R13);
+    m_RegDisplay.pR14 = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R14);
+    m_RegDisplay.pR15 = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R15);
+#endif // TARGET_AMD64  
+                        
+    //                  
+    // scratch regs     
+    //                  
+    m_RegDisplay.pRax = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rax);
+    m_RegDisplay.pRcx = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rcx);
+    m_RegDisplay.pRdx = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, Rdx);
+#ifdef TARGET_AMD64     
+    m_RegDisplay.pR8  = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R8);
+    m_RegDisplay.pR9  = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R9);
+    m_RegDisplay.pR10 = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R10);
+    m_RegDisplay.pR11 = (PTR_UIntNative)PTR_TO_MEMBER(CONTEXT, pCtx, R11);
+#endif // TARGET_AMD64
+#else
+    PORTABILITY_ASSERT("StackFrameIterator::InternalInit");
+#endif // TARGET_ARM
+}
+
 PTR_VOID StackFrameIterator::HandleExCollide(PTR_ExInfo pExInfo)
 {
     STRESS_LOG3(LF_STACKWALK, LL_INFO10000, "   [ ex collide ] kind = %d, pass = %d, idxCurClause = %d\n",
@@ -1290,7 +1407,7 @@ void StackFrameIterator::NextInternal()
 {
 UnwindOutOfCurrentManagedFrame:
     ASSERT(m_dwFlags & MethodStateCalculated);
-    m_dwFlags &= ~(ExCollide|MethodStateCalculated|UnwoundReversePInvoke);
+    m_dwFlags &= ~(ExCollide|MethodStateCalculated|UnwoundReversePInvoke|ActiveStackFrame);
     ASSERT(IsValid());
 
     m_pHijackedReturnValue = NULL;
@@ -1640,6 +1757,12 @@ MethodInfo * StackFrameIterator::GetMethodInfo()
     return &m_methodInfo;
 }
 
+bool StackFrameIterator::IsActiveStackFrame()
+{
+    ASSERT(IsValid());
+    return (m_dwFlags & ActiveStackFrame) != 0;
+}
+
 #ifdef DACCESS_COMPILE
 #define FAILFAST_OR_DAC_RETURN_FALSE(x) if(!(x)) return false;
 #else

src/coreclr/nativeaot/Runtime/StackFrameIterator.h

@@ -36,14 +36,14 @@ class StackFrameIterator
     StackFrameIterator(Thread * pThreadToWalk, PInvokeTransitionFrame* pInitialTransitionFrame);
     StackFrameIterator(Thread * pThreadToWalk, PTR_PAL_LIMITED_CONTEXT pCtx);
 
-
     bool             IsValid();
     void             CalculateCurrentMethodState();
     void             Next();
     PTR_VOID         GetEffectiveSafePointAddress();
     REGDISPLAY *     GetRegisterSet();
     PTR_ICodeManager GetCodeManager();
     MethodInfo *     GetMethodInfo();
+    bool             IsActiveStackFrame();
     bool             GetHijackedReturnValueLocation(PTR_RtuObjectRef * pLocation, GCRefKind * pKind);
     void             SetControlPC(PTR_VOID controlPC);
 
@@ -82,6 +82,8 @@ class StackFrameIterator
 
     void InternalInit(Thread * pThreadToWalk, PTR_PInvokeTransitionFrame pFrame, uint32_t dwFlags); // GC stackwalk
     void InternalInit(Thread * pThreadToWalk, PTR_PAL_LIMITED_CONTEXT pCtx, uint32_t dwFlags);  // EH and hijack stackwalk, and collided unwind
+    void InternalInit(Thread * pThreadToWalk, CONTEXT* pCtx, uint32_t dwFlags);  // GC stackwalk of redirected thread
+
     void InternalInitForEH(Thread * pThreadToWalk, PAL_LIMITED_CONTEXT * pCtx, bool instructionFault); // EH stackwalk
     void InternalInitForStackTrace();  // Environment.StackTrace
 
@@ -137,6 +139,9 @@ class StackFrameIterator
         // This is a state returned by Next() which indicates that we just unwound a reverse pinvoke method
         UnwoundReversePInvoke = 0x20,
 
+        // The thread was interrupted in the current frame at the current IP by a signal, SuspendThread or similar.
+        ActiveStackFrame = 0x40,
+
         GcStackWalkFlags = (CollapseFunclets | RemapHardwareFaultsToSafePoint),
         EHStackWalkFlags = ApplyReturnAddressAdjustment,
         StackTraceStackWalkFlags = GcStackWalkFlags

src/coreclr/nativeaot/Runtime/arm64/GcProbe.asm

@@ -8,7 +8,7 @@
     EXTERN      RhpGcPoll2
     EXTERN      g_fGcStressStarted
 
-PROBE_SAVE_FLAGS_EVERYTHING     equ DEFAULT_FRAME_SAVE_FLAGS + PTFF_SAVE_ALL_SCRATCH
+PROBE_SAVE_FLAGS_EVERYTHING     equ DEFAULT_FRAME_SAVE_FLAGS + PTFF_SAVE_ALL_SCRATCH + PTFF_SAVE_LR
 
     ;; Build a map of symbols representing offsets into the transition frame (see PInvokeTransitionFrame in
     ;; rhbinder.h) and keep these two in sync.

src/coreclr/nativeaot/Runtime/gcrhenv.cpp

@@ -515,7 +515,8 @@ void RedhawkGCInterface::EnumGcRefs(ICodeManager * pCodeManager,
                                     PTR_VOID safePointAddress,
                                     REGDISPLAY * pRegisterSet,
                                     void * pfnEnumCallback,
-                                    void * pvCallbackData)
+                                    void * pvCallbackData,
+                                    bool   isActiveStackFrame)
 {
     EnumGcRefContext ctx;
     ctx.pCallback = EnumGcRefsCallback;
@@ -526,7 +527,8 @@ void RedhawkGCInterface::EnumGcRefs(ICodeManager * pCodeManager,
     pCodeManager->EnumGcRefs(pMethodInfo,
                              safePointAddress,
                              pRegisterSet,
-                             &ctx);
+                             &ctx,
+                             isActiveStackFrame);
 }
 
 // static

src/coreclr/nativeaot/Runtime/gcrhinterface.h

@@ -120,7 +120,8 @@ class RedhawkGCInterface
                            PTR_VOID safePointAddress,
                            REGDISPLAY * pRegisterSet,
                            void * pfnEnumCallback,
-                           void * pvCallbackData);
+                           void * pvCallbackData,
+                           bool   isActiveStackFrame);
 
     static void EnumGcRefsInRegionConservatively(PTR_RtuObjectRef pLowerBound,
                                                  PTR_RtuObjectRef pUpperBound,