Merge #1865

1865: Fix memory leak in host function envs r=MarkMcCaskey a=MarkMcCaskey

TODO: link to issue

This PR contains a number of changes:

1. Make `WasmerEnv: Clone`
2. Store a pointer to the `clone` function when creating a host function (Notably this is a feature that wouldn't work even if we _could_ use a proper trait object because you can't have a `Sized` trait object and `Clone: Sized`).
3. Store a pointer to the `drop` function when creating a host function.
4. Clone the env via pointer every time an `Instance` is made. Therefore each `Instance` gets its own, unique `Env` per host function with `Env`.
5. Add reference counting and drop logic to a sub-field of `wasmer_export::ExportFunction` which frees the original version of the `Env` (the thing that gets cloned each time an `Instance` is made) with the `drop` function pointer.
6. Change some logic in `vm::Instance` from SoA (struct of arrays) to AoS (array of structs): this uses more memory but is a bit less error prone and can be easily changed later.
7. Add logic on this new struct (`vm::ImportEnv`) that contains the function pointers for each import in `Instance` to drop (with the `drop` fn pointer) when the `vm::Instance` is being dropped. This fixes the original memory leak.
8. Add wrapper functions inside the host function creation functions which makes the layout of the user supplied env-pointer the responsibility of each function.  Thus, rather than `drop` being `Env::drop`, it's a function which frees all wrapper types, traverses indirections and frees the internal `Env` with `Env::drop`.  This simplifies code at the cost of making the `host_env` pointer (`vmctx`) not consistent in terms of what it actually points to.  This change fixes another memory leak related to the creation of host functions.

tl;dr: we're leaning into manually doing virtual method dispatch on `WasmerEnv`s and it actually works great! The biggest issue I have with the PR as-is is that the code isn't as clean/readable/robust as I'd ideally like it to be.

Edit (by @Hywan): This PR fixes #1584, #1714, #1865, #1667.

# Review

- [ ] Add a short description of the the change to the CHANGELOG.md file


Co-authored-by: Mark McCaskey <mark@wasmer.io>
Co-authored-by: Mark McCaskey <5770194+MarkMcCaskey@users.noreply.github.com>

examples/imports_function_env.rs

@@ -70,7 +70,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // possible to know the size of the `Env` at compile time (i.e it has to
     // implement the `Sized` trait) and that it implement the `WasmerEnv` trait.
     // We derive a default implementation of `WasmerEnv` here.
-    #[derive(WasmerEnv)]
+    #[derive(WasmerEnv, Clone)]
     struct Env {
         counter: Arc<RefCell<i32>>,
     }

lib/api/src/env.rs

