Make table/memory creation async functions (#11470)

* Make core instance allocation an `async` function

This commit is a step in preparation for #11430, notably core instance
allocation, or `StoreOpaque::allocate_instance` is now an `async fn`.
This function does not actually use the `async`-ness just yet so it's a
noop from that point of view, but this propagates outwards to enough
locations that I wanted to split this off to make future changes more
digestable.

Notably some creation functions here such as making an `Instance`,
`Table`, or `Memory` are refactored internally to use this new `async`
function. Annotations of `assert_ready` or `one_poll` are used as
appropriate as well.

For reference this commit was benchmarked with our `instantiation.rs`
benchmark in the pooling allocator and shows no changes relative to the
original baseline from before-`async`-PRs.

* 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.

* Shuffle `async-trait` dep

* Fix configured build

Author: alexcrichton

crates/wasmtime/Cargo.toml

@@ -196,7 +196,6 @@ cache = ["dep:wasmtime-cache", "std"]
 # `async fn` and calling functions asynchronously.
 async = [
   "dep:wasmtime-fiber",
-  "dep:async-trait",
   "wasmtime-component-macro?/async",
   "runtime",
 ]
@@ -261,6 +260,7 @@ runtime = [
   "dep:windows-sys",
   "pulley-interpreter/interp",
   "dep:wasmtime-unwinder",
+  "dep:async-trait",
 ]
 
 # Enable support for garbage collection-related things.

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

@@ -634,7 +634,7 @@ impl<'a> Instantiator<'a> {
         }
     }
 
-    fn run<T>(&mut self, store: &mut StoreContextMut<'_, T>) -> Result<()> {
+    async fn run<T>(&mut self, store: &mut StoreContextMut<'_, T>) -> Result<()> {
         let env_component = self.component.env_component();
 
         // Before all initializers are processed configure all destructors for
@@ -714,7 +714,7 @@ impl<'a> Instantiator<'a> {
                     // if required.
 
                     let i = unsafe {
-                        crate::Instance::new_started_impl(store, module, imports.as_ref())?
+                        crate::Instance::new_started(store, module, imports.as_ref()).await?
                     };
                     self.instance_mut(store.0).push_instance_id(i.id());
                 }
@@ -991,37 +991,26 @@ impl<T: 'static> InstancePre<T> {
             !store.as_context().async_support(),
             "must use async instantiation when async support is enabled"
         );
-        self.instantiate_impl(store)
+        vm::assert_ready(self._instantiate(store))
     }
     /// Performs the instantiation process into the store specified.
     ///
     /// Exactly like [`Self::instantiate`] except for use on async stores.
     //
     // TODO: needs more docs
     #[cfg(feature = "async")]
