Load generated trampolines into jitdump when profiling (#3344)

* Load generated trampolines into jitdump when profiling

This commit updates the jitdump profiler to generate JIT profiling
records for generated trampolines in a wasm module in addition to the
functions already in a module. It's also updated to learn about
trampolines generated via `Func::new` and friends. These trampolines
were all not previously registered meaning that stack traces with these
pc values would be confusing to see in the profile output. While the
names aren't the best it should at least be more clear than before if a
function is hot!

* Fix more builds

crates/jit/src/instantiate.rs

@@ -397,13 +397,17 @@ impl CompiledModule {
     }
 
     /// Returns the per-signature trampolines for this module.
-    pub fn trampolines(&self) -> impl Iterator<Item = (SignatureIndex, VMTrampoline)> + '_ {
+    pub fn trampolines(&self) -> impl Iterator<Item = (SignatureIndex, VMTrampoline, usize)> + '_ {
         let code = self.code();
         self.trampolines.iter().map(move |info| {
-            (info.signature, unsafe {
+            (
+                info.signature,
+                unsafe {
                     let ptr = &code[info.start as usize];
                     std::mem::transmute::<*const u8, VMTrampoline>(ptr)
-            })
+                },
+                info.length as usize,
+            )
         })
     }
 

crates/jit/src/profiling.rs

@@ -29,7 +29,10 @@ pub use vtune::VTuneAgent;
 /// Common interface for profiling tools.
 pub trait ProfilingAgent: Send + Sync + 'static {
     /// Notify the profiler of a new module loaded into memory
-    fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>) -> ();
+    fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>);
+    /// Notify the profiler that the object file provided contains
+    /// dynamically-generated trampolines which are now being loaded.
+    fn trampoline_load(&self, file: &object::File<'_>);
 }
 
 /// Default agent for unsupported profiling build.
@@ -54,7 +57,8 @@ impl Error for NullProfilerAgentError {
 }
 
 impl ProfilingAgent for NullProfilerAgent {
-    fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) -> () {}
+    fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
+    fn trampoline_load(&self, _file: &object::File<'_>) {}
 }
 
 #[allow(dead_code)]

crates/jit/src/profiling/jitdump_disabled.rs

@@ -20,4 +20,5 @@ impl JitDumpAgent {
 
 impl ProfilingAgent for JitDumpAgent {
     fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
+    fn trampoline_load(&self, _file: &object::File<'_>) {}
 }

crates/jit/src/profiling/jitdump_linux.rs

@@ -21,6 +21,7 @@ use std::ptr;
 use std::sync::Mutex;
 use std::{borrow, mem, process};
 use target_lexicon::Architecture;
+use wasmtime_environ::EntityRef;
 
 use object::elf;
 
@@ -202,6 +203,9 @@ impl ProfilingAgent for JitDumpAgent {
     fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>) {
         self.state.lock().unwrap().module_load(module, dbg_image);
     }
+    fn trampoline_load(&self, file: &object::File<'_>) {
+        self.state.lock().unwrap().trampoline_load(file)
+    }
 }
 
 impl State {
@@ -280,7 +284,7 @@ impl State {
     }
 
     /// Sent when a method is compiled and loaded into memory by the VM.
-    pub fn module_load(&mut self, module: &CompiledModule, dbg_image: Option<&[u8]>) -> () {
+    pub fn module_load(&mut self, module: &CompiledModule, dbg_image: Option<&[u8]>) {
         let pid = process::id();
         let tid = pid; // ThreadId does appear to track underlying thread. Using PID.
 
@@ -299,6 +303,52 @@ impl State {
                 self.dump_code_load_record(&name, addr, len, timestamp, pid, tid);
             }
         }
+        for (idx, func, len) in module.trampolines() {
+            let (addr, len) = (func as usize as *const u8, len);
+            let timestamp = self.get_time_stamp();
+            let name = format!("wasm::trampoline[{}]", idx.index());
+            self.dump_code_load_record(&name, addr, len, timestamp, pid, tid);
+        }
+    }
+
+    fn trampoline_load(&mut self, image: &object::File<'_>) {
+        use object::{ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
+        let pid = process::id();
+        let tid = pid;
+
+        let text_base = match image.sections().find(|s| s.kind() == SectionKind::Text) {
+            Some(section) => match section.data() {
+                Ok(data) => data.as_ptr() as usize,
+                Err(_) => return,
+            },
+            None => return,
+        };
+
+        for sym in image.symbols() {
+            if !sym.is_definition() {
+                continue;
+            }
+            if sym.kind() != SymbolKind::Text {
+                continue;
+            }
+            let address = sym.address();
+            let size = sym.size();
+            if address == 0 || size == 0 {
+                continue;
+            }
+            if let Ok(name) = sym.name() {
+                let addr = text_base + address as usize;
+                let timestamp = self.get_time_stamp();
+                self.dump_code_load_record(
+                    &name,
+                    addr as *const u8,
+                    size as usize,
+                    timestamp,
+                    pid,
+                    tid,
+                );
+            }
+        }
     }
 
     fn dump_code_load_record(
@@ -309,7 +359,7 @@ impl State {
         timestamp: u64,
         pid: u32,
         tid: u32,
-    ) -> () {
+    ) {
         let name_len = method_name.len() + 1;
         let size_limit = mem::size_of::<CodeLoadRecord>();
 

crates/jit/src/profiling/vtune_disabled.rs

@@ -20,4 +20,5 @@ impl VTuneAgent {
 
 impl ProfilingAgent for VTuneAgent {
     fn module_load(&self, _module: &crate::CompiledModule, _dbg_image: Option<&[u8]>) {}
+    fn trampoline_load(&self, _file: &object::File<'_>) {}
 }

crates/jit/src/profiling/vtune_linux.rs

@@ -112,6 +112,9 @@ impl ProfilingAgent for VTuneAgent {
     fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>) {
         self.state.lock().unwrap().module_load(module, dbg_image);
     }
+    fn trampoline_load(&self, _file: &object::File<'_>) {
+        // TODO: needs an implementation
+    }
 }
 
 impl State {

crates/wasmtime/src/module.rs

@@ -500,7 +500,9 @@ impl Module {
         let signatures = Arc::new(SignatureCollection::new_for_module(
             engine.signatures(),
             &types.wasm_signatures,
-            modules.iter().flat_map(|m| m.trampolines()),
+            modules
+                .iter()
+                .flat_map(|m| m.trampolines().map(|(idx, f, _)| (idx, f))),
         ));
 
         let module = modules.remove(main_module);

crates/wasmtime/src/trampoline/func.rs

@@ -86,6 +86,7 @@ where
     // also take care of unwind table registration.
     let mut code_memory = CodeMemory::new(obj);
     let code = code_memory.publish()?;
+    engine.config().profiler.trampoline_load(&code.obj);
 
     // Extract the host/wasm trampolines from the results of compilation since
     // we know their start/length.