Refactor Wasmtime's profiling support

crates/jit/Cargo.toml

@@ -27,6 +27,9 @@ cpp_demangle = "0.3.2"
 log = { workspace = true }
 wasmtime-jit-icache-coherence = { workspace = true }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+rustix = { workspace = true, features = ['thread'] }
+
 [target.'cfg(target_os = "windows")'.dependencies.windows-sys]
 workspace = true
 features = [

crates/jit/src/instantiate.rs

@@ -5,7 +5,7 @@
 
 use crate::code_memory::CodeMemory;
 use crate::debug::create_gdbjit_image;
-use crate::ProfilingAgent;
+use crate::profiling::ProfilingAgent;
 use anyhow::{bail, Context, Error, Result};
 use object::write::{Object, SectionId, StandardSegment, WritableBuffer};
 use object::SectionKind;
@@ -471,17 +471,21 @@ impl CompiledModule {
     }
 
     fn register_debug_and_profiling(&mut self, profiler: &dyn ProfilingAgent) -> Result<()> {
-        // Register GDB JIT images; initialize profiler and load the wasm module.
         if self.meta.native_debug_info_present {
             let text = self.text();
             let bytes = create_gdbjit_image(self.mmap().to_vec(), (text.as_ptr(), text.len()))
                 .context("failed to create jit image for gdb")?;
-            profiler.module_load(self, Some(&bytes));
             let reg = GdbJitImageRegistration::register(bytes);
             self.dbg_jit_registration = Some(reg);
-        } else {
-            profiler.module_load(self, None);
         }
+        profiler.register_module(&self.code_memory, &|addr| {
+            let (idx, _) = self.func_by_text_offset(addr)?;
+            let idx = self.module.func_index(idx);
+            let name = self.func_name(idx)?;
+            let mut demangled = String::new();
+            crate::demangling::demangle_function_name(&mut demangled, name).unwrap();
+            Some(demangled)
+        });
         Ok(())
     }
 
@@ -564,16 +568,6 @@ impl CompiledModule {
         Some(&self.text()[loc.start as usize..][..loc.length as usize])
     }
 
-    /// Returns an iterator over all array-to-Wasm trampolines defined within
-    /// this module, providing both their index and their in-memory body.
-    pub fn array_to_wasm_trampolines(
-        &self,
-    ) -> impl ExactSizeIterator<Item = (DefinedFuncIndex, &[u8])> + '_ {
-        self.funcs
-            .keys()
-            .map(move |i| (i, self.array_to_wasm_trampoline(i).unwrap()))
-    }
-
     /// Get the native-to-Wasm trampoline for the function `index` points to.
     ///
     /// If the function `index` points to does not escape, then `None` is
@@ -586,16 +580,6 @@ impl CompiledModule {
         Some(&self.text()[loc.start as usize..][..loc.length as usize])
     }
 
-    /// Returns an iterator over all native-to-Wasm trampolines defined within
-    /// this module, providing both their index and their in-memory body.
-    pub fn native_to_wasm_trampolines(
-        &self,
-    ) -> impl ExactSizeIterator<Item = (DefinedFuncIndex, &[u8])> + '_ {
-        self.funcs
-            .keys()
-            .map(move |i| (i, self.native_to_wasm_trampoline(i).unwrap()))
-    }
-
     /// Get the Wasm-to-native trampoline for the given signature.
     ///
     /// These trampolines are used for filling in
@@ -610,16 +594,6 @@ impl CompiledModule {
         &self.text()[loc.start as usize..][..loc.length as usize]
     }
 
-    /// Returns an iterator over all native-to-Wasm trampolines defined within
-    /// this module, providing both their index and their in-memory body.
-    pub fn wasm_to_native_trampolines(
-        &self,
-    ) -> impl ExactSizeIterator<Item = (SignatureIndex, &[u8])> + '_ {
-        self.wasm_to_native_trampolines
-            .iter()
-            .map(move |(i, _)| (*i, self.wasm_to_native_trampoline(*i)))
-    }
-
     /// Returns the stack map information for all functions defined in this
     /// module.
     ///

