Fix memory leak in host function envs

lib/api/src/externals/function.rs

@@ -12,6 +12,7 @@ pub use inner::{UnsafeMutableEnv, WithUnsafeMutableEnv};
 
 use std::cmp::max;
 use std::fmt;
+use std::sync::Arc;
 use wasmer_engine::{Export, ExportFunction};
 use wasmer_vm::{
     raise_user_trap, resume_panic, wasmer_call_trampoline, VMCallerCheckedAnyfunc,
@@ -87,10 +88,11 @@ impl Function {
     where
         F: Fn(&[Val]) -> Result<Vec<Val>, RuntimeError> + 'static,
     {
-        let dynamic_ctx = VMDynamicFunctionContext::from_context(VMDynamicFunctionWithoutEnv {
+        let dynamic_ctx: VMDynamicFunctionContext<VMDynamicFunctionWithoutEnv> =
-            func: Box::new(func),
+            VMDynamicFunctionContext::from_context(VMDynamicFunctionWithoutEnv {
-            function_type: ty.clone(),
+                func: Box::new(func),
-        });
+                function_type: ty.clone(),
+            });
         // We don't yet have the address with the Wasm ABI signature.
         // The engine linker will replace the address with one pointing to a
         // generated dynamic trampoline.
@@ -98,20 +100,27 @@ impl Function {
         let vmctx = VMFunctionEnvironment {
             host_env: Box::into_raw(Box::new(dynamic_ctx)) as *mut _,
         };
+        let host_env_drop_fn: Option<fn(*mut std::ffi::c_void)> =
+            Some(|ptr: *mut std::ffi::c_void| {
+                unsafe {
+                    Box::from_raw(ptr as *mut VMDynamicFunctionContext<VMDynamicFunctionWithoutEnv>)
+                };
+            });
 
         Self {
             store: store.clone(),
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: false }),
             exported: ExportFunction {
                 import_init_function_ptr: None,
-                vm_function: VMExportFunction {
+                host_env_drop_fn,
+                vm_function: Arc::new(VMExportFunction {
                     address,
                     kind: VMFunctionKind::Dynamic,
                     vmctx,
                     signature: ty.clone(),
                     call_trampoline: None,
                     instance_allocator: None,
-                },
+                }),
             },
         }
     }
@@ -143,11 +152,12 @@ impl Function {
         F: Fn(&Env, &[Val]) -> Result<Vec<Val>, RuntimeError> + 'static,
         Env: Sized + WasmerEnv + 'static,
     {
-        let dynamic_ctx = VMDynamicFunctionContext::from_context(VMDynamicFunctionWithEnv {
+        let dynamic_ctx: VMDynamicFunctionContext<VMDynamicFunctionWithEnv<Env>> =
-            env: Box::new(env),
+            VMDynamicFunctionContext::from_context(VMDynamicFunctionWithEnv {
-            func: Box::new(func),
+                env: Box::new(env),
-            function_type: ty.clone(),
+                func: Box::new(func),
-        });
+                function_type: ty.clone(),
+            });
         // We don't yet have the address with the Wasm ABI signature.
         // The engine linker will replace the address with one pointing to a
         // generated dynamic trampoline.
@@ -155,26 +165,44 @@ impl Function {
         let vmctx = VMFunctionEnvironment {
             host_env: Box::into_raw(Box::new(dynamic_ctx)) as *mut _,
         };
-        // TODO: look into removing transmute by changing API type signatures
+        let import_init_function_ptr: fn(_, _) -> Result<(), _> =
+            |ptr: *mut std::ffi::c_void, instance: *const std::ffi::c_void| {
+                let ptr = ptr as *mut VMDynamicFunctionContext<VMDynamicFunctionWithEnv<Env>>;
+                unsafe {
+                    let env = &mut *ptr;
+                    let env: &mut Env = &mut *env.ctx.env;
+                    let instance = &*(instance as *const crate::Instance);
+                    Env::init_with_instance(env, instance)
+                }
+            };
         let import_init_function_ptr = Some(unsafe {
             std::mem::transmute::<fn(_, _) -> Result<(), _>, fn(_, _) -> Result<(), _>>(
-                Env::init_with_instance,
+                import_init_function_ptr,
             )
         });
+        let host_env_drop_fn: Option<fn(*mut std::ffi::c_void)> =
+            Some(|ptr: *mut std::ffi::c_void| {
+                unsafe {
+                    Box::from_raw(
+                        ptr as *mut VMDynamicFunctionContext<VMDynamicFunctionWithEnv<Env>>,
+                    )
+                };
+            });
 
         Self {
             store: store.clone(),
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: true }),
             exported: ExportFunction {
                 import_init_function_ptr,
-                vm_function: VMExportFunction {
+                host_env_drop_fn,
+                vm_function: Arc::new(VMExportFunction {
                     address,
                     kind: VMFunctionKind::Dynamic,
                     vmctx,
                     signature: ty.clone(),
                     call_trampoline: None,
                     instance_allocator: None,
-                },
+                }),
             },
         }
     }
@@ -221,14 +249,15 @@ impl Function {
                 // TODO: figure out what's going on in this function: it takes an `Env`
                 // param but also marks itself as not having an env
                 import_init_function_ptr: None,
-                vm_function: VMExportFunction {
+                host_env_drop_fn: None,
+                vm_function: Arc::new(VMExportFunction {
                     address,
                     vmctx,
                     signature,
                     kind: VMFunctionKind::Static,
                     call_trampoline: None,
                     instance_allocator: None,
-                },
+                }),
             },
         }
     }
