Reduce contention on the global module rwlock (#4041)

* Reduce contention on the global module rwlock

This commit intendes to close #4025 by reducing contention on the global
rwlock Wasmtime has for module information during instantiation and
dropping a store. Currently registration of a module into this global
map happens during instantiation, but this can be a hot path as
embeddings may want to, in parallel, instantiate modules.

Instead this switches to a strategy of inserting into the global module
map when a `Module` is created and then removing it from the map when
the `Module` is dropped. Registration in a `Store` now preserves the
entire `Module` within the store as opposed to trying to only save it
piecemeal. In reality the only piece that wasn't saved within a store
was the `TypeTables` which was pretty inconsequential for core wasm
modules anyway.

This means that instantiation should now clone a singluar `Arc` into a
`Store` per `Module` (previously it cloned two) with zero managemnt on
the global rwlock as that happened at `Module` creation time.
Additionally dropping a `Store` again involves zero rwlock management
and only a single `Arc` drop per-instantiated module (previously it was
two).

In the process of doing this I also went ahead and removed the
`Module::new_with_name` API. This has been difficult to support
historically with various variations on the internals of `ModuleInner`
because it involves mutating a `Module` after it's been created. My hope
is that this API is pretty rarely used and/or isn't super important, so
it's ok to remove.

Finally this change removes some internal `Arc` layerings that are no
longer necessary, attempting to use either `T` or `&T` where possible
without dealing with the overhead of an `Arc`.

Closes #4025

* Move back to a `BTreeMap` in `ModuleRegistry`

crates/jit/src/instantiate.rs

@@ -398,7 +398,7 @@ impl CompiledModule {
         info: Option<CompiledModuleInfo>,
         profiler: &dyn ProfilingAgent,
         id_allocator: &CompiledModuleIdAllocator,
-    ) -> Result<Arc<Self>> {
+    ) -> Result<Self> {
         // Transfer ownership of `obj` to a `CodeMemory` object which will
         // manage permissions, such as the executable bit. Once it's located
         // there we also publish it for being able to execute. Note that this
@@ -454,7 +454,7 @@ impl CompiledModule {
         };
         ret.register_debug_and_profiling(profiler)?;
 
-        Ok(Arc::new(ret))
+        Ok(ret)
     }
 
     fn register_debug_and_profiling(&mut self, profiler: &dyn ProfilingAgent) -> Result<()> {

crates/runtime/src/externref.rs

@@ -99,6 +99,7 @@
 //! Examination of Deferred Reference Counting and Cycle Detection* by Quinane:
 //! <https://openresearch-repository.anu.edu.au/bitstream/1885/42030/2/hon-thesis.pdf>
 
+use std::alloc::Layout;
 use std::any::Any;
 use std::cell::UnsafeCell;
 use std::cmp;
@@ -108,7 +109,6 @@ use std::mem;
 use std::ops::Deref;
 use std::ptr::{self, NonNull};
 use std::sync::atomic::{self, AtomicUsize, Ordering};
-use std::{alloc::Layout, sync::Arc};
 use wasmtime_environ::StackMap;
 
 /// An external reference to some opaque data.
@@ -814,7 +814,7 @@ impl VMExternRefActivationsTable {
 /// program counter value.
 pub trait ModuleInfoLookup {
     /// Lookup the module information from a program counter value.
-    fn lookup(&self, pc: usize) -> Option<Arc<dyn ModuleInfo>>;
+    fn lookup(&self, pc: usize) -> Option<&dyn ModuleInfo>;
 }
 
 /// Used by the runtime to query module information.

crates/wasmtime/src/module.rs

@@ -105,9 +105,9 @@ struct ModuleInner {
     /// executed.
     module: Arc<CompiledModule>,
     /// Type information of this module.
-    types: Arc<TypeTables>,
+    types: TypeTables,
     /// Registered shared signature for the module.
-    signatures: Arc<SignatureCollection>,
+    signatures: SignatureCollection,
     /// A set of initialization images for memories, if any.
     ///
     /// Note that this is behind a `OnceCell` to lazily create this image. On
@@ -193,22 +193,6 @@ impl Module {
         Self::from_binary(engine, &bytes)
     }
 
-    /// Creates a new WebAssembly `Module` from the given in-memory `binary`
-    /// data. The provided `name` will be used in traps/backtrace details.
-    ///
-    /// See [`Module::new`] for other details.
-    #[cfg(compiler)]
-    #[cfg_attr(nightlydoc, doc(cfg(feature = "cranelift")))] // see build.rs
-    pub fn new_with_name(engine: &Engine, bytes: impl AsRef<[u8]>, name: &str) -> Result<Module> {
-        let mut module = Self::new(engine, bytes.as_ref())?;
-        Arc::get_mut(&mut Arc::get_mut(&mut module.inner).unwrap().module)
-            .unwrap()
-            .module_mut()
-            .expect("mutable module")
-            .name = Some(name.to_string());
-        Ok(module)
-    }
-
     /// Creates a new WebAssembly `Module` from the contents of the given
     /// `file` on disk.
     ///
@@ -332,7 +316,7 @@ impl Module {
             }
         };
 
-        Self::from_parts(engine, mmap, info, Arc::new(types))
+        Self::from_parts(engine, mmap, info, types)
     }
 
     /// Converts an input binary-encoded WebAssembly module to compilation
@@ -487,23 +471,29 @@ impl Module {
         engine: &Engine,
         mmap: MmapVec,
         info: Option<CompiledModuleInfo>,
-        types: Arc<TypeTables>,
+        types: TypeTables,
     ) -> Result<Self> {
-        let module = CompiledModule::from_artifacts(
+        let module = Arc::new(CompiledModule::from_artifacts(
             mmap,
             info,
             &*engine.config().profiler,
             engine.unique_id_allocator(),
-        )?;
+        )?);
 
         // Validate the module can be used with the current allocator
         engine.allocator().validate(module.module())?;
 
-        let signatures = Arc::new(SignatureCollection::new_for_module(
+        let signatures = SignatureCollection::new_for_module(
             engine.signatures(),
             &types,
             module.trampolines().map(|(idx, f, _)| (idx, f)),
-        ));
+        );
+
+        // We're about to create a `Module` for real now so enter this module
+        // into the global registry of modules so we can resolve traps
+        // appropriately. Note that the corresponding `unregister` happens below
+        // in `Drop for ModuleInner`.
+        registry::register(engine, &module);
 
         Ok(Self {
             inner: Arc::new(ModuleInner {
@@ -564,7 +554,7 @@ impl Module {
         SerializedModule::new(self).to_bytes(&self.inner.engine.config().module_version)
     }
 
-    pub(crate) fn compiled_module(&self) -> &Arc<CompiledModule> {
+    pub(crate) fn compiled_module(&self) -> &CompiledModule {
         &self.inner.module
     }
 
@@ -572,11 +562,11 @@ impl Module {
         self.compiled_module().module()
     }
 
-    pub(crate) fn types(&self) -> &Arc<TypeTables> {
+    pub(crate) fn types(&self) -> &TypeTables {
         &self.inner.types
     }
 
-    pub(crate) fn signatures(&self) -> &Arc<SignatureCollection> {
+    pub(crate) fn signatures(&self) -> &SignatureCollection {
         &self.inner.signatures
     }
 
@@ -599,8 +589,6 @@ impl Module {
     /// let module = Module::new(&engine, "(module)")?;
     /// assert_eq!(module.name(), None);
     ///
-    /// let module = Module::new_with_name(&engine, "(module)", "bar")?;
-    /// assert_eq!(module.name(), Some("bar"));
     /// # Ok(())
     /// # }
     /// ```
@@ -798,6 +786,10 @@ impl Module {
         self.inner.clone()
     }
 
+    pub(crate) fn module_info(&self) -> &dyn wasmtime_runtime::ModuleInfo {
+        &*self.inner
+    }
+
     /// Returns the range of bytes in memory where this module's compilation
     /// image resides.
     ///
@@ -917,6 +909,38 @@ impl wasmtime_runtime::ModuleRuntimeInfo for ModuleInner {
     }
 }
 
+impl wasmtime_runtime::ModuleInfo for ModuleInner {
+    fn lookup_stack_map(&self, pc: usize) -> Option<&wasmtime_environ::StackMap> {
+        let text_offset = pc - self.module.code().as_ptr() as usize;
+        let (index, func_offset) = self.module.func_by_text_offset(text_offset)?;
+        let info = self.module.func_info(index);
+
+        // Do a binary search to find the stack map for the given offset.
+        let index = match info
+            .stack_maps
+            .binary_search_by_key(&func_offset, |i| i.code_offset)
+        {
+            // Found it.
+            Ok(i) => i,
+
+            // No stack map associated with this PC.
+            //
+            // Because we know we are in Wasm code, and we must be at some kind
+            // of call/safepoint, then the Cranelift backend must have avoided
+            // emitting a stack map for this location because no refs were live.
+            Err(_) => return None,
+        };
+
+        Some(&info.stack_maps[index].stack_map)
+    }
+}
+
+impl Drop for ModuleInner {
+    fn drop(&mut self) {
+        registry::unregister(&self.module);
+    }
+}
+
 /// A barebones implementation of ModuleRuntimeInfo that is useful for
 /// cases where a purpose-built environ::Module is used and a full
 /// CompiledModule does not exist (for example, for tests or for the

crates/wasmtime/src/module/registry.rs

@@ -1,11 +1,11 @@
 //! Implements a registry of modules for a store.
 
-use crate::{signatures::SignatureCollection, Module};
+use crate::{Engine, Module};
 use std::{
     collections::BTreeMap,
     sync::{Arc, RwLock},
 };
-use wasmtime_environ::{EntityRef, FilePos, StackMap, TrapCode};
+use wasmtime_environ::{EntityRef, FilePos, TrapCode};
 use wasmtime_jit::CompiledModule;
 use wasmtime_runtime::{ModuleInfo, VMCallerCheckedAnyfunc, VMTrampoline};
 
@@ -15,30 +15,38 @@ lazy_static::lazy_static! {
 
 /// Used for registering modules with a store.
 ///
-/// The map is from the ending (exclusive) address for the module code to
+/// Note that the primary reason for this registry is to ensure that everything
-/// the registered module.
+/// in `Module` is kept alive for the duration of a `Store`. At this time we
-///
+/// need "basically everything" within a `Moudle` to stay alive once it's
-/// The `BTreeMap` is used to quickly locate a module based on a program counter value.
+/// instantiated within a store. While there's some smaller portions that could
+/// theoretically be omitted as they're not needed by the store they're
+/// currently small enough to not worry much about.
 #[derive(Default)]
 pub struct ModuleRegistry {
-    modules_with_code: BTreeMap<usize, Arc<RegisteredModule>>,
+    // Keyed by the end address of the module's code in memory.
-    modules_without_code: Vec<Arc<CompiledModule>>,
+    modules_with_code: BTreeMap<usize, Module>,
+
+    // Preserved for keeping data segments alive or similar
+    modules_without_code: Vec<Module>,
+}
+
+fn start(module: &Module) -> usize {
+    assert!(!module.compiled_module().code().is_empty());
+    module.compiled_module().code().as_ptr() as usize
 }
 
 impl ModuleRegistry {
     /// Fetches information about a registered module given a program counter value.
-    pub fn lookup_module(&self, pc: usize) -> Option<Arc<dyn ModuleInfo>> {
+    pub fn lookup_module(&self, pc: usize) -> Option<&dyn ModuleInfo> {
-        self.module(pc)
+        self.module(pc).map(|m| m.module_info())
-            .map(|m| -> Arc<dyn ModuleInfo> { m.clone() })
     }
 
-    fn module(&self, pc: usize) -> Option<&Arc<RegisteredModule>> {
+    fn module(&self, pc: usize) -> Option<&Module> {
-        let (end, info) = self.modules_with_code.range(pc..).next()?;
+        let (end, module) = self.modules_with_code.range(pc..).next()?;
-        if pc < info.start || *end < pc {
+        if pc < start(module) || *end < pc {
             return None;
         }
-
+        Some(module)
-        Some(info)
     }
 
     /// Registers a new module with the registry.
@@ -52,92 +60,40 @@ impl ModuleRegistry {
         // module in the future. For that reason we continue to register empty
         // modules and retain them.
         if compiled_module.finished_functions().len() == 0 {
-            self.modules_without_code.push(compiled_module.clone());
+            self.modules_without_code.push(module.clone());
             return;
         }
 
         // The module code range is exclusive for end, so make it inclusive as it
         // may be a valid PC value
-        let code = compiled_module.code();
+        let start_addr = start(module);
-        assert!(!code.is_empty());
+        let end_addr = start_addr + compiled_module.code().len() - 1;
-        let start = code.as_ptr() as usize;
-        let end = start + code.len() - 1;
 
         // Ensure the module isn't already present in the registry
         // This is expected when a module is instantiated multiple times in the
         // same store
-        if let Some(m) = self.modules_with_code.get(&end) {
+        if let Some(m) = self.modules_with_code.get(&end_addr) {
-            assert_eq!(m.start, start);
+            assert_eq!(start(m), start_addr);
             return;
         }
 
-        // Assert that this module's code doesn't collide with any other registered modules
+        // Assert that this module's code doesn't collide with any other
-        if let Some((_, prev)) = self.modules_with_code.range(end..).next() {
+        // registered modules
-            assert!(prev.start > end);
+        if let Some((_, prev)) = self.modules_with_code.range(end_addr..).next() {
+            assert!(start(prev) > end_addr);
+        }
+        if let Some((prev_end, _)) = self.modules_with_code.range(..=start_addr).next_back() {
+            assert!(*prev_end < start_addr);
         }
 
-        if let Some((prev_end, _)) = self.modules_with_code.range(..=start).next_back() {
+        let prev = self.modules_with_code.insert(end_addr, module.clone());
-            assert!(*prev_end < start);
-        }
-
-        let prev = self.modules_with_code.insert(
-            end,
-            Arc::new(RegisteredModule {
-                start,
-                module: compiled_module.clone(),
-                signatures: module.signatures().clone(),
-            }),
-        );
         assert!(prev.is_none());
-
-        GLOBAL_MODULES.write().unwrap().register(start, end, module);
     }
 
     /// Looks up a trampoline from an anyfunc.
     pub fn lookup_trampoline(&self, anyfunc: &VMCallerCheckedAnyfunc) -> Option<VMTrampoline> {
         let module = self.module(anyfunc.func_ptr.as_ptr() as usize)?;
-        module.signatures.trampoline(anyfunc.type_index)
+        module.signatures().trampoline(anyfunc.type_index)
-    }
-}
-
-impl Drop for ModuleRegistry {
-    fn drop(&mut self) {
-        let mut info = GLOBAL_MODULES.write().unwrap();
-        for end in self.modules_with_code.keys() {
-            info.unregister(*end);
-        }
-    }
-}
-
-struct RegisteredModule {
-    start: usize,
-    module: Arc<CompiledModule>,
-    signatures: Arc<SignatureCollection>,
-}
-
-impl ModuleInfo for RegisteredModule {
-    fn lookup_stack_map(&self, pc: usize) -> Option<&StackMap> {
-        let text_offset = pc - self.start;
-        let (index, func_offset) = self.module.func_by_text_offset(text_offset)?;
-        let info = self.module.func_info(index);
-
-        // Do a binary search to find the stack map for the given offset.
-        let index = match info
-            .stack_maps
-            .binary_search_by_key(&func_offset, |i| i.code_offset)
-        {
-            // Found it.
-            Ok(i) => i,
-
-            // No stack map associated with this PC.
-            //
-            // Because we know we are in Wasm code, and we must be at some kind
-            // of call/safepoint, then the Cranelift backend must have avoided
-            // emitting a stack map for this location because no refs were live.
-            Err(_) => return None,
-        };
-
-        Some(&info.stack_maps[index].stack_map)
     }
 }
 
@@ -146,11 +102,6 @@ struct GlobalRegisteredModule {
     start: usize,
     module: Arc<CompiledModule>,
     wasm_backtrace_details_env_used: bool,
-    /// Note that modules can be instantiated in many stores, so the purpose of
-    /// this field is to keep track of how many stores have registered a
-    /// module. Information is only removed from the global registry when this
-    /// reference count reaches 0.
-    references: usize,
 }
 
 /// This is the global module registry that stores information for all modules
@@ -173,14 +124,12 @@ impl GlobalModuleRegistry {
     /// Returns whether the `pc`, according to globally registered information,
     /// is a wasm trap or not.
     pub(crate) fn is_wasm_trap_pc(pc: usize) -> bool {
-        let modules = GLOBAL_MODULES.read().unwrap();
+        let (module, text_offset) = match GLOBAL_MODULES.read().unwrap().module(pc) {
+            Some((module, offset)) => (module.module.clone(), offset),
+            None => return false,
+        };
 
-        match modules.module(pc) {
+        wasmtime_environ::lookup_trap_code(module.trap_data(), text_offset).is_some()
-            Some((entry, text_offset)) => {
-                wasmtime_environ::lookup_trap_code(entry.module.trap_data(), text_offset).is_some()
-            }
-            None => false,
-        }
     }
 
     /// Returns, if found, the corresponding module for the `pc` as well as the
@@ -223,36 +172,43 @@ impl GlobalModuleRegistry {
         let (module, offset) = self.module(pc)?;
         wasmtime_environ::lookup_trap_code(module.module.trap_data(), offset)
     }
+}
 
-    /// Registers a new region of code, described by `(start, end)` and with
+/// Registers a new region of code.
-    /// the given function information, with the global information.
+///
-    fn register(&mut self, start: usize, end: usize, module: &Module) {
+/// Must not have been previously registered and must be `unregister`'d to
-        let info = self.0.entry(end).or_insert_with(|| GlobalRegisteredModule {
+/// prevent leaking memory.
-            start,
+///
-            module: module.compiled_module().clone(),
+/// This is required to enable traps to work correctly since the signal handler
-            wasm_backtrace_details_env_used: module
+/// will lookup in the `GLOBAL_MODULES` list to determine which a particular pc
-                .engine()
+/// is a trap or not.
-                .config()
+pub fn register(engine: &Engine, module: &Arc<CompiledModule>) {
-                .wasm_backtrace_details_env_used,
+    let code = module.code();
-            references: 0,
+    if code.is_empty() {
-        });
+        return;
-
-        // Note that ideally we'd debug_assert that the information previously
-        // stored, if any, matches the `functions` we were given, but for now we
-        // just do some simple checks to hope it's the same.
-        assert_eq!(info.start, start);
-        info.references += 1;
     }
+    let start = code.as_ptr() as usize;
+    let end = start + code.len() - 1;
+    let module = GlobalRegisteredModule {
+        start,
+        wasm_backtrace_details_env_used: engine.config().wasm_backtrace_details_env_used,
+        module: module.clone(),
+    };
+    let prev = GLOBAL_MODULES.write().unwrap().0.insert(end, module);
+    assert!(prev.is_none());
+}
 
-    /// Unregisters a region of code (keyed by the `end` address) from the
+/// Unregisters a module from the global map.
-    /// global information.
+///
-    fn unregister(&mut self, end: usize) {
+/// Must hae been previously registered with `register`.
-        let info = self.0.get_mut(&end).unwrap();
+pub fn unregister(module: &Arc<CompiledModule>) {
-        info.references -= 1;
+    let code = module.code();
-        if info.references == 0 {
+    if code.is_empty() {
-            self.0.remove(&end);
+        return;
-        }
     }
+    let end = (code.as_ptr() as usize) + code.len() - 1;
+    let module = GLOBAL_MODULES.write().unwrap().0.remove(&end);
+    assert!(module.is_some());
 }
 
 impl GlobalRegisteredModule {

crates/wasmtime/src/module/serialization.rs

@@ -47,7 +47,6 @@ use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::path::Path;
 use std::str::FromStr;
-use std::sync::Arc;
 use wasmtime_environ::{FlagValue, Tunables, TypeTables};
 use wasmtime_jit::{subslice_range, CompiledModuleInfo};
 use wasmtime_runtime::MmapVec;
@@ -206,7 +205,7 @@ impl<'a> SerializedModule<'a> {
 
     pub fn into_module(self, engine: &Engine) -> Result<Module> {
         let (mmap, info, types) = self.into_parts(engine)?;
-        Module::from_parts(engine, mmap, info, Arc::new(types))
+        Module::from_parts(engine, mmap, info, types)
     }
 
     pub fn into_parts(

crates/wasmtime/src/store.rs

@@ -1960,7 +1960,7 @@ impl Drop for StoreOpaque {
 }
 
 impl wasmtime_runtime::ModuleInfoLookup for ModuleRegistry {
-    fn lookup(&self, pc: usize) -> Option<Arc<dyn ModuleInfo>> {
+    fn lookup(&self, pc: usize) -> Option<&dyn ModuleInfo> {
         self.lookup_module(pc)
     }
 }

tests/all/name.rs

@@ -27,8 +27,5 @@ fn test_module_name() -> anyhow::Result<()> {
     let module = Module::new(&engine, wat)?;
     assert_eq!(module.name(), Some("from_name_section"));
 
-    let module = Module::new_with_name(&engine, wat, "override")?;
-    assert_eq!(module.name(), Some("override"));
-
     Ok(())
 }