Make table/memory creation async functions

This commit is a large-ish refactor which is made possible by the many
previous refactorings to internals w.r.t. async-in-Wasmtime. The end
goal of this change is that table and memory allocation are both `async`
functions. Achieving this, however, required some refactoring to enable
it to work:

* To work with `Send` neither function can close over `dyn VMStore`.
  This required changing their `Option<&mut dyn VMStore>` arugment to
  `Option<&mut StoreResourceLimiter<'_>>`
* Somehow a `StoreResourceLimiter` needed to be acquired from an
  `InstanceAllocationRequest`. Previously the store was stored here as
  an unsafe raw pointer, but I've refactored this now so
  `InstanceAllocationRequest` directly stores `&StoreOpaque` and
  `Option<&mut StoreResourceLimiter>` meaning it's trivial to acquire
  them. This additionally means no more `unsafe` access of the store
  during instance allocation (yay!).
* Now-redundant fields of `InstanceAllocationRequest` were removed since
  they can be safely inferred from `&StoreOpaque`. For example passing
  around `&Tunables` is now all gone.
* Methods upwards from table/memory allocation to the
  `InstanceAllocator` trait needed to be made `async`. This includes new
  `#[async_trait]` methods for example.
* `StoreOpaque::ensure_gc_store` is now an `async` function. This
  internally carries a new `unsafe` block carried over from before with
  the raw point passed around in `InstanceAllocationRequest`. A future
  PR will delete this `unsafe` block, it's just temporary.

I attempted a few times to split this PR up into separate commits but
everything is relatively intertwined here so this is the smallest
"atomic" unit I could manage to land these changes and refactorings.

Author: alexcrichton

crates/wasmtime/src/runtime/instance.rs