@@ -278,6 +307,11 @@ impl Function {
         let vmctx = VMFunctionEnvironment {
             host_env: Box::into_raw(box_env) as *mut _,
         };
+        let host_env_drop_fn: Option<fn(*mut std::ffi::c_void)> =
+            Some(|ptr: *mut std::ffi::c_void| {
+                unsafe { Box::from_raw(ptr as *mut Env) };
+            });
+
         // TODO: look into removing transmute by changing API type signatures
         let import_init_function_ptr = Some(unsafe {
             std::mem::transmute::<fn(_, _) -> Result<(), _>, fn(_, _) -> Result<(), _>>(
@@ -291,14 +325,15 @@ impl Function {
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: true }),
             exported: ExportFunction {
                 import_init_function_ptr,
-                vm_function: VMExportFunction {
+                host_env_drop_fn,
+                vm_function: Arc::new(VMExportFunction {
                     address,
                     kind: VMFunctionKind::Static,
                     vmctx,
                     signature,
                     call_trampoline: None,
                     instance_allocator: None,
-                },
+                }),
             },
         }
     }
@@ -344,14 +379,14 @@ impl Function {
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: true }),
             exported: ExportFunction {
                 import_init_function_ptr,
-                vm_function: VMExportFunction {
+                vm_function: Arc::new(VMExportFunction {
                     address,
                     kind: VMFunctionKind::Static,
                     vmctx,
                     signature,
                     call_trampoline: None,
                     instance_allocator: None,
-                },
+                }),
             },
         }
     }

lib/api/src/native.rs

@@ -15,6 +15,7 @@ use crate::externals::function::{
 };
 use crate::{FromToNativeWasmType, Function, FunctionType, RuntimeError, Store, WasmTypeList};
 use std::panic::{catch_unwind, AssertUnwindSafe};
+use std::sync::Arc;
 use wasmer_engine::ExportFunction;
 use wasmer_types::NativeWasmType;
 use wasmer_vm::{
@@ -88,14 +89,15 @@ where
         Self {
             // TODO:
             import_init_function_ptr: None,
-            vm_function: VMExportFunction {
+            host_env_drop_fn: None,
+            vm_function: Arc::new(VMExportFunction {
                 address: other.address,
                 vmctx: other.vmctx,
                 signature,
                 kind: other.arg_kind,
                 call_trampoline: None,
                 instance_allocator: None,
-            },
+            }),
         }
     }
 }
