From cf3554ab2f627018206987a1cdc2c2ed715867a7 Mon Sep 17 00:00:00 2001
From: Filip Navara <filip.navara@gmail.com>
Date: Fri, 20 Sep 2024 23:01:34 +0200
Subject: [PATCH] Add AOT sampling profiler

---
 .../nativeaot/Runtime/StackFrameIterator.h    |  1 +
 .../nativeaot/Runtime/eventpipe/ep-rt-aot.cpp | 72 ++++++++++++++++++-
 src/coreclr/nativeaot/Runtime/thread.h        |  1 +
 3 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/src/coreclr/nativeaot/Runtime/StackFrameIterator.h b/src/coreclr/nativeaot/Runtime/StackFrameIterator.h
index f174edd4c473b2..c6ecff14506c35 100644
--- a/src/coreclr/nativeaot/Runtime/StackFrameIterator.h
+++ b/src/coreclr/nativeaot/Runtime/StackFrameIterator.h
@@ -53,6 +53,7 @@ class StackFrameIterator
     bool             IsActiveStackFrame();
     bool             GetHijackedReturnValueLocation(PTR_OBJECTREF * pLocation, GCRefKind * pKind);
     void             SetControlPC(PTR_VOID controlPC);
+    PTR_VOID         GetControlPC() { return m_ControlPC; }
 
     static bool     IsValidReturnAddress(PTR_VOID pvAddress);
 
diff --git a/src/coreclr/nativeaot/Runtime/eventpipe/ep-rt-aot.cpp b/src/coreclr/nativeaot/Runtime/eventpipe/ep-rt-aot.cpp
index e54156891593f3..a539cf3b9c98c5 100644
--- a/src/coreclr/nativeaot/Runtime/eventpipe/ep-rt-aot.cpp
+++ b/src/coreclr/nativeaot/Runtime/eventpipe/ep-rt-aot.cpp
@@ -45,8 +45,29 @@ ep_rt_aot_walk_managed_stack_for_thread (
     ep_rt_thread_handle_t thread,
     EventPipeStackContents *stack_contents)
 {
-    // NativeAOT does not support getting the call stack
-    return false;
+    STATIC_CONTRACT_NOTHROW;
+    EP_ASSERT (thread != NULL);
+    EP_ASSERT (stack_contents != NULL);
+
+    StackFrameIterator frameIterator(thread, thread->GetTransitionFrameForSampling()); 
+
+    while (frameIterator.IsValid())
+    {
+        frameIterator.CalculateCurrentMethodState();
+
+        // Get the IP.
+        uintptr_t control_pc = (uintptr_t)frameIterator.GetControlPC();
+
+        if (control_pc != 0)
+        {
+            // Add the IP to the captured stack.
+            ep_stack_contents_append (stack_contents, control_pc, NULL);
+        }
+
+        frameIterator.Next();
+    }
+
+    return true;
 }
 
 bool
@@ -61,6 +82,53 @@ ep_rt_aot_sample_profiler_write_sampling_event_for_threads (
     ep_rt_thread_handle_t sampling_thread,
     EventPipeEvent *sampling_event)
 {
+    STATIC_CONTRACT_NOTHROW;
+    EP_ASSERT (sampling_thread != NULL);
+
+    ThreadStore *thread_store = GetThreadStore ();
+
+    // Check to see if we can suspend managed execution.
+    if (thread_store->GetSuspendingThread () != NULL)
+        return;
+
+    // Actually suspend managed execution.
+    thread_store->LockThreadStore ();
+    thread_store->SuspendAllThreads (false);
+
+    EventPipeStackContents stack_contents;
+    EventPipeStackContents *current_stack_contents;
+    current_stack_contents = ep_stack_contents_init (&stack_contents);
+
+    EP_ASSERT (current_stack_contents != NULL);
+
+    // Walk all managed threads and capture stacks.
+    FOREACH_THREAD (target_thread)
+    {
+        ep_stack_contents_reset (current_stack_contents);
+
+        // 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;
+
+            // Write the sample.
+            ep_write_sample_profile_event (
+                sampling_thread,
+                sampling_event,
+                target_thread,
+                current_stack_contents,
+                (uint8_t *)&payload_data,
+                sizeof (payload_data));
+        }
+    }
+    END_FOREACH_THREAD
+
+    ep_stack_contents_fini (current_stack_contents);
+
+    // Resume managed execution.
+    thread_store->ResumeAllThreads(false);
+    thread_store->UnlockThreadStore();
 }
 
 const ep_char8_t *
diff --git a/src/coreclr/nativeaot/Runtime/thread.h b/src/coreclr/nativeaot/Runtime/thread.h
index 4c0a21e9f9ab7f..d8c4f15a65c7c0 100644
--- a/src/coreclr/nativeaot/Runtime/thread.h
+++ b/src/coreclr/nativeaot/Runtime/thread.h
@@ -266,6 +266,7 @@ class Thread : private RuntimeThreadLocals
     bool                IsCurrentThreadInCooperativeMode();
 
     PInvokeTransitionFrame* GetTransitionFrameForStackTrace();
+    PInvokeTransitionFrame* GetTransitionFrameForSampling() { return GetTransitionFrame(); }
     void *              GetCurrentThreadPInvokeReturnAddress();
 
     //