Use the CallHook::CallingHost and ReturningFromHost with components (#…

…9196)

* Use the CallHook::CallingHost and ReturningFromHost with components

We never implemented calling the CallingHost and ReturningFromHost hooks
for component host calls.

co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

* Add tests

* Run hooks under `catch_unwind_and_longjmp`

* Cleanup the imports

* Remove todo

* Ensure that returning hooks are run

* Appease clippy

* Suggestions from code review

* Reuse infrastructure from the core-wasm call-hook test

* Remove redundant test

* Add a realloc test

* Switch the realloc test to handle returning a string

* Rework the realloc test to check that we're tracking host reallocs

* Use the call hook in the realloc host call

* Unnecessary pub mod

* Add a post-return test

* Remove unnecessary assertions

* Format

* Remove incorrect hook calls for realloc uses

---------

Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

crates/wasmtime/src/runtime/component/func/host.rs

@@ -4,17 +4,17 @@ use crate::component::storage::slice_to_storage_mut;
 use crate::component::{ComponentNamedList, ComponentType, Lift, Lower, Val};
 use crate::prelude::*;
 use crate::runtime::vm::component::{
-    InstanceFlags, VMComponentContext, VMLowering, VMLoweringCallee,
+    ComponentInstance, InstanceFlags, VMComponentContext, VMLowering, VMLoweringCallee,
 };
 use crate::runtime::vm::{VMFuncRef, VMMemoryDefinition, VMOpaqueContext};
