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();
@@ -349,24 +349,7 @@ impl Instance {
             .features()
             .contains(WasmFeatures::BULK_MEMORY);
 
-        if store.async_support() {
-            #[cfg(feature = "async")]
-            store.block_on(|store| {
-                let module = compiled_module.module().clone();
-                Box::pin(
-                    async move { vm::initialize_instance(store, id, &module, bulk_memory).await },
-                )
-            })??;
-            #[cfg(not(feature = "async"))]
-            unreachable!();
-        } else {
-            vm::assert_ready(vm::initialize_instance(
-                store,
-                id,
-                compiled_module.module(),
-                bulk_memory,
-            ))?;
-        }
+        vm::initialize_instance(store, id, compiled_module.module(), bulk_memory).await?;
 
         Ok((instance, compiled_module.module().start_func))
     }

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> {
+        async 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
@@ -2091,6 +2100,7 @@ at https://bytecodealliance.org/security.
 
     /// Retrieve the store's protection key.
     #[inline]
+    #[cfg(feature = "pooling-allocator")]
     pub(crate) fn get_pkey(&self) -> Option<ProtectionKey> {
         self.pkey
     }
@@ -2225,16 +2235,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 {
@@ -2316,82 +2327,6 @@ unsafe impl<T> VMStore for StoreInner<T> {
         )
     }
 
-    fn memory_growing(
-        &mut self,
-        current: usize,
-        desired: usize,
-        maximum: Option<usize>,
-    ) -> Result<bool, anyhow::Error> {
-        match self.limiter {
-            Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
-                limiter(&mut self.data).memory_growing(current, desired, maximum)
-            }
-            #[cfg(feature = "async")]
-            Some(ResourceLimiterInner::Async(_)) => self.block_on(|store| {
-                let limiter = match &mut store.0.limiter {
-                    Some(ResourceLimiterInner::Async(limiter)) => limiter,
-                    _ => unreachable!(),
-                };
-                limiter(&mut store.0.data).memory_growing(current, desired, maximum)
-            })?,
-            None => Ok(true),
-        }
-    }
-
-    fn memory_grow_failed(&mut self, error: anyhow::Error) -> Result<()> {
-        match self.limiter {
-            Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
-                limiter(&mut self.data).memory_grow_failed(error)
-            }
-            #[cfg(feature = "async")]
-            Some(ResourceLimiterInner::Async(ref mut limiter)) => {
-                limiter(&mut self.data).memory_grow_failed(error)
-            }
-            None => {
-                log::debug!("ignoring memory growth failure error: {error:?}");
-                Ok(())
-            }
-        }
-    }
-
-    fn table_growing(
-        &mut self,
-        current: usize,
-        desired: usize,
-        maximum: Option<usize>,
-    ) -> Result<bool, anyhow::Error> {
-        match self.limiter {
-            Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
-                limiter(&mut self.data).table_growing(current, desired, maximum)
-            }
-            #[cfg(feature = "async")]
-            Some(ResourceLimiterInner::Async(_)) => self.block_on(|store| {
-                let limiter = match &mut store.0.limiter {
-                    Some(ResourceLimiterInner::Async(limiter)) => limiter,
-                    _ => unreachable!(),
-                };
-                limiter(&mut store.0.data).table_growing(current, desired, maximum)
-            })?,
-            None => Ok(true),
-        }
-    }
-
-    fn table_grow_failed(&mut self, error: anyhow::Error) -> Result<()> {
-        match self.limiter {
-            Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
-                limiter(&mut self.data).table_grow_failed(error)
-            }
-            #[cfg(feature = "async")]
-            Some(ResourceLimiterInner::Async(ref mut limiter)) => {
-                limiter(&mut self.data).table_grow_failed(error)
-            }
-            None => {
-                log::debug!("ignoring table growth failure: {error:?}");
-                Ok(())
-            }
-        }
-    }
-
     fn out_of_gas(&mut self) -> Result<()> {
         if !self.refuel() {
             return Err(Trap::OutOfFuel.into());

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>>() {
@@ -208,7 +208,7 @@ impl StoreOpaque {
     where
         T: Send + Sync + 'static,
     {
-        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(