From 9b088756b3c0971b231a231a7f9e077c84c694a9 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Tue, 27 Jul 2021 16:17:45 -0500 Subject: [PATCH] Implement `Linker::module_async` (#3121) This implements and adds the async counterpart of the `Linker::module` method. Closes #3077 --- crates/wasmtime/src/linker.rs | 146 ++++++++++++++++++++++++++-------- tests/all/async_functions.rs | 82 +++++++++++++++++++ 2 files changed, 197 insertions(+), 31 deletions(-) diff --git a/crates/wasmtime/src/linker.rs b/crates/wasmtime/src/linker.rs index 22ddead971..69c596b31b 100644 --- a/crates/wasmtime/src/linker.rs +++ b/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,8 +596,46 @@ impl Linker { 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::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. + let instance = instance_pre.instantiate(&mut caller)?; + + // `unwrap()` everything here because we know the instance contains a + // function export with the given name and signature because we're + // iterating over the module it was instantiated from. + let command_results = instance + .get_export(&mut caller, &export_name) + .unwrap() + .into_func() + .unwrap() + .call(&mut caller, params) + .map_err(|error| error.downcast::().unwrap())?; + + // Copy the return values into the output slice. + for (result, command_result) in + results.iter_mut().zip(command_results.into_vec()) + { + *result = command_result; + } + + Ok(()) + }, + ) + }, + ) + } ModuleKind::Reactor => { let instance = self.instantiate(&mut store, &module)?; @@ -614,47 +652,93 @@ impl Linker { } } - fn command( + /// 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, 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::().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, + module_name: &str, + module: &Module, + mk_func: impl Fn(&mut StoreContextMut, &FuncType, String, InstancePre) -> 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 = Func::new( - &mut store, - func_ty.clone(), - move |mut caller, params, results| { - // Create a new instance for this command execution. - let instance = instance_pre.instantiate(&mut caller)?; - - // `unwrap()` everything here because we know the instance contains a - // function export with the given name and signature because we're - // iterating over the module it was instantiated from. - let command_results = instance - .get_export(&mut caller, &export_name) - .unwrap() - .into_func() - .unwrap() - .call(&mut caller, params) - .map_err(|error| error.downcast::().unwrap())?; - - // Copy the return values into the output slice. - for (result, command_result) in - results.iter_mut().zip(command_results.into_vec()) - { - *result = command_result; - } - - Ok(()) - }, - ); + 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() { diff --git a/tests/all/async_functions.rs b/tests/all/async_functions.rs index 6246e2631a..20349c5a58 100644 --- a/tests/all/async_functions.rs +++ b/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(()) + }) +}