Keep frame info registered until internal instance is gone (#1514)

This commit fixes an issue where the global registration of frame data
goes away once the `wasmtime::Module` has been dropped. Even after this
has been dropped, though, there may still be `wasmtime::Func` instances
which reference the original module, so it's only once the underlying
`wasmtime_runtime::Instance` has gone away that we can drop everything.

Closes #1479

crates/api/src/instance.rs

@@ -3,6 +3,7 @@ use crate::module::Module;
 use crate::runtime::{Config, Store};
 use crate::trap::Trap;
 use anyhow::{bail, Error, Result};
+use std::any::Any;
 use wasmtime_jit::{CompiledModule, Resolver};
 use wasmtime_runtime::{Export, InstanceHandle, InstantiationError, SignatureRegistry};
 
@@ -23,6 +24,7 @@ fn instantiate(
     compiled_module: &CompiledModule,
     imports: &[Extern],
     sig_registry: &SignatureRegistry,
+    host: Box<dyn Any>,
 ) -> Result<InstanceHandle, Error> {
     let mut resolver = SimpleResolver { imports };
     unsafe {
@@ -32,6 +34,7 @@ fn instantiate(
                 &mut resolver,
                 sig_registry,
                 config.memory_creator.as_ref().map(|a| a as _),
+                host,
             )
             .map_err(|e| -> Error {
                 match e {
@@ -132,13 +135,14 @@ impl Instance {
             );
         }
 
-        module.register_frame_info();
+        let info = module.register_frame_info();
         let config = store.engine().config();
         let instance_handle = instantiate(
             config,
             module.compiled_module(),
             imports,
             store.compiler().signatures(),
+            Box::new(info),
         )?;
 
         let mut exports = Vec::with_capacity(module.exports().len());

crates/api/src/module.rs

@@ -137,7 +137,7 @@ struct ModuleInner {
     imports: Box<[ImportType]>,
     exports: Box<[ExportType]>,
     compiled: CompiledModule,
-    frame_info_registration: Mutex<Option<Option<GlobalFrameInfoRegistration>>>,
+    frame_info_registration: Mutex<Option<Option<Arc<GlobalFrameInfoRegistration>>>>,
 }
 
 impl Module {
@@ -665,11 +665,13 @@ and for re-adding support for interface types you can see this issue:
     /// 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_frame_info(&self) {
+    pub(crate) fn register_frame_info(&self) -> Option<Arc<GlobalFrameInfoRegistration>> {
         let mut info = self.inner.frame_info_registration.lock().unwrap();
-        if info.is_some() {
+        if let Some(info) = &*info {
-            return;
+            return info.clone();
         }
-        *info = Some(super::frame_info::register(&self.inner.compiled));
+        let ret = super::frame_info::register(&self.inner.compiled).map(Arc::new);
+        *info = Some(ret.clone());
+        return ret;
     }
 }

crates/api/tests/traps.rs

@@ -411,3 +411,25 @@ wasm backtrace:
     );
     Ok(())
 }
+
+#[test]
+fn present_after_module_drop() -> Result<()> {
+    let store = Store::default();
+    let module = Module::new(&store, r#"(func (export "foo") unreachable)"#)?;
+    let instance = Instance::new(&module, &[])?;
+    let func = instance.exports()[0].func().unwrap().clone();
+
+    println!("asserting before we drop modules");
+    assert_trap(func.call(&[]).unwrap_err().downcast()?);
+    drop((instance, module));
+
+    println!("asserting after drop");
+    assert_trap(func.call(&[]).unwrap_err().downcast()?);
+    return Ok(());
+
+    fn assert_trap(t: Trap) {
+        println!("{}", t);
+        assert_eq!(t.trace().len(), 1);
+        assert_eq!(t.trace()[0].func_index(), 0);
+    }
+}

crates/jit/src/instantiate.rs

@@ -7,6 +7,7 @@ use crate::compiler::Compiler;
 use crate::imports::resolve_imports;
 use crate::link::link_module;
 use crate::resolver::Resolver;
+use std::any::Any;
 use std::collections::HashMap;
 use std::io::Write;
 use std::rc::Rc;
@@ -202,6 +203,7 @@ impl CompiledModule {
         resolver: &mut dyn Resolver,
         sig_registry: &SignatureRegistry,
         mem_creator: Option<&dyn RuntimeMemoryCreator>,
+        host_state: Box<dyn Any>,
     ) -> Result<InstanceHandle, InstantiationError> {
         let data_initializers = self
             .data_initializers
@@ -222,7 +224,7 @@ impl CompiledModule {
             self.signatures.clone(),
             self.dbg_jit_registration.as_ref().map(|r| Rc::clone(&r)),
             is_bulk_memory,
-            Box::new(()),
+            host_state,
         )
     }
 
@@ -275,29 +277,3 @@ impl OwnedDataInitializer {
         }
     }
 }
-
-/// Create a new wasm instance by compiling the wasm module in `data` and instatiating it.
-///
-/// This is equivalent to creating a `CompiledModule` and calling `instantiate()` on it,
-/// but avoids creating an intermediate copy of the data initializers.
-///
-/// # Unsafety
-///
-/// See `InstanceHandle::new`
-#[allow(clippy::implicit_hasher)]
-pub unsafe fn instantiate(
-    compiler: &mut Compiler,
-    data: &[u8],
-    resolver: &mut dyn Resolver,
-    is_bulk_memory: bool,
-    profiler: &dyn ProfilingAgent,
-    mem_creator: Option<&dyn RuntimeMemoryCreator>,
-) -> Result<InstanceHandle, SetupError> {
-    let instance = CompiledModule::new(compiler, data, profiler)?.instantiate(
-        is_bulk_memory,
-        resolver,
-        compiler.signatures(),
-        mem_creator,
-    )?;
-    Ok(instance)
-}

crates/jit/src/lib.rs

@@ -34,7 +34,7 @@ pub mod trampoline;
 
 pub use crate::code_memory::CodeMemory;
 pub use crate::compiler::{make_trampoline, Compilation, CompilationStrategy, Compiler};
-pub use crate::instantiate::{instantiate, CompiledModule, SetupError};
+pub use crate::instantiate::{CompiledModule, SetupError};
 pub use crate::link::link_module;
 pub use crate::resolver::{NullResolver, Resolver};