From b5e31a5c33725ab8a7dfbe8505c56b5cf282ffed Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Mon, 8 Jul 2024 13:53:16 -0500 Subject: [PATCH] Refactor internal profiling APIs to be safer (#8914) Pass around `&[u8]` instead of `*const u8` and `usize` to avoid the need for raw unsafe abstractions. Closes #8905 --- crates/jit-debug/src/perf_jitdump.rs | 16 ++++++---------- crates/wasmtime/src/profiling_agent.rs | 12 ++++++------ .../wasmtime/src/profiling_agent/jitdump.rs | 5 ++--- .../wasmtime/src/profiling_agent/perfmap.rs | 19 ++++++++++--------- crates/wasmtime/src/profiling_agent/vtune.rs | 15 ++++++--------- 5 files changed, 30 insertions(+), 37 deletions(-) diff --git a/crates/jit-debug/src/perf_jitdump.rs b/crates/jit-debug/src/perf_jitdump.rs index 6d6856e8d35e..4592318c0bc8 100644 --- a/crates/jit-debug/src/perf_jitdump.rs +++ b/crates/jit-debug/src/perf_jitdump.rs @@ -252,8 +252,7 @@ impl JitDumpFile { pub fn dump_code_load_record( &mut self, method_name: &str, - addr: *const u8, - len: usize, + code: &[u8], timestamp: u64, pid: u32, tid: u32, @@ -263,7 +262,7 @@ impl JitDumpFile { let rh = RecordHeader { id: RecordId::JitCodeLoad as u32, - record_size: size_limit as u32 + name_len as u32 + len as u32, + record_size: size_limit as u32 + name_len as u32 + code.len() as u32, timestamp, }; @@ -271,16 +270,13 @@ impl JitDumpFile { header: rh, pid, tid, - virtual_address: addr as u64, - address: addr as u64, - size: len as u64, + virtual_address: code.as_ptr() as u64, + address: code.as_ptr() as u64, + size: code.len() as u64, index: self.next_code_index(), }; - unsafe { - let code_buffer: &[u8] = std::slice::from_raw_parts(addr, len); - self.write_code_load_record(method_name, clr, code_buffer) - } + self.write_code_load_record(method_name, clr, code) } } diff --git a/crates/wasmtime/src/profiling_agent.rs b/crates/wasmtime/src/profiling_agent.rs index b32b846193f4..50480a62450a 100644 --- a/crates/wasmtime/src/profiling_agent.rs +++ b/crates/wasmtime/src/profiling_agent.rs @@ -45,7 +45,7 @@ cfg_if::cfg_if! { /// Common interface for profiling tools. pub trait ProfilingAgent: Send + Sync + 'static { - fn register_function(&self, name: &str, addr: *const u8, size: usize); + fn register_function(&self, name: &str, code: &[u8]); fn register_module(&self, code: &[u8], custom_name: &dyn Fn(usize) -> Option) { use object::{File, Object as _, ObjectSection, ObjectSymbol, SectionKind, SymbolKind}; @@ -55,9 +55,9 @@ pub trait ProfilingAgent: Send + Sync + 'static { Err(_) => return, }; - let text_base = match image.sections().find(|s| s.kind() == SectionKind::Text) { + let text = match image.sections().find(|s| s.kind() == SectionKind::Text) { Some(section) => match section.data() { - Ok(data) => data.as_ptr() as usize, + Ok(data) => data, Err(_) => return, }, None => return, @@ -76,7 +76,6 @@ pub trait ProfilingAgent: Send + Sync + 'static { 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) => { @@ -85,7 +84,8 @@ pub trait ProfilingAgent: Send + Sync + 'static { } None => name, }; - self.register_function(name, addr as *const u8, size as usize); + let code = &text[address as usize..][..size as usize]; + self.register_function(name, code) } } } @@ -99,6 +99,6 @@ pub fn new_null() -> Box { struct NullProfilerAgent; impl ProfilingAgent for NullProfilerAgent { - fn register_function(&self, _name: &str, _addr: *const u8, _size: usize) {} + fn register_function(&self, _name: &str, _code: &[u8]) {} fn register_module(&self, _code: &[u8], _custom_name: &dyn Fn(usize) -> Option) {} } diff --git a/crates/wasmtime/src/profiling_agent/jitdump.rs b/crates/wasmtime/src/profiling_agent/jitdump.rs index 5de53984f848..cb8a49ece91f 100644 --- a/crates/wasmtime/src/profiling_agent/jitdump.rs +++ b/crates/wasmtime/src/profiling_agent/jitdump.rs @@ -50,14 +50,13 @@ pub fn new() -> Result> { } impl ProfilingAgent for JitDumpAgent { - fn register_function(&self, name: &str, addr: *const u8, size: usize) { + fn register_function(&self, name: &str, code: &[u8]) { 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) + if let Err(err) = jitdump_file.dump_code_load_record(&name, code, timestamp, self.pid, tid) { println!("Jitdump: write_code_load_failed_record failed: {:?}\n", err); } diff --git a/crates/wasmtime/src/profiling_agent/perfmap.rs b/crates/wasmtime/src/profiling_agent/perfmap.rs index c60a284401db..8ed7fc8924fe 100644 --- a/crates/wasmtime/src/profiling_agent/perfmap.rs +++ b/crates/wasmtime/src/profiling_agent/perfmap.rs @@ -21,26 +21,27 @@ pub fn new() -> Result> { } impl PerfMapAgent { - fn make_line( - writer: &mut dyn Write, - name: &str, - addr: *const u8, - len: usize, - ) -> io::Result<()> { + fn make_line(writer: &mut dyn Write, name: &str, code: &[u8]) -> io::Result<()> { // Format is documented here: https://github.com/torvalds/linux/blob/master/tools/perf/Documentation/jit-interface.txt // Try our best to sanitize the name, since wasm allows for any utf8 string in there. let sanitized_name = name.replace('\n', "_").replace('\r', "_"); - write!(writer, "{:x} {:x} {}\n", addr as usize, len, sanitized_name)?; + write!( + writer, + "{:p} {:x} {}\n", + code.as_ptr(), + code.len(), + sanitized_name + )?; writer.flush()?; Ok(()) } } impl ProfilingAgent for PerfMapAgent { - fn register_function(&self, name: &str, addr: *const u8, size: usize) { + fn register_function(&self, name: &str, code: &[u8]) { let mut file = PERFMAP_FILE.lock().unwrap(); let file = file.as_mut().unwrap(); - if let Err(err) = Self::make_line(file, name, addr, size) { + if let Err(err) = Self::make_line(file, name, code) { eprintln!("Error when writing import trampoline info to the perf map file: {err}"); } } diff --git a/crates/wasmtime/src/profiling_agent/vtune.rs b/crates/wasmtime/src/profiling_agent/vtune.rs index 159792401261..cbf7f5d7b287 100644 --- a/crates/wasmtime/src/profiling_agent/vtune.rs +++ b/crates/wasmtime/src/profiling_agent/vtune.rs @@ -47,10 +47,10 @@ impl Drop for VTuneAgent { impl State { /// Notify vtune about a newly tracked code region. - fn notify_code(&mut self, module_name: &str, method_name: &str, addr: *const u8, len: usize) { + fn notify_code(&mut self, module_name: &str, method_name: &str, code: &[u8]) { self.vtune .load_method( - MethodLoadBuilder::new(method_name.to_owned(), addr, len) + MethodLoadBuilder::new(method_name.to_owned(), code.as_ptr(), code.len()) .class_file_name(module_name.to_owned()) .source_file_name("".to_owned()), ) @@ -65,16 +65,13 @@ impl State { } impl ProfilingAgent for VTuneAgent { - fn register_function(&self, name: &str, addr: *const u8, size: usize) { - self.state - .lock() - .unwrap() - .register_function(name, addr, size); + fn register_function(&self, name: &str, code: &[u8]) { + self.state.lock().unwrap().register_function(name, code); } } impl State { - fn register_function(&mut self, name: &str, addr: *const u8, size: usize) { - self.notify_code("wasmtime", name, addr, size); + fn register_function(&mut self, name: &str, code: &[u8]) { + self.notify_code("wasmtime", name, code); } }