diff --git a/crates/wasmtime/src/instance.rs b/crates/wasmtime/src/instance.rs index 32a77e1f9f..26a1776cd4 100644 --- a/crates/wasmtime/src/instance.rs +++ b/crates/wasmtime/src/instance.rs @@ -16,6 +16,11 @@ fn instantiate( imports: Imports<'_>, host: Box, ) -> Result { + // Register the module just before instantiation to ensure we have a + // trampoline registered for every signature and to preserve the module's + // compiled JIT code within the `Store`. + store.register_module(compiled_module); + let config = store.engine().config(); let instance = unsafe { let instance = compiled_module.instantiate( @@ -98,11 +103,6 @@ fn instantiate( #[derive(Clone)] pub struct Instance { pub(crate) handle: StoreInstanceHandle, - // Note that this is required to keep the module's code memory alive while - // we have a handle to this `Instance`. We may eventually want to shrink - // this to only hold onto the bare minimum each instance needs to allow - // deallocating some `Module` resources early, but until then we just hold - // on to everything. module: Module, } @@ -165,7 +165,6 @@ impl Instance { bail!("cross-`Engine` instantiation is not currently supported"); } - store.register_module(module.compiled_module()); let handle = with_imports(store, module.compiled_module(), imports, |imports| { instantiate(store, module.compiled_module(), imports, Box::new(())) })?; diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index de21842dbc..5317631cc4 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -5,6 +5,7 @@ use crate::Engine; use anyhow::{bail, Result}; use std::any::Any; use std::cell::RefCell; +use std::collections::HashSet; use std::fmt; use std::hash::{Hash, Hasher}; use std::rc::{Rc, Weak}; @@ -63,9 +64,9 @@ pub(crate) struct StoreInner { /// Information about JIT code which allows us to test if a program counter /// is in JIT code, lookup trap information, etc. frame_info: RefCell, - /// List of all compiled modules that we're holding a strong reference to + /// Set of all compiled modules that we're holding a strong reference to /// the module's code for. This includes JIT functions, trampolines, etc. - modules: RefCell>>, + modules: RefCell>, } struct HostInfoKey(VMExternRef); @@ -169,6 +170,15 @@ impl Store { // a `Func` wrapper for any function in the module, which requires that // we know about the signature and trampoline for all instances. self.register_signatures(module); + + // And finally with a module being instantiated into this `Store` we + // need to preserve its jit-code. References to this module's code and + // trampolines are not owning-references so it's our responsibility to + // keep it all alive within the `Store`. + self.inner + .modules + .borrow_mut() + .insert(ArcModuleCode(module.code().clone())); } fn register_jit_code(&self, module: &CompiledModule) { @@ -180,7 +190,6 @@ impl Store { // Only register this module if it hasn't already been registered. if !self.is_wasm_code(first_pc) { self.inner.frame_info.borrow_mut().register(module); - self.inner.modules.borrow_mut().push(module.code().clone()); } } @@ -435,3 +444,21 @@ impl InterruptHandle { self.interrupts.interrupt() } } + +// Wrapper struct to implement hash/equality based on the pointer value of the +// `Arc` in question. +struct ArcModuleCode(Arc); + +impl PartialEq for ArcModuleCode { + fn eq(&self, other: &ArcModuleCode) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } +} + +impl Eq for ArcModuleCode {} + +impl Hash for ArcModuleCode { + fn hash(&self, hasher: &mut H) { + Arc::as_ptr(&self.0).hash(hasher) + } +} diff --git a/tests/all/func.rs b/tests/all/func.rs index fca6aed814..c626d977bf 100644 --- a/tests/all/func.rs +++ b/tests/all/func.rs @@ -522,3 +522,30 @@ fn externref_signature_no_reference_types() -> anyhow::Result<()> { ); Ok(()) } + +#[test] +fn trampolines_always_valid() -> anyhow::Result<()> { + let func = { + // Compile two modules up front + let store = Store::default(); + let module1 = Module::new(store.engine(), "(module (import \"\" \"\" (func)))")?; + let module2 = Module::new(store.engine(), "(module (func (export \"\")))")?; + // Start instantiating the first module, but this will fail. + // Historically this registered the module's trampolines with `Store` + // before the failure, but then after the failure the `Store` didn't + // hold onto the trampoline. + drop(Instance::new(&store, &module1, &[])); + drop(module1); + + // Then instantiate another module which has the same function type (no + // parameters or results) which tries to use the trampoline defined in + // the previous module. Then we extract the function and, in another + // scope where everything is dropped, we call the func. + let i = Instance::new(&store, &module2, &[])?; + i.get_func("").unwrap() + }; + + // ... and no segfaults! right? right? ... + func.call(&[])?; + Ok(()) +}