fix(vm) Fix memory leak of InstanceHandle

lib/api/src/instance.rs

@@ -4,6 +4,7 @@ use crate::module::Module;
 use crate::store::Store;
 use crate::InstantiationError;
 use std::fmt;
+use std::sync::{Arc, Mutex};
 use wasmer_engine::Resolver;
 use wasmer_vm::{InstanceHandle, VMContext};
 
@@ -17,7 +18,7 @@ use wasmer_vm::{InstanceHandle, VMContext};
 /// Spec: https://webassembly.github.io/spec/core/exec/runtime.html#module-instances
 #[derive(Clone)]
 pub struct Instance {
-    handle: InstanceHandle,
+    handle: Arc<Mutex<InstanceHandle>>,
     module: Module,
     /// The exports for an instance.
     pub exports: Exports,
@@ -30,6 +31,7 @@ mod send_test {
     fn is_send<T: Send>() -> bool {
         true
     }
+
     #[test]
     fn instance_is_send() {
         assert!(is_send::<Instance>());
@@ -72,9 +74,7 @@ impl Instance {
     ///  * Runtime errors that happen when running the module `start` function.
     pub fn new(module: &Module, resolver: &dyn Resolver) -> Result<Self, InstantiationError> {
         let store = module.store();
-
         let handle = module.instantiate(resolver)?;
-
         let exports = module
             .exports()
             .map(|export| {
@@ -86,7 +86,7 @@ impl Instance {
             .collect::<Exports>();
 
         Ok(Self {
-            handle,
+            handle: Arc::new(Mutex::new(handle)),
             module: module.clone(),
             exports,
         })
@@ -104,7 +104,7 @@ impl Instance {
 
     #[doc(hidden)]
     pub fn vmctx_ptr(&self) -> *mut VMContext {
-        self.handle.vmctx_ptr()
+        self.handle.lock().unwrap().vmctx_ptr()
     }
 }
 

lib/compiler-cranelift/src/compiler.rs

@@ -55,7 +55,7 @@ impl Compiler for CraneliftCompiler {
         &self,
         target: &Target,
         compile_info: &CompileModuleInfo,
-        module_translation: &ModuleTranslationState,
+        module_translation_state: &ModuleTranslationState,
         function_body_inputs: PrimaryMap<LocalFunctionIndex, FunctionBodyData<'_>>,
     ) -> Result<Compilation, CompileError> {
         let isa = self.config().isa(target);
@@ -115,7 +115,7 @@ impl Compiler for CraneliftCompiler {
                 // }
 
                 func_translator.translate(
-                    module_translation,
+                    module_translation_state,
                     input.data,
                     input.module_offset,
                     &mut context.func,

lib/compiler/src/translator/environ.rs

@@ -44,7 +44,7 @@ pub struct ModuleInfoTranslation<'data> {
     pub data_initializers: Vec<DataInitializer<'data>>,
 
     /// The decoded Wasm types for the module.
-    pub module_translation: Option<ModuleTranslationState>,
+    pub module_translation_state: Option<ModuleTranslationState>,
 }
 
 /// Object containing the standalone environment information.
@@ -62,7 +62,7 @@ impl<'data> ModuleEnvironment<'data> {
                 module: ModuleInfo::new(),
                 function_body_inputs: PrimaryMap::new(),
                 data_initializers: Vec::new(),
-                module_translation: None,
+                module_translation_state: None,
             },
             imports: 0,
         }
@@ -71,9 +71,9 @@ impl<'data> ModuleEnvironment<'data> {
     /// Translate a wasm module using this environment. This consumes the
     /// `ModuleEnvironment` and produces a `ModuleInfoTranslation`.
     pub fn translate(mut self, data: &'data [u8]) -> WasmResult<ModuleInfoTranslation<'data>> {
-        assert!(self.result.module_translation.is_none());
-        let module_translation = translate_module(data, &mut self)?;
-        self.result.module_translation = Some(module_translation);
+        assert!(self.result.module_translation_state.is_none());
+        let module_translation_state = translate_module(data, &mut self)?;
+        self.result.module_translation_state = Some(module_translation_state);
         Ok(self.result)
     }
 
@@ -370,7 +370,7 @@ impl<'data> ModuleEnvironment<'data> {
 
     pub(crate) fn define_function_body(
         &mut self,
-        _module_translation: &ModuleTranslationState,
+        _module_translation_state: &ModuleTranslationState,
         body_bytes: &'data [u8],
         body_offset: usize,
     ) -> WasmResult<()> {

lib/engine-jit/src/artifact.rs

@@ -83,7 +83,10 @@ impl JITArtifact {
         let compilation = compiler.compile_module(
             &jit.target(),
             &compile_info,
-            translation.module_translation.as_ref().unwrap(),
+            // SAFETY: Calling `unwrap` is correct since
+            // `environ.translate()` above will write some data into
+            // `module_translation_state`.
+            translation.module_translation_state.as_ref().unwrap(),
             translation.function_body_inputs,
         )?;
         let function_call_trampolines = compilation.get_function_call_trampolines();

lib/engine-native/src/artifact.rs

@@ -133,7 +133,7 @@ impl NativeArtifact {
             compile_info,
             translation.function_body_inputs,
             translation.data_initializers,
-            translation.module_translation,
+            translation.module_translation_state,
         ))
     }
 

lib/engine-object-file/src/artifact.rs

@@ -128,7 +128,7 @@ impl ObjectFileArtifact {
             compile_info,
             translation.function_body_inputs,
             translation.data_initializers,
-            translation.module_translation,
+            translation.module_translation_state,
         ))
     }
 

lib/vm/src/instance.rs

@@ -777,6 +777,13 @@ impl Instance {
 #[derive(Hash, PartialEq, Eq)]
 pub struct InstanceHandle {
     instance: *mut Instance,
+
+    /// Whether `Self` owns `self.instance`. It's not always the case;
+    /// e.g. when `Self` is built with `Self::from_vmctx`.
+    ///
+    /// This information is necessary to know whether `Self` should
+    /// deallocate `self.instance`.
+    owned_instance: bool,
 }
 
 /// # Safety
@@ -878,8 +885,10 @@ impl InstanceHandle {
     /// safety.
     ///
     /// However the following must be taken care of before calling this function:
-    /// - `instance_ptr` must point to valid memory sufficiently large for there
-    ///    `Instance`.
+    /// - `instance_ptr` must point to valid memory sufficiently large
+    ///    for the `Instance`. `instance_ptr` will be owned by
+    ///    `InstanceHandle`, see `InstanceHandle::owned_instance` to
+    ///    learn more.
     /// - The memory at `instance.tables_ptr()` must be initialized with data for
     ///   all the local tables.
     /// - The memory at `instance.memories_ptr()` must be initialized with data for
@@ -924,8 +933,10 @@ impl InstanceHandle {
                 vmctx: VMContext {},
             };
             ptr::write(instance_ptr, instance);
+
             Self {
                 instance: instance_ptr,
+                owned_instance: true,
             }
         };
         let instance = handle.instance();
@@ -1012,6 +1023,9 @@ impl InstanceHandle {
 
         Self {
             instance: instance as *const Instance as *mut Instance,
+
+            // We consider `vmctx` owns the `Instance`, not `Self`.
+            owned_instance: false,
         }
     }
 
@@ -1140,23 +1154,14 @@ impl InstanceHandle {
     }
 }
 
-impl Clone for InstanceHandle {
-    fn clone(&self) -> Self {
-        Self {
-            instance: self.instance,
+impl Drop for InstanceHandle {
+    fn drop(&mut self) {
+        if self.owned_instance {
+            unsafe { self.dealloc() }
         }
     }
 }
 
-// TODO: uncomment this, as we need to store the handles
-// in the store, and once the store is dropped, then the instances
-// will too.
-// impl Drop for InstanceHandle {
-//     fn drop(&mut self) {
-//         unsafe { self.dealloc() }
-//     }
-// }
-
 fn check_table_init_bounds(instance: &Instance) -> Result<(), Trap> {
     let module = Arc::clone(&instance.module);
     for init in &module.table_initializers {