avoid code duplication with jitdump_linux

crates/jit/src/profiling.rs

@@ -30,10 +30,10 @@ pub use vtune::VTuneAgent;
 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]>);
-    /// Notify the profiler that the object file provided contains
-    /// dynamically-generated trampolines (for host functions) which are now
-    /// being loaded.
-    fn trampoline_load(&self, file: &object::File<'_>);
+
+    /// Notify the profiler about a single dynamically-generated trampoline (for host function)
+    /// that is being loaded now.`
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32);
 }
 
 /// Default agent for unsupported profiling build.
@@ -59,7 +59,15 @@ impl Error for NullProfilerAgentError {
 
 impl ProfilingAgent for NullProfilerAgent {
     fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn trampoline_load(&self, _file: &object::File<'_>) {}
+    fn load_single_trampoline(
+        &self,
+        _name: &str,
+        _addr: *const u8,
+        _size: usize,
+        _pid: u32,
+        _tid: u32,
+    ) {
+    }
 }
 
 #[allow(dead_code)]

crates/jit/src/profiling/jitdump_disabled.rs

@@ -20,5 +20,13 @@ impl JitDumpAgent {
 
 impl ProfilingAgent for JitDumpAgent {
     fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn trampoline_load(&self, _file: &object::File<'_>) {}
+    fn load_single_trampoline(
+        &self,
+        _name: &str,
+        _addr: *const u8,
+        _size: usize,
+        __pid: u32,
+        _tid: u32,
+    ) {
+    }
 }

crates/jit/src/profiling/jitdump_linux.rs

@@ -203,8 +203,11 @@ 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)
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32) {
+        self.state
+            .lock()
+            .unwrap()
+            .load_single_trampoline(name, addr, size, pid, tid);
     }
 }
 
@@ -313,44 +316,16 @@ impl State {
         }
     }
 
-    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 load_single_trampoline(
+        &mut self,
+        name: &str,
+        addr: *const u8,
+        size: usize,
+        pid: u32,
+        tid: u32,
+    ) {
+        let timestamp = self.get_time_stamp();
+        self.dump_code_load_record(name, addr, size, timestamp, pid, tid);
     }
 
     fn dump_code_load_record(

crates/jit/src/profiling/vtune_disabled.rs

@@ -20,5 +20,13 @@ impl VTuneAgent {
 
 impl ProfilingAgent for VTuneAgent {
     fn module_load(&self, _module: &crate::CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn trampoline_load(&self, _file: &object::File<'_>) {}
+    fn load_single_trampoline(
+        &self,
+        _name: &str,
+        _addr: *const u8,
+        _size: usize,
+        __pid: u32,
+        _tid: u32,
+    ) {
+    }
 }

crates/jit/src/profiling/vtune_linux.rs

@@ -20,7 +20,6 @@ use crate::{CompiledModule, ProfilingAgent};
 use anyhow::Result;
 use core::ptr;
 use ittapi_rs::*;
-use object::Object as _;
 use std::ffi::CString;
 use std::sync::{atomic, Mutex};
 use wasmtime_environ::EntityRef;
@@ -109,8 +108,11 @@ 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, image: &object::File<'_>) {
-        self.state.lock().unwrap().trampoline_load(image);
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32) {
+        self.state
+            .lock()
+            .unwrap()
+            .load_single_trampoline(name, addr, size, pid, tid);
     }
 }
 
@@ -157,40 +159,15 @@ impl State {
         }
     }
 
-    fn trampoline_load(&mut self, image: &object::File<'_>) {
-        use object::{ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
-
-        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 method_id = self.get_method_id();
-                self.event_load(
-                    method_id,
-                    "wasm trampoline for Func::new",
-                    &name,
-                    addr as *const u8,
-                    size as usize,
-                );
-            }
-        }
+    fn load_single_trampoline(
+        &mut self,
+        name: &str,
+        addr: *const u8,
+        size: usize,
+        _pid: u32,
+        _tid: u32,
+    ) {
+        let method_id = self.get_method_id();
+        self.event_load(method_id, "wasm trampoline for Func::new", name, addr, size);
     }
 }

crates/wasmtime/src/trampoline/func.rs

@@ -6,7 +6,7 @@ use std::any::Any;
 use std::panic::{self, AssertUnwindSafe};
 use std::sync::Arc;
 use wasmtime_environ::{EntityIndex, Module, ModuleType, PrimaryMap, SignatureIndex};
-use wasmtime_jit::{CodeMemory, MmapVec};
+use wasmtime_jit::{CodeMemory, MmapVec, ProfilingAgent};
 use wasmtime_runtime::{
     Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
     OnDemandInstanceAllocator, StorePtr, VMContext, VMFunctionBody, VMSharedSignatureIndex,
@@ -66,6 +66,39 @@ unsafe extern "C" fn stub_fn<F>(
     }
 }
 
+#[cfg(compiler)]
+fn register_trampolines(profiler: &dyn ProfilingAgent, image: &object::File<'_>) {
+    use object::{Object as _, ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
+    let pid = std::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;
+            profiler.load_single_trampoline(name, addr as *const u8, size as usize, pid, tid);
+        }
+    }
+}
+
 #[cfg(compiler)]
 pub fn create_function<F>(
     ft: &FuncType,
@@ -87,7 +120,8 @@ 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);
+
+    register_trampolines(engine.config().profiler.as_ref(), &code.obj);
 
     // Extract the host/wasm trampolines from the results of compilation since
     // we know their start/length.