Fix call hooks with yields and concurrent execution

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::<Self>().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)
             })
         }
     }

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
         };
 

crates/wasmtime/src/runtime/store.rs

@@ -2311,6 +2311,11 @@ unsafe impl<T> VMStore for StoreInner<T> {
         &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<StoreResourceLimiter<'_>>, &mut StoreOpaque) {

crates/wasmtime/src/runtime/vm.rs

@@ -213,6 +213,11 @@ pub unsafe trait VMStore: 'static {
         &mut self,
     ) -> (Option<StoreResourceLimiter<'_>>, &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.

crates/wasmtime/src/runtime/vm/traphandlers.rs

@@ -245,7 +245,38 @@ where
         store: &mut dyn VMStore,
         f: impl FnOnce(&mut dyn VMStore) -> Result<T, E>,
     ) -> (T::Abi, Option<UnwindReason>) {
-        // 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) {

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(())
+}

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(())
+}