From 9003f9884327628a3e3d873c3589b0c5e22dc7d9 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Mon, 13 Jul 2026 09:07:26 -0700 Subject: [PATCH] Fix call hooks with yields and concurrent execution This commit fixes two (known) issues with call hooks. The first is a long-standing issue where when Wasmtime yielded due to fuel or epochs the call hooks were not invoked. If hosts were using call hooks to track time executed in a guest, for example, then that would lead to inaccuracies. The second is with component-model-async fiber switches weren't correctly accounted for in terms of the call hook machinery which could lead results to get mixed up. The fix for both here is the same which is to sink the invocation of call hooks far down into the stack such that they cannot be forgotten when the guest exits to the host. This moves manual invocations, for example, on host functions to instead one central location that handles both host functions and libcalls. Closes #13870 --- .../src/runtime/component/func/host.rs | 9 +- crates/wasmtime/src/runtime/func.rs | 8 -- crates/wasmtime/src/runtime/store.rs | 5 + crates/wasmtime/src/runtime/vm.rs | 5 + .../wasmtime/src/runtime/vm/traphandlers.rs | 33 +++++- tests/all/call_hook.rs | 35 ++++++ tests/all/component_model/call_hook.rs | 102 ++++++++++++++++++ 7 files changed, 181 insertions(+), 16 deletions(-) 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(()) +}