-use crate::{AsContextMut, StoreContextMut, ValRaw};
+use crate::{AsContextMut, CallHook, StoreContextMut, ValRaw};
 use alloc::sync::Arc;
 use core::any::Any;
 use core::mem::{self, MaybeUninit};
 use core::ptr::NonNull;
 use wasmtime_environ::component::{
-    CanonicalAbiInfo, InterfaceType, StringEncoding, TypeFuncIndex, MAX_FLAT_PARAMS,
-    MAX_FLAT_RESULTS,
+    CanonicalAbiInfo, ComponentTypes, InterfaceType, StringEncoding, TypeFuncIndex,
+    MAX_FLAT_PARAMS, MAX_FLAT_RESULTS,
 };
 
 pub struct HostFunc {
@@ -55,9 +55,11 @@ impl HostFunc {
     {
         let data = data as *const F;
         unsafe {
-            handle_result(|| {
+            call_host_and_handle_result::<T>(cx, |instance, types, store| {
                 call_host::<_, _, _, _>(
-                    cx,
+                    instance,
+                    types,
+                    store,
                     ty,
                     flags,
                     memory,
@@ -131,7 +133,9 @@ where
 /// must be upheld. Generally that's done by ensuring this is only called from
 /// the select few places it's intended to be called from.
 unsafe fn call_host<T, Params, Return, F>(
-    cx: *mut VMOpaqueContext,
+    instance: *mut ComponentInstance,
+    types: &Arc<ComponentTypes>,
+    mut cx: StoreContextMut<'_, T>,
     ty: TypeFuncIndex,
     mut flags: InstanceFlags,
     memory: *mut VMMemoryDefinition,
@@ -163,10 +167,6 @@ where
         ret: U,
     }
 
-    let cx = VMComponentContext::from_opaque(cx);
-    let instance = (*cx).instance();
-    let mut cx = StoreContextMut::from_raw((*instance).store());
-
     let options = Options::new(
         cx.0.id(),
         NonNull::new(memory),
@@ -181,7 +181,6 @@ where
         bail!("cannot leave component instance");
     }
 
-    let types = (*instance).component_types();
     let ty = &types[ty];
     let param_tys = InterfaceType::Tuple(ty.params);
     let result_tys = InterfaceType::Tuple(ty.results);
@@ -289,15 +288,37 @@ fn validate_inbounds<T: ComponentType>(memory: &[u8], ptr: &ValRaw) -> Result<us
     Ok(ptr)
 }
 
-unsafe fn handle_result(func: impl FnOnce() -> Result<()>) {
-    match crate::runtime::vm::catch_unwind_and_longjmp(func) {
+unsafe fn call_host_and_handle_result<T>(
+    cx: *mut VMOpaqueContext,
+    func: impl FnOnce(
+        *mut ComponentInstance,
+        &Arc<ComponentTypes>,
+        StoreContextMut<'_, T>,
+    ) -> Result<()>,
+) {
+    let cx = VMComponentContext::from_opaque(cx);
+    let instance = (*cx).instance();
+    let types = (*instance).component_types();
+    let raw_store = (*instance).store();
+    let mut store = StoreContextMut::from_raw(raw_store);
+
+    let res = crate::runtime::vm::catch_unwind_and_longjmp(|| {
+        store.0.call_hook(CallHook::CallingHost)?;
+        let res = func(instance, types, store.as_context_mut());
+        store.0.call_hook(CallHook::ReturningFromHost)?;
+        res
+    });
+
+    match res {
         Ok(()) => {}
         Err(e) => crate::trap::raise(e),
     }
 }
 
 unsafe fn call_host_dynamic<T, F>(
-    cx: *mut VMOpaqueContext,
+    instance: *mut ComponentInstance,
+    types: &Arc<ComponentTypes>,
+    mut store: StoreContextMut<'_, T>,
     ty: TypeFuncIndex,
     mut flags: InstanceFlags,
     memory: *mut VMMemoryDefinition,
@@ -309,10 +330,6 @@ unsafe fn call_host_dynamic<T, F>(
 where
     F: FnOnce(StoreContextMut<'_, T>, &[Val], &mut [Val]) -> Result<()>,
 {
-    let cx = VMComponentContext::from_opaque(cx);
-    let instance = (*cx).instance();
-    let mut store = StoreContextMut::from_raw((*instance).store());
-    let types = (*instance).component_types();
     let options = Options::new(
         store.0.id(),
         NonNull::new(memory),
@@ -423,9 +440,11 @@ extern "C" fn dynamic_entrypoint<T, F>(
 {
     let data = data as *const F;
     unsafe {
-        handle_result(|| {
+        call_host_and_handle_result(cx, |instance, types, store| {
             call_host_dynamic::<T, _>(
-                cx,
+                instance,
+                types,
+                store,
                 ty,
                 flags,
                 memory,

tests/all/call_hook.rs

@@ -822,23 +822,23 @@ async fn drop_suspended_async_hook() -> Result<(), Error> {
 }
 
 #[derive(Debug, PartialEq, Eq)]
-enum Context {
+pub enum Context {
     Host,
     Wasm,
 }
 
-struct State {
-    context: Vec<Context>,
+pub struct State {
+    pub context: Vec<Context>,
 
-    calls_into_host: usize,
-    returns_from_host: usize,
-    calls_into_wasm: usize,
-    returns_from_wasm: usize,
+    pub calls_into_host: usize,
+    pub returns_from_host: usize,
+    pub calls_into_wasm: usize,
+    pub returns_from_wasm: usize,
 
-    trap_next_call_host: bool,
-    trap_next_return_host: bool,
-    trap_next_call_wasm: bool,
-    trap_next_return_wasm: bool,
+    pub trap_next_call_host: bool,
+    pub trap_next_return_host: bool,
+    pub trap_next_call_wasm: bool,
+    pub trap_next_return_wasm: bool,
 }
 
 impl Default for State {
@@ -906,6 +906,6 @@ impl State {
     }
 }
 
-fn sync_call_hook(mut ctx: StoreContextMut<'_, State>, transition: CallHook) -> Result<()> {
+pub fn sync_call_hook(mut ctx: StoreContextMut<'_, State>, transition: CallHook) -> Result<()> {
     ctx.data_mut().call_hook(transition)
 }

tests/all/component_model.rs

@@ -8,6 +8,7 @@ use wasmtime_component_util::REALLOC_AND_FREE;
 mod aot;
 mod r#async;
 mod bindgen;
+mod call_hook;
 mod dynamic;
 mod func;
 mod import;