Remove global state for trap registration (#909)

* Remove global state for trap registration

There's a number of changes brought about in this commit, motivated by a
few things. One motivation was to remove an instance of using
`lazy_static!` in an effort to remove global state and encapsulate it
wherever possible. A second motivation came when investigating a
slowly-compiling wasm module (a bit too slowly) where a good chunk of
time was spent in managing trap registrations.

The specific change made here is that `TrapRegistry` is now stored
inside of a `Compiler` instead of inside a global. Additionally traps
are "bulk registered" for a module rather than one-by-one. This form of
bulk-registration allows optimizing the locks used here, where a lock is
only held for a module at-a-time instead of once-per-function.

With these changes the "unregister" logic has also been tweaked a bit
here and there to continue to work. As a nice side effect the `Compiler`
type now has one fewer field that requires actual mutability and has
been updated for multi-threaded compilation, nudging us closer to a
world where we can support multi-threaded compilation. Yay!

In terms of performance improvements, a local wasm test file that
previously took 3 seconds to compile is now 10% faster to compile,
taking ~2.7 seconds now.

* Perform trap resolution after unwinding

This avoids taking locks in signal handlers which feels a bit iffy...

* Remove `TrapRegistration::dummy()`

Avoid an case where you're trying to lookup trap information from a
dummy module for something that happened in a different module.

* Tweak some comments

crates/api/src/externals.rs

@@ -181,7 +181,7 @@ impl Global {
         if val.ty() != *ty.content() {
             bail!("value provided does not match the type of this global");
         }
-        let (wasmtime_export, wasmtime_state) = generate_global_export(&ty, val)?;
+        let (wasmtime_export, wasmtime_state) = generate_global_export(store, &ty, val)?;
         Ok(Global {
             inner: Rc::new(GlobalInner {
                 _store: store.clone(),
@@ -321,7 +321,7 @@ impl Table {
     /// Returns an error if `init` does not match the element type of the table.
     pub fn new(store: &Store, ty: TableType, init: Val) -> Result<Table> {
         let item = into_checked_anyfunc(init, store)?;
-        let (mut wasmtime_handle, wasmtime_export) = generate_table_export(&ty)?;
+        let (mut wasmtime_handle, wasmtime_export) = generate_table_export(store, &ty)?;
 
         // Initialize entries with the init value.
         match wasmtime_export {
@@ -473,7 +473,7 @@ impl Memory {
     /// type's configuration. All WebAssembly memory is initialized to zero.
     pub fn new(store: &Store, ty: MemoryType) -> Memory {
         let (wasmtime_handle, wasmtime_export) =
-            generate_memory_export(&ty).expect("generated memory");
+            generate_memory_export(store, &ty).expect("generated memory");
         Memory {
             _store: store.clone(),
             ty,

crates/api/src/trampoline/create_handle.rs

@@ -12,7 +12,7 @@ use wasmtime_runtime::{Imports, InstanceHandle, VMFunctionBody};
 
 pub(crate) fn create_handle(
     module: Module,
-    store: Option<&Store>,
+    store: &Store,
     finished_functions: PrimaryMap<DefinedFuncIndex, *const VMFunctionBody>,
     state: Box<dyn Any>,
 ) -> Result<InstanceHandle> {
@@ -26,19 +26,16 @@ pub(crate) fn create_handle(
     let data_initializers = Vec::new();
 
     // Compute indices into the shared signature table.
-    let signatures = store
+    let signatures = module
-        .map(|store| {
-            module
         .signatures
         .values()
         .map(|sig| store.compiler().signatures().register(sig))
-                .collect::<PrimaryMap<_, _>>()
+        .collect::<PrimaryMap<_, _>>();
-        })
-        .unwrap_or_else(PrimaryMap::new);
 
     unsafe {
         Ok(InstanceHandle::new(
             Arc::new(module),
+            store.compiler().trap_registry().register_traps(Vec::new()),
             finished_functions.into_boxed_slice(),
             imports,
             &data_initializers,

crates/api/src/trampoline/func.rs

@@ -287,7 +287,7 @@ pub fn create_handle_with_function(
 
     create_handle(
         module,
-        Some(store),
+        store,
         finished_functions,
         Box::new(trampoline_state),
     )
@@ -322,5 +322,5 @@ pub unsafe fn create_handle_with_raw_function(
         .insert("trampoline".to_string(), Export::Function(func_id));
     finished_functions.push(func);
 
-    create_handle(module, Some(store), finished_functions, state)
+    create_handle(module, store, finished_functions, state)
 }

crates/api/src/trampoline/global.rs

@@ -1,4 +1,5 @@
 use super::create_handle::create_handle;
+use crate::Store;
 use crate::{GlobalType, Mutability, Val};
 use anyhow::{bail, Result};
 use wasmtime_environ::entity::PrimaryMap;
@@ -11,7 +12,11 @@ pub struct GlobalState {
     handle: InstanceHandle,
 }
 
-pub fn create_global(gt: &GlobalType, val: Val) -> Result<(wasmtime_runtime::Export, GlobalState)> {
+pub fn create_global(
+    store: &Store,
+    gt: &GlobalType,
+    val: Val,
+) -> Result<(wasmtime_runtime::Export, GlobalState)> {
     let mut definition = Box::new(VMGlobalDefinition::new());
     unsafe {
         match val {
@@ -35,7 +40,7 @@ pub fn create_global(gt: &GlobalType, val: Val) -> Result<(wasmtime_runtime::Exp
         initializer: wasm::GlobalInit::Import, // TODO is it right?
     };
     let handle =
-        create_handle(Module::new(), None, PrimaryMap::new(), Box::new(())).expect("handle");
+        create_handle(Module::new(), store, PrimaryMap::new(), Box::new(())).expect("handle");
     Ok((
         wasmtime_runtime::Export::Global {
             definition: definition.as_mut(),

crates/api/src/trampoline/memory.rs

@@ -1,13 +1,12 @@
 use super::create_handle::create_handle;
 use crate::MemoryType;
+use crate::Store;
 use anyhow::Result;
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::{wasm, Module};
 use wasmtime_runtime::InstanceHandle;
 
-#[allow(dead_code)]
+pub fn create_handle_with_memory(store: &Store, memory: &MemoryType) -> Result<InstanceHandle> {
-
-pub fn create_handle_with_memory(memory: &MemoryType) -> Result<InstanceHandle> {
     let mut module = Module::new();
 
     let memory = wasm::Memory {
@@ -24,5 +23,5 @@ pub fn create_handle_with_memory(memory: &MemoryType) -> Result<InstanceHandle>
         wasmtime_environ::Export::Memory(memory_id),
     );
 
-    create_handle(module, None, PrimaryMap::new(), Box::new(()))
+    create_handle(module, store, PrimaryMap::new(), Box::new(()))
 }

crates/api/src/trampoline/mod.rs

@@ -43,24 +43,27 @@ pub unsafe fn generate_raw_func_export(
 }
 
 pub fn generate_global_export(
+    store: &Store,
     gt: &GlobalType,
     val: Val,
 ) -> Result<(wasmtime_runtime::Export, GlobalState)> {
-    create_global(gt, val)
+    create_global(store, gt, val)
 }
 
 pub fn generate_memory_export(
+    store: &Store,
     m: &MemoryType,
 ) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
-    let instance = create_handle_with_memory(m)?;
+    let instance = create_handle_with_memory(store, m)?;
     let export = instance.lookup("memory").expect("memory export");
     Ok((instance, export))
 }
 
 pub fn generate_table_export(
+    store: &Store,
     t: &TableType,
 ) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
-    let instance = create_handle_with_table(t)?;
+    let instance = create_handle_with_table(store, t)?;
     let export = instance.lookup("table").expect("table export");
     Ok((instance, export))
 }

crates/api/src/trampoline/table.rs

@@ -1,11 +1,12 @@
 use super::create_handle::create_handle;
+use crate::Store;
 use crate::{TableType, ValType};
 use anyhow::{bail, Result};
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::{wasm, Module};
 use wasmtime_runtime::InstanceHandle;
 
-pub fn create_handle_with_table(table: &TableType) -> Result<InstanceHandle> {
+pub fn create_handle_with_table(store: &Store, table: &TableType) -> Result<InstanceHandle> {
     let mut module = Module::new();
 
     let table = wasm::Table {
@@ -25,5 +26,5 @@ pub fn create_handle_with_table(table: &TableType) -> Result<InstanceHandle> {
         wasmtime_environ::Export::Table(table_id),
     );
 
-    create_handle(module, None, PrimaryMap::new(), Box::new(()))
+    create_handle(module, store, PrimaryMap::new(), Box::new(()))
 }

crates/api/tests/traps.rs

@@ -63,7 +63,11 @@ fn test_trap_trace() -> Result<()> {
     assert_eq!(trace[1].module_name().unwrap(), "hello_mod");
     assert_eq!(trace[1].func_index(), 0);
     assert_eq!(trace[1].func_name(), None);
-    assert!(e.message().contains("unreachable"));
+    assert!(
+        e.message().contains("unreachable"),
+        "wrong message: {}",
+        e.message()
+    );
 
     Ok(())
 }

crates/jit/src/compiler.rs

@@ -21,8 +21,8 @@ use wasmtime_environ::{
     VMOffsets,
 };
 use wasmtime_runtime::{
-    get_mut_trap_registry, jit_function_registry, InstantiationError, SignatureRegistry,
+    jit_function_registry, InstantiationError, SignatureRegistry, TrapRegistration, TrapRegistry,
-    TrapRegistrationGuard, VMFunctionBody,
+    VMFunctionBody,
 };
 
 /// Select which kind of compilation to use.
@@ -51,8 +51,8 @@ pub struct Compiler {
     isa: Box<dyn TargetIsa>,
 
     code_memory: CodeMemory,
-    trap_registration_guards: Vec<TrapRegistrationGuard>,
     jit_function_ranges: Vec<(usize, usize)>,
+    trap_registry: TrapRegistry,
     trampoline_park: HashMap<*const VMFunctionBody, *const VMFunctionBody>,
     signatures: SignatureRegistry,
     strategy: CompilationStrategy,
@@ -67,28 +67,18 @@ impl Compiler {
         Self {
             isa,
             code_memory: CodeMemory::new(),
-            trap_registration_guards: Vec::new(),
             jit_function_ranges: Vec::new(),
             trampoline_park: HashMap::new(),
             signatures: SignatureRegistry::new(),
             fn_builder_ctx: FunctionBuilderContext::new(),
             strategy,
+            trap_registry: TrapRegistry::default(),
         }
     }
 }
 
 impl Drop for Compiler {
     fn drop(&mut self) {
-        // We must deregister traps before freeing the code memory.
-        // Otherwise, we have a race:
-        // - Compiler #1 dropped code memory, but hasn't deregistered the trap yet
-        // - Compiler #2 allocated code memory and tries to register a trap,
-        //   but the trap at certain address happens to be already registered,
-        //   since Compiler #1 hasn't deregistered it yet => assertion in trap registry fails.
-        // Having a custom drop implementation we are independent from the field order
-        // in the struct what reduces potential human error.
-        self.trap_registration_guards.clear();
-
         for (start, end) in self.jit_function_ranges.iter() {
             jit_function_registry::unregister(*start, *end);
         }
@@ -119,6 +109,7 @@ impl Compiler {
             PrimaryMap<DefinedFuncIndex, ir::JumpTableOffsets>,
             Relocations,
             Option<Vec<u8>>,
+            TrapRegistration,
         ),
         SetupError,
     > {
@@ -156,11 +147,7 @@ impl Compiler {
                 )))
             })?;
 
-        register_traps(
+        let trap_registration = register_traps(&allocated_functions, &traps, &self.trap_registry);
-            &allocated_functions,
-            &traps,
-            &mut self.trap_registration_guards,
-        );
 
         for (i, allocated) in allocated_functions.iter() {
             let ptr = (*allocated) as *const VMFunctionBody;
@@ -217,7 +204,13 @@ impl Compiler {
 
         let jt_offsets = compilation.get_jt_offsets();
 
-        Ok((allocated_functions, jt_offsets, relocations, dbg))
+        Ok((
+            allocated_functions,
+            jt_offsets,
+            relocations,
+            dbg,
+            trap_registration,
+        ))
     }
 
     /// Create a trampoline for invoking a function.
@@ -267,6 +260,11 @@ impl Compiler {
     pub fn signatures(&self) -> &SignatureRegistry {
         &self.signatures
     }
+
+    /// Shared registration of trap information
+    pub fn trap_registry(&self) -> &TrapRegistry {
+        &self.trap_registry
+    }
 }
 
 /// Create a trampoline for invoking a function.
@@ -408,19 +406,21 @@ fn allocate_functions(
 fn register_traps(
     allocated_functions: &PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
     traps: &Traps,
-    trap_registration_guards: &mut Vec<TrapRegistrationGuard>,
+    registry: &TrapRegistry,
-) {
+) -> TrapRegistration {
-    let mut trap_registry = get_mut_trap_registry();
+    let traps =
-    for (func_addr, func_traps) in allocated_functions.values().zip(traps.values()) {
+        allocated_functions
-        for trap_desc in func_traps.iter() {
+            .values()
+            .zip(traps.values())
+            .flat_map(|(func_addr, func_traps)| {
+                func_traps.iter().map(move |trap_desc| {
                     let func_addr = *func_addr as *const u8 as usize;
                     let offset = usize::try_from(trap_desc.code_offset).unwrap();
                     let trap_addr = func_addr + offset;
-            let guard =
+                    (trap_addr, trap_desc.source_loc, trap_desc.trap_code)
-                trap_registry.register_trap(trap_addr, trap_desc.source_loc, trap_desc.trap_code);
+                })
-            trap_registration_guards.push(guard);
+            });
-        }
+    registry.register_traps(traps)
-    }
 }
 
 /// We don't expect trampoline compilation to produce any relocations, so

crates/jit/src/instantiate.rs

@@ -18,7 +18,7 @@ use wasmtime_environ::{
     CompileError, DataInitializer, DataInitializerLocation, Module, ModuleEnvironment,
 };
 use wasmtime_runtime::{
-    GdbJitImageRegistration, InstanceHandle, InstantiationError, VMFunctionBody,
+    GdbJitImageRegistration, InstanceHandle, InstantiationError, TrapRegistration, VMFunctionBody,
     VMSharedSignatureIndex,
 };
 
@@ -52,6 +52,7 @@ struct RawCompiledModule<'data> {
     data_initializers: Box<[DataInitializer<'data>]>,
     signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
     dbg_jit_registration: Option<GdbJitImageRegistration>,
+    trap_registration: TrapRegistration,
 }
 
 impl<'data> RawCompiledModule<'data> {
@@ -73,7 +74,8 @@ impl<'data> RawCompiledModule<'data> {
             None
         };
 
-        let (allocated_functions, jt_offsets, relocations, dbg_image) = compiler.compile(
+        let (allocated_functions, jt_offsets, relocations, dbg_image, trap_registration) = compiler
+            .compile(
                 &translation.module,
                 translation.module_translation.as_ref().unwrap(),
                 translation.function_body_inputs,
@@ -127,6 +129,7 @@ impl<'data> RawCompiledModule<'data> {
             data_initializers: translation.data_initializers.into_boxed_slice(),
             signatures: signatures.into_boxed_slice(),
             dbg_jit_registration,
+            trap_registration,
         })
     }
 }
@@ -138,6 +141,7 @@ pub struct CompiledModule {
     data_initializers: Box<[OwnedDataInitializer]>,
     signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
     dbg_jit_registration: Option<Rc<GdbJitImageRegistration>>,
+    trap_registration: TrapRegistration,
 }
 
 impl CompiledModule {
@@ -159,6 +163,7 @@ impl CompiledModule {
                 .into_boxed_slice(),
             raw.signatures.clone(),
             raw.dbg_jit_registration,
+            raw.trap_registration,
         ))
     }
 
@@ -169,6 +174,7 @@ impl CompiledModule {
         data_initializers: Box<[OwnedDataInitializer]>,
         signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
         dbg_jit_registration: Option<GdbJitImageRegistration>,
+        trap_registration: TrapRegistration,
     ) -> Self {
         Self {
             module: Arc::new(module),
@@ -176,6 +182,7 @@ impl CompiledModule {
             data_initializers,
             signatures,
             dbg_jit_registration: dbg_jit_registration.map(Rc::new),
+            trap_registration,
         }
     }
 
@@ -203,6 +210,7 @@ impl CompiledModule {
         let imports = resolve_imports(&self.module, resolver)?;
         InstanceHandle::new(
             Arc::clone(&self.module),
+            self.trap_registration.clone(),
             self.finished_functions.clone(),
             imports,
             &data_initializers,

crates/runtime/src/instance.rs

@@ -15,6 +15,7 @@ use crate::vmcontext::{
     VMGlobalDefinition, VMGlobalImport, VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex,
     VMTableDefinition, VMTableImport,
 };
+use crate::TrapRegistration;
 use memoffset::offset_of;
 use more_asserts::assert_lt;
 use std::any::Any;
@@ -100,6 +101,10 @@ pub(crate) struct Instance {
     /// Handler run when `SIGBUS`, `SIGFPE`, `SIGILL`, or `SIGSEGV` are caught by the instance thread.
     pub(crate) signal_handler: Cell<Option<Box<SignalHandler>>>,
 
+    /// Handle to our registration of traps so signals know what trap to return
+    /// when a segfault/sigill happens.
+    pub(crate) trap_registration: TrapRegistration,
+
     /// Additional context used by compiled wasm code. This field is last, and
     /// represents a dynamically-sized array that extends beyond the nominal
     /// end of the struct (similar to a flexible array member).
@@ -534,6 +539,7 @@ impl InstanceHandle {
     /// safety.
     pub unsafe fn new(
         module: Arc<Module>,
+        trap_registration: TrapRegistration,
         finished_functions: BoxedSlice<DefinedFuncIndex, *const VMFunctionBody>,
         imports: Imports,
         data_initializers: &[DataInitializer<'_>],
@@ -582,6 +588,7 @@ impl InstanceHandle {
                 dbg_jit_registration,
                 host_state,
                 signal_handler: Cell::new(None),
+                trap_registration,
                 vmctx: VMContext {},
             };
             ptr::write(instance_ptr, instance);

crates/runtime/src/lib.rs

@@ -43,8 +43,7 @@ pub use crate::instance::{InstanceHandle, InstantiationError, LinkError};
 pub use crate::jit_int::GdbJitImageRegistration;
 pub use crate::mmap::Mmap;
 pub use crate::sig_registry::SignatureRegistry;
-pub use crate::trap_registry::TrapDescription;
+pub use crate::trap_registry::{TrapDescription, TrapRegistration, TrapRegistry};
-pub use crate::trap_registry::{get_mut_trap_registry, get_trap_registry, TrapRegistrationGuard};
 pub use crate::traphandlers::resume_panic;
 pub use crate::traphandlers::{raise_user_trap, wasmtime_call, wasmtime_call_trampoline, Trap};
 pub use crate::vmcontext::{

crates/runtime/src/trap_registry.rs

@@ -1,19 +1,64 @@
-use lazy_static::lazy_static;
+use std::collections::{BTreeMap, HashMap};
-use std::collections::HashMap;
 use std::fmt;
-use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+use std::sync::{Arc, RwLock};
 use wasmtime_environ::ir;
 
-lazy_static! {
-    static ref REGISTRY: RwLock<TrapRegistry> = RwLock::new(TrapRegistry::default());
-}
-
 /// The registry maintains descriptions of traps in currently allocated functions.
 #[derive(Default)]
 pub struct TrapRegistry {
+    // This data structure is intended to be safe to use across many threads
+    // since this is stored inside of a `Compiler` which, eventually, will be
+    // used across many threads. To that end this is internally use an `Arc`
+    // plus an `RwLock`.
+    //
+    // The problem that this data structure is solving is that when a
+    // segfault/illegal instruction happens we need to answer "given this
+    // hardware program counter what is the wasm reason this trap is being
+    // raised"?
+    //
+    // The way this is answered here is done to minimize the amount of
+    // synchronization (in theory) and have something like so:
+    //
+    // * Each module bulk-registers a list of in-memory pc addresses that have
+    //   traps. We assume that the range of traps for each module are always
+    //   disjoint.
+    // * Each key in this `BTreeMap` is the highest trapping address and the
+    //   value contains the lowest address as well as all the individual
+    //   addresses in their own `HashMap`.
+    // * Registration then looks by calculating the start/end and inserting
+    //   into this map (with some assertions about disjointed-ness)
+    // * Lookup is done in two layers. First we find the corresponding entry
+    //   in the map and verify that a program counter falls in the start/end
+    //   range. Next we look up the address in the `traps` hash map below.
+    //
+    // The `register_traps` function works by returning an RAII guard that owns
+    // a handle to this `Arc` as well, and when that type is dropped it will
+    // automatically remove all trap information from this `ranges` list.
+    ranges: Arc<RwLock<BTreeMap<usize, TrapGroup>>>,
+}
+
+#[derive(Debug)]
+struct TrapGroup {
+    /// The lowest key in the `trap` field.
+    ///
+    /// This represents the start of the range of this group of traps, and the
+    /// end of the range for this group of traps is stored as the key in the
+    /// `ranges` struct above in `TrapRegistry`.
+    start: usize,
+
+    /// All known traps in this group, mapped from program counter to the
+    /// description of the trap itself.
     traps: HashMap<usize, TrapDescription>,
 }
 
+/// RAII structure returned from `TrapRegistry::register_trap` to unregister
+/// trap information on drop.
+#[derive(Clone)]
+pub struct TrapRegistration {
+    ranges: Arc<RwLock<BTreeMap<usize, TrapGroup>>>,
+    end: Option<usize>,
+}
+
 /// Description of a trap.
 #[derive(Clone, Copy, PartialEq, Debug)]
 pub struct TrapDescription {
@@ -52,50 +97,77 @@ fn trap_code_to_expected_string(trap_code: ir::TrapCode) -> String {
     }
 }
 
-/// RAII guard for deregistering traps
-pub struct TrapRegistrationGuard(usize);
-
 impl TrapRegistry {
-    /// Registers a new trap.
+    /// Registers a list of traps.
-    /// Returns a RAII guard that deregisters the trap when dropped.
+    ///
-    pub fn register_trap(
+    /// Returns a RAII guard that deregisters all traps when dropped.
-        &mut self,
+    pub fn register_traps(
-        address: usize,
+        &self,
-        source_loc: ir::SourceLoc,
+        list: impl IntoIterator<Item = (usize, ir::SourceLoc, ir::TrapCode)>,
-        trap_code: ir::TrapCode,
+    ) -> TrapRegistration {
-    ) -> TrapRegistrationGuard {
+        let mut start = usize::max_value();
-        let entry = TrapDescription {
+        let mut end = 0;
+        let mut traps = HashMap::new();
+        for (addr, source_loc, trap_code) in list.into_iter() {
+            traps.insert(
+                addr,
+                TrapDescription {
                     source_loc,
                     trap_code,
+                },
+            );
+            if addr < start {
+                start = addr;
+            }
+            if addr > end {
+                end = addr;
+            }
+        }
+        if traps.len() == 0 {
+            return TrapRegistration {
+                ranges: self.ranges.clone(),
+                end: None,
             };
-        let previous_trap = self.traps.insert(address, entry);
+        }
-        assert!(previous_trap.is_none());
+        let mut ranges = self.ranges.write().unwrap();
-        TrapRegistrationGuard(address)
+
+        // Sanity check that no other group of traps overlaps with our
+        // registration...
+        if let Some((_, prev)) = ranges.range(end..).next() {
+            assert!(prev.start > end);
+        }
+        if let Some((prev_end, _)) = ranges.range(..=start).next_back() {
+            assert!(*prev_end < start);
         }
 
-    fn deregister_trap(&mut self, address: usize) {
+        // ... and then register ourselves
-        assert!(self.traps.remove(&address).is_some());
+        assert!(ranges.insert(end, TrapGroup { start, traps }).is_none());
+        TrapRegistration {
+            ranges: self.ranges.clone(),
+            end: Some(end),
+        }
+    }
 }
 
+impl TrapRegistration {
     /// Gets a trap description at given address.
     pub fn get_trap(&self, address: usize) -> Option<TrapDescription> {
-        self.traps.get(&address).copied()
+        let ranges = self.ranges.read().ok()?;
+        let (end, group) = ranges.range(address..).next()?;
+        if group.start <= address && address <= *end {
+            group.traps.get(&address).copied()
+        } else {
+            None
+        }
     }
 }
 
-impl Drop for TrapRegistrationGuard {
+impl Drop for TrapRegistration {
     fn drop(&mut self) {
-        let mut registry = get_mut_trap_registry();
+        if let Some(end) = self.end {
-        registry.deregister_trap(self.0);
+            if let Ok(mut ranges) = self.ranges.write() {
+                ranges.remove(&end);
             }
         }
-
-/// Gets guarded writable reference to traps registry
-pub fn get_mut_trap_registry() -> RwLockWriteGuard<'static, TrapRegistry> {
-    REGISTRY.write().expect("trap registry lock got poisoned")
     }
-
-/// Gets guarded readable reference to traps registry
-pub fn get_trap_registry() -> RwLockReadGuard<'static, TrapRegistry> {
-    REGISTRY.read().expect("trap registry lock got poisoned")
 }

crates/runtime/src/traphandlers.rs

@@ -2,7 +2,6 @@
 //! signalhandling mechanisms.
 
 use crate::instance::{InstanceHandle, SignalHandler};
-use crate::trap_registry::get_trap_registry;
 use crate::trap_registry::TrapDescription;
 use crate::vmcontext::{VMContext, VMFunctionBody};
 use backtrace::Backtrace;
@@ -78,8 +77,7 @@ cfg_if::cfg_if! {
 /// Only safe to call when wasm code is on the stack, aka `wasmtime_call` or
 /// `wasmtime_call_trampoline` must have been previously called.
 pub unsafe fn raise_user_trap(data: Box<dyn Error + Send + Sync>) -> ! {
-    let trap = Trap::User(data);
+    tls::with(|info| info.unwrap().unwind_with(UnwindReason::UserTrap(data)))
-    tls::with(|info| info.unwrap().unwind_with(UnwindReason::Trap(trap)))
 }
 
 /// Carries a Rust panic across wasm code and resumes the panic on the other
@@ -187,7 +185,8 @@ pub struct CallThreadState {
 enum UnwindReason {
     None,
     Panic(Box<dyn Any + Send>),
-    Trap(Trap),
+    UserTrap(Box<dyn Error + Send + Sync>),
+    Trap { backtrace: Backtrace, pc: usize },
 }
 
 impl CallThreadState {
@@ -210,9 +209,25 @@ impl CallThreadState {
                     debug_assert_eq!(ret, 1);
                     Ok(())
                 }
-                UnwindReason::Trap(trap) => {
+                UnwindReason::UserTrap(data) => {
                     debug_assert_eq!(ret, 0);
-                    Err(trap)
+                    Err(Trap::User(data))
+                }
+                UnwindReason::Trap { backtrace, pc } => {
+                    debug_assert_eq!(ret, 0);
+                    let instance = unsafe { InstanceHandle::from_vmctx(self.vmctx) };
+
+                    Err(Trap::Wasm {
+                        desc: instance
+                            .instance()
+                            .trap_registration
+                            .get_trap(pc)
+                            .unwrap_or_else(|| TrapDescription {
+                                source_loc: ir::SourceLoc::default(),
+                                trap_code: ir::TrapCode::StackOverflow,
+                            }),
+                        backtrace,
+                    })
                 }
                 UnwindReason::Panic(panic) => {
                     debug_assert_eq!(ret, 0);
@@ -289,20 +304,12 @@ impl CallThreadState {
         if self.jmp_buf.get().is_null() {
             return ptr::null();
         }
-
+        let backtrace = Backtrace::new_unresolved();
-        let registry = get_trap_registry();
-        let trap = Trap::Wasm {
-            desc: registry
-                .get_trap(pc as usize)
-                .unwrap_or_else(|| TrapDescription {
-                    source_loc: ir::SourceLoc::default(),
-                    trap_code: ir::TrapCode::StackOverflow,
-                }),
-            backtrace: Backtrace::new_unresolved(),
-        };
-
         self.reset_guard_page.set(reset_guard_page);
-        self.unwind.replace(UnwindReason::Trap(trap));
+        self.unwind.replace(UnwindReason::Trap {
+            backtrace,
+            pc: pc as usize,
+        });
         self.jmp_buf.get()
     }
 }