From d821bf2fc1b5f6fd10dda81a695431a653edd3e8 Mon Sep 17 00:00:00 2001
From: Markus Stange <markus.stange@gmail.com>
Date: Mon, 15 Jul 2024 17:49:23 -0400
Subject: [PATCH] Add support for SourceLoad events.

Now we get proper JS stacks out of regular Chrome builds.
---
 samply/src/shared/jit_category_manager.rs |  1 +
 samply/src/windows/etw_gecko.rs           | 30 +++++++++++-------
 samply/src/windows/profile_context.rs     | 37 ++++++++++++++++++++++-
 3 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/samply/src/shared/jit_category_manager.rs b/samply/src/shared/jit_category_manager.rs
index 66b5ce76..ccb174a2 100644
--- a/samply/src/shared/jit_category_manager.rs
+++ b/samply/src/shared/jit_category_manager.rs
@@ -33,6 +33,7 @@ impl JitCategoryManager {
         ("JS:^", "Baseline", CategoryColor::Blue, true),
         ("JS:+", "Maglev", CategoryColor::Green, true),
         ("JS:*", "Turbofan", CategoryColor::Green, true),
+        ("JS:?", "JavaScript", CategoryColor::Blue, true),
         ("Builtin:", "Builtin", CategoryColor::Brown, false),
         ("BytecodeHandler:", "Interpreter", CategoryColor::Red, false),
         ("Interpreter: ", "Interpreter", CategoryColor::Red, true),
diff --git a/samply/src/windows/etw_gecko.rs b/samply/src/windows/etw_gecko.rs
index 44c2c2c2..2e5caab0 100644
--- a/samply/src/windows/etw_gecko.rs
+++ b/samply/src/windows/etw_gecko.rs
@@ -332,9 +332,20 @@ fn process_trace(
                 // these events can give us the unblocking stack
                 let _thread_id: u32 = parser.parse("TThreadId");
             }
+            "V8.js/SourceLoad/Start"
+            | "Microsoft-JScript/ScriptContextRuntime/SourceLoad"
+            | "Microsoft-JScript/ScriptContextRundown/SourceDCStart" => {
+                let pid = s.process_id();
+                if !context.has_process_at_time(pid, timestamp_raw) {
+                    return;
+                }
+                let source_id: u64 = parser.parse("SourceID");
+                let url: String = parser.parse("Url");
+                context.handle_js_source_load(timestamp_raw, pid, source_id, url);
+            }
             "V8.js/MethodLoad/Start"
-            | "Microsoft-JScript/MethodRuntime/MethodDCStart"
-            | "Microsoft-JScript/MethodRuntime/MethodLoad" => {
+            | "Microsoft-JScript/MethodRuntime/MethodLoad"
+            | "Microsoft-JScript/MethodRundown/MethodDCStart" => {
                 let pid = s.process_id();
                 if !context.has_process_at_time(pid, timestamp_raw) {
                     return;
@@ -342,23 +353,20 @@ fn process_trace(
                 let method_name: String = parser.parse("MethodName");
                 let method_start_address: Address = parser.parse("MethodStartAddress");
                 let method_size: u64 = parser.parse("MethodSize");
-                // let source_id: u64 = parser.parse("SourceID");
+                let source_id: u64 = parser.parse("SourceID");
+                let line: u32 = parser.parse("Line");
+                let column: u32 = parser.parse("Column");
                 context.handle_js_method_load(
                     timestamp_raw,
                     pid,
                     method_name,
                     method_start_address.as_u64(),
                     method_size as u32,
+                    source_id,
+                    line,
+                    column,
                 );
             }
-            /*"V8.js/SourceLoad/" |
-            "Microsoft-JScript/MethodRuntime/MethodDCStart" |
-            "Microsoft-JScript/MethodRuntime/MethodLoad" => {
-                let source_id: u64 = parser.parse("SourceID");
-                let url: String = parser.parse("Url");
-                jscript_sources.insert(source_id, url);
-                //dbg!(s.process_id(), jscript_symbols.keys());
-            }*/
             "Microsoft-Windows-Direct3D11/ID3D11VideoContext_SubmitDecoderBuffers/win:Start" => {
                 if !context.is_in_time_range(timestamp_raw) {
                     return;
diff --git a/samply/src/windows/profile_context.rs b/samply/src/windows/profile_context.rs
index c702a48e..00298461 100644
--- a/samply/src/windows/profile_context.rs
+++ b/samply/src/windows/profile_context.rs
@@ -300,6 +300,7 @@ pub struct Process {
     pub parent_id: u32,
     pub thread_recycler: Option<ThreadRecycler>,
     pub jit_function_recycler: Option<JitFunctionRecycler>,
+    pub js_sources: HashMap<u64, String>,
 }
 
 impl Process {
@@ -329,6 +330,7 @@ impl Process {
             parent_id,
             thread_recycler,
             jit_function_recycler,
+            js_sources: HashMap::new(),
         }
     }
 
@@ -1551,18 +1553,51 @@ impl ProfileContext {
         }
     }
 
+    pub fn handle_js_source_load(
+        &mut self,
+        timestamp_raw: u64,
+        pid: u32,
+        source_id: u64,
+        url: String,
+    ) {
+        let Some(process) = self.processes.get_by_pid_and_timestamp(pid, timestamp_raw) else {
+            return;
+        };
+
+        process.js_sources.insert(source_id, url);
+    }
+
+    #[allow(clippy::too_many_arguments)]
     pub fn handle_js_method_load(
         &mut self,
         timestamp_raw: u64,
         pid: u32,
-        method_name: String,
+        mut method_name: String,
         method_start_address: u64,
         method_size: u32,
+        source_id: u64,
+        line: u32,
+        column: u32,
     ) {
         let Some(process) = self.processes.get_by_pid_and_timestamp(pid, timestamp_raw) else {
             return;
         };
 
+        if let Some(url) = process.js_sources.get(&source_id) {
+            if !method_name.starts_with("JS:") {
+                method_name.insert_str(0, "JS:?");
+                method_name.push(' ');
+                method_name.push_str(url);
+                if line != 0 {
+                    use std::fmt::Write;
+                    write!(&mut method_name, ":{line}").unwrap();
+                    if column != 0 {
+                        write!(&mut method_name, ":{column}").unwrap();
+                    }
+                }
+            }
+        }
+
         let lib = &mut self.js_jit_lib;
         let (category, js_frame) = self
             .js_category_manager