T0b1/wasmtime
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add support for async call hooks (#3876)

Browse Source 
* Instead of simply panicking, return an error when we attempt to resume on a dying fiber.

This situation should never occur in the existing code base, but can be
triggered if support for running outside async code in a call hook.

* Shift `async_cx()` to return an `Option`, reflecting if the fiber is dying.

This should never happen in the existing code base, but is a nice
forward-looking guard. The current implementations simply lift the
trap that would eventually be produced by such an operation into
a `Trap` (or similar) at the invocation of `async_cx()`.

* Add support for using `async` call hooks.

This retains the ability to do non-async hooks. Hooks end up being
implemented as an async trait with a handler call, to get around some
issues passing around async closures. This change requires some of
the prior changes to handle picking up blocked tasks during fiber
shutdown, to avoid some panics during timeouts and other such events.

* More fully specify a doc link, to avoid a doc-building error.

* Revert the use of catchable traps on cancellation of a fiber; turn them into expect()/unwrap().

The justification for this revert is that (a) these events shouldn't
happen, and (b) they wouldn't be catchable by wasm anyways.

* Replace a duplicated check in `async` hook evaluation with a single check.

This also moves the checks inside of their respective Async variants,
meaning that if you're using an async-enabled version of wasmtime but
using the synchronous versions of the callbacks, you won't pay any
penalty for validating the async context.

* Use `match &mut ...` insead of `ref mut`.

* Add some documentation on why/when `async_cx` can return None.

* Add two simple test cases for async call hooks.

* Fix async_cx() to check both the box and the value for current_poll_cx.

In the prior version, we only checked that the box had not been cleared,
but had not ensured that there was an actual context for us to use. This
updates the check to validate both, returning None if the inner context
is missing. This allows us to skip a validation check inside `block_on`,
since all callers will have run through the `async_cx` check prior to
arrival.

* Tweak the timeout test to address PR suggestions.

* Add a test about dropping async hooks while suspended

Should help exercise that the check for `None` is properly handled in a
few more locations.

Co-authored-by: Alex Crichton <alex@alexcrichton.com>
This commit is contained in:
Adam Wick
2022-03-23 08:43:34 -07:00
committed by GitHub
parent 923faaff4f
commit 6a60e8363f
5 changed files with 369 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

272
tests/all/call_hook.rs
View File

@@ -1,4 +1,7 @@
use anyhow::Error;
use std::future::Future;
use std::pin::Pin;
use std::task::{self, Poll};
use wasmtime::*;
// Crate a synchronous Func, call it directly:
@@ -551,6 +554,275 @@ fn trapping() -> Result<(), Error> {
Ok(())
}
#[tokio::test]
async fn basic_async_hook() -> Result<(), Error> {
struct HandlerR;
#[async_trait::async_trait]
impl CallHookHandler<State> for HandlerR {
async fn handle_call_event(
&self,
obj: &mut State,
ch: CallHook,
) -> Result<(), wasmtime::Trap> {
State::call_hook(obj, ch)
}
}
let mut config = Config::new();
config.async_support(true);
let engine = Engine::new(&config)?;
let mut store = Store::new(&engine, State::default());
store.call_hook_async(HandlerR {});
assert_eq!(store.data().calls_into_host, 0);
assert_eq!(store.data().returns_from_host, 0);
assert_eq!(store.data().calls_into_wasm, 0);
assert_eq!(store.data().returns_from_wasm, 0);
let mut linker = Linker::new(&engine);
linker.func_wrap(
"host",
"f",
|caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
// Calling this func will switch context into wasm, then back to host:
assert_eq!(caller.data().context, vec![Context::Wasm, Context::Host]);
assert_eq!(
caller.data().calls_into_host,
caller.data().returns_from_host + 1
);
assert_eq!(
caller.data().calls_into_wasm,
caller.data().returns_from_wasm + 1
);
assert_eq!(a, 1);
assert_eq!(b, 2);
assert_eq!(c, 3.0);
assert_eq!(d, 4.0);
},
)?;
let wat = r#"
(module
(import "host" "f"
(func $f (param i32) (param i64) (param f32) (param f64)))
(func (export "export")
(call $f (i32.const 1) (i64.const 2) (f32.const 3.0) (f64.const 4.0)))
)
"#;
let module = Module::new(&engine, wat)?;
let inst = linker.instantiate(&mut store, &module)?;
let export = inst
.get_export(&mut store, "export")
.expect("get export")
.into_func()
.expect("export is func");
export.call_async(&mut store, &[], &mut []).await?;
// One switch from vm to host to call f, another in return from f.
assert_eq!(store.data().calls_into_host, 1);
assert_eq!(store.data().returns_from_host, 1);
assert_eq!(store.data().calls_into_wasm, 1);
assert_eq!(store.data().returns_from_wasm, 1);
Ok(())
}
#[tokio::test]
async fn timeout_async_hook() -> Result<(), Error> {
struct HandlerR;
#[async_trait::async_trait]
impl CallHookHandler<State> for HandlerR {
async fn handle_call_event(
&self,
obj: &mut State,
ch: CallHook,
) -> Result<(), wasmtime::Trap> {
if obj.calls_into_host > 200 {
return Err(wasmtime::Trap::new("timeout"));
}
match ch {
CallHook::CallingHost => obj.calls_into_host += 1,
CallHook::CallingWasm => obj.calls_into_wasm += 1,
CallHook::ReturningFromHost => obj.returns_from_host += 1,
CallHook::ReturningFromWasm => obj.returns_from_wasm += 1,
}
Ok(())
}
}
let mut config = Config::new();
config.async_support(true);
let engine = Engine::new(&config)?;
let mut store = Store::new(&engine, State::default());
store.call_hook_async(HandlerR {});
assert_eq!(store.data().calls_into_host, 0);
assert_eq!(store.data().returns_from_host, 0);
assert_eq!(store.data().calls_into_wasm, 0);
assert_eq!(store.data().returns_from_wasm, 0);
let mut linker = Linker::new(&engine);
linker.func_wrap(
"host",
"f",
|_caller: Caller<State>, a: i32, b: i64, c: f32, d: f64| {
assert_eq!(a, 1);
assert_eq!(b, 2);
assert_eq!(c, 3.0);
assert_eq!(d, 4.0);
},
)?;
let wat = r#"
(module
(import "host" "f"
(func $f (param i32) (param i64) (param f32) (param f64)))
(func (export "export")
(loop $start
(call $f (i32.const 1) (i64.const 2) (f32.const 3.0) (f64.const 4.0))
(br $start)))
)
"#;
let module = Module::new(&engine, wat)?;
let inst = linker.instantiate(&mut store, &module)?;
let export = inst
.get_typed_func::<(), (), _>(&mut store, "export")
.expect("export is func");
store.set_epoch_deadline(1);
store.epoch_deadline_async_yield_and_update(1);
assert!(export.call_async(&mut store, ()).await.is_err());
// One switch from vm to host to call f, another in return from f.
assert!(store.data().calls_into_host > 1);
assert!(store.data().returns_from_host > 1);
assert_eq!(store.data().calls_into_wasm, 1);
assert_eq!(store.data().returns_from_wasm, 0);
Ok(())
}
#[tokio::test]
async fn drop_suspended_async_hook() -> Result<(), Error> {
struct Handler;
#[async_trait::async_trait]
impl CallHookHandler<u32> for Handler {
async fn handle_call_event(
&self,
state: &mut u32,
_ch: CallHook,
) -> Result<(), wasmtime::Trap> {
assert_eq!(*state, 0);
*state += 1;
let _dec = Decrement(state);
// Simulate some sort of event which takes a number of yields
for _ in 0..500 {
tokio::task::yield_now().await;
}
Ok(())
}
}
let mut config = Config::new();
config.async_support(true);
let engine = Engine::new(&config)?;
let mut store = Store::new(&engine, 0);
store.call_hook_async(Handler);
let mut linker = Linker::new(&engine);
// Simulate a host function that has lots of yields with an infinite loop.
linker.func_wrap0_async("host", "f", |mut cx| {
Box::new(async move {
let state = cx.data_mut();
assert_eq!(*state, 0);
*state += 1;
let _dec = Decrement(state);
loop {
tokio::task::yield_now().await;
}
})
})?;
let wat = r#"
(module
(import "host" "f" (func $f))
(func (export "") call $f)
)
"#;
let module = Module::new(&engine, wat)?;
let inst = linker.instantiate(&mut store, &module)?;
assert_eq!(*store.data(), 0);
let export = inst
.get_typed_func::<(), (), _>(&mut store, "")
.expect("export is func");
// First test that if we drop in the middle of an async hook that everything
// is alright.
PollNTimes {
future: Box::pin(export.call_async(&mut store, ())),
times: 200,
}
.await;
assert_eq!(*store.data(), 0); // double-check user dtors ran
// Next test that if we drop while in a host async function that everything
// is also alright.
PollNTimes {
future: Box::pin(export.call_async(&mut store, ())),
times: 1_000,
}
.await;
assert_eq!(*store.data(), 0); // double-check user dtors ran
return Ok(());
// A helper struct to poll an inner `future` N `times` and then resolve.
// This is used above to test that when futures are dropped while they're
// pending everything works and is cleaned up on the Wasmtime side of
// things.
struct PollNTimes<F> {
future: F,
times: u32,
}
impl<F: Future + Unpin> Future for PollNTimes<F> {
type Output = ();
fn poll(mut self: Pin<&mut Self>, task: &mut task::Context<'_>) -> Poll<()> {
for _ in 0..self.times {
match Pin::new(&mut self.future).poll(task) {
Poll::Ready(_) => panic!("future should not be ready"),
Poll::Pending => {}
}
}
Poll::Ready(())
}
}
// helper struct to decrement a counter on drop
struct Decrement<'a>(&'a mut u32);
impl Drop for Decrement<'_> {
fn drop(&mut self) {
*self.0 -= 1;
}
}
}
#[derive(Debug, PartialEq, Eq)]
enum Context {
Host,