wasmtime: Remove GET_WASM_TRAP indirection (bytecodealliance#8949)

In the past, the wasmtime-runtime crate couldn't directly call `get_wasm_trap` because the registry of loaded modules was in the wasmtime crate, so it called through a global function pointer registered with `init_traps` instead. Since the two crates were merged in bytecodealliance#8501, we no longer need this indirection. While I'm here, I've also split the former `get_wasm_trap` function into two parts: `lookup_code` finds a loaded module that had been previously registered with `register_code`, and the `lookup_trap_code` step is now done by a helper on `CodeMemory`. This makes the module registry more broadly useful. I also simplified the code lookup step in two ways: - I removed a redundant check from the code lookup. `BTreeMap::range` will only return entries where `end >= pc`, so the `end < pc` condition is always false. - I used checked_sub instead of writing both the comparison and subtraction explicitly.
ahaoboy · Jul 15, 2024 · 66dfa36 · 66dfa36
1 parent 91e059d
commit 66dfa36
Show file tree

Hide file tree

Showing 8 changed files with 35 additions and 52 deletions.
diff --git a/crates/wasmtime/src/engine.rs b/crates/wasmtime/src/engine.rs
@@ -93,10 +93,7 @@ impl Engine {
             // configured. This is the per-program initialization required for
             // handling traps, such as configuring signals, vectored exception
             // handlers, etc.
-            crate::runtime::vm::init_traps(
-                crate::module::get_wasm_trap,
-                config.macos_use_mach_ports,
-            );
+            crate::runtime::vm::init_traps(config.macos_use_mach_ports);
             #[cfg(feature = "debug-builtins")]
             crate::runtime::vm::debug_builtins::ensure_exported();
         }

diff --git a/crates/wasmtime/src/runtime/code_memory.rs b/crates/wasmtime/src/runtime/code_memory.rs
@@ -7,7 +7,7 @@ use core::ops::Range;
 use object::endian::NativeEndian;
 use object::read::{elf::ElfFile64, Object, ObjectSection};
 use object::ObjectSymbol;
-use wasmtime_environ::obj;
+use wasmtime_environ::{lookup_trap_code, obj, Trap};
 use wasmtime_jit_icache_coherence as icache_coherence;
 
 /// Management of executable memory within a `MmapVec`
@@ -317,6 +317,12 @@ impl CodeMemory {
         *self.unwind_registration = Some(registration);
         Ok(())
     }
+
+    /// Looks up the given offset within this module's text section and returns
+    /// the trap code associated with that instruction, if there is one.
+    pub fn lookup_trap_code(&self, text_offset: usize) -> Option<Trap> {
+        lookup_trap_code(self.trap_data(), text_offset)
+    }
 }
 
 /// Returns the range of `inner` within `outer`, such that `outer[range]` is the

