Add allocator to store so that it can be used by host functions

This commit is contained in:
Isaac Clayton 2022-06-09 12:44:10 +02:00
parent 47520f0ca1
commit 1f5903d16d
2 changed files with 109 additions and 58 deletions

View file

@ -6,7 +6,10 @@ use anyhow::{anyhow, Error};
use serde::{de::DeserializeOwned, Serialize};
use wasi_common::{dir, file};
use wasmtime::{Caller, Config, Engine, Instance, Linker, Module, Store, TypedFunc};
use wasmtime::{
AsContext, AsContextMut, Caller, Config, Engine, Extern, Instance, Linker, Module, Store,
StoreContext, StoreContextMut, Trap, TypedFunc,
};
use wasmtime::{IntoFunc, Memory};
use wasmtime_wasi::{Dir, WasiCtx, WasiCtxBuilder};
@ -35,15 +38,6 @@ impl<A: Serialize, R: DeserializeOwned> Clone for WasiFn<A, R> {
// }
// }
pub struct Wasi {
engine: Engine,
module: Module,
store: Store<WasiCtx>,
instance: Instance,
alloc_buffer: TypedFunc<u32, u32>,
// free_buffer: TypedFunc<(u32, u32), ()>,
}
// type signature derived from:
// https://docs.rs/wasmtime/latest/wasmtime/struct.Linker.html#method.func_wrap2_async
// macro_rules! dynHostFunction {
@ -68,14 +62,10 @@ pub struct Wasi {
// };
// }
// This type signature goodness gracious
pub type HostFunction = Box<dyn IntoFunc<WasiCtx, (u32, u32), u32>>;
pub struct WasiPluginBuilder {
// host_functions: HashMap<String, Box<dyn Fn(&str, &mut Linker<WasiCtx>) -> Result<(), Error>>>,
wasi_ctx: WasiCtx,
engine: Engine,
linker: Linker<WasiCtx>,
linker: Linker<WasiCtxAlloc>,
}
impl WasiPluginBuilder {
@ -101,7 +91,7 @@ impl WasiPluginBuilder {
Self::new(wasi_ctx)
}
pub fn host_function<A: Serialize, R: DeserializeOwned>(
pub fn host_function<A: DeserializeOwned, R: Serialize>(
mut self,
name: &str,
function: impl Fn(A) -> R + Send + Sync + 'static,
@ -109,10 +99,23 @@ impl WasiPluginBuilder {
self.linker.func_wrap(
"env",
name,
move |ctx: Caller<'_, WasiCtx>, ptr: u32, len: u32| {
// TODO: insert serialization code
function(todo!());
7u32
move |mut caller: Caller<'_, WasiCtxAlloc>, ptr: u32, len: u32| {
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 args = Wasi::deserialize_from_buffer(&mut plugin_memory, &caller, ptr, len)?;
let result = function(args);
let buffer = Wasi::serialize_to_buffer(
caller.data().alloc_buffer(),
&mut plugin_memory,
&mut caller,
result,
)
.await;
Ok(7u32)
},
)?;
Ok(self)
@ -123,14 +126,50 @@ impl WasiPluginBuilder {
}
}
// TODO: remove
/// Represents a to-be-initialized plugin.
/// Please use [`WasiPluginBuilder`], don't use this directly.
pub struct WasiPlugin {
pub module: Vec<u8>,
pub wasi_ctx: WasiCtx,
pub host_functions:
HashMap<String, Box<dyn Fn(&str, &mut Linker<WasiCtx>) -> Result<(), Error>>>,
// // TODO: remove
// /// Represents a to-be-initialized plugin.
// /// Please use [`WasiPluginBuilder`], don't use this directly.
// pub struct WasiPlugin {
// pub module: Vec<u8>,
// pub wasi_ctx: WasiCtx,
// pub host_functions:
// HashMap<String, Box<dyn Fn(&str, &mut Linker<WasiCtx>) -> Result<(), Error>>>,
// }
#[derive(Copy, Clone)]
struct WasiAlloc {
alloc_buffer: TypedFunc<u32, u32>,
free_buffer: TypedFunc<u32, u32>,
}
struct WasiCtxAlloc {
wasi_ctx: WasiCtx,
alloc: Option<WasiAlloc>,
}
impl WasiCtxAlloc {
fn alloc_buffer(&self) -> TypedFunc<u32, u32> {
self.alloc
.expect("allocator has been not initialized, cannot allocate buffer!")
.alloc_buffer
}
fn free_buffer(&self) -> TypedFunc<u32, u32> {
self.alloc
.expect("allocator has been not initialized, cannot free buffer!")
.free_buffer
}
fn init_alloc(&mut self, alloc: WasiAlloc) {
self.alloc = Some(alloc)
}
}
pub struct Wasi {
engine: Engine,
module: Module,
store: Store<WasiCtxAlloc>,
instance: Instance,
}
impl Wasi {
@ -154,30 +193,40 @@ impl Wasi {
impl Wasi {
async fn init(module: Vec<u8>, plugin: WasiPluginBuilder) -> Result<Self, Error> {
// initialize the WebAssembly System Interface context
let engine = plugin.engine;
let mut linker = plugin.linker;
wasmtime_wasi::add_to_linker(&mut linker, |s| &mut s.wasi_ctx)?;
linker.func_wrap("env", "__hello", |x: u32| x * 2).unwrap();
linker.func_wrap("env", "__bye", |x: u32| x / 2).unwrap();
wasmtime_wasi::add_to_linker(&mut linker, |s| s)?;
let mut store: Store<_> = Store::new(&engine, plugin.wasi_ctx);
// create a store, note that we can't initialize the allocator,
// because we can't grab the functions until initialized.
let mut store: Store<WasiCtxAlloc> = Store::new(
&engine,
WasiCtxAlloc {
wasi_ctx: plugin.wasi_ctx,
alloc: None,
},
);
let module = Module::new(&engine, module)?;
// load the provided module into the asynchronous runtime
linker.module_async(&mut store, "", &module).await?;
let instance = linker.instantiate_async(&mut store, &module).await?;
// now that the module is initialized,
// we can initialize the store's allocator
let alloc_buffer = instance.get_typed_func(&mut store, "__alloc_buffer")?;
// let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
let free_buffer = instance.get_typed_func(&mut store, "__free_buffer")?;
store.data_mut().init_alloc(WasiAlloc {
alloc_buffer,
free_buffer,
});
Ok(Wasi {
engine,
module,
store,
instance,
alloc_buffer,
// free_buffer,
})
}
@ -194,13 +243,14 @@ impl Wasi {
let dir = Box::new(wasmtime_wasi::dir::Dir::from_cap_std(dir));
// grab an empty file descriptor, specify capabilities
let fd = ctx.table().push(Box::new(()))?;
let fd = ctx.wasi_ctx.table().push(Box::new(()))?;
let caps = dir::DirCaps::all();
let file_caps = file::FileCaps::all();
// insert the directory at the given fd,
// return a handle to the resource
ctx.insert_dir(fd, dir, caps, file_caps, path.as_ref().to_path_buf());
ctx.wasi_ctx
.insert_dir(fd, dir, caps, file_caps, path.as_ref().to_path_buf());
Ok(WasiResource(fd))
}
@ -208,6 +258,7 @@ impl Wasi {
pub fn remove_resource(&mut self, resource: WasiResource) -> Result<(), Error> {
self.store
.data_mut()
.wasi_ctx
.table()
.delete(resource.0)
.ok_or_else(|| anyhow!("Resource did not exist, but a valid handle was passed in"))?;
@ -269,11 +320,11 @@ impl Wasi {
/// Takes an item, allocates a buffer, serializes the argument to that buffer,
/// and returns a (ptr, len) pair to that buffer.
async fn serialize_to_buffer<T: Serialize>(
async fn serialize_to_buffer<A: Serialize>(
alloc_buffer: TypedFunc<u32, u32>,
plugin_memory: &mut Memory,
mut store: &mut Store<WasiCtx>,
item: T,
mut store: impl AsContextMut<Data = WasiCtxAlloc>,
item: A,
) -> Result<(u32, u32), Error> {
// serialize the argument using bincode
let item = bincode::serialize(&item)?;
@ -288,7 +339,7 @@ impl Wasi {
/// Takes `ptr to a `(ptr, len)` pair, and returns `(ptr, len)`.
fn deref_buffer(
plugin_memory: &mut Memory,
store: &mut Store<WasiCtx>,
store: impl AsContext<Data = WasiCtxAlloc>,
buffer: u32,
) -> Result<(u32, u32), Error> {
// create a buffer to read the (ptr, length) pair into
@ -308,7 +359,7 @@ impl Wasi {
/// Takes a `(ptr, len)` pair and returns the corresponding deserialized buffer.
fn deserialize_from_buffer<R: DeserializeOwned>(
plugin_memory: &mut Memory,
store: &mut Store<WasiCtx>,
store: impl AsContext<Data = WasiCtxAlloc>,
buffer_ptr: u32,
buffer_len: u32,
) -> Result<R, Error> {
@ -318,10 +369,10 @@ impl Wasi {
// 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)[buffer_ptr..buffer_end];
let result = &plugin_memory.data(store.as_context())[buffer_ptr..buffer_end];
let result = bincode::deserialize(result)?;
// TODO: this is handled wasm-side, but I'd like to double-check
// TODO: this is handled wasm-side
// // deallocate the argument buffer
// self.free_buffer.call(&mut self.store, arg_buffer);
@ -358,9 +409,13 @@ impl Wasi {
// write the argument to linear memory
// this returns a (ptr, lentgh) pair
let arg_buffer =
Self::serialize_to_buffer(self.alloc_buffer, &mut plugin_memory, &mut self.store, arg)
.await?;
let arg_buffer = Self::serialize_to_buffer(
self.store.data().alloc_buffer(),
&mut plugin_memory,
&mut self.store,
arg,
)
.await?;
// call the function, passing in the buffer and its length
// this returns a ptr to a (ptr, lentgh) pair

View file

@ -1,16 +1,12 @@
use super::installation::{npm_install_packages, npm_package_latest_version};
use anyhow::{anyhow, Context, Result};
use anyhow::{anyhow, Result};
use client::http::HttpClient;
use futures::lock::Mutex;
use futures::{future::BoxFuture, FutureExt, StreamExt};
use gpui::executor::{self, Background};
use isahc::http::version;
use futures::{future::BoxFuture, FutureExt};
use gpui::executor::Background;
use language::{LanguageServerName, LspAdapter};
use plugin_runtime::{Wasi, WasiFn, WasiPlugin, WasiPluginBuilder};
use serde_json::json;
use std::fs;
use plugin_runtime::{Wasi, WasiFn, WasiPluginBuilder};
use std::{any::Any, path::PathBuf, sync::Arc};
use util::{ResultExt, TryFutureExt};
use util::ResultExt;
pub async fn new_json(executor: Arc<Background>) -> Result<PluginLspAdapter> {
let plugin = WasiPluginBuilder::new_with_default_ctx()?