Add support for concurrent call recording

arjunr2 · arjunr2 · commit 875a4b1f73a7 · 2025-11-20T17:59:43.000-05:00
diff --git a/crates/wasmtime/src/runtime/component/concurrent.rs b/crates/wasmtime/src/runtime/component/concurrent.rs
@@ -53,6 +53,7 @@
 use crate::component::func::{self, Func};
 use crate::component::{HasData, HasSelf, Instance, Resource, ResourceTable, ResourceTableError};
 use crate::fiber::{self, StoreFiber, StoreFiberYield};
+use crate::rr::{RRWasmFuncType, component_hooks};
 use crate::store::{Store, StoreId, StoreInner, StoreOpaque, StoreToken};
 use crate::vm::component::{CallContext, ComponentInstance, InstanceFlags, ResourceTables};
 use crate::vm::{AlwaysMut, SendSyncPtr, VMFuncRef, VMMemoryDefinition, VMStore};
@@ -1715,7 +1716,8 @@ impl Instance {
     ///
     /// SAFETY: The raw pointer arguments must be valid references to guest
     /// functions (with the appropriate signatures) when the closures queued by
-    /// this function are called.
+    /// this function are called. For RR, the rr_handle must be a valid `Func`
+    /// corresponding to `callee` if provided
     unsafe fn queue_call<T: 'static>(
         self,
         mut store: StoreContextMut<T>,
@@ -1727,6 +1729,7 @@ impl Instance {
         async_: bool,
         callback: Option<SendSyncPtr<VMFuncRef>>,
         post_return: Option<SendSyncPtr<VMFuncRef>>,
+        rr_handle: Option<Func>,
     ) -> Result<()> {
         /// Return a closure which will call the specified function in the scope
         /// of the specified task.
@@ -1741,14 +1744,16 @@ impl Instance {
         /// nothing.
         ///
         /// SAFETY: `callee` must be a valid `*mut VMFuncRef` at the time when
-        /// the returned closure is called.
+        /// the returned closure is called. For RR, the handle must be a valid `Func`
+        /// corresponding to `callee` if provided
         unsafe fn make_call<T: 'static>(
             store: StoreContextMut<T>,
             guest_thread: QualifiedThreadId,
             callee: SendSyncPtr<VMFuncRef>,
             param_count: usize,
             result_count: usize,
             flags: Option<InstanceFlags>,
+            rr_handle: Option<Func>,
         ) -> impl FnOnce(&mut dyn VMStore) -> Result<[MaybeUninit<ValRaw>; MAX_FLAT_PARAMS]>
         + Send
         + Sync
@@ -1766,6 +1771,14 @@ impl Instance {
                 let may_enter_after_call = task.call_post_return_automatically();
                 let lower = task.lower_params.take().unwrap();
 
+                if let Some(func) = rr_handle {
+                    component_hooks::record_wasm_func_begin(
+                        func.instance().id().instance(),
+                        func.index(),
+                        store.store_opaque_mut(),
+                    )?;
+                }
+
                 lower(store, &mut storage[..param_count])?;
 
                 let mut store = token.as_context_mut(store);
@@ -1776,15 +1789,31 @@ impl Instance {
                     if let Some(mut flags) = flags {
                         flags.set_may_enter(false);
                     }
-                    crate::Func::call_unchecked_raw(
+
+                    let rr_type = if let Some(func) = rr_handle {
+                        let type_idx = func.abi_info(store.0).2;
+                        let types = func
+                            .instance()
+                            .id()
+                            .get(store.0)
+                            .component()
+                            .types()
+                            .clone();
+                        RRWasmFuncType::Component { type_idx, types }
+                    } else {
+                        RRWasmFuncType::None
+                    };
+                    crate::Func::call_unchecked_raw_with_rr(
                         &mut store,
                         callee.as_non_null(),
                         NonNull::new(
                             &mut storage[..param_count.max(result_count)]
                                 as *mut [MaybeUninit<ValRaw>] as _,
                         )
                         .unwrap(),
+                        rr_type,
                     )?;
+
                     if let Some(mut flags) = flags {
                         flags.set_may_enter(may_enter_after_call);
                     }
@@ -1805,6 +1834,7 @@ impl Instance {
                 param_count,
                 result_count,
                 flags,
+                rr_handle,
             )
         };
 
@@ -2330,6 +2360,7 @@ impl Instance {
                 (flags & START_FLAG_ASYNC_CALLEE) != 0,
                 NonNull::new(callback).map(SendSyncPtr::new),
                 NonNull::new(post_return).map(SendSyncPtr::new),
+                None,
             )?;
         }
 
@@ -5047,6 +5078,7 @@ fn queue_call0<T: 'static>(
             is_concurrent,
             callback,
             post_return.map(SendSyncPtr::new),
+            Some(handle),
         )
     }
 }
diff --git a/crates/wasmtime/src/runtime/component/func.rs b/crates/wasmtime/src/runtime/component/func.rs
@@ -4,7 +4,7 @@ use crate::component::storage::storage_as_slice;
 use crate::component::types::ComponentFunc;
 use crate::component::values::Val;
 use crate::prelude::*;
