Host function record/replay hooks split up

arjunr2 · arjunr2 · commit 3fed3d9c52bc · 2025-11-13T20:10:12.000-05:00
diff --git a/crates/wasmtime/src/runtime/component/func/host.rs b/crates/wasmtime/src/runtime/component/func/host.rs
@@ -257,12 +257,13 @@ where
 
     let types = vminstance.component().types().clone();
 
-    rr::component_hooks::record_and_replay_validate_host_func_entry(storage, &types, &ty, store.0)?;
-
     let func_ty = &types[ty];
     let param_tys = InterfaceType::Tuple(func_ty.params);
     let result_tys = InterfaceType::Tuple(func_ty.results);
 
+    rr::component_hooks::record_validate_host_func_entry(storage, &types, &ty, store.0)?;
+    rr::component_hooks::replay_validate_host_func_entry(storage, &types, &ty, store.0)?;
+
     let storage_type = if async_ {
         #[cfg(feature = "component-model-async")]
         {
@@ -373,7 +374,9 @@ where
         {
             match storage_type {
                 #[cfg(feature = "component-model-async")]
-                StorageType::Async(_) => unreachable!("`rr` should not be configurable with async"),
+                StorageType::Async(_) => {
+                    unreachable!("`rr` should not be configurable with `component-model-async`")
+                }
                 StorageType::Sync(mut storage) => {
                     unsafe {
                         flags.set_may_leave(false);
@@ -813,13 +816,12 @@ where
 
     let types = instance.id().get(store.0).component().types().clone();
 
-    rr::component_hooks::record_and_replay_validate_host_func_entry(storage, &types, &ty, store.0)?;
-
     let func_ty = &types[ty];
     let param_tys = &types[func_ty.params];
     let result_tys = &types[func_ty.results];
 
     if !store.0.replay_enabled() {
+        rr::component_hooks::record_validate_host_func_entry(storage, &types, &ty, store.0)?;
         let mut params_and_results = Vec::new();
         let mut lift = &mut LiftContext::new(store.0.store_opaque_mut(), &options, instance);
         lift.enter_call();
@@ -967,9 +969,10 @@ where
             cx.exit_call()?;
         }
     } else {
+        rr::component_hooks::replay_validate_host_func_entry(storage, &types, &ty, store.0)?;
         #[cfg(feature = "rr-component")]
         if async_ {
-            unreachable!("`rr` should not be configurable with async");
+            unreachable!("`rr` should not be configurable with `component-model-async`");
         } else {
             unsafe {
                 flags.set_may_leave(false);
diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs
@@ -1287,19 +1287,17 @@ impl Func {
         let nparams = ty.params().len();
         val_vec.reserve(nparams + ty.results().len());
 
-        // Recording host function entry
         let flat_params = ty
             .params()
             .into_iter()
             .map(|x| x.to_wasm_type().byte_size());
 
-        rr::core_hooks::record_and_replay_validate_host_func_entry(
-            values_vec,
-            flat_params,
-            &mut caller.store.0,
-        )?;
-
         if !caller.store.0.replay_enabled() {
+            rr::core_hooks::record_validate_host_func_entry(
+                values_vec,
+                flat_params,
+                &mut caller.store.0,
+            )?;
             for (i, ty) in ty.params().enumerate() {
                 val_vec.push(unsafe { Val::from_raw(&mut caller.store, values_vec[i], ty) })
             }
@@ -1324,6 +1322,11 @@ impl Func {
                 .map(|x| x.to_wasm_type().byte_size());
             rr::core_hooks::record_host_func_return(values_vec, flat_results, &mut caller.store.0)?;
         } else {
+            rr::core_hooks::replay_validate_host_func_entry(
+                values_vec,
+                flat_params,
+                &mut caller.store.0,
+            )?;
             rr::core_hooks::replay_host_func_return(values_vec, &mut caller.store.0)?;
         }
 
@@ -2436,15 +2439,13 @@ impl HostContext {
                 wasm_func_type.returns().into_iter().map(|x| x.byte_size()),
             );
 
-            // Record/replay(validation) of the raw parameter arguments
-            // Don't need auto-assert GC store here since we aren't using P, just raw args
-            rr::core_hooks::record_and_replay_validate_host_func_entry(
-                unsafe { &args.as_ref()[..num_params] },
-                flat_size_params,
-                caller.store.0,
-            )?;
-
             if !caller.store.0.replay_enabled() {
+                // Don't need auto-assert GC store here since we aren't using P, just raw args for recording
+                rr::core_hooks::record_validate_host_func_entry(
+                    unsafe { &args.as_ref()[..num_params] },
+                    flat_size_params,
+                    caller.store.0,
+                )?;
                 let ret = 'ret: {
                     if let Err(trap) = caller.store.0.call_hook(CallHook::CallingHost) {
                         break 'ret R::fallible_from_error(trap);
@@ -2489,6 +2490,12 @@ impl HostContext {
                     caller.store.0,
                 )?;
             } else {
+                rr::core_hooks::replay_validate_host_func_entry(
+                    unsafe { &args.as_ref()[..num_params] },
+                    flat_size_params,
+                    caller.store.0,
+                )?;
+
                 // Replay the return values
                 rr::core_hooks::replay_host_func_return(
                     unsafe { &mut args.as_mut()[..num_results] },
diff --git a/crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs b/crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs
@@ -3,6 +3,8 @@ use crate::ValRaw;
 use crate::component::func::LowerContext;
 #[cfg(feature = "rr-component")]
 use crate::rr::ResultEvent;
+#[cfg(all(feature = "rr-component", feature = "rr-validate"))]
+use crate::rr::component_events::HostFuncEntryEvent;
 #[cfg(feature = "rr-component")]
 use crate::rr::component_events::{
     HostFuncReturnEvent, LowerFlatReturnEvent, LowerMemoryReturnEvent, WasmFuncEntryEvent,
@@ -40,15 +42,15 @@ where
 {
     let _ = (args, type_idx, &types);
     #[cfg(feature = "rr-component")]
-    {
+    store.0.record_event(|| {
         let flat_params = types.flat_types_storage(
             &InterfaceType::Tuple(types[type_idx].params),
             MAX_FLAT_PARAMS,
         );
-        store.0.record_event(|| WasmFuncEntryEvent {
+        WasmFuncEntryEvent {
             args: RRFuncArgVals::from_flat_storage(args, flat_params),
-        })?;
-    }
+        }
+    })?;
     let result = wasm_call(store);
     #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
     {
@@ -72,29 +74,32 @@ where
     return result;
 }
 
-/// Record/replay hook operation for host function entry events
+/// Record hook operation for host function entry events
 #[inline]
-pub fn record_and_replay_validate_host_func_entry(
+pub fn record_validate_host_func_entry(
     args: &mut [MaybeUninit<ValRaw>],
     types: &Arc<ComponentTypes>,
     type_idx: &TypeFuncIndex,
     store: &mut StoreOpaque,
 ) -> Result<()> {
     #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
-    {
-        use crate::rr::component_events::HostFuncEntryEvent;
-        let event = || {
-            let flat_params = types.flat_types_storage(
-                &InterfaceType::Tuple(types[*type_idx].params),
-                MAX_FLAT_PARAMS,
-            );
-            HostFuncEntryEvent {
-                args: RRFuncArgVals::from_flat_storage(args, flat_params),
-            }
-        };
-        store.record_event_validation(|| event())?;
-        store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| event())?;
-    }
+    store.record_event_validation(|| create_host_func_entry_event(args, types, type_idx))?;
+    let _ = (args, types, type_idx, store);
+    Ok(())
+}
+
+/// Replay hook operation for host function entry events
+#[inline]
+pub fn replay_validate_host_func_entry(
+    args: &mut [MaybeUninit<ValRaw>],
+    types: &Arc<ComponentTypes>,
+    type_idx: &TypeFuncIndex,
+    store: &mut StoreOpaque,
+) -> Result<()> {
+    #[cfg(all(feature = "rr-component", feature = "rr-validate"))]
+    store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| {
+        create_host_func_entry_event(args, types, type_idx)
+    })?;
     let _ = (args, types, type_idx, store);
     Ok(())
 }
@@ -168,3 +173,19 @@ where
         .record_event(|| LowerFlatReturnEvent(ResultEvent::from_anyhow_result(&lower_result)))?;
     lower_result
 }
+
+#[cfg(all(feature = "rr-component", feature = "rr-validate"))]
+#[inline(always)]
+fn create_host_func_entry_event(
+    args: &mut [MaybeUninit<ValRaw>],
+    types: &Arc<ComponentTypes>,
+    type_idx: &TypeFuncIndex,
+) -> HostFuncEntryEvent {
+    let flat_params = types.flat_types_storage(
+        &InterfaceType::Tuple(types[*type_idx].params),
+        MAX_FLAT_PARAMS,
+    );
+    HostFuncEntryEvent {
+        args: RRFuncArgVals::from_flat_storage(args, flat_params),
+    }
+}
diff --git a/crates/wasmtime/src/runtime/rr/hooks/core_hooks.rs b/crates/wasmtime/src/runtime/rr/hooks/core_hooks.rs
@@ -65,9 +65,9 @@ where
     return result;
 }
 
-/// Record and replay hook operation for host function entry events
+/// Record hook operation for host function entry events
 #[inline]
-pub fn record_and_replay_validate_host_func_entry<T>(
+pub fn record_validate_host_func_entry<T>(
     args: &[T],
     flat: impl Iterator<Item = u8>,
     store: &mut StoreOpaque,
@@ -77,20 +77,9 @@ where
 {
     let _ = (args, &flat, &store);
     #[cfg(all(feature = "rr", feature = "rr-validate"))]
-    {
-        // Record/replay the raw parameter args
-        if store.replay_enabled() {
-            store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| {
-                HostFuncEntryEvent {
-                    args: RRFuncArgVals::from_flat_iter(args, flat),
-                }
-            })?;
-        } else {
-            store.record_event_validation(|| HostFuncEntryEvent {
-                args: RRFuncArgVals::from_flat_iter(args, flat),
-            })?;
-        }
-    }
+    store.record_event_validation(|| HostFuncEntryEvent {
+        args: RRFuncArgVals::from_flat_iter(args, flat),
+    })?;
     Ok(())
 }
 
@@ -113,6 +102,24 @@ where
     Ok(())
 }
 
+/// Replay hook operation for host function entry events
+#[inline]
+pub fn replay_validate_host_func_entry<T>(
+    args: &[T],
+    flat: impl Iterator<Item = u8>,
+    store: &mut StoreOpaque,
+) -> Result<()>
+where
+    T: FlatBytes,
+{
+    let _ = (args, &flat, &store);
+    #[cfg(all(feature = "rr", feature = "rr-validate"))]
+    store.next_replay_event_validation::<HostFuncEntryEvent, _, _>(|| HostFuncEntryEvent {
+        args: RRFuncArgVals::from_flat_iter(args, flat),
+    })?;
+    Ok(())
+}
+
 /// Replay hook operation for host function return events
 #[inline]
 pub fn replay_host_func_return<T>(args: &mut [T], store: &mut StoreOpaque) -> Result<()>