@@ -30,12 +30,12 @@ impl From<ExportError> for HostEnvInitError {
 /// ```
 /// use wasmer::{WasmerEnv, LazyInit, Memory, NativeFunc};
 ///
-/// #[derive(WasmerEnv)]
+/// #[derive(WasmerEnv, Clone)]
 /// pub struct MyEnvWithNoInstanceData {
 ///     non_instance_data: u8,
 /// }
 ///
-/// #[derive(WasmerEnv)]
+/// #[derive(WasmerEnv, Clone)]
 /// pub struct MyEnvWithInstanceData {
 ///     non_instance_data: u8,
 ///     #[wasmer(export)]
@@ -54,6 +54,7 @@ impl From<ExportError> for HostEnvInitError {
 /// This trait can also be implemented manually:
 /// ```
 /// # use wasmer::{WasmerEnv, LazyInit, Memory, Instance, HostEnvInitError};
+/// #[derive(Clone)]
 /// pub struct MyEnv {
 ///    memory: LazyInit<Memory>,
 /// }
@@ -66,7 +67,7 @@ impl From<ExportError> for HostEnvInitError {
 ///     }
 /// }
 /// ```
-pub trait WasmerEnv {
+pub trait WasmerEnv: Clone {
     /// The function that Wasmer will call on your type to let it finish
     /// setting up the environment with data from the `Instance`.
     ///
@@ -94,26 +95,26 @@ impl WasmerEnv for isize {}
 impl WasmerEnv for char {}
 impl WasmerEnv for bool {}
 impl WasmerEnv for String {}
-impl WasmerEnv for ::std::sync::atomic::AtomicBool {}
-impl WasmerEnv for ::std::sync::atomic::AtomicI8 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicU8 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicI16 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicU16 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicI32 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicU32 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicI64 {}
-impl WasmerEnv for ::std::sync::atomic::AtomicUsize {}
-impl WasmerEnv for ::std::sync::atomic::AtomicIsize {}
-impl WasmerEnv for dyn ::std::any::Any {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicBool {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicI8 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicU8 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicI16 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicU16 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicI32 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicU32 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicI64 {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicUsize {}
+impl<'a> WasmerEnv for &'a ::std::sync::atomic::AtomicIsize {}
 impl<T: WasmerEnv> WasmerEnv for Box<T> {
     fn init_with_instance(&mut self, instance: &Instance) -> Result<(), HostEnvInitError> {
         (&mut **self).init_with_instance(instance)
     }
 }
 
-impl<T: WasmerEnv> WasmerEnv for &'static mut T {
+impl<T: WasmerEnv> WasmerEnv for ::std::sync::Arc<::std::sync::Mutex<T>> {
     fn init_with_instance(&mut self, instance: &Instance) -> Result<(), HostEnvInitError> {
-        (*self).init_with_instance(instance)
+        let mut guard = self.lock().unwrap();
+        guard.init_with_instance(instance)
     }
 }
 

lib/api/src/externals/function.rs

@@ -12,7 +12,8 @@ pub use inner::{UnsafeMutableEnv, WithUnsafeMutableEnv};
 
 use std::cmp::max;
 use std::fmt;
-use wasmer_engine::{Export, ExportFunction};
+use std::sync::Arc;
+use wasmer_engine::{Export, ExportFunction, ExportFunctionMetadata};
 use wasmer_vm::{
     raise_user_trap, resume_panic, wasmer_call_trampoline, VMCallerCheckedAnyfunc,
     VMDynamicFunctionContext, VMExportFunction, VMFunctionBody, VMFunctionEnvironment,
@@ -106,23 +107,41 @@ impl Function {
         F: Fn(&[Val]) -> Result<Vec<Val>, RuntimeError> + 'static,
     {
         let ty: FunctionType = ty.into();
-        let dynamic_ctx = VMDynamicFunctionContext::from_context(VMDynamicFunctionWithoutEnv {
-            func: Box::new(func),
-            function_type: ty.clone(),
-        });
+        let dynamic_ctx: VMDynamicFunctionContext<DynamicFunctionWithoutEnv> =
+            VMDynamicFunctionContext::from_context(DynamicFunctionWithoutEnv {
+                func: Arc::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.
         let address = std::ptr::null() as *const VMFunctionBody;
-        let vmctx = VMFunctionEnvironment {
-            host_env: Box::into_raw(Box::new(dynamic_ctx)) as *mut _,
+        let host_env = Box::into_raw(Box::new(dynamic_ctx)) as *mut _;
+        let vmctx = VMFunctionEnvironment { host_env };
+        let host_env_clone_fn: fn(*mut std::ffi::c_void) -> *mut std::ffi::c_void = |ptr| {
+            let duped_env: VMDynamicFunctionContext<DynamicFunctionWithoutEnv> = unsafe {
+                let ptr: *mut VMDynamicFunctionContext<DynamicFunctionWithoutEnv> = ptr as _;
+                let item: &VMDynamicFunctionContext<DynamicFunctionWithoutEnv> = &*ptr;
+                item.clone()
+            };
+            Box::into_raw(Box::new(duped_env)) as _
+        };
+        let host_env_drop_fn: fn(*mut std::ffi::c_void) = |ptr: *mut std::ffi::c_void| {
+            unsafe {
+                Box::from_raw(ptr as *mut VMDynamicFunctionContext<DynamicFunctionWithoutEnv>)
+            };
         };
 
         Self {
             store: store.clone(),
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: false }),
             exported: ExportFunction {
-                import_init_function_ptr: None,
+                metadata: Some(Arc::new(ExportFunctionMetadata::new(
+                    host_env,
+                    None,
+                    host_env_clone_fn,
+                    host_env_drop_fn,
+                ))),
                 vm_function: VMExportFunction {
                     address,
                     kind: VMFunctionKind::Dynamic,
@@ -187,30 +206,58 @@ impl Function {
         Env: Sized + WasmerEnv + 'static,
     {
         let ty: FunctionType = ty.into();
-        let dynamic_ctx = VMDynamicFunctionContext::from_context(VMDynamicFunctionWithEnv {
-            env: Box::new(env),
-            func: Box::new(func),
-            function_type: ty.clone(),
-        });
+        let dynamic_ctx: VMDynamicFunctionContext<DynamicFunctionWithEnv<Env>> =
+            VMDynamicFunctionContext::from_context(DynamicFunctionWithEnv {
+                env: Box::new(env),
+                func: Arc::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.
         let address = std::ptr::null() as *const VMFunctionBody;
-        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 host_env = Box::into_raw(Box::new(dynamic_ctx)) as *mut _;
+        let vmctx = VMFunctionEnvironment { host_env };
+        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<DynamicFunctionWithEnv<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_clone_fn: fn(*mut std::ffi::c_void) -> *mut std::ffi::c_void = |ptr| {
+            let duped_env: VMDynamicFunctionContext<DynamicFunctionWithEnv<Env>> = unsafe {
+                let ptr: *mut VMDynamicFunctionContext<DynamicFunctionWithEnv<Env>> = ptr as _;
+                let item: &VMDynamicFunctionContext<DynamicFunctionWithEnv<Env>> = &*ptr;
+                item.clone()
+            };
+            Box::into_raw(Box::new(duped_env)) as _
+        };
+        let host_env_drop_fn: fn(*mut std::ffi::c_void) = |ptr: *mut std::ffi::c_void| {
+            unsafe {
+                Box::from_raw(ptr as *mut VMDynamicFunctionContext<DynamicFunctionWithEnv<Env>>)
+            };
+        };
 
         Self {
             store: store.clone(),
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: true }),
             exported: ExportFunction {
-                import_init_function_ptr,
+                metadata: Some(Arc::new(ExportFunctionMetadata::new(
+                    host_env,
+                    import_init_function_ptr,
+                    host_env_clone_fn,
+                    host_env_drop_fn,
+                ))),
                 vm_function: VMExportFunction {
                     address,
                     kind: VMFunctionKind::Dynamic,
@@ -264,7 +311,7 @@ impl Function {
             exported: ExportFunction {
                 // 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,
+                metadata: None,
                 vm_function: VMExportFunction {
                     address,
                     vmctx,
@@ -318,10 +365,20 @@ impl Function {
         // Wasm-defined functions have a `VMContext`.
         // In the case of Host-defined functions `VMContext` is whatever environment
         // the user want to attach to the function.
-        let box_env = Box::new(env);
-        let vmctx = VMFunctionEnvironment {
-            host_env: Box::into_raw(box_env) as *mut _,
+        let host_env = Box::into_raw(Box::new(env)) as *mut _;
+        let vmctx = VMFunctionEnvironment { host_env };
+        let host_env_clone_fn: fn(*mut std::ffi::c_void) -> *mut std::ffi::c_void = |ptr| {
+            let duped_env = unsafe {
+                let ptr: *mut Env = ptr as _;
+                let item: &Env = &*ptr;
+                item.clone()
+            };
+            Box::into_raw(Box::new(duped_env)) as _
         };
+        let host_env_drop_fn: fn(*mut std::ffi::c_void) = |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<(), _>>(
@@ -334,7 +391,12 @@ impl Function {
             store: store.clone(),
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: true }),
             exported: ExportFunction {
-                import_init_function_ptr,
+                metadata: Some(Arc::new(ExportFunctionMetadata::new(
+                    host_env,
+                    import_init_function_ptr,
+                    host_env_clone_fn,
+                    host_env_drop_fn,
+                ))),
                 vm_function: VMExportFunction {
                     address,
                     kind: VMFunctionKind::Static,
@@ -373,9 +435,20 @@ impl Function {
         let address = function.address();
 
         let box_env = Box::new(env);
-        let vmctx = VMFunctionEnvironment {
-            host_env: Box::into_raw(box_env) as *mut _,
+        let host_env = Box::into_raw(box_env) as *mut _;
+        let vmctx = VMFunctionEnvironment { host_env };
+        let host_env_clone_fn: fn(*mut std::ffi::c_void) -> *mut std::ffi::c_void = |ptr| {
+            let duped_env: Env = {
+                let ptr: *mut Env = ptr as _;
+                let item: &Env = &*ptr;
+                item.clone()
+            };
+            Box::into_raw(Box::new(duped_env)) as _
         };
+        let host_env_drop_fn: fn(*mut std::ffi::c_void) = |ptr: *mut std::ffi::c_void| {
+            Box::from_raw(ptr as *mut Env);
+        };
+
         let signature = function.ty();
         // TODO: look into removing transmute by changing API type signatures
         let import_init_function_ptr = Some(std::mem::transmute::<
@@ -387,7 +460,12 @@ impl Function {
             store: store.clone(),
             definition: FunctionDefinition::Host(HostFunctionDefinition { has_env: true }),
             exported: ExportFunction {
-                import_init_function_ptr,
+                metadata: Some(Arc::new(ExportFunctionMetadata {
+                    host_env,
+                    host_env_clone_fn,
+                    host_env_drop_fn,
+                    import_init_function_ptr,
+                })),
                 vm_function: VMExportFunction {
                     address,
                     kind: VMFunctionKind::Static,
@@ -756,13 +834,14 @@ pub(crate) trait VMDynamicFunction {
     fn function_type(&self) -> &FunctionType;
 }
 
-pub(crate) struct VMDynamicFunctionWithoutEnv {
+#[derive(Clone)]
+pub(crate) struct DynamicFunctionWithoutEnv {
     #[allow(clippy::type_complexity)]
-    func: Box<dyn Fn(&[Val]) -> Result<Vec<Val>, RuntimeError> + 'static>,
+    func: Arc<dyn Fn(&[Val]) -> Result<Vec<Val>, RuntimeError> + 'static>,
     function_type: FunctionType,
 }
 
-impl VMDynamicFunction for VMDynamicFunctionWithoutEnv {
+impl VMDynamicFunction for DynamicFunctionWithoutEnv {
     fn call(&self, args: &[Val]) -> Result<Vec<Val>, RuntimeError> {
         (*self.func)(&args)
     }
@@ -771,19 +850,27 @@ impl VMDynamicFunction for VMDynamicFunctionWithoutEnv {
     }
 }
 
-#[repr(C)]
-pub(crate) struct VMDynamicFunctionWithEnv<Env>
+pub(crate) struct DynamicFunctionWithEnv<Env>
 where
     Env: Sized + 'static,
 {
-    // This field _must_ come first in this struct.
-    env: Box<Env>,
     function_type: FunctionType,
     #[allow(clippy::type_complexity)]
-    func: Box<dyn Fn(&Env, &[Val]) -> Result<Vec<Val>, RuntimeError> + 'static>,
+    func: Arc<dyn Fn(&Env, &[Val]) -> Result<Vec<Val>, RuntimeError> + 'static>,
+    env: Box<Env>,
+}
+
+impl<Env: Sized + Clone + 'static> Clone for DynamicFunctionWithEnv<Env> {
+    fn clone(&self) -> Self {
+        Self {
+            env: self.env.clone(),
+            function_type: self.function_type.clone(),
+            func: self.func.clone(),
+        }
+    }
 }
 
-impl<Env> VMDynamicFunction for VMDynamicFunctionWithEnv<Env>
+impl<Env> VMDynamicFunction for DynamicFunctionWithEnv<Env>
 where
     Env: Sized + 'static,
 {

lib/api/src/native.rs

@@ -10,8 +10,8 @@
 use std::marker::PhantomData;
 
 use crate::externals::function::{
-    FunctionDefinition, HostFunctionDefinition, VMDynamicFunction, VMDynamicFunctionWithEnv,
-    VMDynamicFunctionWithoutEnv, WasmFunctionDefinition,
+    DynamicFunctionWithEnv, DynamicFunctionWithoutEnv, FunctionDefinition, HostFunctionDefinition,
+    VMDynamicFunction, WasmFunctionDefinition,
 };
 use crate::{FromToNativeWasmType, Function, RuntimeError, Store, WasmTypeList};
 use std::panic::{catch_unwind, AssertUnwindSafe};
@@ -21,6 +21,7 @@ use wasmer_vm::{VMDynamicFunctionContext, VMFunctionBody, VMFunctionEnvironment,
 
 /// A WebAssembly function that can be called natively
 /// (using the Native ABI).
+#[derive(Clone)]
 pub struct NativeFunc<Args = (), Rets = ()> {
     definition: FunctionDefinition,
     store: Store,
@@ -183,13 +184,13 @@ macro_rules! impl_native_traits {
                             VMFunctionKind::Dynamic => {
                                 let params_list = [ $( $x.to_native().to_value() ),* ];
                                 let results = if !has_env {
-                                    type VMContextWithoutEnv = VMDynamicFunctionContext<VMDynamicFunctionWithoutEnv>;
+                                    type VMContextWithoutEnv = VMDynamicFunctionContext<DynamicFunctionWithoutEnv>;
                                     unsafe {
                                         let ctx = self.vmctx().host_env as *mut VMContextWithoutEnv;
                                         (*ctx).ctx.call(&params_list)?
                                     }
                                 } else {
-                                    type VMContextWithEnv = VMDynamicFunctionContext<VMDynamicFunctionWithEnv<std::ffi::c_void>>;
+                                    type VMContextWithEnv = VMDynamicFunctionContext<DynamicFunctionWithEnv<std::ffi::c_void>>;
                                     unsafe {
                                         let ctx = self.vmctx().host_env as *mut VMContextWithEnv;
                                         (*ctx).ctx.call(&params_list)?

lib/api/src/types.rs

@@ -76,7 +76,7 @@ impl ValFuncRef for Val {
         let export = wasmer_engine::ExportFunction {
             // TODO:
             // figure out if we ever need a value here: need testing with complicated import patterns
-            import_init_function_ptr: None,
+            metadata: None,
             vm_function: wasmer_vm::VMExportFunction {
                 address: item.func_ptr,
                 signature,