Remove the jit_function_registry global state (bytecodealliance#915)

* Remove the `jit_function_registry` global state

This commit removes on the final pieces of global state in wasmtime
today, the `jit_function_registry` module. The purpose of this module is
to help translate a native backtrace with native program counters into a
wasm backtrace with module names, function names, and wasm module
indices. To that end this module retained a global map of function
ranges to this metadata information for each compiled function.

It turns out that we already had a `NAMES` global in the `wasmtime`
crate for symbolicating backtrace addresses, so this commit moves that
global into its own file and restructures the internals to account for
program counter ranges as well. The general set of changes here are:

* Remove `jit_function_registry`
* Remove `NAMES`
* Create a new `frame_info` module which has a singleton global
  registering compiled module's frame information.
* Update traps to use the `frame_info` module to symbolicate pcs,
  directly extracting a `FrameInfo` from the module.
* Register and unregister information on a module level instead of on a
  per-function level (at least in terms of locking granluarity).

This commit leaves the new `FRAME_INFO` global variable as the only
remaining "critical" global variable in `wasmtime`, which only exists
due to the API of `Trap` where it doesn't take in any extra context when
capturing a stack trace through which we could hang off frame
information. I'm thinking though that this is ok, and we can always
tweak the API of `Trap` in the future if necessary if we truly need to
accomodate this.

* Remove a lazy_static dep

* Add some comments and restructure

crates/api/src/frame_info.rs

@@ -0,0 +1,183 @@
+use crate::module::Names;
+use std::collections::BTreeMap;
+use std::sync::{Arc, RwLock};
+use wasmtime_environ::entity::EntityRef;
+use wasmtime_environ::wasm::FuncIndex;
+use wasmtime_jit::CompiledModule;
+
+lazy_static::lazy_static! {
+    /// This is a global cache of backtrace frame information for all active
+    ///
+    /// This global cache is used during `Trap` creation to symbolicate frames.
+    /// This is populated on module compilation, and it is cleared out whenever
+    /// all references to a module are dropped.
+    pub static ref FRAME_INFO: GlobalFrameInfo = GlobalFrameInfo::default();
+}
+
+#[derive(Default)]
+pub struct GlobalFrameInfo {
+    /// An internal map that keeps track of backtrace frame information for
+    /// each module.
+    ///
+    /// This map is morally a map of ranges to a map of information for that
+    /// module. Each module is expected to reside in a disjoint section of
+    /// contiguous memory. No modules can overlap.
+    ///
+    /// The key of this map is the highest address in the module and the value
+    /// is the module's information, which also contains the start address.
+    ranges: RwLock<BTreeMap<usize, ModuleFrameInfo>>,
+}
+
+/// An RAII structure used to unregister a module's frame information when the
+/// module is destroyed.
+pub struct GlobalFrameInfoRegistration {
+    /// The key that will be removed from the global `ranges` map when this is
+    /// dropped.
+    key: usize,
+}
+
+struct ModuleFrameInfo {
+    start: usize,
+    functions: BTreeMap<usize, (usize, FuncIndex)>,
+    names: Arc<Names>,
+}
+
+impl GlobalFrameInfo {
+    /// Registers a new compiled module's frame information.
+    ///
+    /// This function will register the `names` information for all of the
+    /// compiled functions within `module`. If the `module` has no functions
+    /// then `None` will be returned. Otherwise the returned object, when
+    /// dropped, will be used to unregister all name information from this map.
+    pub fn register(
+        &self,
+        names: &Arc<Names>,
+        module: &CompiledModule,
+    ) -> Option<GlobalFrameInfoRegistration> {
+        let mut min = usize::max_value();
+        let mut max = 0;
+        let mut functions = BTreeMap::new();
+        for (i, allocated) in module.finished_functions() {
+            let (start, end) = unsafe {
+                let ptr = (**allocated).as_ptr();
+                let len = (**allocated).len();
+                (ptr as usize, ptr as usize + len)
+            };
+            if start < min {
+                min = start;
+            }
+            if end > max {
+                max = end;
+            }
+            let func_index = module.module().func_index(i);
+            assert!(functions.insert(end, (start, func_index)).is_none());
+        }
+        if functions.len() == 0 {
+            return None;
+        }
+
+        let mut ranges = self.ranges.write().unwrap();
+        // First up assert that our chunk of jit functions doesn't collide with
+        // any other known chunks of jit functions...
+        if let Some((_, prev)) = ranges.range(max..).next() {
+            assert!(prev.start > max);
+        }
+        if let Some((prev_end, _)) = ranges.range(..=min).next_back() {
+            assert!(*prev_end < min);
+        }
+
+        // ... then insert our range and assert nothing was there previously
+        let prev = ranges.insert(
+            max,
+            ModuleFrameInfo {
+                start: min,
+                functions,
+                names: names.clone(),
+            },
+        );
+        assert!(prev.is_none());
+        Some(GlobalFrameInfoRegistration { key: max })
+    }
+
+    /// Fetches information about a program counter in a backtrace.
+    ///
+    /// Returns an object if this `pc` is known to some previously registered
+    /// module, or returns `None` if no information can be found.
+    pub fn lookup(&self, pc: usize) -> Option<FrameInfo> {
+        let ranges = self.ranges.read().ok()?;
+        let (end, info) = ranges.range(pc..).next()?;
+        if pc < info.start || *end < pc {
+            return None;
+        }
+        let (end, (start, func_index)) = info.functions.range(pc..).next()?;
+        if pc < *start || *end < pc {
+            return None;
+        }
+        Some(FrameInfo {
+            module_name: info.names.module_name.clone(),
+            func_index: func_index.index() as u32,
+            func_name: info.names.module.func_names.get(func_index).cloned(),
+        })
+    }
+}
+
+impl Drop for GlobalFrameInfoRegistration {
+    fn drop(&mut self) {
+        if let Ok(mut map) = FRAME_INFO.ranges.write() {
+            map.remove(&self.key);
+        }
+    }
+}
+
+/// Description of a frame in a backtrace for a [`Trap`].
+///
+/// Whenever a WebAssembly trap occurs an instance of [`Trap`] is created. Each
+/// [`Trap`] has a backtrace of the WebAssembly frames that led to the trap, and
+/// each frame is described by this structure.
+#[derive(Debug)]
+pub struct FrameInfo {
+    module_name: Option<String>,
+    func_index: u32,
+    func_name: Option<String>,
+}
+
+impl FrameInfo {
+    /// Returns the WebAssembly function index for this frame.
+    ///
+    /// This function index is the index in the function index space of the
+    /// WebAssembly module that this frame comes from.
+    pub fn func_index(&self) -> u32 {
+        self.func_index
+    }
+
+    /// Returns the identifer of the module that this frame is for.
+    ///
+    /// Module identifiers are present in the `name` section of a WebAssembly
+    /// binary, but this may not return the exact item in the `name` section.
+    /// Module names can be overwritten at construction time or perhaps inferred
+    /// from file names. The primary purpose of this function is to assist in
+    /// debugging and therefore may be tweaked over time.
+    ///
+    /// This function returns `None` when no name can be found or inferred.
+    pub fn module_name(&self) -> Option<&str> {
+        self.module_name.as_deref()
+    }
+
+    /// Returns a descriptive name of the function for this frame, if one is
+    /// available.
+    ///
+    /// The name of this function may come from the `name` section of the
+    /// WebAssembly binary, or wasmtime may try to infer a better name for it if
+    /// not available, for example the name of the export if it's exported.
+    ///
+    /// This return value is primarily used for debugging and human-readable
+    /// purposes for things like traps. Note that the exact return value may be
+    /// tweaked over time here and isn't guaranteed to be something in
+    /// particular about a wasm module due to its primary purpose of assisting
+    /// in debugging.
+    ///
+    /// This function returns `None` when no name could be inferred.
+    pub fn func_name(&self) -> Option<&str> {
+        self.func_name.as_deref()
+    }
+}

crates/api/src/func.rs

@@ -73,7 +73,10 @@ macro_rules! wrappers {
             unsafe {
                 let (instance, export) = crate::trampoline::generate_raw_func_export(
                     &ty,
-                    shim::<F, $($args,)* R> as *const _,
+                    std::slice::from_raw_parts_mut(
+                        shim::<F, $($args,)* R> as *mut _,
+                        0,
+                    ),
                     store,
                     Box::new(func),
                 )

crates/api/src/instance.rs

@@ -120,7 +120,7 @@ impl Instance {
             }
             exports.into_boxed_slice()
         };
-        module.register_names();
+        module.register_frame_info();
         Ok(Instance {
             instance_handle,
             module: module.clone(),

crates/api/src/lib.rs

@@ -10,6 +10,7 @@
 
 mod callable;
 mod externals;
+mod frame_info;
 mod func;
 mod instance;
 mod module;
@@ -22,12 +23,13 @@ mod values;
 
 pub use crate::callable::Callable;
 pub use crate::externals::*;
+pub use crate::frame_info::FrameInfo;
 pub use crate::func::{Func, WasmArg, WasmRet};
 pub use crate::instance::Instance;
 pub use crate::module::Module;
 pub use crate::r#ref::{AnyRef, HostInfo, HostRef};
 pub use crate::runtime::{Config, Engine, OptLevel, Store, Strategy};
-pub use crate::trap::{FrameInfo, Trap};
+pub use crate::trap::Trap;
 pub use crate::types::*;
 pub use crate::values::*;
 

crates/api/src/module.rs

@@ -1,15 +1,12 @@
+use crate::frame_info::{GlobalFrameInfoRegistration, FRAME_INFO};
 use crate::runtime::Store;
 use crate::types::{
     ExportType, ExternType, FuncType, GlobalType, ImportType, Limits, MemoryType, Mutability,
     TableType, ValType,
 };
 use anyhow::{bail, Error, Result};
-use lazy_static::lazy_static;
-use std::cell::Cell;
-use std::collections::HashMap;
 use std::path::Path;
-use std::rc::Rc;
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, Mutex};
 use wasmparser::{
     validate, CustomSectionKind, ExternalKind, ImportSectionEntryType, ModuleReader, Name,
     OperatorValidatorConfig, SectionCode, ValidatingParserConfig,
@@ -83,16 +80,15 @@ fn into_table_type(tt: wasmparser::TableType) -> TableType {
 /// In other words it's a shallow copy, not a deep copy.
 #[derive(Clone)]
 pub struct Module {
-    // FIXME(#777) should be `Arc` and this type should be thread-safe
-    inner: Rc<ModuleInner>,
+    inner: Arc<ModuleInner>,
 }
 
 struct ModuleInner {
     store: Store,
     imports: Box<[ImportType]>,
     exports: Box<[ExportType]>,
     compiled: CompiledModule,
-    registered_names: Cell<bool>,
+    frame_info_registration: Mutex<Option<Option<GlobalFrameInfoRegistration>>>,
     names: Arc<Names>,
 }
 
@@ -101,17 +97,6 @@ pub struct Names {
     pub module_name: Option<String>,
 }
 
-lazy_static! {
-    /// This is a global cache of names known for all compiled modules in this
-    /// process.
-    ///
-    /// This global cache is used during `Trap` creation to symbolicate frames.
-    /// This is populated on module compilation, and it is cleared out whenever
-    /// all references to a module are dropped, aka the `Drop for ModuleInner`
-    /// below.
-    pub static ref NAMES: RwLock<HashMap<usize, Arc<Names>>> = RwLock::default();
-}
-
 impl Module {
     /// Creates a new WebAssembly `Module` from the given in-memory `bytes`.
     ///
@@ -170,7 +155,7 @@ impl Module {
     /// See [`Module::new`] for other details.
     pub fn new_with_name(store: &Store, bytes: impl AsRef<[u8]>, name: &str) -> Result<Module> {
         let mut module = Module::new(store, bytes.as_ref())?;
-        let inner = Rc::get_mut(&mut module.inner).unwrap();
+        let inner = Arc::get_mut(&mut module.inner).unwrap();
         Arc::get_mut(&mut inner.names).unwrap().module_name = Some(name.to_string());
         Ok(module)
     }
@@ -199,9 +184,9 @@ impl Module {
     /// used instead.
     pub fn from_binary(store: &Store, binary: &[u8]) -> Result<Module> {
         Module::validate(store, binary)?;
-        // Note that the call to `validate` here should be ok because we
-        // previously validated the binary, meaning we're guaranteed to pass a
-        // valid binary for `store`.
+        // Note that the call to `from_binary_unchecked` here should be ok
+        // because we previously validated the binary, meaning we're guaranteed
+        // to pass a valid binary for `store`.
         unsafe { Module::from_binary_unchecked(store, binary) }
     }
 
@@ -284,13 +269,13 @@ impl Module {
             module: compiled.module().clone(),
         });
         Ok(Module {
-            inner: Rc::new(ModuleInner {
+            inner: Arc::new(ModuleInner {
                 store: store.clone(),
                 imports: Box::new([]),
                 exports: Box::new([]),
                 names,
                 compiled,
-                registered_names: Cell::new(false),
+                frame_info_registration: Mutex::new(None),
             }),
         })
     }
@@ -323,7 +308,7 @@ impl Module {
     }
 
     fn read_imports_and_exports(&mut self, binary: &[u8]) -> Result<()> {
-        let inner = Rc::get_mut(&mut self.inner).unwrap();
+        let inner = Arc::get_mut(&mut self.inner).unwrap();
         let mut reader = ModuleReader::new(binary)?;
         let mut imports = Vec::new();
         let mut exports = Vec::new();
@@ -452,23 +437,14 @@ impl Module {
         Ok(())
     }
 
-    /// Register this module's names in the global map of module names.
+    /// Register this module's stack frame information into the global scope.
     ///
     /// This is required to ensure that any traps can be properly symbolicated.
-    pub(crate) fn register_names(&self) {
-        if self.inner.registered_names.get() {
+    pub(crate) fn register_frame_info(&self) {
+        let mut info = self.inner.frame_info_registration.lock().unwrap();
+        if info.is_some() {
             return;
         }
-        let names = self.inner.names.clone();
-        NAMES.write().unwrap().insert(names.module.id, names);
-        self.inner.registered_names.set(true);
-    }
-}
-
-impl Drop for ModuleInner {
-    fn drop(&mut self) {
-        if self.registered_names.get() {
-            NAMES.write().unwrap().remove(&self.names.module.id);
-        }
+        *info = Some(FRAME_INFO.register(&self.inner.names, &self.inner.compiled));
     }
 }

crates/api/src/trampoline/create_handle.rs

@@ -13,7 +13,7 @@ use wasmtime_runtime::{Imports, InstanceHandle, VMFunctionBody};
 pub(crate) fn create_handle(
     module: Module,
     store: &Store,
-    finished_functions: PrimaryMap<DefinedFuncIndex, *const VMFunctionBody>,
+    finished_functions: PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
     state: Box<dyn Any>,
 ) -> Result<InstanceHandle> {
     let imports = Imports::new(