@@ -295,7 +295,7 @@ impl Instance {
 
         // Allocate the GC heap, if necessary.
         if module.env_module().needs_gc_heap {
-            store.ensure_gc_store()?;
+            store.ensure_gc_store().await?;
         }
 
         let compiled_module = module.compiled_module();

crates/wasmtime/src/runtime/store.rs

@@ -93,8 +93,7 @@ use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::{
     self, GcStore, Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
     Interpreter, InterpreterRef, ModuleRuntimeInfo, OnDemandInstanceAllocator, SendSyncPtr,
-    SignalHandler, StoreBox, StorePtr, Unwind, VMContext, VMFuncRef, VMGcRef, VMStore,
-    VMStoreContext,
+    SignalHandler, StoreBox, Unwind, VMContext, VMFuncRef, VMGcRef, VMStore, VMStoreContext,
 };
 use crate::trampoline::VMHostGlobalContext;
 use crate::{Engine, Module, Trap, Val, ValRaw, module::ModuleRegistry};
@@ -470,6 +469,20 @@ pub struct StoreOpaque {
     executor: Executor,
 }
 
+/// Self-pointer to `StoreInner<T>` from within a `StoreOpaque` which is chiefly
+/// used to copy into instances during instantiation.
+///
+/// FIXME: ideally this type would get deleted and Wasmtime's reliance on it
+/// would go away.
+struct StorePtr(Option<NonNull<dyn VMStore>>);
+
+// We can't make `VMStore: Send + Sync` because that requires making all of
+// Wastime's internals generic over the `Store`'s `T`. So instead, we take care
+// in the whole VM layer to only use the `VMStore` in ways that are `Send`- and
+// `Sync`-safe and we have to have these unsafe impls.
+unsafe impl Send for StorePtr {}
+unsafe impl Sync for StorePtr {}
+
 /// Executor state within `StoreOpaque`.
 ///
 /// Effectively stores Pulley interpreter state and handles conditional support
@@ -646,7 +659,7 @@ impl<T> Store<T> {
             fuel_reserve: 0,
             fuel_yield_interval: None,
             store_data,
-            traitobj: StorePtr::empty(),
+            traitobj: StorePtr(None),
             default_caller_vmctx: SendSyncPtr::new(NonNull::dangling()),
             hostcall_val_storage: Vec::new(),
             wasm_val_raw_storage: Vec::new(),
@@ -670,7 +683,7 @@ impl<T> Store<T> {
             data: ManuallyDrop::new(data),
         });
 
-        inner.traitobj = StorePtr::new(NonNull::from(&mut *inner));
+        inner.traitobj = StorePtr(Some(NonNull::from(&mut *inner)));
 
         // Wasmtime uses the callee argument to host functions to learn about
         // the original pointer to the `Store` itself, allowing it to
@@ -1489,15 +1502,15 @@ impl StoreOpaque {
     /// `ResourceLimiterAsync` which means that this should only be executed
     /// in a fiber context at this time.
     #[inline]
-    pub(crate) fn ensure_gc_store(&mut self) -> Result<&mut GcStore> {
+    pub(crate) async fn ensure_gc_store(&mut self) -> Result<&mut GcStore> {
         if self.gc_store.is_some() {
             return Ok(self.gc_store.as_mut().unwrap());
         }
-        self.allocate_gc_store()
+        self.allocate_gc_store().await
     }
 
     #[inline(never)]
-    fn allocate_gc_store(&mut self) -> Result<&mut GcStore> {
+    async fn allocate_gc_store(&mut self) -> Result<&mut GcStore> {
         log::trace!("allocating GC heap for store {:?}", self.id());
 
         assert!(self.gc_store.is_none());
@@ -1507,19 +1520,24 @@ impl StoreOpaque {
         );
         assert_eq!(self.vm_store_context.gc_heap.current_length(), 0);
 
-        let vmstore = self.traitobj();
-        let gc_store = allocate_gc_store(self.engine(), vmstore, self.get_pkey())?;
+        let gc_store = allocate_gc_store(self).await?;
         self.vm_store_context.gc_heap = gc_store.vmmemory_definition();
         return Ok(self.gc_store.insert(gc_store));
 
         #[cfg(feature = "gc")]
-        fn allocate_gc_store(
-            engine: &Engine,
-            vmstore: NonNull<dyn VMStore>,
-            pkey: Option<ProtectionKey>,
-        ) -> Result<GcStore> {
+        async fn allocate_gc_store(store: &mut StoreOpaque) -> Result<GcStore> {
             use wasmtime_environ::packed_option::ReservedValue;
 
+            // FIXME(#11409) this is not a sound widening borrow
+            let (mut limiter, store) = unsafe {
+                store
+                    .traitobj()
+                    .as_mut()
+                    .resource_limiter_and_store_opaque()
+            };
+
+            let engine = store.engine();
+            let mem_ty = engine.tunables().gc_heap_memory_type();
             ensure!(
                 engine.features().gc_types(),
                 "cannot allocate a GC store when GC is disabled at configuration time"
@@ -1532,19 +1550,14 @@ impl StoreOpaque {
                     wasmtime_environ::Module::default(),
                 )),
                 imports: vm::Imports::default(),
-                store: StorePtr::new(vmstore),
-                #[cfg(feature = "wmemcheck")]
-                wmemcheck: false,
-                pkey,
-                tunables: engine.tunables(),
+                store,
+                limiter: limiter.as_mut(),
             };
-            let mem_ty = engine.tunables().gc_heap_memory_type();
-            let tunables = engine.tunables();
 
-            let (mem_alloc_index, mem) =
-                engine
-                    .allocator()
-                    .allocate_memory(&mut request, &mem_ty, tunables, None)?;
+            let (mem_alloc_index, mem) = engine
+                .allocator()
+                .allocate_memory(&mut request, &mem_ty, None)
+                .await?;
 
             // Then, allocate the actual GC heap, passing in that memory
             // storage.
@@ -1560,11 +1573,7 @@ impl StoreOpaque {
         }
 
         #[cfg(not(feature = "gc"))]
-        fn allocate_gc_store(
-            _engine: &Engine,
-            _vmstore: NonNull<dyn VMStore>,
-            _pkey: Option<ProtectionKey>,
-        ) -> Result<GcStore> {
+        fn allocate_gc_store(_: &mut StoreOpaque) -> Result<GcStore> {
             bail!("cannot allocate a GC store: the `gc` feature was disabled at compile time")
         }
     }
@@ -1944,7 +1953,7 @@ impl StoreOpaque {
 
     #[inline]
     pub fn traitobj(&self) -> NonNull<dyn VMStore> {
-        self.traitobj.as_raw().unwrap()
+        self.traitobj.0.unwrap()
     }
 
     /// Takes the cached `Vec<Val>` stored internally across hostcalls to get
@@ -2225,16 +2234,17 @@ at https://bytecodealliance.org/security.
         // SAFETY: this function's own contract is the same as
         // `allocate_module`, namely the imports provided are valid.
         let handle = unsafe {
-            allocator.allocate_module(InstanceAllocationRequest {
-                id,
-                runtime_info,
-                imports,
-                store: StorePtr::new(self.traitobj()),
-                #[cfg(feature = "wmemcheck")]
-                wmemcheck: self.engine().config().wmemcheck,
-                pkey: self.get_pkey(),
-                tunables: self.engine().tunables(),
-            })?
+            // FIXME(#11409) this is not a sound widening borrow
+            let (mut limiter, store) = self.traitobj().as_mut().resource_limiter_and_store_opaque();
+            allocator
+                .allocate_module(InstanceAllocationRequest {
+                    id,
+                    runtime_info,
+                    imports,
+                    store,
+                    limiter: limiter.as_mut(),
+                })
+                .await?
         };
 
         let actual = match kind {

crates/wasmtime/src/runtime/store/gc.rs

@@ -179,7 +179,7 @@ impl StoreOpaque {
             !self.async_support(),
             "use the `*_async` versions of methods when async is configured"
         );
-        self.ensure_gc_store()?;
+        vm::assert_ready(self.ensure_gc_store())?;
         match alloc_func(self, value) {
             Ok(x) => Ok(x),
             Err(e) => match e.downcast::<crate::GcHeapOutOfMemory<T>>() {
@@ -212,7 +212,16 @@ impl StoreOpaque {
     where
         T: Send + Sync + 'static,
     {
-        self.ensure_gc_store()?;
+        if self.async_support() {
+            self.block_on(|store| {
+                Box::pin(async move {
+                    store.ensure_gc_store().await?;
+                    anyhow::Ok(())
+                })
+            })??;
+        } else {
+            vm::assert_ready(self.ensure_gc_store())?;
+        }
         match alloc_func(self, value) {
             Ok(x) => Ok(x),
             Err(e) => match e.downcast::<crate::GcHeapOutOfMemory<T>>() {
@@ -262,7 +271,7 @@ impl StoreOpaque {
             self.async_support(),
             "you must configure async to use the `*_async` versions of methods"
         );
-        self.ensure_gc_store()?;
+        self.ensure_gc_store().await?;
         match alloc_func(self, value) {
             Ok(x) => Ok(x),
             Err(e) => match e.downcast::<crate::GcHeapOutOfMemory<T>>() {

crates/wasmtime/src/runtime/trampoline/memory.rs

@@ -124,6 +124,7 @@ struct SingleMemoryInstance<'a> {
     ondemand: OnDemandInstanceAllocator,
 }
 
+#[async_trait::async_trait]
 unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
     #[cfg(feature = "component-model")]
     fn validate_component<'a>(
@@ -167,11 +168,10 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         self.ondemand.decrement_core_instance_count();
     }
 
-    fn allocate_memory(
+    async fn allocate_memory(
         &self,
-        request: &mut InstanceAllocationRequest,
+        request: &mut InstanceAllocationRequest<'_>,
         ty: &wasmtime_environ::Memory,
-        tunables: &Tunables,
         memory_index: Option<DefinedMemoryIndex>,
     ) -> Result<(MemoryAllocationIndex, Memory)> {
         if cfg!(debug_assertions) {
@@ -186,9 +186,11 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
                 MemoryAllocationIndex::default(),
                 shared_memory.clone().as_memory(),
             )),
-            None => self
-                .ondemand
-                .allocate_memory(request, ty, tunables, memory_index),
+            None => {
+                self.ondemand
+                    .allocate_memory(request, ty, memory_index)
+                    .await
+            }
         }
     }
 
@@ -204,14 +206,13 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         }
     }
 
-    fn allocate_table(
+    async fn allocate_table(
         &self,
-        req: &mut InstanceAllocationRequest,
+        req: &mut InstanceAllocationRequest<'_>,
         ty: &wasmtime_environ::Table,
-        tunables: &Tunables,
         table_index: DefinedTableIndex,
     ) -> Result<(TableAllocationIndex, Table)> {
-        self.ondemand.allocate_table(req, ty, tunables, table_index)
+        self.ondemand.allocate_table(req, ty, table_index).await
     }
 
     unsafe fn deallocate_table(

crates/wasmtime/src/runtime/vm.rs

@@ -99,8 +99,7 @@ pub use crate::runtime::vm::gc::*;
 pub use crate::runtime::vm::imports::Imports;
 pub use crate::runtime::vm::instance::{
     GcHeapAllocationIndex, Instance, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
-    MemoryAllocationIndex, OnDemandInstanceAllocator, StorePtr, TableAllocationIndex,
-    initialize_instance,
+    MemoryAllocationIndex, OnDemandInstanceAllocator, TableAllocationIndex, initialize_instance,
 };
 #[cfg(feature = "pooling-allocator")]
 pub use crate::runtime::vm::instance::{

crates/wasmtime/src/runtime/vm/instance.rs

@@ -487,32 +487,22 @@ impl Instance {
         unsafe { self.vmctx_plus_offset_mut(offset) }
     }
 
-    pub(crate) unsafe fn set_store(mut self: Pin<&mut Self>, store: Option<NonNull<dyn VMStore>>) {
+    pub(crate) unsafe fn set_store(mut self: Pin<&mut Self>, store: &StoreOpaque) {
         // FIXME: should be more targeted ideally with the `unsafe` than just
         // throwing this entire function in a large `unsafe` block.
         unsafe {
-            *self.as_mut().store_mut() = store.map(VMStoreRawPtr);
-            if let Some(mut store) = store {
-                let store = store.as_mut();
-                self.vm_store_context()
-                    .write(Some(store.vm_store_context_ptr().into()));
-                #[cfg(target_has_atomic = "64")]
-                {
-                    *self.as_mut().epoch_ptr() =
-                        Some(NonNull::from(store.engine().epoch_counter()).into());
-                }
+            *self.as_mut().store_mut() = Some(VMStoreRawPtr(store.traitobj()));
+            self.vm_store_context()
+                .write(Some(store.vm_store_context_ptr().into()));
+            #[cfg(target_has_atomic = "64")]
+            {
+                *self.as_mut().epoch_ptr() =
+                    Some(NonNull::from(store.engine().epoch_counter()).into());
+            }
 
-                if self.env_module().needs_gc_heap {
-                    self.as_mut().set_gc_heap(Some(store.unwrap_gc_store()));
-                } else {
-                    self.as_mut().set_gc_heap(None);
-                }
+            if self.env_module().needs_gc_heap {
+                self.as_mut().set_gc_heap(Some(store.unwrap_gc_store()));
             } else {
-                self.vm_store_context().write(None);
-                #[cfg(target_has_atomic = "64")]
-                {
-                    *self.as_mut().epoch_ptr() = None;
-                }
                 self.as_mut().set_gc_heap(None);
             }
         }
@@ -1246,7 +1236,7 @@ impl Instance {
         mut self: Pin<&mut Self>,
         module: &Module,
         offsets: &VMOffsets<HostPtr>,
-        store: StorePtr,
+        store: &StoreOpaque,
         imports: Imports,
     ) {
         assert!(ptr::eq(module, self.env_module().as_ref()));
@@ -1260,7 +1250,7 @@ impl Instance {
 
         // SAFETY: it's up to the caller to provide a valid store pointer here.
         unsafe {
-            self.as_mut().set_store(store.as_raw());
+            self.as_mut().set_store(store);
         }
 
         // Initialize shared types