component model: async host function & embedding support (#5055)

* func_wrap_async typechecks

* func call async

* instantiate_async

* fixes

* async engine creation for tests

* start adding a component model test for async

* fix wrong check for async support, factor out Instance::new_started to an unchecked impl

* tests: wibbles

* component::Linker::func_wrap: replace IntoComponentFunc with directly accepting a closure

We find that this makes the Linker::func_wrap type signature much easier
to read. The IntoComponentFunc abstraction was adding a lot of weight to
"splat" a set of arguments from a tuple of types into individual
arguments to the closure. Additionally, making the StoreContextMut
argument optional, or the Result<return> optional, wasn't very
worthwhile.

* Fixes for the new style of closure required by component::Linker::func_wrap

* future of result of return

* add Linker::instantiate_async and {Typed}Func::post_return_async

* fix fuzzing generator

* note optimisation opportunity

* simplify test

crates/misc/component-test-util/src/lib.rs

@@ -65,6 +65,12 @@ pub fn engine() -> Engine {
     Engine::new(&config()).unwrap()
 }
 
+pub fn async_engine() -> Engine {
+    let mut config = config();
+    config.async_support(true);
+    Engine::new(&config).unwrap()
+}
+
 /// Newtype wrapper for `f32` whose `PartialEq` impl considers NaNs equal to each other.
 #[derive(Copy, Clone, Debug, Arbitrary)]
 pub struct Float32(pub f32);

crates/wasmtime/src/component/func.rs

@@ -277,11 +277,59 @@ impl Func {
     /// The `params` here must match the type signature of this `Func`, or this will return an error. If a trap
     /// occurs while executing this function, then an error will also be returned.
     // TODO: say more -- most of the docs for `TypedFunc::call` apply here, too
+    //
+    // # Panics
+    //
+    // Panics if this is called on a function in an asyncronous store. This only works
+    // with functions defined within a synchronous store. Also panics if `store`
+    // does not own this function.
     pub fn call(
         &self,
         mut store: impl AsContextMut,
         params: &[Val],
         results: &mut [Val],
+    ) -> Result<()> {
+        let mut store = store.as_context_mut();
+        assert!(
+            !store.0.async_support(),
+            "must use `call_async` when async support is enabled on the config"
+        );
+        self.call_impl(&mut store.as_context_mut(), params, results)
+    }
+
+    /// Exactly like [`Self::call`] except for use on async stores.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this is called on a function in a synchronous store. This only works
+    /// with functions defined within an asynchronous store. Also panics if `store`
+    /// does not own this function.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub async fn call_async<T>(
+        &self,
+        mut store: impl AsContextMut<Data = T>,
+        params: &[Val],
+        results: &mut [Val],
+    ) -> Result<()>
+    where
+        T: Send,
+    {
+        let mut store = store.as_context_mut();
+        assert!(
+            store.0.async_support(),
+            "cannot use `call_async` without enabling async support in the config"
+        );
+        store
+            .on_fiber(|store| self.call_impl(store, params, results))
+            .await?
+    }
+
+    fn call_impl(
+        &self,
+        mut store: impl AsContextMut,
+        params: &[Val],
+        results: &mut [Val],
     ) -> Result<()> {
         let store = &mut store.as_context_mut();
 
@@ -490,7 +538,42 @@ impl Func {
     /// called, then it will panic. If a different [`Func`] for the same
     /// component instance was invoked then this function will also panic
     /// because the `post-return` needs to happen for the other function.
+    ///
+    /// Panics if this is called on a function in an asynchronous store.
+    /// This only works with functions defined within a synchronous store.
     pub fn post_return(&self, mut store: impl AsContextMut) -> Result<()> {
+        let store = store.as_context_mut();
+        assert!(
+            !store.0.async_support(),
+            "must use `post_return_async` when async support is enabled on the config"
+        );
+        self.post_return_impl(store)
+    }
+
+    /// Exactly like [`Self::post_return`] except for use on async stores.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this is called on a function in a synchronous store. This
+    /// only works with functions defined within an asynchronous store.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub async fn post_return_async<T: Send>(
+        &self,
+        mut store: impl AsContextMut<Data = T>,
+    ) -> Result<()> {
+        let mut store = store.as_context_mut();
+        assert!(
+            store.0.async_support(),
+            "cannot use `call_async` without enabling async support in the config"
+        );
+        // Future optimization opportunity: conditionally use a fiber here since
+        // some func's post_return will not need the async context (i.e. end up
+        // calling async host functionality)
+        store.on_fiber(|store| self.post_return_impl(store)).await?
+    }
+
+    fn post_return_impl(&self, mut store: impl AsContextMut) -> Result<()> {
         let mut store = store.as_context_mut();
         let data = &mut store.0[self.0];
         let instance = data.instance;

crates/wasmtime/src/component/func/typed.rs

@@ -145,8 +145,47 @@ where
     ///
     /// # Panics
     ///
-    /// This function will panic if `store` does not own this function.
+    /// Panics if this is called on a function in an asynchronous store. This
-    pub fn call(&self, mut store: impl AsContextMut, params: Params) -> Result<Return> {
+    /// only works with functions defined within a synchonous store. Also
+    /// panics if `store` does not own this function.
+    pub fn call(&self, store: impl AsContextMut, params: Params) -> Result<Return> {
+        assert!(
+            !store.as_context().async_support(),
+            "must use `call_async` when async support is enabled on the config"
+        );
+        self.call_impl(store, params)
+    }
+
+    /// Exactly like [`Self::call`], except for use on asynchronous stores.
+    ///
+    /// # Panics
+    ///
+    /// Panics if this is called on a function in a synchronous store. This
+    /// only works with functions defined within an asynchronous store. Also
+    /// panics if `store` does not own this function.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub async fn call_async<T>(
+        &self,
+        mut store: impl AsContextMut<Data = T>,
+        params: Params,
+    ) -> Result<Return>
+    where
+        T: Send,
+        Params: Send + Sync,
+        Return: Send + Sync,
+    {
+        let mut store = store.as_context_mut();
+        assert!(
+            store.0.async_support(),
+            "cannot use `call_async` when async support is not enabled on the config"
+        );
+        store
+            .on_fiber(|store| self.call_impl(store, params))
+            .await?
+    }
+
+    fn call_impl(&self, mut store: impl AsContextMut, params: Params) -> Result<Return> {
         let store = &mut store.as_context_mut();
         // Note that this is in theory simpler than it might read at this time.
         // Here we're doing a runtime dispatch on the `flatten_count` for the
@@ -286,6 +325,16 @@ where
     pub fn post_return(&self, store: impl AsContextMut) -> Result<()> {
         self.func.post_return(store)
     }
+
+    /// See [`Func::post_return_async`]
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub async fn post_return_async<T: Send>(
+        &self,
+        store: impl AsContextMut<Data = T>,
+    ) -> Result<()> {
+        self.func.post_return_async(store).await
+    }
 }
 
 /// A trait representing a static list of named types that can be passed to or

crates/wasmtime/src/component/instance.rs

@@ -259,10 +259,16 @@ impl<'a> Instantiator<'a> {
 
                     // Note that the unsafety here should be ok because the
                     // validity of the component means that type-checks have
-                    // already been performed. This maens that the unsafety due
+                    // already been performed. This means that the unsafety due
                     // to imports having the wrong type should not happen here.
-                    let i =
+                    //
-                        unsafe { crate::Instance::new_started(store, module, imports.as_ref())? };
+                    // Also note we are calling new_started_impl because we have
+                    // already checked for asyncness and are running on a fiber
+                    // if required.
+
+                    let i = unsafe {
+                        crate::Instance::new_started_impl(store, module, imports.as_ref())?
+                    };
                     self.data.instances.push(i);
                 }
 
@@ -484,7 +490,36 @@ impl<T> InstancePre<T> {
     /// Performs the instantiation process into the store specified.
     //
     // TODO: needs more docs
-    pub fn instantiate(&self, mut store: impl AsContextMut<Data = T>) -> Result<Instance> {
+    pub fn instantiate(&self, store: impl AsContextMut<Data = T>) -> Result<Instance> {
+        assert!(
+            !store.as_context().async_support(),
+            "must use async instantiation when async support is enabled"
+        );
+        self.instantiate_impl(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")]
+    #[cfg_attr(nightlydoc, doc(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?
+    }
+
+    fn instantiate_impl(&self, mut store: impl AsContextMut<Data = T>) -> Result<Instance> {
         let mut store = store.as_context_mut();
         let mut i = Instantiator::new(&self.component, store.0, &self.imports);
         i.run(&mut store)?;

crates/wasmtime/src/component/linker.rs

@@ -5,7 +5,9 @@ use crate::component::{Component, ComponentNamedList, Instance, InstancePre, Lif
 use crate::{AsContextMut, Engine, Module, StoreContextMut};
 use anyhow::{anyhow, bail, Context, Result};
 use std::collections::hash_map::{Entry, HashMap};
+use std::future::Future;
 use std::marker;
+use std::pin::Pin;
 use std::sync::Arc;
 use wasmtime_environ::component::TypeDef;
 use wasmtime_environ::PrimaryMap;
@@ -36,6 +38,7 @@ pub struct Strings {
 /// a "bag of named items", so each [`LinkerInstance`] can further define items
 /// internally.
 pub struct LinkerInstance<'a, T> {
+    engine: Engine,
     strings: &'a mut Strings,
     map: &'a mut NameMap,
     allow_shadowing: bool,
@@ -82,6 +85,7 @@ impl<T> Linker<T> {
     /// the root namespace.
     pub fn root(&mut self) -> LinkerInstance<'_, T> {
         LinkerInstance {
+            engine: self.engine.clone(),
             strings: &mut self.strings,
             map: &mut self.map,
             allow_shadowing: self.allow_shadowing,
@@ -187,13 +191,47 @@ impl<T> Linker<T> {
         store: impl AsContextMut<Data = T>,
         component: &Component,
     ) -> Result<Instance> {
+        assert!(
+            !store.as_context().async_support(),
+            "must use async instantiation when async support is enabled"
+        );
         self.instantiate_pre(component)?.instantiate(store)
     }
+
+    /// Instantiates the [`Component`] provided into the `store` specified.
+    ///
+    /// This is exactly like [`Linker::instantiate`] except for async stores.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if this [`Linker`] doesn't define an import that
+    /// `component` requires or if it is of the wrong type. Additionally this
+    /// can return an error if something goes wrong during instantiation such as
+    /// a runtime trap or a runtime limit being exceeded.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub async fn instantiate_async(
+        &self,
+        store: impl AsContextMut<Data = T>,
+        component: &Component,
+    ) -> Result<Instance>
+    where
+        T: Send,
+    {
+        assert!(
+            store.as_context().async_support(),
+            "must use sync instantiation when async support is disabled"
+        );
+        self.instantiate_pre(component)?
+            .instantiate_async(store)
+            .await
+    }
 }
 
 impl<T> LinkerInstance<'_, T> {
     fn as_mut(&mut self) -> LinkerInstance<'_, T> {
         LinkerInstance {
+            engine: self.engine.clone(),
             strings: self.strings,
             map: self.map,
             allow_shadowing: self.allow_shadowing,
@@ -229,6 +267,36 @@ impl<T> LinkerInstance<'_, T> {
         self.insert(name, Definition::Func(HostFunc::from_closure(func)))
     }
 
+    /// Defines a new host-provided async function into this [`Linker`].
+    ///
+    /// This is exactly like [`Self::func_wrap`] except it takes an async
+    /// host function.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub fn func_wrap_async<Params, Return, F>(&mut self, name: &str, f: F) -> Result<()>
+    where
+        F: for<'a> Fn(
+                StoreContextMut<'a, T>,
+                Params,
+            ) -> Box<dyn Future<Output = Result<Return>> + Send + 'a>
+            + Send
+            + Sync
+            + 'static,
+        Params: ComponentNamedList + Lift + 'static,
+        Return: ComponentNamedList + Lower + 'static,
+    {
+        assert!(
+            self.engine.config().async_support,
+            "cannot use `func_wrap_async` without enabling async support in the config"
+        );
+        let ff = move |mut store: StoreContextMut<'_, T>, params: Params| -> Result<Return> {
+            let async_cx = store.as_context_mut().0.async_cx().expect("async cx");
+            let mut future = Pin::from(f(store.as_context_mut(), params));
+            unsafe { async_cx.block_on(future.as_mut()) }?
+        };
+        self.func_wrap(name, ff)
+    }
+
     /// Define a new host-provided function using dynamic types.
     ///
     /// `name` must refer to a function type import in `component`.  If and when
@@ -260,6 +328,8 @@ impl<T> LinkerInstance<'_, T> {
         Err(anyhow!("import `{name}` not found"))
     }
 
+    // TODO: define func_new_async
+
     /// Defines a [`Module`] within this instance.
     ///
     /// This can be used to provide a core wasm [`Module`] as an import to a

crates/wasmtime/src/instance.rs

@@ -174,7 +174,18 @@ impl Instance {
             !store.0.async_support(),
             "must use async instantiation when async support is enabled",
         );
+        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>(
+        store: &mut StoreContextMut<'_, T>,
+        module: &Module,
+        imports: Imports<'_>,
+    ) -> Result<Instance> {
         let (instance, start) = Instance::new_raw(store.0, module, imports)?;
         if let Some(start) = start {
             instance.start_raw(store, start)?;
@@ -194,22 +205,13 @@ impl Instance {
     where
         T: Send,
     {
-        // Note that the body of this function is intentionally quite similar
-        // to the `new_started` function, and it's intended that the two bodies
-        // here are small enough to be ok duplicating.
         assert!(
             store.0.async_support(),
             "must use sync instantiation when async support is disabled",
         );
 
         store
-            .on_fiber(|store| {
+            .on_fiber(|store| Self::new_started_impl(store, module, imports))
-                let (instance, start) = Instance::new_raw(store.0, module, imports)?;
-                if let Some(start) = start {
-                    instance.start_raw(store, start)?;
-                }
-                Ok(instance)
-            })
             .await?
     }
 

tests/all/component_model.rs

@@ -1,10 +1,11 @@
 use anyhow::Result;
-use component_test_util::{engine, TypedFuncExt};
+use component_test_util::{async_engine, engine, TypedFuncExt};
 use std::fmt::Write;
 use std::iter;
 use wasmtime::component::Component;
 use wasmtime_component_util::REALLOC_AND_FREE;
 
+mod r#async;
 mod dynamic;
 mod func;
 mod import;

tests/all/component_model/async.rs

@@ -0,0 +1,88 @@
+use anyhow::Result;
+use wasmtime::component::*;
+use wasmtime::{Store, StoreContextMut, Trap, TrapCode};
+
+/// This is super::func::thunks, except with an async store.
+#[tokio::test]
+async fn smoke() -> Result<()> {
+    let component = r#"
+        (component
+            (core module $m
+                (func (export "thunk"))
+                (func (export "thunk-trap") unreachable)
+            )
+            (core instance $i (instantiate $m))
+            (func (export "thunk")
+                (canon lift (core func $i "thunk"))
+            )
+            (func (export "thunk-trap")
+                (canon lift (core func $i "thunk-trap"))
+            )
+        )
+    "#;
+
+    let engine = super::async_engine();
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    let instance = Linker::new(&engine)
+        .instantiate_async(&mut store, &component)
+        .await?;
+
+    let thunk = instance.get_typed_func::<(), (), _>(&mut store, "thunk")?;
+
+    thunk.call_async(&mut store, ()).await?;
+    thunk.post_return_async(&mut store).await?;
+
+    let err = instance
+        .get_typed_func::<(), (), _>(&mut store, "thunk-trap")?
+        .call_async(&mut store, ())
+        .await
+        .unwrap_err();
+    assert!(err.downcast::<Trap>()?.trap_code() == Some(TrapCode::UnreachableCodeReached));
+
+    Ok(())
+}
+
+/// Handle an import function, created using component::Linker::func_wrap_async.
+#[tokio::test]
+async fn smoke_func_wrap() -> Result<()> {
+    let component = r#"
+        (component
+            (type $f (func))
+            (import "i" (func $f))
+
+            (core module $m
+                (import "imports" "i" (func $i))
+                (func (export "thunk") call $i)
+            )
+
+            (core func $f (canon lower (func $f)))
+            (core instance $i (instantiate $m
+                (with "imports" (instance
+                    (export "i" (func $f))
+                ))
+             ))
+            (func (export "thunk")
+                (canon lift (core func $i "thunk"))
+            )
+        )
+    "#;
+
+    let engine = super::async_engine();
+    let component = Component::new(&engine, component)?;
+    let mut store = Store::new(&engine, ());
+    let mut linker = Linker::new(&engine);
+    let mut root = linker.root();
+    root.func_wrap_async("i", |_: StoreContextMut<()>, _: ()| {
+        Box::new(async { Ok(()) })
+    })?;
+
+    let instance = linker.instantiate_async(&mut store, &component).await?;
+
+    let thunk = instance.get_typed_func::<(), (), _>(&mut store, "thunk")?;
+
+    thunk.call_async(&mut store, ()).await?;
+    thunk.post_return_async(&mut store).await?;
+
+    Ok(())
+}