Add TSF_SuspensionTrapped flag to NativeAOT for accurate sample profiler reporting

src/coreclr/nativeaot/Runtime/eventpipe/ep-rt-aot.cpp

@@ -26,6 +26,7 @@
 #include "threadstore.inl"
 #include "eventtrace_context.h"
 #include "eventtracebase.h"
+#include "thread.inl"
 
 // Uses _rt_aot_lock_internal_t that has CrstStatic as a field
 // This is initialized at the beginning and EventPipe library requires the lock handle to be maintained by the runtime
@@ -122,9 +123,11 @@ ep_rt_aot_sample_profiler_write_sampling_event_for_threads (
 
         // Walk the stack and write it out as an event.
         if (ep_rt_aot_walk_managed_stack_for_thread (target_thread, current_stack_contents) && !ep_stack_contents_is_empty (current_stack_contents)) {
-            // Set the payload.
-            // TODO: We can actually detect whether we are in managed or external code but does it matter?!
-            uint32_t payload_data = EP_SAMPLE_PROFILER_SAMPLE_TYPE_EXTERNAL;
+            // Set the payload. If the thread is trapped for suspension, it was in cooperative mode
+            // (managed code). Otherwise, it was in preemptive mode (external code).
+            uint32_t payload_data = target_thread->IsSuspensionTrapped()
+                ? EP_SAMPLE_PROFILER_SAMPLE_TYPE_MANAGED
+                : EP_SAMPLE_PROFILER_SAMPLE_TYPE_EXTERNAL;
 
             // Write the sample.
             ep_write_sample_profile_event (

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -83,6 +83,10 @@ void Thread::WaitForGC(PInvokeTransitionFrame* pTransitionFrame)
     // restored after the wait operation;
     int32_t lastErrorOnEntry = PalGetLastError();
 
+    // Mark that this thread is trapped for suspension.
+    // Used by the sample profiler to determine this thread was in managed code.
+    SetState(TSF_SuspensionTrapped);
+
     do
     {
         // set preemptive mode
@@ -102,6 +106,8 @@ void Thread::WaitForGC(PInvokeTransitionFrame* pTransitionFrame)
     }
     while (ThreadStore::IsTrapThreadsRequested());
 
+    ClearState(TSF_SuspensionTrapped);
+
     // Restore the saved error
     PalSetLastError(lastErrorOnEntry);
 }

src/coreclr/nativeaot/Runtime/thread.h

@@ -214,6 +214,9 @@ class Thread : private RuntimeThreadLocals
                                                     // On Unix this is an optimization to not queue up more signals when one is
                                                     // still being processed.
         TSF_Interrupted         = 0x00000200,       // Set to indicate Thread.Interrupt() has been called on this thread
+
+        TSF_SuspensionTrapped   = 0x00000400,       // Set when thread is trapped waiting for suspension to complete
+                                                    // (was in managed code).
     };
 private:
 
@@ -306,6 +309,8 @@ class Thread : private RuntimeThreadLocals
     bool                IsDetached();
     void                SetDetached();
 
+    bool                IsSuspensionTrapped();
+
     PTR_VOID            GetThreadStressLog() const;
 #ifndef DACCESS_COMPILE
     void                SetThreadStressLog(void * ptsl);

src/coreclr/nativeaot/Runtime/thread.inl

@@ -146,6 +146,11 @@ inline bool Thread::IsDoNotTriggerGcSet()
     return IsStateSet(TSF_DoNotTriggerGc);
 }
 
+inline bool Thread::IsSuspensionTrapped()
+{
+    return IsStateSet(TSF_SuspensionTrapped);
+}
+
 inline bool Thread::IsCurrentThreadInCooperativeMode()
 {
 #ifndef DACCESS_COMPILE

src/tests/tracing/eventpipe/eventsvalidation/SampleProfilerSampleType.cs

@@ -22,7 +22,6 @@ public class SampleProfilerSampleType
         private const uint SampleTypeExternal = 1;
         private const uint SampleTypeManaged = 2;
 
-        [ActiveIssue("https://github.com/dotnet/runtime/issues/125217", typeof(Utilities), nameof(Utilities.IsNativeAot))]
         [Fact]
         public static int TestEntryPoint()
         {