Make core instance allocation an async function

This commit is a step in preparation for bytecodealliance#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.

Author: alexcrichton

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].
@@ -197,7 +198,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(
@@ -239,62 +240,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.
     ///
@@ -310,7 +280,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<'_>,
@@ -338,11 +308,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
@@ -885,7 +857,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
@@ -919,7 +894,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

@@ -1,5 +1,6 @@
 use crate::Trap;
 use crate::prelude::*;
+use crate::runtime::vm;
 use crate::store::{StoreInstanceId, StoreOpaque};
 use crate::trampoline::generate_memory_export;
 use crate::{AsContext, AsContextMut, Engine, MemoryType, StoreContext, StoreContextMut};
@@ -259,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
@@ -272,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.
@@ -1007,7 +1004,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)
     }

crates/wasmtime/src/runtime/store.rs

@@ -666,15 +666,17 @@ impl<T> Store<T> {
             .unwrap();
 
         unsafe {
-            let id = inner
-                .allocate_instance(
-                    AllocateInstanceKind::Dummy {
-                        allocator: &allocator,
-                    },
-                    &shim,
-                    Default::default(),
-                )
-                .expect("failed to allocate default callee");
+            // Note that this dummy instance doesn't allocate tables or memories
+            // so it won't have an async await point meaning that it should be
+            // ok to assert the future is always ready.
+            let id = vm::assert_ready(inner.allocate_instance(
+                AllocateInstanceKind::Dummy {
+                    allocator: &allocator,
+                },
+                &shim,
+                Default::default(),
+            ))
+            .expect("failed to allocate default callee");
             let default_caller_vmctx = inner.instance(id).vmctx();
             inner.default_caller_vmctx = default_caller_vmctx.into();
         }
@@ -2173,7 +2175,7 @@ at https://bytecodealliance.org/security.
     ///
     /// The `imports` provided must be correctly sized/typed for the module
     /// being allocated.
-    pub(crate) unsafe fn allocate_instance(
+    pub(crate) async unsafe fn allocate_instance(
         &mut self,
         kind: AllocateInstanceKind<'_>,
         runtime_info: &ModuleRuntimeInfo,

crates/wasmtime/src/runtime/trampoline.rs

@@ -19,21 +19,21 @@ use crate::store::StoreOpaque;
 use crate::{MemoryType, TableType, TagType};
 use wasmtime_environ::{MemoryIndex, TableIndex, TagIndex};
 
-pub fn generate_memory_export(
+pub async fn generate_memory_export(
     store: &mut StoreOpaque,
     m: &MemoryType,
     preallocation: Option<&SharedMemory>,
 ) -> Result<crate::Memory> {
     let id = store.id();
-    let instance = create_memory(store, m, preallocation)?;
+    let instance = create_memory(store, m, preallocation).await?;
     Ok(store
         .instance_mut(instance)
         .get_exported_memory(id, MemoryIndex::from_u32(0)))
 }
 
-pub fn generate_table_export(store: &mut StoreOpaque, t: &TableType) -> Result<crate::Table> {
+pub async fn generate_table_export(store: &mut StoreOpaque, t: &TableType) -> Result<crate::Table> {
     let id = store.id();
-    let instance = create_table(store, t)?;
+    let instance = create_table(store, t).await?;
     Ok(store
         .instance_mut(instance)
         .get_exported_table(id, TableIndex::from_u32(0)))

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

@@ -24,7 +24,7 @@ use wasmtime_environ::{
 /// This separate instance is necessary because Wasm objects in Wasmtime must be
 /// attached to instances (versus the store, e.g.) and some objects exist
 /// outside: a host-provided memory import, shared memory.
-pub fn create_memory(
+pub async fn create_memory(
     store: &mut StoreOpaque,
     memory_ty: &MemoryType,
     preallocation: Option<&SharedMemory>,
@@ -52,13 +52,15 @@ pub fn create_memory(
         ondemand: OnDemandInstanceAllocator::default(),
     };
     unsafe {
-        store.allocate_instance(
-            AllocateInstanceKind::Dummy {
-                allocator: &allocator,
-            },
-            &ModuleRuntimeInfo::bare(Arc::new(module)),
-            Default::default(),
-        )
+        store
+            .allocate_instance(
+                AllocateInstanceKind::Dummy {
+                    allocator: &allocator,
+                },
+                &ModuleRuntimeInfo::bare(Arc::new(module)),
+                Default::default(),
+            )
+            .await
     }
 }