Revamp memory management of InstanceHandle (#1624)

* Revamp memory management of `InstanceHandle`

This commit fixes a known but in Wasmtime where an instance could still
be used after it was freed. Unfortunately the fix here is a bit of a
hammer, but it's the best that we can do for now. The changes made in
this commit are:

* A `Store` now stores all `InstanceHandle` objects it ever creates.
  This keeps all instances alive unconditionally (along with all host
  functions and such) until the `Store` is itself dropped. Note that a
  `Store` is reference counted so basically everything has to be dropped
  to drop anything, there's no longer any partial deallocation of instances.

* The `InstanceHandle` type's own reference counting has been removed.
  This is largely redundant with what's already happening in `Store`, so
  there's no need to manage two reference counts.

* Each `InstanceHandle` no longer tracks its dependencies in terms of
  instance handles. This set was actually inaccurate due to dynamic
  updates to tables and such, so we needed to revamp it anyway.

* Initialization of an `InstanceHandle` is now deferred until after
  `InstanceHandle::new`. This allows storing the `InstanceHandle` before
  side-effectful initialization, such as copying element segments or
  running the start function, to ensure that regardless of the result of
  instantiation the underlying `InstanceHandle` is still available to
  persist in storage.

Overall this should fix a known possible way to safely segfault Wasmtime
today (yay!) and it should also fix some flaikness I've seen on CI.
Turns out one of the spec tests
(bulk-memory-operations/partial-init-table-segment.wast) exercises this
functionality and we were hitting sporating use-after-free, but only on
Windows.

* Shuffle some APIs around

* Comment weak cycle

crates/api/src/trampoline/create_handle.rs

@@ -1,9 +1,10 @@
 //! Support for a calling of an imported function.
 
-use crate::runtime::Store;
+use crate::trampoline::StoreInstanceHandle;
+use crate::Store;
 use anyhow::Result;
 use std::any::Any;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::sync::Arc;
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::wasm::DefinedFuncIndex;
@@ -18,15 +19,13 @@ pub(crate) fn create_handle(
     finished_functions: PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
     trampolines: HashMap<VMSharedSignatureIndex, VMTrampoline>,
     state: Box<dyn Any>,
-) -> Result<InstanceHandle> {
+) -> Result<StoreInstanceHandle> {
     let imports = Imports::new(
-        HashSet::new(),
         PrimaryMap::new(),
         PrimaryMap::new(),
         PrimaryMap::new(),
         PrimaryMap::new(),
     );
-    let data_initializers = Vec::new();
 
     // Compute indices into the shared signature table.
     let signatures = module
@@ -37,24 +36,17 @@ pub(crate) fn create_handle(
         .collect::<PrimaryMap<_, _>>();
 
     unsafe {
-        Ok(InstanceHandle::new(
+        let handle = InstanceHandle::new(
             Arc::new(module),
             finished_functions.into_boxed_slice(),
             trampolines,
             imports,
             store.memory_creator(),
-            &data_initializers,
             signatures.into_boxed_slice(),
             None,
-            store
-                .engine()
-                .config()
-                .validating_config
-                .operator_config
-                .enable_bulk_memory,
             state,
             store.compiler().interrupts().clone(),
-            store.engine().config().max_wasm_stack,
-        )?)
+        )?;
+        Ok(store.add_instance(handle))
     }
 }

crates/api/src/trampoline/func.rs

@@ -1,6 +1,7 @@
 //! Support for a calling of an imported function.
 
 use super::create_handle::create_handle;
+use crate::trampoline::StoreInstanceHandle;
 use crate::{FuncType, Store, Trap};
 use anyhow::{bail, Result};
 use std::any::Any;
@@ -203,7 +204,7 @@ pub fn create_handle_with_function(
     ft: &FuncType,
     func: Box<dyn Fn(*mut VMContext, *mut u128) -> Result<(), Trap>>,
     store: &Store,
-) -> Result<(InstanceHandle, VMTrampoline)> {
+) -> Result<(StoreInstanceHandle, VMTrampoline)> {
     let isa = {
         let isa_builder = native::builder();
         let flag_builder = settings::builder();
@@ -267,7 +268,7 @@ pub unsafe fn create_handle_with_raw_function(
     trampoline: VMTrampoline,
     store: &Store,
     state: Box<dyn Any>,
-) -> Result<InstanceHandle> {
+) -> Result<StoreInstanceHandle> {
     let isa = {
         let isa_builder = native::builder();
         let flag_builder = settings::builder();

crates/api/src/trampoline/global.rs

@@ -1,12 +1,12 @@
 use super::create_handle::create_handle;
+use crate::trampoline::StoreInstanceHandle;
 use crate::Store;
 use crate::{GlobalType, Mutability, Val};
 use anyhow::{bail, Result};
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::{wasm, EntityIndex, Module};
-use wasmtime_runtime::InstanceHandle;
 
-pub fn create_global(store: &Store, gt: &GlobalType, val: Val) -> Result<InstanceHandle> {
+pub fn create_global(store: &Store, gt: &GlobalType, val: Val) -> Result<StoreInstanceHandle> {
     let global = wasm::Global {
         ty: match gt.content().get_wasmtime_type() {
             Some(t) => t,

crates/api/src/trampoline/memory.rs

@@ -1,17 +1,19 @@
 use super::create_handle::create_handle;
 use crate::externals::{LinearMemory, MemoryCreator};
+use crate::trampoline::StoreInstanceHandle;
 use crate::Store;
 use crate::{Limits, MemoryType};
 use anyhow::Result;
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::{wasm, EntityIndex, MemoryPlan, Module, WASM_PAGE_SIZE};
-use wasmtime_runtime::{
-    InstanceHandle, RuntimeLinearMemory, RuntimeMemoryCreator, VMMemoryDefinition,
-};
+use wasmtime_runtime::{RuntimeLinearMemory, RuntimeMemoryCreator, VMMemoryDefinition};
 
 use std::sync::Arc;
 
-pub fn create_handle_with_memory(store: &Store, memory: &MemoryType) -> Result<InstanceHandle> {
+pub fn create_handle_with_memory(
+    store: &Store,
+    memory: &MemoryType,
+) -> Result<StoreInstanceHandle> {
     let mut module = Module::new();
 
     let memory = wasm::Memory {

crates/api/src/trampoline/mod.rs

@@ -15,14 +15,42 @@ use self::table::create_handle_with_table;
 use crate::{FuncType, GlobalType, MemoryType, Store, TableType, Trap, Val};
 use anyhow::Result;
 use std::any::Any;
-use wasmtime_runtime::{VMContext, VMFunctionBody, VMTrampoline};
+use std::ops::Deref;
+use wasmtime_runtime::{InstanceHandle, VMContext, VMFunctionBody, VMTrampoline};
+
+/// A wrapper around `wasmtime_runtime::InstanceHandle` which pairs it with the
+/// `Store` that it's rooted within. The instance is deallocated when `Store` is
+/// deallocated, so this is a safe handle in terms of memory management for the
+/// `Store`.
+pub struct StoreInstanceHandle {
+    pub store: Store,
+    pub handle: InstanceHandle,
+}
+
+impl Clone for StoreInstanceHandle {
+    fn clone(&self) -> StoreInstanceHandle {
+        StoreInstanceHandle {
+            store: self.store.clone(),
+            // Note should be safe because the lifetime of the instance handle
+            // is tied to the `Store` which this is paired with.
+            handle: unsafe { self.handle.clone() },
+        }
+    }
+}
+
+impl Deref for StoreInstanceHandle {
+    type Target = InstanceHandle;
+    fn deref(&self) -> &InstanceHandle {
+        &self.handle
+    }
+}
 
 pub fn generate_func_export(
     ft: &FuncType,
     func: Box<dyn Fn(*mut VMContext, *mut u128) -> Result<(), Trap>>,
     store: &Store,
 ) -> Result<(
-    wasmtime_runtime::InstanceHandle,
+    StoreInstanceHandle,
     wasmtime_runtime::ExportFunction,
     VMTrampoline,
 )> {
@@ -42,10 +70,7 @@ pub unsafe fn generate_raw_func_export(
     trampoline: VMTrampoline,
     store: &Store,
     state: Box<dyn Any>,
-) -> Result<(
-    wasmtime_runtime::InstanceHandle,
-    wasmtime_runtime::ExportFunction,
-)> {
+) -> Result<(StoreInstanceHandle, wasmtime_runtime::ExportFunction)> {
     let instance = func::create_handle_with_raw_function(ft, func, trampoline, store, state)?;
     match instance.lookup("trampoline").expect("trampoline export") {
         wasmtime_runtime::Export::Function(f) => Ok((instance, f)),
@@ -57,10 +82,7 @@ pub fn generate_global_export(
     store: &Store,
     gt: &GlobalType,
     val: Val,
-) -> Result<(
-    wasmtime_runtime::InstanceHandle,
-    wasmtime_runtime::ExportGlobal,
-)> {
+) -> Result<(StoreInstanceHandle, wasmtime_runtime::ExportGlobal)> {
     let instance = create_global(store, gt, val)?;
     match instance.lookup("global").expect("global export") {
         wasmtime_runtime::Export::Global(g) => Ok((instance, g)),
@@ -71,10 +93,7 @@ pub fn generate_global_export(
 pub fn generate_memory_export(
     store: &Store,
     m: &MemoryType,
-) -> Result<(
-    wasmtime_runtime::InstanceHandle,
-    wasmtime_runtime::ExportMemory,
-)> {
+) -> Result<(StoreInstanceHandle, wasmtime_runtime::ExportMemory)> {
     let instance = create_handle_with_memory(store, m)?;
     match instance.lookup("memory").expect("memory export") {
         wasmtime_runtime::Export::Memory(m) => Ok((instance, m)),
@@ -85,10 +104,7 @@ pub fn generate_memory_export(
 pub fn generate_table_export(
     store: &Store,
     t: &TableType,
-) -> Result<(
-    wasmtime_runtime::InstanceHandle,
-    wasmtime_runtime::ExportTable,
-)> {
+) -> Result<(StoreInstanceHandle, wasmtime_runtime::ExportTable)> {
     let instance = create_handle_with_table(store, t)?;
     match instance.lookup("table").expect("table export") {
         wasmtime_runtime::Export::Table(t) => Ok((instance, t)),

crates/api/src/trampoline/table.rs

@@ -1,12 +1,12 @@
 use super::create_handle::create_handle;
+use crate::trampoline::StoreInstanceHandle;
 use crate::Store;
 use crate::{TableType, ValType};
 use anyhow::{bail, Result};
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::{wasm, EntityIndex, Module};
-use wasmtime_runtime::InstanceHandle;
 
-pub fn create_handle_with_table(store: &Store, table: &TableType) -> Result<InstanceHandle> {
+pub fn create_handle_with_table(store: &Store, table: &TableType) -> Result<StoreInstanceHandle> {
     let mut module = Module::new();
 
     let table = wasm::Table {