Implement Linker::module_async (#3121)

This implements and adds the async counterpart of the `Linker::module`
method.

Closes #3077

crates/wasmtime/src/linker.rs

@@ -3,7 +3,7 @@ use crate::instance::{InstanceData, InstancePre};
 use crate::store::StoreOpaque;
 use crate::{
     AsContextMut, Caller, Engine, Extern, ExternType, Func, FuncType, ImportType, Instance,
-    IntoFunc, Module, Trap, Val,
+    IntoFunc, Module, StoreContextMut, Trap, Val,
 };
 use anyhow::{anyhow, bail, Context, Error, Result};
 use log::warn;
@@ -596,39 +596,17 @@ impl<T> Linker<T> {
     where
         T: 'static,
     {
+        // NB: this is intended to function the same as `Linker::module_async`,
+        // they should be kept in sync.
         match ModuleKind::categorize(module)? {
-            ModuleKind::Command => self.command(store, module_name, module),
-            ModuleKind::Reactor => {
-                let instance = self.instantiate(&mut store, &module)?;
-
-                if let Some(export) = instance.get_export(&mut store, "_initialize") {
-                    if let Extern::Func(func) = export {
-                        func.typed::<(), (), _>(&store)
-                            .and_then(|f| f.call(&mut store, ()).map_err(Into::into))
-                            .context("calling the Reactor initialization function")?;
-                    }
-                }
-
-                self.instance(store, module_name, instance)
-            }
-        }
-    }
-
-    fn command(
-        &mut self,
-        mut store: impl AsContextMut<Data = T>,
-        module_name: &str,
-        module: &Module,
-    ) -> Result<&mut Self>
-    where
-        T: 'static,
-    {
-        for export in module.exports() {
-            if let Some(func_ty) = export.ty().func() {
-                let instance_pre = self.instantiate_pre(&mut store, module)?;
-                let export_name = export.name().to_owned();
-                let func = Func::new(
-                    &mut store,
+            ModuleKind::Command => {
+                self.command(
+                    store,
+                    module_name,
+                    module,
+                    |store, func_ty, export_name, instance_pre| {
+                        Func::new(
+                            store,
                             func_ty.clone(),
                             move |mut caller, params, results| {
                                 // Create a new instance for this command execution.
@@ -654,7 +632,113 @@ impl<T> Linker<T> {
 
                                 Ok(())
                             },
-                );
+                        )
+                    },
+                )
+            }
+            ModuleKind::Reactor => {
+                let instance = self.instantiate(&mut store, &module)?;
+
+                if let Some(export) = instance.get_export(&mut store, "_initialize") {
+                    if let Extern::Func(func) = export {
+                        func.typed::<(), (), _>(&store)
+                            .and_then(|f| f.call(&mut store, ()).map_err(Into::into))
+                            .context("calling the Reactor initialization function")?;
+                    }
+                }
+
+                self.instance(store, module_name, instance)
+            }
+        }
+    }
+
+    /// Define automatic instantiations of a [`Module`] in this linker.
+    ///
+    /// This is the same as [`Linker::module`], except for async `Store`s.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub async fn module_async(
+        &mut self,
+        mut store: impl AsContextMut<Data = T>,
+        module_name: &str,
+        module: &Module,
+    ) -> Result<&mut Self>
+    where
+        T: Send + 'static,
+    {
+        // NB: this is intended to function the same as `Linker::module`, they
+        // should be kept in sync.
+        match ModuleKind::categorize(module)? {
+            ModuleKind::Command => self.command(
+                store,
+                module_name,
+                module,
+                |store, func_ty, export_name, instance_pre| {
+                    let upvars = Arc::new((instance_pre, export_name));
+                    Func::new_async(
+                        store,
+                        func_ty.clone(),
+                        move |mut caller, params, results| {
+                            let upvars = upvars.clone();
+                            Box::new(async move {
+                                let (instance_pre, export_name) = &*upvars;
+                                let instance = instance_pre.instantiate_async(&mut caller).await?;
+
+                                let command_results = instance
+                                    .get_export(&mut caller, &export_name)
+                                    .unwrap()
+                                    .into_func()
+                                    .unwrap()
+                                    .call_async(&mut caller, params)
+                                    .await
+                                    .map_err(|error| error.downcast::<Trap>().unwrap())?;
+
+                                for (result, command_result) in
+                                    results.iter_mut().zip(command_results.into_vec())
+                                {
+                                    *result = command_result;
+                                }
+                                Ok(())
+                            })
+                        },
+                    )
+                },
+            ),
+            ModuleKind::Reactor => {
+                let instance = self.instantiate_async(&mut store, &module).await?;
+
+                if let Some(export) = instance.get_export(&mut store, "_initialize") {
+                    if let Extern::Func(func) = export {
+                        let func = func
+                            .typed::<(), (), _>(&store)
+                            .context("loading the Reactor initialization function")?;
+                        func.call_async(&mut store, ())
+                            .await
+                            .context("calling the Reactor initialization function")?;
+                    }
+                }
+
+                self.instance(store, module_name, instance)
+            }
+        }
+    }
+
+    fn command(
+        &mut self,
+        mut store: impl AsContextMut<Data = T>,
+        module_name: &str,
+        module: &Module,
+        mk_func: impl Fn(&mut StoreContextMut<T>, &FuncType, String, InstancePre<T>) -> Func,
+    ) -> Result<&mut Self>
+    where
+        T: 'static,
+    {
+        let mut store = store.as_context_mut();
+        for export in module.exports() {
+            if let Some(func_ty) = export.ty().func() {
+                let instance_pre = self.instantiate_pre(&mut store, module)?;
+                let export_name = export.name().to_owned();
+                let func = mk_func(&mut store, func_ty, export_name, instance_pre);
                 let key = self.import_key(module_name, Some(export.name()));
                 self.insert(key, Definition::Extern(func.into()))?;
             } else if export.name() == "memory" && export.ty().memory().is_some() {

tests/all/async_functions.rs

@@ -639,3 +639,85 @@ fn recursive_async() -> Result<()> {
     run(f2.call_async(&mut store, &[]))?;
     Ok(())
 }
+
+#[test]
+fn linker_module_command() -> Result<()> {
+    run(async {
+        let mut store = async_store();
+        let mut linker = Linker::new(store.engine());
+        let module1 = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (global $g (mut i32) (i32.const 0))
+
+                    (func (export "_start"))
+
+                    (func (export "g") (result i32)
+                        global.get $g
+                        i32.const 1
+                        global.set $g)
+                )
+            "#,
+        )?;
+        let module2 = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (import "" "g" (func (result i32)))
+
+                    (func (export "get") (result i32)
+                        call 0)
+                )
+            "#,
+        )?;
+
+        linker.module_async(&mut store, "", &module1).await?;
+        let instance = linker.instantiate_async(&mut store, &module2).await?;
+        let f = instance.get_typed_func::<(), i32, _>(&mut store, "get")?;
+        assert_eq!(f.call_async(&mut store, ()).await?, 0);
+        assert_eq!(f.call_async(&mut store, ()).await?, 0);
+
+        Ok(())
+    })
+}
+
+#[test]
+fn linker_module_reactor() -> Result<()> {
+    run(async {
+        let mut store = async_store();
+        let mut linker = Linker::new(store.engine());
+        let module1 = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (global $g (mut i32) (i32.const 0))
+
+                    (func (export "g") (result i32)
+                        global.get $g
+                        i32.const 1
+                        global.set $g)
+                )
+            "#,
+        )?;
+        let module2 = Module::new(
+            store.engine(),
+            r#"
+                (module
+                    (import "" "g" (func (result i32)))
+
+                    (func (export "get") (result i32)
+                        call 0)
+                )
+            "#,
+        )?;
+
+        linker.module_async(&mut store, "", &module1).await?;
+        let instance = linker.instantiate_async(&mut store, &module2).await?;
+        let f = instance.get_typed_func::<(), i32, _>(&mut store, "get")?;
+        assert_eq!(f.call_async(&mut store, ()).await?, 0);
+        assert_eq!(f.call_async(&mut store, ()).await?, 1);
+
+        Ok(())
+    })
+}