Added post return event to rr

Author: arjunr2

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

@@ -704,7 +704,12 @@ impl Func {
     /// This only works with functions defined within a synchronous store.
     #[inline]
     pub fn post_return(&self, mut store: impl AsContextMut) -> Result<()> {
-        let store = store.as_context_mut();
+        let mut store = store.as_context_mut();
+        component_hooks::record_wasm_func_post_return(
+            self.instance.id().instance(),
+            self.index,
+            &mut store,
+        )?;
         assert!(
             !store.0.async_support(),
             "must use `post_return_async` when async support is enabled on the config"

crates/wasmtime/src/runtime/rr/core.rs

@@ -156,6 +156,8 @@ rr_event! {
     ComponentMemorySliceWrite(__component_events::MemorySliceWriteEvent),
     /// Return from a component builtin
     ComponentBuiltinReturn(__component_events::BuiltinReturnEvent),
+    /// Call to `post_return` (after the function call)
+    ComponentPostReturn(__component_events::PostReturnEvent),
 
     // OPTIONAL events for replay validation (Component)
 

crates/wasmtime/src/runtime/rr/core/events/component_events.rs

@@ -26,6 +26,15 @@ pub struct InstantiationEvent {
     pub instance: ComponentInstanceId,
 }
 
+/// A call to `post_return` (after the function call)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PostReturnEvent {
+    /// Instance ID for the component instance
+    pub instance: ComponentInstanceId,
+    /// Export index for the function on which post_return is invoked
+    pub func_idx: ExportIndex,
+}
+
 /// A call event from Host into a Wasm component function
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct WasmFuncEntryEvent {

crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs

@@ -11,7 +11,7 @@ use crate::rr::common_events::{HostFuncEntryEvent, WasmFuncReturnEvent};
 #[cfg(feature = "rr-component")]
 use crate::rr::component_events::{
     LowerFlatEntryEvent, LowerFlatReturnEvent, LowerMemoryEntryEvent, LowerMemoryReturnEvent,
-    WasmFuncBeginEvent, WasmFuncEntryEvent,
+    PostReturnEvent, WasmFuncBeginEvent, WasmFuncEntryEvent,
 };
 #[cfg(feature = "rr-component")]
 use crate::rr::{RRFuncArgVals, ResultEvent, common_events::HostFuncReturnEvent};
@@ -47,6 +47,21 @@ pub fn record_wasm_func_begin(
     Ok(())
 }
 
+/// Record hook for wasm component function post_return call
+#[inline]
+pub fn record_wasm_func_post_return<T>(
+    instance: ComponentInstanceId,
+    func_idx: ExportIndex,
+    store: &mut StoreContextMut<'_, T>,
+) -> Result<()> {
+    #[cfg(feature = "rr-component")]
+    store
+        .0
+        .record_event(|| PostReturnEvent { 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]

crates/wasmtime/src/runtime/rr/replay_driver.rs

@@ -277,6 +277,28 @@ impl<T: 'static> ReplayInstance<T> {
                     );
                 }
             }
+            RREvent::ComponentPostReturn(event) => {
+                #[cfg(feature = "rr-component")]
+                {
+                    // Grab the correct component instance
+                    let key = event.instance;
+                    let instance = self
+                        .component_instances
+                        .get_mut(&key)
+                        .ok_or(ReplayError::MissingComponentInstance(key.as_u32()))?;
+
+                    let func = component::Func::from_lifted_func(*instance, event.func_idx);
+                    let mut store = self.store.as_context_mut();
+
+                    func.post_return(&mut store)?;
+                }
+                #[cfg(not(feature = "rr-component"))]
+                {
+                    bail!(
+                        "Cannot parse ComponentPostReturn replay event without rr-component feature enabled"
+                    );
+                }
+            }
             RREvent::CoreWasmInstantiation(event) => {
                 // Find matching module from environment to instantiate
                 let module = self
@@ -442,6 +464,28 @@ impl<T: 'static> ReplayInstance<T> {
                     );
                 }
             }
+            RREvent::ComponentPostReturn(event) => {
+                #[cfg(feature = "rr-component")]
+                {
+                    // Grab the correct component instance
+                    let key = event.instance;
+                    let instance = self
+                        .component_instances
+                        .get_mut(&key)
+                        .ok_or(ReplayError::MissingComponentInstance(key.as_u32()))?;
+
+                    let func = component::Func::from_lifted_func(*instance, event.func_idx);
+                    let mut store = self.store.as_context_mut();
+
+                    func.post_return_async(&mut store).await?;
+                }
+                #[cfg(not(feature = "rr-component"))]
+                {
+                    bail!(
+                        "Cannot parse ComponentPostReturn replay event without rr-component feature enabled"
+                    );
+                }
+            }
             RREvent::CoreWasmInstantiation(event) => {
                 // Find matching module from environment to instantiate
                 let module = self