@@ -113,14 +115,15 @@ where
             exported: ExportFunction {
                 // TODO:
                 import_init_function_ptr: None,
-                vm_function: VMExportFunction {
+                host_env_drop_fn: None,
+                vm_function: Arc::new(VMExportFunction {
                     address: other.address,
                     vmctx: other.vmctx,
                     signature,
                     kind: other.arg_kind,
                     call_trampoline: None,
                     instance_allocator: None,
-                },
+                }),
             },
         }
     }

lib/api/src/types.rs

@@ -77,7 +77,8 @@ impl ValFuncRef for Val {
             // TODO:
             // figure out if we ever need a value here: need testing with complicated import patterns
             import_init_function_ptr: None,
-            vm_function: wasmer_vm::VMExportFunction {
+            host_env_drop_fn: None,
+            vm_function: std::sync::Arc::new(wasmer_vm::VMExportFunction {
                 address: item.func_ptr,
                 signature,
                 // All functions in tables are already Static (as dynamic functions
@@ -86,7 +87,7 @@ impl ValFuncRef for Val {
                 vmctx: item.vmctx,
                 call_trampoline: None,
                 instance_allocator: None,
-            },
+            }),
         };
         let f = Function::from_vm_export(store, export);
         Self::FuncRef(f)

lib/engine/src/artifact.rs

@@ -96,7 +96,7 @@ pub trait Artifact: Send + Sync + Upcastable {
 
         let module = self.module();
         let (instance_ptr, offsets) = InstanceHandle::allocate_instance(&module);
-        let (imports, import_initializers) = {
+        let (imports, import_envs) = {
             let mut imports = resolve_imports(
                 &module,
                 resolver,
@@ -108,9 +108,9 @@ pub trait Artifact: Send + Sync + Upcastable {
 
             // Get the `WasmerEnv::init_with_instance` function pointers and the pointers
             // to the envs to call it on.
-            let import_initializers: Vec<(_, _)> = imports.get_import_initializers();
+            let import_envs = imports.get_import_initializers();
 
-            (imports, import_initializers)
+            (imports, import_envs)
         };
 
         // Get pointers to where metadata about local memories should live in VM memory.
@@ -146,7 +146,7 @@ pub trait Artifact: Send + Sync + Upcastable {
             imports,
             self.signatures().clone(),
             host_state,
-            import_initializers,
+            import_envs,
         )
         .map_err(|trap| InstantiationError::Start(RuntimeError::from_trap(trap)))?;
         Ok(handle)

lib/engine/src/export.rs

@@ -3,6 +3,8 @@ use wasmer_vm::{
     VMExportTable,
 };
 
+use std::sync::Arc;
+
 /// The value of an export passed from one instance to another.
 #[derive(Debug, Clone)]
 pub enum Export {
@@ -36,6 +38,7 @@ impl From<VMExport> for Export {
             VMExport::Function(vm_function) => Export::Function(ExportFunction {
                 vm_function,
                 import_init_function_ptr: None,
+                host_env_drop_fn: None,
             }),
             VMExport::Memory(vm_memory) => Export::Memory(ExportMemory { vm_memory }),
             VMExport::Table(vm_table) => Export::Table(ExportTable { vm_table }),
@@ -49,12 +52,14 @@ impl From<VMExport> for Export {
 #[derive(Debug, Clone, PartialEq)]
 pub struct ExportFunction {
     /// The VM function, containing most of the data.
-    pub vm_function: VMExportFunction,
+    pub vm_function: Arc<VMExportFunction>,
     /// Function pointer to `WasmerEnv::init_with_instance(&mut self, instance: &Instance)`.
     ///
     /// This function is called to finish setting up the environment after
     /// we create the `api::Instance`.
     pub import_init_function_ptr: Option<ImportInitializerFuncPtr>,
+    /// The destructor to free the host environment.
+    pub host_env_drop_fn: Option<fn(*mut std::ffi::c_void)>,
 }
 
 impl From<ExportFunction> for Export {

lib/engine/src/resolver.rs

@@ -3,13 +3,14 @@
 
 use crate::{Export, ImportError, LinkError};
 use more_asserts::assert_ge;
+use std::sync::Arc;
 use wasmer_types::entity::{BoxedSlice, EntityRef, PrimaryMap};
 use wasmer_types::{ExternType, FunctionIndex, ImportIndex, MemoryIndex, TableIndex};
 
 use wasmer_vm::{
-    FunctionBodyPtr, Imports, MemoryStyle, ModuleInfo, TableStyle, VMDynamicFunctionContext,
+    FunctionBodyPtr, ImportEnv, Imports, MemoryStyle, ModuleInfo, TableStyle,
-    VMFunctionBody, VMFunctionImport, VMFunctionKind, VMGlobalImport, VMMemoryImport,
+    VMDynamicFunctionContext, VMFunctionBody, VMFunctionImport, VMFunctionKind, VMGlobalImport,
-    VMTableImport,
+    VMMemoryImport, VMTableImport,
 };
 
 /// Import resolver connects imports with available exported values.
@@ -155,36 +156,13 @@ pub fn resolve_imports(
         }
         match resolved {
             Export::Function(ref f) => {
-                let (address, env_ptr) = match f.vm_function.kind {
+                let address = match f.vm_function.kind {
                     VMFunctionKind::Dynamic => {
                         // If this is a dynamic imported function,
                         // the address of the function is the address of the
                         // reverse trampoline.
                         let index = FunctionIndex::new(function_imports.len());
-                        let address = finished_dynamic_function_trampolines[index].0
+                        finished_dynamic_function_trampolines[index].0 as *mut VMFunctionBody as _
-                            as *mut VMFunctionBody as _;
-                        let env_ptr = if f.import_init_function_ptr.is_some() {
-                            // Our function env looks like:
-                            // Box<VMDynamicFunctionContext<VMDynamicFunctionWithEnv<Env>>>
-                            // Which we can interpret as `*const <field offset> *const Env` (due to
-                            // the precise layout of these types via `repr(C)`)
-                            // We extract the `*const Env`:
-                            unsafe {
-                                // Box<VMDynamicFunctionContext<...>>
-                                let dyn_func_ctx_ptr = f.vm_function.vmctx.host_env
-                                    as *mut VMDynamicFunctionContext<*mut std::ffi::c_void>;
-                                // maybe report error here if it's null?
-                                // invariants of these types are not enforced.
-
-                                // &VMDynamicFunctionContext<...>
-                                let dyn_func_ctx = &*dyn_func_ctx_ptr;
-                                dyn_func_ctx.ctx
-                            }
-                        } else {
-                            std::ptr::null_mut()
-                        };
-
-                        (address, env_ptr)
 
                         // TODO: We should check that the f.vmctx actually matches
                         // the shape of `VMDynamicFunctionImportContext`
@@ -198,9 +176,7 @@ pub fn resolve_imports(
                             assert!(num_params < 9, "Only native functions with less than 9 arguments are allowed in Apple Silicon (for now). Received {} in the import {}.{}", num_params, module_name, field);
                         }
 
-                        (f.vm_function.address, unsafe {
+                        f.vm_function.address
-                            f.vm_function.vmctx.host_env
-                        })
                     }
                 };
                 function_imports.push(VMFunctionImport {
@@ -208,7 +184,11 @@ pub fn resolve_imports(
                     environment: f.vm_function.vmctx,
                 });
 
-                host_function_env_initializers.push((f.import_init_function_ptr, env_ptr));
+                host_function_env_initializers.push(Arc::new(ImportEnv {
+                    env: unsafe { f.vm_function.vmctx.host_env },
+                    initializer: f.import_init_function_ptr,
+                    destructor: f.host_env_drop_fn,
+                }));
             }
             Export::Table(ref t) => {
                 table_imports.push(VMTableImport {

lib/vm/src/export.rs

@@ -10,10 +10,10 @@ use std::sync::Arc;
 use wasmer_types::{FunctionType, MemoryType, TableType};
 
 /// The value of an export passed from one instance to another.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub enum VMExport {
     /// A function export value.
-    Function(VMExportFunction),
+    Function(Arc<VMExportFunction>),
 
     /// A table export value.
     Table(VMExportTable),
@@ -26,7 +26,7 @@ pub enum VMExport {
 }
 
 /// A function export value.
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, PartialEq)]
 pub struct VMExportFunction {
     /// The address of the native-code function.
     pub address: *const VMFunctionBody,
@@ -63,7 +63,7 @@ unsafe impl Sync for VMExportFunction {}
 
 impl From<VMExportFunction> for VMExport {
     fn from(func: VMExportFunction) -> Self {
-        Self::Function(func)
+        Self::Function(Arc::new(func))
     }
 }
 

lib/vm/src/imports.rs

@@ -1,8 +1,9 @@
 // This file contains code from external sources.
 // Attributions: https://github.com/wasmerio/wasmer/blob/master/ATTRIBUTIONS.md
 
-use crate::instance::{ImportInitializerFuncPtr, ImportInitializerThunks};
+use crate::instance::ImportEnv;
 use crate::vmcontext::{VMFunctionImport, VMGlobalImport, VMMemoryImport, VMTableImport};
+use std::sync::Arc;
 use wasmer_types::entity::{BoxedSlice, PrimaryMap};
 use wasmer_types::{FunctionIndex, GlobalIndex, MemoryIndex, TableIndex};
 
@@ -17,9 +18,7 @@ pub struct Imports {
     /// space may affect Wasm runtime performance due to increased cache pressure.
     ///
     /// We make it optional so that we can free the data after use.
-    pub host_function_env_initializers: Option<
+    pub host_function_env_initializers: Option<BoxedSlice<FunctionIndex, Arc<ImportEnv>>>,
-        BoxedSlice<FunctionIndex, (Option<ImportInitializerFuncPtr>, *mut std::ffi::c_void)>,
-    >,
 
     /// Resolved addresses for imported tables.
     pub tables: BoxedSlice<TableIndex, VMTableImport>,
@@ -35,10 +34,7 @@ impl Imports {
     /// Construct a new `Imports` instance.
     pub fn new(
         function_imports: PrimaryMap<FunctionIndex, VMFunctionImport>,
-        host_function_env_initializers: PrimaryMap<
+        host_function_env_initializers: PrimaryMap<FunctionIndex, Arc<ImportEnv>>,
-            FunctionIndex,
-            (Option<ImportInitializerFuncPtr>, *mut std::ffi::c_void),
-        >,
         table_imports: PrimaryMap<TableIndex, VMTableImport>,
         memory_imports: PrimaryMap<MemoryIndex, VMMemoryImport>,
         global_imports: PrimaryMap<GlobalIndex, VMGlobalImport>,
@@ -56,7 +52,7 @@ impl Imports {
     pub fn none() -> Self {
         Self {
             functions: PrimaryMap::new().into_boxed_slice(),
-            host_function_env_initializers: Some(PrimaryMap::new().into_boxed_slice()),
+            host_function_env_initializers: None,
             tables: PrimaryMap::new().into_boxed_slice(),
             memories: PrimaryMap::new().into_boxed_slice(),
             globals: PrimaryMap::new().into_boxed_slice(),
@@ -68,25 +64,9 @@ impl Imports {
     ///
     /// This function can only be called once, it deletes the data it returns after
     /// returning it to ensure that it's not called more than once.
-    pub fn get_import_initializers(&mut self) -> ImportInitializerThunks {
+    pub fn get_import_initializers(&mut self) -> BoxedSlice<FunctionIndex, Arc<ImportEnv>> {
-        let result = if let Some(inner) = &self.host_function_env_initializers {
+        self.host_function_env_initializers
-            inner
+            .take()
-                .values()
+            .unwrap_or_else(|| PrimaryMap::new().into_boxed_slice())
-                .cloned()
-                .map(|(func_init, env_ptr)| {
-                    let host_env = if func_init.is_some() {
-                        env_ptr
-                    } else {
-                        std::ptr::null_mut()
-                    };
-                    (func_init, host_env)
-                })
-                .collect()
-        } else {
-            vec![]
-        };
-        // ensure we only call these functions once and free this now useless memory.
-        self.host_function_env_initializers = None;
-        result
     }
 }

lib/vm/src/instance.rs

@@ -100,7 +100,8 @@ pub(crate) struct Instance {
     ///
     /// TODO: Be sure to test with serialize/deserialize and imported
     /// functions from other Wasm modules.
-    import_initializers: ImportInitializerThunks,
+    /// TODO: update this comment
+    import_envs: BoxedSlice<FunctionIndex, Arc<ImportEnv>>,
 
     /// Additional context used by compiled WebAssembly code. This
     /// field is last, and represents a dynamically-sized array that
@@ -109,6 +110,27 @@ pub(crate) struct Instance {
     vmctx: VMContext,
 }
 
+/// TODO: figure out what to do with Clone here
+///       We can probably remove the Arc... just need to figure that out.
+/// TODO:
+#[derive(Debug)]
+pub struct ImportEnv {
+    /// TODO:
+    pub env: *mut std::ffi::c_void,
+    /// TODO:
+    pub initializer: Option<ImportInitializerFuncPtr>,
+    /// TODO:
+    pub destructor: Option<fn(*mut std::ffi::c_void)>,
+}
+
+impl Drop for ImportEnv {
+    fn drop(&mut self) {
+        if let Some(destructor) = self.destructor {
+            (destructor)(self.env);
+        }
+    }
+}
+
 impl fmt::Debug for Instance {
     fn fmt(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
         formatter.debug_struct("Instance").finish()
@@ -155,7 +177,7 @@ impl Instance {
         &self,
         index: FunctionIndex,
     ) -> Option<ImportInitializerFuncPtr> {
-        self.import_initializers[index.as_u32() as usize].0
+        self.import_envs[index].initializer
     }
 
     /// Return a pointer to the `VMFunctionImport`s.
@@ -1007,7 +1029,7 @@ impl InstanceHandle {
         imports: Imports,
         vmshared_signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
         host_state: Box<dyn Any>,
-        import_initializers: ImportInitializerThunks,
+        import_envs: BoxedSlice<FunctionIndex, Arc<ImportEnv>>,
     ) -> Result<Self, Trap> {
         // `NonNull<u8>` here actually means `NonNull<Instance>`. See
         // `Self::allocate_instance` to understand why.
@@ -1035,7 +1057,7 @@ impl InstanceHandle {
                 passive_data,
                 host_state,
                 signal_handler: Cell::new(None),
-                import_initializers,
+                import_envs,
                 vmctx: VMContext {},
             };
 
@@ -1411,18 +1433,20 @@ impl InstanceHandle {
     ) -> Result<(), Err> {
         let instance_ref = self.instance.as_mut();
 
-        for (func, env) in instance_ref.import_initializers.drain(..) {
+        for ImportEnv {
-            if let Some(ref f) = func {
+            env, initializer, ..
+        } in instance_ref.import_envs.values().map(|v| &**v)
+        {
+            if let Some(ref f) = initializer {
                 // transmute our function pointer into one with the correct error type
                 let f = mem::transmute::<
                     &ImportInitializerFuncPtr,
                     &fn(*mut ffi::c_void, *const ffi::c_void) -> Result<(), Err>,
                 >(f);
-                f(env, instance_ptr)?;
+                f(*env, instance_ptr)?;
+                //*initializer = None;
             }
         }
-        // free memory now that it's empty.
-        instance_ref.import_initializers.shrink_to_fit();
 
         Ok(())
     }

lib/vm/src/lib.rs

@@ -39,7 +39,7 @@ pub mod libcalls;
 pub use crate::export::*;
 pub use crate::global::*;
 pub use crate::imports::Imports;
-pub use crate::instance::{ImportInitializerFuncPtr, InstanceHandle};
+pub use crate::instance::{ImportEnv, ImportInitializerFuncPtr, InstanceHandle};
 pub use crate::memory::{LinearMemory, Memory, MemoryError, MemoryStyle};
 pub use crate::mmap::Mmap;
 pub use crate::module::{ExportsIterator, ImportsIterator, ModuleInfo};