diff --git a/crates/wasmtime/src/runtime/module.rs b/crates/wasmtime/src/runtime/module.rs
@@ -24,7 +24,7 @@ use wasmtime_environ::{
 mod registry;
 
 pub use registry::{
-    get_wasm_trap, register_code, unregister_code, ModuleRegistry, RegisteredModuleId,
+    lookup_code, register_code, unregister_code, ModuleRegistry, RegisteredModuleId,
 };
 
 /// A compiled WebAssembly module, ready to be instantiated.

diff --git a/crates/wasmtime/src/runtime/module/registry.rs b/crates/wasmtime/src/runtime/module/registry.rs
@@ -6,7 +6,7 @@ use crate::component::Component;
 use crate::prelude::*;
 use crate::runtime::vm::VMWasmCallFunction;
 use crate::sync::{OnceLock, RwLock};
-use crate::{code_memory::CodeMemory, FrameInfo, Module, Trap};
+use crate::{code_memory::CodeMemory, FrameInfo, Module};
 use alloc::collections::btree_map::{BTreeMap, Entry};
 use alloc::sync::Arc;
 use core::ptr::NonNull;
@@ -254,23 +254,13 @@ fn global_code() -> &'static RwLock<GlobalRegistry> {
 
 type GlobalRegistry = BTreeMap<usize, (usize, Arc<CodeMemory>)>;
 
-/// Returns whether the `pc`, according to globally registered information,
-/// is a wasm trap or not.
-pub fn get_wasm_trap(pc: usize) -> Option<Trap> {
-    let (code, text_offset) = {
-        let all_modules = global_code().read();
-
-        let (end, (start, module)) = match all_modules.range(pc..).next() {
-            Some(info) => info,
-            None => return None,
-        };
-        if pc < *start || *end < pc {
-            return None;
-        }
-        (module.clone(), pc - *start)
-    };
-
-    wasmtime_environ::lookup_trap_code(code.trap_data(), text_offset)
+/// Find which registered region of code contains the given program counter, and
+/// what offset that PC is within that module's code.
+pub fn lookup_code(pc: usize) -> Option<(Arc<CodeMemory>, usize)> {
+    let all_modules = global_code().read();
+    let (_end, (start, module)) = all_modules.range(pc..).next()?;
+    let text_offset = pc.checked_sub(*start)?;
+    Some((module.clone(), text_offset))
 }
 
 /// Registers a new region of code.

diff --git a/crates/wasmtime/src/runtime/vm/sys/custom/traphandlers.rs b/crates/wasmtime/src/runtime/vm/sys/custom/traphandlers.rs
@@ -20,10 +20,8 @@ pub unsafe fn wasmtime_setjmp(
     capi::wasmtime_setjmp(jmp_buf, callback, payload, callee.cast())
 }
 
-pub fn platform_init(_macos_use_mach_ports: bool) {
-    unsafe {
-        capi::wasmtime_init_traps(handle_trap);
-    }
+pub unsafe fn platform_init(_macos_use_mach_ports: bool) {
+    capi::wasmtime_init_traps(handle_trap);
 }
 
 extern "C" fn handle_trap(ip: usize, fp: usize, has_faulting_addr: bool, faulting_addr: usize) {

diff --git a/crates/wasmtime/src/runtime/vm/sys/miri/traphandlers.rs b/crates/wasmtime/src/runtime/vm/sys/miri/traphandlers.rs
@@ -37,7 +37,7 @@ pub fn wasmtime_longjmp(_jmp_buf: *const u8) -> ! {
 #[allow(missing_docs)]
 pub type SignalHandler<'a> = dyn Fn() + Send + Sync + 'a;
 
-pub fn platform_init(_macos_use_mach_ports: bool) {}
+pub unsafe fn platform_init(_macos_use_mach_ports: bool) {}
 
 pub fn lazy_per_thread_init() {}
 

diff --git a/crates/wasmtime/src/runtime/vm/sys/unix/machports.rs b/crates/wasmtime/src/runtime/vm/sys/unix/machports.rs
@@ -33,6 +33,7 @@
 
 #![allow(non_snake_case, clippy::cast_sign_loss)]
 
+use crate::runtime::module::lookup_code;
 use crate::runtime::vm::sys::traphandlers::wasmtime_longjmp;
 use crate::runtime::vm::traphandlers::tls;
 use mach2::exc::*;
@@ -325,9 +326,12 @@ unsafe fn handle_exception(request: &mut ExceptionRequest) -> bool {
     // pointer value and if `MAP` changes happen after we read our entry that's
     // ok since they won't invalidate our entry.
     let (pc, fp) = get_pc_and_fp(&thread_state);
-    let trap = match crate::runtime::vm::traphandlers::GET_WASM_TRAP(pc as usize) {
-        Some(trap) => trap,
-        None => return false,
+    let Some((code, text_offset)) = lookup_code(pc as usize) else {
+        return false;
+    };
+
+    let Some(trap) = code.lookup_trap_code(text_offset) else {
+        return false;
     };
 
     // We have determined that this is a wasm trap and we need to actually

diff --git a/crates/wasmtime/src/runtime/vm/traphandlers.rs b/crates/wasmtime/src/runtime/vm/traphandlers.rs
@@ -11,6 +11,7 @@ mod coredump;
 mod coredump;
 
 use crate::prelude::*;
+use crate::runtime::module::lookup_code;
 use crate::runtime::vm::sys::traphandlers;
 use crate::runtime::vm::{Instance, VMContext, VMRuntimeLimits};
 use crate::sync::OnceLock;
@@ -24,37 +25,21 @@ pub use self::tls::{tls_eager_initialize, AsyncWasmCallState, PreviousAsyncWasmC
 
 pub use traphandlers::SignalHandler;
 
-/// Globally-set callback to determine whether a program counter is actually a
-/// wasm trap.
-///
-/// This is initialized during `init_traps` below. The definition lives within
-/// `wasmtime` currently.
-pub(crate) static mut GET_WASM_TRAP: fn(usize) -> Option<wasmtime_environ::Trap> = |_| None;
-
 /// This function is required to be called before any WebAssembly is entered.
 /// This will configure global state such as signal handlers to prepare the
 /// process to receive wasm traps.
 ///
-/// The `get_wasm_trap` argument is used when a trap happens to determine if a
-/// program counter is the pc of an actual wasm trap or not. This is then used
-/// to disambiguate faults that happen due to wasm and faults that happen due to
-/// bugs in Rust or elsewhere.
-///
 /// # Panics
 ///
 /// This function will panic on macOS if it is called twice or more times with
 /// different values of `macos_use_mach_ports`.
 ///
 /// This function will also panic if the `std` feature is disabled and it's
 /// called concurrently.
-pub fn init_traps(
-    get_wasm_trap: fn(usize) -> Option<wasmtime_environ::Trap>,
-    macos_use_mach_ports: bool,
-) {
+pub fn init_traps(macos_use_mach_ports: bool) {
     static INIT: OnceLock<()> = OnceLock::new();
 
     INIT.get_or_init(|| unsafe {
-        GET_WASM_TRAP = get_wasm_trap;
         traphandlers::platform_init(macos_use_mach_ports);
     });
 
@@ -460,9 +445,12 @@ impl CallThreadState {
         }
 
         // If this fault wasn't in wasm code, then it's not our problem
-        let trap = match unsafe { GET_WASM_TRAP(pc as usize) } {
-            Some(trap) => trap,
-            None => return TrapTest::NotWasm,
+        let Some((code, text_offset)) = lookup_code(pc as usize) else {
+            return TrapTest::NotWasm;
+        };
+
+        let Some(trap) = code.lookup_trap_code(text_offset) else {
+            return TrapTest::NotWasm;
         };
 
         // If all that passed then this is indeed a wasm trap, so return the