crates/jit/src/lib.rs

@@ -24,7 +24,7 @@ mod code_memory;
 mod debug;
 mod demangling;
 mod instantiate;
-mod profiling;
+pub mod profiling;
 mod unwind;
 
 pub use crate::code_memory::CodeMemory;
@@ -33,7 +33,6 @@ pub use crate::instantiate::{
     SymbolizeContext,
 };
 pub use demangling::*;
-pub use profiling::*;
 
 /// Version number of this crate.
 pub const VERSION: &str = env!("CARGO_PKG_VERSION");

crates/jit/src/profiling.rs

@@ -1,74 +1,108 @@
-use crate::{demangling::demangle_function_name_or_index, CompiledModule};
-use wasmtime_environ::{DefinedFuncIndex, EntityRef};
+#![allow(missing_docs)]
+
+use crate::CodeMemory;
+#[allow(unused_imports)]
+use anyhow::{bail, Result};
 
 cfg_if::cfg_if! {
     if #[cfg(all(feature = "jitdump", target_os = "linux"))] {
-        #[path = "profiling/jitdump_linux.rs"]
         mod jitdump;
+        pub use jitdump::new as new_jitdump;
     } else {
-        #[path = "profiling/jitdump_disabled.rs"]
-        mod jitdump;
+        pub fn new_jitdump() -> Result<Box<dyn ProfilingAgent>> {
+            if cfg!(feature = "jitdump") {
+                bail!("jitdump is not supported on this platform");
+            } else {
+                bail!("jitdump support disabled at compile time");
+            }
+        }
     }
 }
 
 cfg_if::cfg_if! {
     if #[cfg(target_os = "linux")] {
-        #[path = "profiling/perfmap_linux.rs"]
         mod perfmap;
+        pub use perfmap::new as new_perfmap;
     } else {
-        #[path = "profiling/perfmap_disabled.rs"]
-        mod perfmap;
+        pub fn new_perfmap() -> Result<Box<dyn ProfilingAgent>> {
+            bail!("perfmap support not supported on this platform");
+        }
     }
 }
 
 cfg_if::cfg_if! {
     // Note: VTune support is disabled on windows mingw because the ittapi crate doesn't compile
     // there; see also https://github.com/bytecodealliance/wasmtime/pull/4003 for rationale.
     if #[cfg(all(feature = "vtune", target_arch = "x86_64", not(all(target_os = "windows", target_env = "gnu"))))] {
-        #[path = "profiling/vtune.rs"]
         mod vtune;
+        pub use vtune::new as new_vtune;
     } else {
-        #[path = "profiling/vtune_disabled.rs"]
-        mod vtune;
+        pub fn new_vtune() -> Result<Box<dyn ProfilingAgent>> {
+            if cfg!(feature = "vtune") {
+                bail!("VTune is not supported on this platform.");
+            } else {
+                bail!("VTune support disabled at compile time.");
+            }
+        }
     }
 }
 
-pub use jitdump::JitDumpAgent;
-pub use perfmap::PerfMapAgent;
-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 register_function(&self, name: &str, addr: *const u8, size: usize);
 
-    /// 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);
-}
+    fn register_module(&self, code: &CodeMemory, custom_name: &dyn Fn(usize) -> Option<String>) {
+        use object::{File, Object as _, ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
 
-/// Default agent for unsupported profiling build.
-#[derive(Debug, Default, Clone, Copy)]
-pub struct NullProfilerAgent;
+        let image = match File::parse(&code.mmap()[..]) {
+            Ok(image) => image,
+            Err(_) => return,
+        };
 
-impl ProfilingAgent for NullProfilerAgent {
-    fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn load_single_trampoline(
-        &self,
-        _name: &str,
-        _addr: *const u8,
-        _size: usize,
-        _pid: u32,
-        _tid: u32,
-    ) {
+        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 owned;
+                let name = match custom_name(address as usize) {
+                    Some(name) => {
+                        owned = name;
+                        &owned
+                    }
+                    None => name,
+                };
+                self.register_function(name, addr as *const u8, size as usize);
+            }
+        }
     }
 }
 