-    pub async fn instantiate_async(
-        &self,
-        mut store: impl AsContextMut<Data = T>,
-    ) -> Result<Instance>
-    where
-        T: Send,
-    {
-        let mut store = store.as_context_mut();
-        assert!(
-            store.0.async_support(),
-            "must use sync instantiation when async support is disabled"
-        );
-        store.on_fiber(|store| self.instantiate_impl(store)).await?
+    pub async fn instantiate_async(&self, store: impl AsContextMut<Data = T>) -> Result<Instance> {
+        self._instantiate(store).await
     }
 
-    fn instantiate_impl(&self, mut store: impl AsContextMut<Data = T>) -> Result<Instance> {
+    async fn _instantiate(&self, mut store: impl AsContextMut<Data = T>) -> Result<Instance> {
         let mut store = store.as_context_mut();
         store
             .engine()
             .allocator()
             .increment_component_instance_count()?;
         let mut instantiator = Instantiator::new(&self.component, store.0, &self.imports);
-        instantiator.run(&mut store).map_err(|e| {
+        instantiator.run(&mut store).await.map_err(|e| {
             store
                 .engine()
                 .allocator()

crates/wasmtime/src/runtime/externals/table.rs

@@ -94,7 +94,8 @@ impl Table {
     /// # }
     /// ```
     pub fn new(mut store: impl AsContextMut, ty: TableType, init: Ref) -> Result<Table> {
-        Table::_new(store.as_context_mut().0, ty, init)
+        vm::one_poll(Table::_new(store.as_context_mut().0, ty, init))
+            .expect("must use `new_async` when async resource limiters are in use")
     }
 
     /// Async variant of [`Table::new`]. You must use this variant with
@@ -111,18 +112,11 @@ impl Table {
         ty: TableType,
         init: Ref,
     ) -> Result<Table> {
-        let mut store = store.as_context_mut();
-        assert!(
-            store.0.async_support(),
-            "cannot use `new_async` without enabling async support on the config"
-        );
-        store
-            .on_fiber(|store| Table::_new(store.0, ty, init))
-            .await?
+        Table::_new(store.as_context_mut().0, ty, init).await
     }
 
-    fn _new(store: &mut StoreOpaque, ty: TableType, init: Ref) -> Result<Table> {
-        let table = generate_table_export(store, &ty)?;
+    async fn _new(store: &mut StoreOpaque, ty: TableType, init: Ref) -> Result<Table> {
+        let table = generate_table_export(store, &ty).await?;
         table._fill(store, 0, init, ty.minimum())?;
         Ok(table)
     }

crates/wasmtime/src/runtime/instance.rs

@@ -117,7 +117,8 @@ impl Instance {
         // Note that the unsafety here should be satisfied by the call to
         // `typecheck_externs` above which satisfies the condition that all
         // the imports are valid for this module.
-        unsafe { Instance::new_started(&mut store, module, imports.as_ref()) }
+        assert!(!store.0.async_support());
+        vm::assert_ready(unsafe { Instance::new_started(&mut store, module, imports.as_ref()) })
     }
 
     /// Same as [`Instance::new`], except for usage in [asynchronous stores].
@@ -200,7 +201,7 @@ impl Instance {
         let mut store = store.as_context_mut();
         let imports = Instance::typecheck_externs(store.0, module, imports)?;
         // See `new` for notes on this unsafety
-        unsafe { Instance::new_started_async(&mut store, module, imports.as_ref()).await }
+        unsafe { Instance::new_started(&mut store, module, imports.as_ref()).await }
     }
 
     fn typecheck_externs(
@@ -242,62 +243,31 @@ impl Instance {
     /// Internal function to create an instance and run the start function.
     ///
     /// This function's unsafety is the same as `Instance::new_raw`.
-    pub(crate) unsafe fn new_started<T>(
-        store: &mut StoreContextMut<'_, T>,
-        module: &Module,
-        imports: Imports<'_>,
-    ) -> Result<Instance> {
-        assert!(
-            !store.0.async_support(),
-            "must use async instantiation when async support is enabled",
-        );
-
-        // SAFETY: the safety contract of `new_started_impl` is the same as this
-        // function.
-        unsafe { Self::new_started_impl(store, module, imports) }
-    }
-
-    /// Internal function to create an instance and run the start function.
-    ///
-    /// ONLY CALL THIS IF YOU HAVE ALREADY CHECKED FOR ASYNCNESS AND HANDLED
-    /// THE FIBER NONSENSE
-    pub(crate) unsafe fn new_started_impl<T>(
+    pub(crate) async unsafe fn new_started<T>(
         store: &mut StoreContextMut<'_, T>,
         module: &Module,
         imports: Imports<'_>,
     ) -> Result<Instance> {
         // SAFETY: the safety contract of `new_raw` is the same as this
         // function.
-        let (instance, start) = unsafe { Instance::new_raw(store.0, module, imports)? };
+        let (instance, start) = unsafe { Instance::new_raw(store.0, module, imports).await? };
         if let Some(start) = start {
-            instance.start_raw(store, start)?;
+            if store.0.async_support() {
+                #[cfg(feature = "async")]
+                {
+                    store
+                        .on_fiber(|store| instance.start_raw(store, start))
+                        .await??;
+                }
+                #[cfg(not(feature = "async"))]
+                unreachable!();
+            } else {
+                instance.start_raw(store, start)?;
+            }
         }
         Ok(instance)
     }
 
-    /// Internal function to create an instance and run the start function.
-    ///
-    /// This function's unsafety is the same as `Instance::new_raw`.
-    #[cfg(feature = "async")]
-    async unsafe fn new_started_async<T>(
-        store: &mut StoreContextMut<'_, T>,
-        module: &Module,
-        imports: Imports<'_>,
-    ) -> Result<Instance> {
-        assert!(
-            store.0.async_support(),
-            "must use sync instantiation when async support is disabled",
-        );
-
-        store
-            .on_fiber(|store| {
-                // SAFETY: the unsafe contract of `new_started_impl` is the same
-                // as this function.
-                unsafe { Self::new_started_impl(store, module, imports) }
-            })
-            .await?
-    }
-
     /// Internal function to create an instance which doesn't have its `start`
     /// function run yet.
     ///
@@ -313,7 +283,7 @@ impl Instance {
     /// This method is unsafe because it does not type-check the `imports`
     /// provided. The `imports` provided must be suitable for the module
     /// provided as well.
-    unsafe fn new_raw(
+    async unsafe fn new_raw(
         store: &mut StoreOpaque,
         module: &Module,
         imports: Imports<'_>,
@@ -325,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();
@@ -341,11 +311,13 @@ impl Instance {
         // SAFETY: this module, by construction, was already validated within
         // the store.
         let id = unsafe {
-            store.allocate_instance(
-                AllocateInstanceKind::Module(module_id),
-                &ModuleRuntimeInfo::Module(module.clone()),
-                imports,
-            )?
+            store
+                .allocate_instance(
+                    AllocateInstanceKind::Module(module_id),
+                    &ModuleRuntimeInfo::Module(module.clone()),
+                    imports,
+                )
+                .await?
         };
 
         // Additionally, before we start doing fallible instantiation, we
@@ -377,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))
     }
@@ -905,7 +860,10 @@ impl<T: 'static> InstancePre<T> {
         // This unsafety should be handled by the type-checking performed by the
         // constructor of `InstancePre` to assert that all the imports we're passing
         // in match the module we're instantiating.
-        unsafe { Instance::new_started(&mut store, &self.module, imports.as_ref()) }
+        assert!(!store.0.async_support());
+        vm::assert_ready(unsafe {
+            Instance::new_started(&mut store, &self.module, imports.as_ref())
+        })
     }
 
     /// Creates a new instance, running the start function asynchronously
@@ -935,7 +893,7 @@ impl<T: 'static> InstancePre<T> {
         // This unsafety should be handled by the type-checking performed by the
         // constructor of `InstancePre` to assert that all the imports we're passing
         // in match the module we're instantiating.
-        unsafe { Instance::new_started_async(&mut store, &self.module, imports.as_ref()).await }
+        unsafe { Instance::new_started(&mut store, &self.module, imports.as_ref()).await }
     }
 }
 

crates/wasmtime/src/runtime/memory.rs

@@ -260,7 +260,8 @@ impl Memory {
     /// # }
     /// ```
     pub fn new(mut store: impl AsContextMut, ty: MemoryType) -> Result<Memory> {
-        Self::_new(store.as_context_mut().0, ty)
+        vm::one_poll(Self::_new(store.as_context_mut().0, ty))
+            .expect("must use `new_async` when async resource limiters are in use")
     }
 
     /// Async variant of [`Memory::new`]. You must use this variant with
@@ -273,17 +274,12 @@ impl Memory {
     /// [`Store`](`crate::Store`).
     #[cfg(feature = "async")]
     pub async fn new_async(mut store: impl AsContextMut, ty: MemoryType) -> Result<Memory> {
-        let mut store = store.as_context_mut();
-        assert!(
-            store.0.async_support(),
-            "cannot use `new_async` without enabling async support on the config"
-        );
-        store.on_fiber(|store| Self::_new(store.0, ty)).await?
+        Self::_new(store.as_context_mut().0, ty).await
     }
 
     /// Helper function for attaching the memory to a "frankenstein" instance
-    fn _new(store: &mut StoreOpaque, ty: MemoryType) -> Result<Memory> {
-        generate_memory_export(store, &ty, None)
+    async fn _new(store: &mut StoreOpaque, ty: MemoryType) -> Result<Memory> {
+        generate_memory_export(store, &ty, None).await
     }
 
     /// Returns the underlying type of this memory.
@@ -1005,7 +1001,10 @@ impl SharedMemory {
     /// Construct a single-memory instance to provide a way to import
     /// [`SharedMemory`] into other modules.
     pub(crate) fn vmimport(&self, store: &mut StoreOpaque) -> crate::runtime::vm::VMMemoryImport {
-        generate_memory_export(store, &self.ty(), Some(&self.vm))
+        // Note `vm::assert_ready` shouldn't panic here because this isn't
+        // actually allocating any new memory so resource limiting shouldn't
+        // kick in.
+        vm::assert_ready(generate_memory_export(store, &self.ty(), Some(&self.vm)))
             .unwrap()
             .vmimport(store)
     }