From 31e3a4d208b340e850b5e2c33b8aa9e1f751e9f5 Mon Sep 17 00:00:00 2001 From: Isaac Clayton Date: Mon, 13 Jun 2022 12:52:11 +0200 Subject: [PATCH] WIP: wrap async closures host-side --- crates/plugin_runtime/src/lib.rs | 8 ++ crates/plugin_runtime/src/wasi.rs | 167 +++++++++++++++++++++++++++++- 2 files changed, 173 insertions(+), 2 deletions(-) diff --git a/crates/plugin_runtime/src/lib.rs b/crates/plugin_runtime/src/lib.rs index cbcfb540a1..afbcfd9278 100644 --- a/crates/plugin_runtime/src/lib.rs +++ b/crates/plugin_runtime/src/lib.rs @@ -20,6 +20,12 @@ mod tests { imports: WasiFn, } + async fn half(a: u32) -> u32 { + a / 2 + } + + let x = half; + async { let mut runtime = WasiPluginBuilder::new_with_default_ctx() .unwrap() @@ -31,6 +37,8 @@ mod tests { .unwrap() .host_function("import_swap", |(a, b): (u32, u32)| (b, a)) .unwrap() + // .host_function_async("import_half", half) + // .unwrap() .init(include_bytes!("../../../plugins/bin/test_plugin.wasm")) .await .unwrap(); diff --git a/crates/plugin_runtime/src/wasi.rs b/crates/plugin_runtime/src/wasi.rs index 19acf7a6b6..595162c7fd 100644 --- a/crates/plugin_runtime/src/wasi.rs +++ b/crates/plugin_runtime/src/wasi.rs @@ -79,11 +79,153 @@ impl WasiPluginBuilder { Self::new(wasi_ctx) } - pub fn host_function( + // pub fn host_function_async( + // mut self, + // name: &str, + // function: impl Fn(A) -> Pin + Send + Sync>> + Sync + Send + 'static, + // ) -> Result + // where + // A: DeserializeOwned + Send, + // R: Serialize + Send, + // { + // self.linker.func_wrap1_async( + // "env", + // &format!("__{}", name), + // move |caller: Caller<'_, WasiCtxAlloc>, packed_buffer: u64| { + // // let function = &function; + // Box::new(async move { + // // grab a handle to the memory + // let mut plugin_memory = match caller.get_export("memory") { + // Some(Extern::Memory(mem)) => mem, + // _ => return Err(Trap::new("Could not grab slice of plugin memory"))?, + // }; + + // let buffer = WasiBuffer::from_u64(packed_buffer); + + // // get the args passed from Guest + // let args = Wasi::buffer_to_type(&mut plugin_memory, &mut caller, &buffer)?; + + // // Call the Host-side function + // let result: R = function(args).await; + + // // Serialize the result back to guest + // let result = Wasi::serialize_to_bytes(result).map_err(|_| { + // Trap::new("Could not serialize value returned from function") + // })?; + + // // Ok((buffer, plugin_memory, result)) + // Wasi::buffer_to_free(caller.data().free_buffer(), &mut caller, buffer).await?; + + // let buffer = Wasi::bytes_to_buffer( + // caller.data().alloc_buffer(), + // &mut plugin_memory, + // &mut caller, + // result, + // ) + // .await?; + + // Ok(buffer.into_u64()) + // }) + // }, + // )?; + // Ok(self) + // } + + // pub fn host_function_async(mut self, name: &str, function: F) -> Result + // where + // F: Fn(u64) -> Pin + Send + Sync + 'static>> + // + Send + // + Sync + // + 'static, + // { + // self.linker.func_wrap1_async( + // "env", + // &format!("__{}", name), + // move |_: Caller<'_, WasiCtxAlloc>, _: u64| { + // // let function = &function; + // Box::new(async { + // let function = function; + // // Call the Host-side function + // let result: u64 = function(7).await; + // Ok(result) + // }) + // }, + // )?; + // Ok(self) + // } + + // pub fn host_function_async(mut self, name: &str, function: F) -> Result + // where + // F: Fn(A) -> Pin + Send + 'static>> + Send + Sync + 'static, + // A: DeserializeOwned + Send, + // R: Serialize + Send + Sync, + // { + // self.linker.func_wrap1_async( + // "env", + // &format!("__{}", name), + // move |mut caller: Caller<'_, WasiCtxAlloc>, packed_buffer: u64| { + // let function = |args: Vec| { + // let args = args; + // let args: A = Wasi::deserialize_to_type(&args)?; + // Ok(async { + // let result = function(args); + // Wasi::serialize_to_bytes(result.await).map_err(|_| { + // Trap::new("Could not serialize value returned from function").into() + // }) + // }) + // }; + + // // TODO: use try block once avaliable + // let result: Result<(WasiBuffer, Memory, _), Trap> = (|| { + // // grab a handle to the memory + // let mut plugin_memory = match caller.get_export("memory") { + // Some(Extern::Memory(mem)) => mem, + // _ => return Err(Trap::new("Could not grab slice of plugin memory"))?, + // }; + + // let buffer = WasiBuffer::from_u64(packed_buffer); + + // // get the args passed from Guest + // let args = Wasi::buffer_to_bytes(&mut plugin_memory, &mut caller, &buffer)?; + + // // Call the Host-side function + // let result = function(args); + + // Ok((buffer, plugin_memory, result)) + // })(); + + // Box::new(async move { + // let (buffer, mut plugin_memory, thingo) = result?; + // let thingo: Result<_, Error> = thingo; + // let result: Result, Error> = thingo?.await; + + // // Wasi::buffer_to_free(caller.data().free_buffer(), &mut caller, buffer).await?; + + // // let buffer = Wasi::bytes_to_buffer( + // // caller.data().alloc_buffer(), + // // &mut plugin_memory, + // // &mut caller, + // // result, + // // ) + // // .await?; + + // // Ok(buffer.into_u64()) + // Ok(27) + // }) + // }, + // )?; + // Ok(self) + // } + + pub fn host_function( mut self, name: &str, function: impl Fn(A) -> R + Send + Sync + 'static, - ) -> Result { + ) -> Result + where + A: DeserializeOwned + Send, + R: Serialize + Send + Sync, + { self.linker.func_wrap1_async( "env", &format!("__{}", name), @@ -335,6 +477,12 @@ impl Wasi { Ok(bytes) } + fn deserialize_to_type(bytes: &[u8]) -> Result { + // serialize the argument using bincode + let bytes = bincode::deserialize(bytes)?; + Ok(bytes) + } + // fn deserialize( // plugin_memory: &mut Memory, // mut store: impl AsContextMut, @@ -382,6 +530,21 @@ impl Wasi { Ok(result) } + /// Takes a `(ptr, len)` pair and returns the corresponding deserialized buffer. + fn buffer_to_bytes<'a>( + plugin_memory: &'a Memory, + store: impl AsContext + 'a, + buffer: &WasiBuffer, + ) -> Result, Error> { + let buffer_start = buffer.ptr as usize; + let buffer_end = buffer_start + buffer.len as usize; + + // read the buffer at this point into a byte array + // deserialize the byte array into the provided serde type + let result = plugin_memory.data(store.as_context())[buffer_start..buffer_end].to_vec(); + Ok(result) + } + async fn buffer_to_free( free_buffer: TypedFunc, mut store: impl AsContextMut,