-use crate::rr::component_hooks;
+use crate::rr::{RRWasmFuncType, component_hooks};
 use crate::runtime::vm::component::{ComponentInstance, InstanceFlags, ResourceTables};
 use crate::runtime::vm::{Export, VMFuncRef};
 use crate::store::StoreOpaque;
@@ -581,16 +581,12 @@ impl Func {
         LowerParams: Copy,
         LowerReturn: Copy,
     {
-        #[cfg(feature = "rr-component")]
-        {
-            use crate::rr::component_events::WasmFuncBeginEvent;
+        component_hooks::record_wasm_func_begin(
+            self.instance.id().instance(),
+            self.index,
+            store.0,
+        )?;
 
-            let instance = self.instance.id().instance();
-            let func_idx = self.index;
-            store
-                .0
-                .record_event(|| WasmFuncBeginEvent { instance, func_idx })?;
-        }
         let export = self.lifted_core_func(store.0);
 
         #[repr(C)]
@@ -635,12 +631,13 @@ impl Func {
             ))
             .unwrap();
 
-            component_hooks::record_and_replay_validate_wasm_func(
-                |store| crate::Func::call_unchecked_raw(store, export, params_and_returns),
-                params_and_returns.as_ref(),
-                self.abi_info(store.0).2,
-                self.instance.id().get(store.0).component().types().clone(),
+            let type_idx = self.abi_info(store.0).2;
+            let types = self.instance.id().get(store.0).component().types().clone();
+            crate::Func::call_unchecked_raw_with_rr(
                 &mut store,
+                export,
+                params_and_returns,
+                RRWasmFuncType::Component { type_idx, types },
             )?;
         }
 
diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs
@@ -1,5 +1,5 @@
 use crate::prelude::*;
-use crate::rr;
+use crate::rr::{self, RRWasmFuncType};
 use crate::runtime::Uninhabited;
 use crate::runtime::vm::{
     self, InterpreterRef, SendSyncPtr, StoreBox, VMArrayCallHostFuncContext,
@@ -1043,17 +1043,69 @@ impl Func {
         let func_ref = self.vm_func_ref(store.0);
         let params_and_returns = NonNull::new(params_and_returns).unwrap_or(NonNull::from(&mut []));
 
-        rr::core_hooks::record_and_replay_validate_wasm_func(
-            |mut store| {
-                // SAFETY: the safety of this function call is the same as the contract
-                // of this function.
-                unsafe { Self::call_unchecked_raw(&mut store, func_ref, params_and_returns) }
+        unsafe {
+            let ty = &self.ty(&store);
+            let origin = self.origin.expand();
+            Self::call_unchecked_raw_with_rr(
+                &mut store,
+                func_ref,
+                params_and_returns,
+                RRWasmFuncType::Core { ty, origin },
+            )
+        }
+    }
+
+    /// Same as [`Func::call_unchecked_raw`] but enables recording and replaying
+    /// hooks for func entry/exit based on whether the intended call was from a component or
+    /// core wasm module.
+    ///
+    /// This method is essentially a wrapper over [`Func::call_unchecked_raw`] and operates exactly like it
+    /// when recording/replay is not enabled, so use this for all calls that will potentially need
+    /// to be recorded/replayed.
+    ///
+    /// `RRWasmFuncType::None` can be used to disable recording/replaying and passthrough to `call_unchecked_raw`.
+    /// Eventually we can replace all occurences of `call_unchecked_raw` with this method with `RRWasmFuncType::None``
+    pub(crate) unsafe fn call_unchecked_raw_with_rr<T>(
+        mut store: &mut StoreContextMut<'_, T>,
+        func_ref: NonNull<VMFuncRef>,
+        params_and_returns: NonNull<[ValRaw]>,
+        rr: RRWasmFuncType,
+    ) -> Result<()> {
+        // SAFETY: the safety of this function call is the same as the contract
+        // of this function.
+        match rr {
+            RRWasmFuncType::Core { ty, origin } => {
+                rr::core_hooks::record_and_replay_validate_wasm_func(
+                    |mut store| {
+                        // SAFETY: the safety of this function call is the same as the contract
+                        // of this function.
+                        unsafe {
+                            Self::call_unchecked_raw(&mut store, func_ref, params_and_returns)
+                        }
+                    },
+                    unsafe { params_and_returns.as_ref() },
+                    ty,
+                    origin,
+                    &mut store,
+                )
+            }
+            #[cfg(feature = "component-model")]
+            RRWasmFuncType::Component { type_idx, types } => {
+                rr::component_hooks::record_and_replay_validate_wasm_func(
+                    |mut store| unsafe {
+                        Self::call_unchecked_raw(&mut store, func_ref, params_and_returns)
+                    },
+                    unsafe { params_and_returns.as_ref() },
+                    type_idx,
+                    types,
+                    &mut store,
+                )
+            }
+            // Passthrough
+            RRWasmFuncType::None => unsafe {
+                Self::call_unchecked_raw(&mut store, func_ref, params_and_returns)
             },
-            unsafe { params_and_returns.as_ref() },
-            &self.ty(&store),
-            self.origin.expand(),
-            &mut store,
-        )
+        }
     }
 
     pub(crate) unsafe fn call_unchecked_raw<T>(
diff --git a/crates/wasmtime/src/runtime/func/typed.rs b/crates/wasmtime/src/runtime/func/typed.rs
@@ -211,6 +211,8 @@ where
         // For component mode, Realloc uses this method `TypedFunc::call_raw`, but Realloc is its
         // own separate event for record/replay purposes. For now, we use the should_record flag to
         // distinguish but this could be removed in the future by folding it into the main function call event.
+        //
+        // Note: This can't use [`crate::Func::call_unchecked_raw_with_rr`] directly because of the closure capture
         rr::core_hooks::record_and_replay_validate_wasm_func(
             |store| {
                 invoke_wasm_and_catch_traps(store, |caller, vm| {
diff --git a/crates/wasmtime/src/runtime/rr.rs b/crates/wasmtime/src/runtime/rr.rs
@@ -43,7 +43,7 @@ impl FlatBytes for MaybeUninit<ValRaw> {
 
 /// Convenience method hooks for injecting event recording/replaying in the rest of the engine
 mod hooks;
-pub(crate) use hooks::core_hooks;
+pub(crate) use hooks::{RRWasmFuncType, core_hooks};
 #[cfg(feature = "component-model")]
 pub(crate) use hooks::{
     component_hooks, component_hooks::ConstMemorySliceCell, component_hooks::MemorySliceCell,
diff --git a/crates/wasmtime/src/runtime/rr/hooks.rs b/crates/wasmtime/src/runtime/rr/hooks.rs
@@ -3,3 +3,26 @@
 pub mod component_hooks;
 /// Core RR hooks
 pub mod core_hooks;
+
+use crate::{FuncType, WasmFuncOrigin};
+#[cfg(feature = "component-model")]
+use alloc::sync::Arc;
+#[cfg(feature = "component-model")]
+use wasmtime_environ::component::{ComponentTypes, TypeFuncIndex};
+
+/// Wasm function type information for RR hooks
+pub enum RRWasmFuncType<'a> {
+    /// No RR hooks to be performed
+    None,
+    /// Core RR hooks to be performed
+    Core {
+        ty: &'a FuncType,
+        origin: Option<WasmFuncOrigin>,
+    },
+    /// Component RR hooks to be performed
+    #[cfg(feature = "component-model")]
+    Component {
+        type_idx: TypeFuncIndex,
+        types: Arc<ComponentTypes>,
+    },
+}
diff --git a/crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs b/crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs
@@ -1,4 +1,5 @@
 use crate::ValRaw;
+use crate::component::ComponentInstanceId;
 #[cfg(feature = "rr-component")]
 use crate::rr::{RecordBuffer, Recorder, component_events::MemorySliceWriteEvent};
 
@@ -10,15 +11,15 @@ use crate::rr::common_events::{HostFuncEntryEvent, WasmFuncReturnEvent};
 #[cfg(feature = "rr-component")]
 use crate::rr::component_events::{
     LowerFlatEntryEvent, LowerFlatReturnEvent, LowerMemoryEntryEvent, LowerMemoryReturnEvent,
-    WasmFuncEntryEvent,
+    WasmFuncBeginEvent, WasmFuncEntryEvent,
 };
 #[cfg(feature = "rr-component")]
 use crate::rr::{RRFuncArgVals, ResultEvent, common_events::HostFuncReturnEvent};
 use crate::store::StoreOpaque;
 use crate::{StoreContextMut, prelude::*};
 use alloc::sync::Arc;
 use core::mem::MaybeUninit;
-use wasmtime_environ::component::{ComponentTypes, InterfaceType, TypeFuncIndex};
+use wasmtime_environ::component::{ComponentTypes, ExportIndex, InterfaceType, TypeFuncIndex};
 #[cfg(all(feature = "rr-component"))]
 use wasmtime_environ::component::{MAX_FLAT_PARAMS, MAX_FLAT_RESULTS};
 
@@ -30,6 +31,22 @@ pub enum ReplayLoweringPhase {
     HostFuncReturn,
 }
 
+/// Record hook for initiating wasm component function call
+///
+/// This differs from WasmFuncEntryEvent since this is pre-lowering, and
+/// WasmFuncEntryEvent is post-lowering
+#[inline]
+pub fn record_wasm_func_begin(
+    instance: ComponentInstanceId,
+    func_idx: ExportIndex,
+    store: &mut StoreOpaque,
+) -> Result<()> {
+    #[cfg(feature = "rr-component")]
+    store.record_event(|| WasmFuncBeginEvent { instance, func_idx })?;
+    let _ = (instance, func_idx, store);
+    Ok(())
+}
+
 /// Record hook wrapping a wasm component export function invocation and replay
 /// validation of return value
 #[inline]
