Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Factor out buffer code

Browse source
This commit is contained in:
Isaac Clayton 2022-06-07 10:10:03 +02:00
parent fbaff615a3
commit e9b87f3dc3
4 changed files with 109 additions and 118 deletions
Download patch file Download diff file Expand all files Collapse all files

104
crates/plugin_runtime/src/wasi.rs
View file

@ -14,8 +14,8 @@ pub struct Wasi {
module: Module,
store: Store<WasiCtx>,
instance: Instance,
alloc_buffer: TypedFunc<i32, i32>,
// free_buffer: TypedFunc<(i32, i32), ()>,
alloc_buffer: TypedFunc<u32, u32>,
// free_buffer: TypedFunc<(u32, u32), ()>,
}
pub struct WasiPlugin {
@ -54,27 +54,20 @@ impl Wasi {
let engine = Engine::default();
let mut linker = Linker::new(&engine);
linker.func_wrap("env", "hello", |x: u32| x * 2).unwrap();
linker.func_wrap("env", "bye", |x: u32| x / 2).unwrap();
linker.func_wrap("env", "__hello", |x: u32| x * 2).unwrap();
linker.func_wrap("env", "__bye", |x: u32| x / 2).unwrap();
println!("linking");
wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
println!("linked");
let mut store: Store<_> = Store::new(&engine, plugin.wasi_ctx);
println!("moduling");
let module = Module::new(&engine, plugin.module)?;
println!("moduled");
linker.module(&mut store, "", &module)?;
println!("linked again");
let instance = linker.instantiate(&mut store, &module)?;
println!("instantiated");
let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
// let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
println!("can alloc");
Ok(Wasi {
engine,
module,
@ -99,7 +92,6 @@ impl Wasi {
// grab an empty file descriptor, specify capabilities
let fd = ctx.table().push(Box::new(()))?;
dbg!(fd);
let caps = dir::DirCaps::all();
let file_caps = file::FileCaps::all();
@ -172,62 +164,78 @@ impl Wasi {
// This isn't a problem because Wasm stops executing after the function returns,
// so the heap is still valid for our inspection when we want to pull things out.
// TODO: dont' use as for conversions
pub fn call<A: Serialize, R: DeserializeOwned>(
&mut self,
handle: &str,
arg: A,
) -> Result<R, Error> {
dbg!(&handle);
// dbg!(serde_json::to_string(&arg)).unwrap();
/// Takes an item, allocates a buffer, serializes the argument to that buffer,
/// and returns a (ptr, len) pair to that buffer.
fn serialize_to_buffer<T: Serialize>(&mut self, item: T) -> Result<(u32, u32), Error> {
// serialize the argument using bincode
let arg = bincode::serialize(&arg)?;
let arg_buffer_len = arg.len();
let item = bincode::serialize(&item)?;
let buffer_len = item.len() as u32;
// allocate a buffer and write the argument to that buffer
let arg_buffer_ptr = self
.alloc_buffer
.call(&mut self.store, arg_buffer_len as i32)?;
let buffer_ptr = self.alloc_buffer.call(&mut self.store, buffer_len)?;
let plugin_memory = self
.instance
.get_memory(&mut self.store, "memory")
.ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
plugin_memory.write(&mut self.store, arg_buffer_ptr as usize, &arg)?;
// get the webassembly function we want to actually call
// TODO: precompute handle
let fun_name = format!("__{}", handle);
let fun = self
.instance
.get_typed_func::<(i32, i32), i32, _>(&mut self.store, &fun_name)?;
// call the function, passing in the buffer and its length
// this should return a pointer to a (ptr, lentgh) pair
let arg_buffer = (arg_buffer_ptr, arg_buffer_len as i32);
let result_buffer = fun.call(&mut self.store, arg_buffer)?;
plugin_memory.write(&mut self.store, buffer_ptr as usize, &item)?;
Ok((buffer_ptr, buffer_len))
}
/// Takes a ptr to a (ptr, len) pair and returns the corresponding deserialized buffer
fn deserialize_from_buffer<R: DeserializeOwned>(&mut self, buffer: u32) -> Result<R, Error> {
// create a buffer to read the (ptr, length) pair into
// this is a total of 4 + 4 = 8 bytes.
let buffer = &mut [0; 8];
plugin_memory.read(&mut self.store, result_buffer as usize, buffer)?;
let raw_buffer = &mut [0; 8];
let plugin_memory = self
.instance
.get_memory(&mut self.store, "memory")
.ok_or_else(|| anyhow!("Could not grab slice of plugin memory"))?;
plugin_memory.read(&mut self.store, buffer as usize, raw_buffer)?;
// use these bytes (wasm stores things little-endian)
// to get a pointer to the buffer and its length
let b = buffer;
let result_buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
let result_buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
let result_buffer_end = result_buffer_ptr + result_buffer_len;
let b = raw_buffer;
let buffer_ptr = u32::from_le_bytes([b[0], b[1], b[2], b[3]]) as usize;
let buffer_len = u32::from_le_bytes([b[4], b[5], b[6], b[7]]) as usize;
let buffer_end = buffer_ptr + buffer_len;
// read the buffer at this point into a byte array
// deserialize the byte array into the provided serde type
let result = &plugin_memory.data(&mut self.store)[result_buffer_ptr..result_buffer_end];
let result = &plugin_memory.data(&mut self.store)[buffer_ptr..buffer_end];
let result = bincode::deserialize(result)?;
// TODO: this is handled wasm-side, but I'd like to double-check
// // deallocate the argument buffer
// self.free_buffer.call(&mut self.store, arg_buffer);
return Ok(result);
Ok(result)
}
// TODO: dont' use as for conversions
pub fn call<A: Serialize, R: DeserializeOwned>(
&mut self,
handle: &str,
arg: A,
) -> Result<R, Error> {
let start = std::time::Instant::now();
dbg!(&handle);
// dbg!(serde_json::to_string(&arg)).unwrap();
// write the argument to linear memory
// this returns a (ptr, lentgh) pair
let arg_buffer = self.serialize_to_buffer(arg)?;
// get the webassembly function we want to actually call
// TODO: precompute handle
let fun_name = format!("__{}", handle);
let fun = self
.instance
.get_typed_func::<(u32, u32), u32, _>(&mut self.store, &fun_name)?;
// call the function, passing in the buffer and its length
// this returns a ptr to a (ptr, lentgh) pair
let result_buffer = fun.call(&mut self.store, arg_buffer)?;
self.deserialize_from_buffer(result_buffer)
}
}