-#[allow(dead_code)]
-fn debug_name(module: &CompiledModule, index: DefinedFuncIndex) -> String {
-    let index = module.module().func_index(index);
-    let mut debug_name = String::new();
-    demangle_function_name_or_index(&mut debug_name, module.func_name(index), index.index())
-        .unwrap();
-    debug_name
+pub fn new_null() -> Box<dyn ProfilingAgent> {
+    Box::new(NullProfilerAgent)
+}
+
+#[derive(Debug, Default, Clone, Copy)]
+struct NullProfilerAgent;
+
+impl ProfilingAgent for NullProfilerAgent {
+    fn register_function(&self, _name: &str, _addr: *const u8, _size: usize) {}
+    fn register_module(&self, _code: &CodeMemory, _custom_name: &dyn Fn(usize) -> Option<String>) {}
 }

crates/jit/src/profiling/jitdump.rs

@@ -0,0 +1,66 @@
+//! Support for jitdump files which can be used by perf for profiling jitted code.
+//! Spec definitions for the output format is as described here:
+//! <https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/tools/perf/Documentation/jitdump-specification.txt>
+//!
+//! Usage Example:
+//!     Record
+//!         sudo perf record -k 1 -e instructions:u target/debug/wasmtime -g --profile=jitdump test.wasm
+//!     Combine
+//!         sudo perf inject -v -j -i perf.data -o perf.jit.data
+//!     Report
+//!         sudo perf report -i perf.jit.data -F+period,srcline
+//! Note: For descriptive results, the WASM file being executed should contain dwarf debug data
+
+use crate::profiling::ProfilingAgent;
+use anyhow::Result;
+use std::process;
+use std::sync::Mutex;
+use target_lexicon::Architecture;
+use wasmtime_jit_debug::perf_jitdump::*;
+
+use object::elf;
+
+/// Interface for driving the creation of jitdump files
+struct JitDumpAgent {
+    pid: u32,
+}
+
+/// Process-wide JIT dump file. Perf only accepts a unique file per process, in the injection step.
+static JITDUMP_FILE: Mutex<Option<JitDumpFile>> = Mutex::new(None);
+
+/// Intialize a JitDumpAgent and write out the header.
+pub fn new() -> Result<Box<dyn ProfilingAgent>> {
+    let mut jitdump_file = JITDUMP_FILE.lock().unwrap();
+
+    if jitdump_file.is_none() {
+        let filename = format!("./jit-{}.dump", process::id());
+        let e_machine = match target_lexicon::HOST.architecture {
+            Architecture::X86_64 => elf::EM_X86_64 as u32,
+            Architecture::X86_32(_) => elf::EM_386 as u32,
+            Architecture::Arm(_) => elf::EM_ARM as u32,
+            Architecture::Aarch64(_) => elf::EM_AARCH64 as u32,
+            Architecture::S390x => elf::EM_S390 as u32,
+            _ => unimplemented!("unrecognized architecture"),
+        };
+        *jitdump_file = Some(JitDumpFile::new(filename, e_machine)?);
+    }
+
+    Ok(Box::new(JitDumpAgent {
+        pid: std::process::id(),
+    }))
+}
+
+impl ProfilingAgent for JitDumpAgent {
+    fn register_function(&self, name: &str, addr: *const u8, size: usize) {
+        let mut jitdump_file = JITDUMP_FILE.lock().unwrap();
+        let jitdump_file = jitdump_file.as_mut().unwrap();
+        let timestamp = jitdump_file.get_time_stamp();
+        #[allow(trivial_numeric_casts)]
+        let tid = rustix::thread::gettid().as_raw_nonzero().get() as u32;
+        if let Err(err) =
+            jitdump_file.dump_code_load_record(&name, addr, size, timestamp, self.pid, tid)
+        {
+            println!("Jitdump: write_code_load_failed_record failed: {:?}\n", err);
+        }
+    }
+}