diff --git a/crates/wasmtime/src/runtime/component/func/host.rs b/crates/wasmtime/src/runtime/component/func/host.rs index b1d8a931f8bd..9eacb9093db0 100644 --- a/crates/wasmtime/src/runtime/component/func/host.rs +++ b/crates/wasmtime/src/runtime/component/func/host.rs @@ -15,7 +15,7 @@ use crate::runtime::vm::component::{ }; use crate::runtime::vm::{VMOpaqueContext, VMStore}; use crate::store::Asyncness; -use crate::{AsContextMut, CallHook, StoreContextMut, ValRaw}; +use crate::{AsContextMut, StoreContextMut, ValRaw}; use alloc::sync::Arc; use core::any::Any; use core::mem::{self, MaybeUninit}; @@ -347,12 +347,7 @@ where let options = OptionsIndex::from_u32(options); let storage = NonNull::slice_from_raw_parts(storage, storage_len).as_mut(); let data = data.cast::().as_ref(); - - store.0.call_hook(CallHook::CallingHost)?; - let res = data.entrypoint(store.as_context_mut(), instance, ty, options, storage); - store.0.call_hook(CallHook::ReturningFromHost)?; - - res + data.entrypoint(store.as_context_mut(), instance, ty, options, storage) }) } } diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs index 0f0a60a67e57..f6b9d99f1a8b 100644 --- a/crates/wasmtime/src/runtime/func.rs +++ b/crates/wasmtime/src/runtime/func.rs @@ -2373,10 +2373,6 @@ impl HostFunc { // this up. let mut store = unsafe { store.unchecked_context_mut() }; - // Handle the entry call hook, with a corresponding exit call hook - // below. - store.0.call_hook(CallHook::CallingHost)?; - // SAFETY: this function itself requires that the `vmctx` is // valid to use here. let state = unsafe { @@ -2414,10 +2410,6 @@ impl HostFunc { store.0.exit_gc_lifo_scope(gc_lifo_scope); - // Note that if this returns a trap then `ret` is discarded - // entirely. - store.0.call_hook(CallHook::ReturningFromHost)?; - ret }; diff --git a/crates/wasmtime/src/runtime/store.rs b/crates/wasmtime/src/runtime/store.rs index fae724cbbaaa..3873612fef18 100644 --- a/crates/wasmtime/src/runtime/store.rs +++ b/crates/wasmtime/src/runtime/store.rs @@ -2311,6 +2311,11 @@ unsafe impl VMStore for StoreInner { &mut self.inner } + #[cfg(feature = "call-hook")] + fn call_hook(&mut self, s: CallHook) -> Result<()> { + StoreInner::call_hook(self, s) + } + fn resource_limiter_and_store_opaque( &mut self, ) -> (Option>, &mut StoreOpaque) { diff --git a/crates/wasmtime/src/runtime/vm.rs b/crates/wasmtime/src/runtime/vm.rs index b663b63f9b88..b9f468166c65 100644 --- a/crates/wasmtime/src/runtime/vm.rs +++ b/crates/wasmtime/src/runtime/vm.rs @@ -213,6 +213,11 @@ pub unsafe trait VMStore: 'static { &mut self, ) -> (Option>, &mut StoreOpaque); + /// Invoke this store's configured call hook, if any, to notify the + /// embedder of a transition between the host and WebAssembly. + #[cfg(feature = "call-hook")] + fn call_hook(&mut self, s: crate::CallHook) -> Result<()>; + /// Callback invoked whenever an instance observes a new epoch /// number. Cannot fail; cooperative epoch-based yielding is /// completely semantically transparent. Returns the new deadline. diff --git a/crates/wasmtime/src/runtime/vm/traphandlers.rs b/crates/wasmtime/src/runtime/vm/traphandlers.rs index acbe5c356ad0..7ea34b1473f1 100644 --- a/crates/wasmtime/src/runtime/vm/traphandlers.rs +++ b/crates/wasmtime/src/runtime/vm/traphandlers.rs @@ -245,7 +245,38 @@ where store: &mut dyn VMStore, f: impl FnOnce(&mut dyn VMStore) -> Result, ) -> (T::Abi, Option) { - // First prepare the closure `f` as something that'll be invoked to + // First wrap `f` in call hooks if that feature is enabled. This is used + // as a "pretty far down in the stack" mechanism of ensuring that hooks + // aren't forgotten. + // + // Note that by being placed here this is handling: + // + // * libcalls + // * host functions + // * component versions of the above + // + // This specifically is NOT handling libcalls where the result can't + // carry a result. This should be safe as anything which doesn't return + // a `Result` sort of has to be simple enough to not allow recursion so + // it's just a brief exit from the guest to the host. + // + // Also note that this only happens with the `call-hook` feature because + // this otherwise imposes a dynamic dispatch on the `store` trait object + // which otherwise can't be optimized away. + #[cfg(feature = "call-hook")] + let f = move |store: &mut dyn VMStore| { + store.call_hook(crate::CallHook::CallingHost)?; + + let res = f(store); + + // Note that if this returns a trap then `ret` is discarded + // entirely. + store.call_hook(crate::CallHook::ReturningFromHost)?; + + res.map_err(|e| e.into()) + }; + + // Next prepare the closure `f` as something that'll be invoked to // generate the return value of this function. This is the // conditionally, below, passed to `catch_unwind`. let f = move || match f(store) { diff --git a/tests/all/call_hook.rs b/tests/all/call_hook.rs index 532ba73420d6..1041952e42c0 100644 --- a/tests/all/call_hook.rs +++ b/tests/all/call_hook.rs @@ -897,3 +897,38 @@ impl State { pub fn sync_call_hook(mut ctx: StoreContextMut<'_, State>, transition: CallHook) -> Result<()> { ctx.data_mut().call_hook(transition) } + +#[tokio::test] +async fn hooks_used_at_yield_points() -> Result<(), Error> { + let mut config = Config::new(); + config.consume_fuel(true); + let engine = Engine::new(&config)?; + let mut store = Store::new(&engine, State::default()); + store.call_hook(sync_call_hook); + store.set_fuel(100)?; + store.fuel_async_yield_interval(Some(1))?; + + 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 wat = r#" + (module + (func $start) + (start $start) + ) + "#; + let module = Module::new(&engine, wat)?; + + Linker::new(&engine) + .instantiate_async(&mut store, &module) + .await?; + + assert!(store.data().calls_into_host > 0); + assert!(store.data().returns_from_host > 0); + assert_eq!(store.data().calls_into_wasm, 1); + assert_eq!(store.data().returns_from_wasm, 1); + + Ok(()) +} diff --git a/tests/all/component_model/call_hook.rs b/tests/all/component_model/call_hook.rs index 4d5ca9d524cd..b0af1322b3ca 100644 --- a/tests/all/component_model/call_hook.rs +++ b/tests/all/component_model/call_hook.rs @@ -618,3 +618,105 @@ async fn drop_suspended_async_hook() -> Result<()> { } } } + +#[tokio::test] +async fn concurrent_behavior_balanced() -> Result<()> { + let engine = Engine::default(); + + let mut store = Store::new(&engine, State::default()); + store.call_hook(sync_call_hook); + + // A composition of two components where `$lowerer.run` makes an + // async-lowered call to the async-lifted (with callback) `$lifter.foo`. + // Starting `foo` forces the fiber running `run` to suspend while it waits + // for `foo` to start, and `foo`'s `task.return` call reenters wasm to run + // the fused adapter's `return_` function, both of which must be reflected + // in the call hooks fired in order to keep the embedder-visible + // transitions balanced. + let component = Component::new( + &engine, + r#"(component + (component $lifter + (core module $m + (import "" "task.return" (func $task-return (param i32))) + (func (export "callback") (param i32 i32 i32) (result i32) + unreachable) + (func (export "foo") (param i32) (result i32) + (call $task-return (i32.add (local.get 0) (i32.const 1))) + i32.const 0 (; EXIT ;) + ) + ) + (core func $task-return (canon task.return (result u32))) + (core instance $i (instantiate $m + (with "" (instance (export "task.return" (func $task-return)))) + )) + (func (export "foo") async (param "p1" u32) (result u32) + (canon lift (core func $i "foo") async (callback (func $i "callback"))) + ) + ) + + (component $lowerer + (import "a" (func $foo async (param "p1" u32) (result u32))) + (core module $libc (memory (export "memory") 1)) + (core instance $libc (instantiate $libc)) + (core func $foo (canon lower (func $foo) async (memory $libc "memory"))) + (core module $m + (import "libc" "memory" (memory 1)) + (import "" "foo" (func $foo (param i32 i32) (result i32))) + (func (export "run") (result i32) + ;; call `foo`, asserting that it returned immediately + block + (call $foo (i32.const 41) (i32.const 1204)) + i32.const 0xf + i32.and + i32.const 2 (; RETURNED ;) + i32.eq + br_if 0 + unreachable + end + (i32.load offset=0 (i32.const 1204)) + ) + ) + (core instance $i (instantiate $m + (with "libc" (instance $libc)) + (with "" (instance (export "foo" (func $foo)))) + )) + (func (export "run") (result u32) (canon lift (core func $i "run"))) + ) + + (instance $lifter (instantiate $lifter)) + (instance $lowerer (instantiate $lowerer (with "a" (func $lifter "foo")))) + (export "run" (func $lowerer "run")) + )"#, + )?; + + let instance = Linker::new(&engine) + .instantiate_async(&mut store, &component) + .await?; + let run = instance.get_typed_func::<(), (u32,)>(&mut store, "run")?; + let (result,) = store + .run_concurrent(async |store| run.call_concurrent(store, ()).await) + .await??; + assert_eq!(result, 42); + + // There are three guest exits (currently) in the above component: + // + // 1. Guest-to-guest async transitions pay an exit to prepare a task to be + // run. + // 2. Guest-to-guest async transitions pay an exit to then run the initial + // turn of the task created. + // 3. The `task.return` call exits the guest component. + assert_eq!(store.data().calls_into_host, 3); + assert_eq!(store.data().returns_from_host, 3); + + // There are four guest entries (currently) in the above component: + // + // 1. The root invocation of `run`. + // 2. The synthesized adapter which translates arguments to `foo`. + // 3. The initial invocation of `foo` in `$lifter`. + // 4. The synthesized adapter which translates `task.return` values. + assert_eq!(store.data().calls_into_wasm, 4); + assert_eq!(store.data().returns_from_wasm, 4); + + Ok(()) +}