Fix a memory leak with command modules (#2017)

This commit fixes a memory leak that can happen with `Linker::module`
when the provided module is a command. This function creates a closure
but the closure closed over a strong reference to `Store` (and
transitively through any imports provided). Unfortunately a `Store`
keeps everything alive, including `Func`, so this meant that `Store` was
inserted into a cycle which caused the leak.

The cycle here is manually broken by closing over the raw value of each
external value rather than the external value itself (which has a
strong reference to `Store`).

crates/wasmtime/src/externals.rs

@@ -160,6 +160,17 @@ impl Extern {
             Extern::Module(_) => "module",
         }
     }
+
+    pub(crate) fn wasmtime_export(&self) -> wasmtime_runtime::Export {
+        match self {
+            Extern::Func(f) => f.wasmtime_export().clone().into(),
+            Extern::Global(f) => f.wasmtime_export().clone().into(),
+            Extern::Table(f) => f.wasmtime_export().clone().into(),
+            Extern::Memory(f) => f.wasmtime_export().clone().into(),
+            Extern::Instance(f) => wasmtime_runtime::Export::Instance(f.wasmtime_export().clone()),
+            Extern::Module(f) => wasmtime_runtime::Export::Module(Box::new(f.clone())),
+        }
+    }
 }
 
 impl From<Func> for Extern {

crates/wasmtime/src/instance.rs

@@ -121,6 +121,10 @@ impl Instance {
         ty
     }
 
+    pub(crate) fn wasmtime_export(&self) -> &RuntimeInstance {
+        &self.items
+    }
+
     /// Returns the associated [`Store`] that this `Instance` is compiled into.
     ///
     /// This is the [`Store`] that generally serves as a sort of global cache

crates/wasmtime/src/linker.rs

@@ -402,11 +402,29 @@ impl Linker {
     fn command(&mut self, module_name: &str, module: &Module) -> Result<&mut Self> {
         for export in module.exports() {
             if let Some(func_ty) = export.ty().func() {
-                let imports = self.compute_imports(module)?;
+                let imports = self
-                let store = self.store.clone();
+                    .compute_imports(module)?
+                    .into_iter()
+                    .map(|e| e.wasmtime_export())
+                    .collect::<Vec<_>>();
                 let module = module.clone();
                 let export_name = export.name().to_owned();
-                let func = Func::new(&self.store, func_ty.clone(), move |_, params, results| {
+                let func = Func::new(
+                    &self.store,
+                    func_ty.clone(),
+                    move |caller, params, results| {
+                        let store = caller.store();
+
+                        // Note that the unsafety here is due to the validity of
+                        // `i` and the validity of `i` within `store`. For our
+                        // case though these items all come from `imports` above
+                        // so they're all valid. They're also all kept alive by
+                        // the store itself used here so this should be safe.
+                        let imports = imports
+                            .iter()
+                            .map(|i| unsafe { Extern::from_wasmtime_export(&i, &store) })
+                            .collect::<Vec<_>>();
+
                         // Create a new instance for this command execution.
                         let instance = Instance::new(&store, &module, &imports)?;
 
@@ -429,7 +447,8 @@ impl Linker {
                         }
 
                         Ok(())
-                });
+                    },
+                );
                 self.insert(module_name, export.name(), func.into())?;
             } else if export.name() == "memory" && export.ty().memory().is_some() {
                 // Allow an exported "memory" memory for now.

tests/all/linker.rs

@@ -1,4 +1,6 @@
 use anyhow::Result;
+use std::cell::Cell;
+use std::rc::Rc;
 use wasmtime::*;
 
 #[test]
@@ -160,3 +162,64 @@ fn module_interposition() -> Result<()> {
     assert_eq!(func()?, 112);
     Ok(())
 }
+
+#[test]
+fn no_leak() -> Result<()> {
+    struct DropMe(Rc<Cell<bool>>);
+
+    impl Drop for DropMe {
+        fn drop(&mut self) {
+            self.0.set(true);
+        }
+    }
+
+    let flag = Rc::new(Cell::new(false));
+    {
+        let store = Store::default();
+        let mut linker = Linker::new(&store);
+        let drop_me = DropMe(flag.clone());
+        linker.func("", "", move || drop(&drop_me))?;
+        let module = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (func (export "_start"))
+                )
+            "#,
+        )?;
+        linker.module("a", &module)?;
+    }
+    assert!(flag.get(), "store was leaked");
+    Ok(())
+}
+
+#[test]
+fn no_leak_with_imports() -> Result<()> {
+    struct DropMe(Rc<Cell<bool>>);
+
+    impl Drop for DropMe {
+        fn drop(&mut self) {
+            self.0.set(true);
+        }
+    }
+
+    let flag = Rc::new(Cell::new(false));
+    {
+        let store = Store::default();
+        let mut linker = Linker::new(&store);
+        let drop_me = DropMe(flag.clone());
+        linker.func("", "", move || drop(&drop_me))?;
+        let module = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (import "" "" (func))
+                    (func (export "_start"))
+                )
+            "#,
+        )?;
+        linker.module("a", &module)?;
+    }
+    assert!(flag.get(), "store was leaked");
+    Ok(())
+}