From ae9b4c4123dfc013bf87b4401d976115ae1d779d Mon Sep 17 00:00:00 2001
From: Peter Huene <peter@huene.dev>
Date: Mon, 14 Dec 2020 17:12:54 -0800
Subject: [PATCH] Ensure default allocator is used for instance deallocation.

Handles created with `create_handle` need to be deallocated with the default
(on-demand) instance allocator.

This commit changes Store such that handles can be added with a flag that is
used to force deallocation via the default instance allocator when the Store is
dropped.
---
 crates/wasmtime/src/instance.rs               |  2 +-
 crates/wasmtime/src/store.rs                  | 37 +++++++++++++++----
 .../wasmtime/src/trampoline/create_handle.rs  |  2 +-
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/crates/wasmtime/src/instance.rs b/crates/wasmtime/src/instance.rs
index 146541d2aa7f..fb5f1abbadf5 100644
--- a/crates/wasmtime/src/instance.rs
+++ b/crates/wasmtime/src/instance.rs
@@ -460,7 +460,7 @@ impl<'a> Instantiator<'a> {
             // initializers may have run which placed elements into other instance's
             // tables. This means that from this point on, regardless of whether
             // initialization is successful, we need to keep the instance alive.
-            let instance = self.store.add_instance(instance);
+            let instance = self.store.add_instance(instance, false);
             allocator
                 .initialize(
                     &instance.handle,
diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs
index 6a9621880f50..ee0c89b67ff0 100644
--- a/crates/wasmtime/src/store.rs
+++ b/crates/wasmtime/src/store.rs
@@ -13,10 +13,19 @@ use std::sync::Arc;
 use wasmtime_environ::wasm;
 use wasmtime_jit::{CompiledModule, ModuleCode, TypeTables};
 use wasmtime_runtime::{
-    InstanceHandle, SignalHandler, StackMapRegistry, TrapInfo, VMContext, VMExternRef,
-    VMExternRefActivationsTable, VMInterrupts, VMSharedSignatureIndex,
+    InstanceAllocator, InstanceHandle, SignalHandler, StackMapRegistry, TrapInfo, VMContext,
+    VMExternRef, VMExternRefActivationsTable, VMInterrupts, VMSharedSignatureIndex,
 };
 
+/// Used to associate instances with the store.
+///
+/// This is needed to track if the instance was allocated expliclty with the default
+/// instance allocator.
+struct StoreInstance {
+    handle: InstanceHandle,
+    use_default_allocator: bool,
+}
+
 /// A `Store` is a collection of WebAssembly instances and host-defined items.
 ///
 /// All WebAssembly instances and items will be attached to and refer to a
@@ -57,7 +66,7 @@ pub(crate) struct StoreInner {
     engine: Engine,
     interrupts: Arc<VMInterrupts>,
     signatures: RefCell<SignatureRegistry>,
-    instances: RefCell<Vec<InstanceHandle>>,
+    instances: RefCell<Vec<StoreInstance>>,
     signal_handler: RefCell<Option<Box<SignalHandler<'static>>>>,
     externref_activations_table: VMExternRefActivationsTable,
     stack_map_registry: StackMapRegistry,
@@ -216,8 +225,15 @@ impl Store {
         Ok(())
     }
 
-    pub(crate) unsafe fn add_instance(&self, handle: InstanceHandle) -> StoreInstanceHandle {
-        self.inner.instances.borrow_mut().push(handle.clone());
+    pub(crate) unsafe fn add_instance(
+        &self,
+        handle: InstanceHandle,
+        use_default_allocator: bool,
+    ) -> StoreInstanceHandle {
+        self.inner.instances.borrow_mut().push(StoreInstance {
+            handle: handle.clone(),
+            use_default_allocator,
+        });
         StoreInstanceHandle {
             store: self.clone(),
             handle,
@@ -230,7 +246,7 @@ impl Store {
             .instances
             .borrow()
             .iter()
-            .any(|i| i.vmctx_ptr() == handle.vmctx_ptr()));
+            .any(|i| i.handle.vmctx_ptr() == handle.vmctx_ptr()));
         StoreInstanceHandle {
             store: self.clone(),
             handle,
@@ -425,7 +441,14 @@ impl Drop for StoreInner {
         let allocator = self.engine.config().instance_allocator();
         for instance in self.instances.borrow().iter() {
             unsafe {
-                allocator.deallocate(instance);
+                if instance.use_default_allocator {
+                    self.engine
+                        .config()
+                        .default_instance_allocator
+                        .deallocate(&instance.handle);
+                } else {
+                    allocator.deallocate(&instance.handle);
+                }
             }
         }
     }
diff --git a/crates/wasmtime/src/trampoline/create_handle.rs b/crates/wasmtime/src/trampoline/create_handle.rs
index ff088b9aa17d..71595ff7296a 100644
--- a/crates/wasmtime/src/trampoline/create_handle.rs
+++ b/crates/wasmtime/src/trampoline/create_handle.rs
@@ -44,6 +44,6 @@ pub(crate) fn create_handle(
                 stack_map_registry: store.stack_map_registry() as *const StackMapRegistry as *mut _,
             })?;
 
-        Ok(store.add_instance(handle))
+        Ok(store.add_instance(handle, true))
     }
 }