Initial factoring out of replay driver

Things left to complete:
* Core wasm support
* Improved error handling
* Wasmtime CLI integration

Author: arjunr2

crates/cli-flags/src/lib.rs

@@ -1027,7 +1027,7 @@ impl CommonOptions {
                     _v => (),
                     _ => err,
                 }
-                config.recording(true);
+                config.rr(wasmtime::RRConfig::Recording);
             },
             _ => err,
         }

crates/wasmtime/src/config.rs

@@ -110,6 +110,17 @@ impl ModuleVersionStrategy {
     }
 }
 
+/// Configuration for record/replay
+#[derive(Clone)]
+pub enum RRConfig {
+    /// Recording on store is enabled
+    Recording,
+    /// Replaying on store is enabled
+    Replaying,
+    /// No record/replay is enabled
+    None,
+}
+
 /// Global configuration options used to create an [`Engine`](crate::Engine)
 /// and customize its behavior.
 ///
@@ -175,9 +186,7 @@ pub struct Config {
     pub(crate) macos_use_mach_ports: bool,
     pub(crate) detect_host_feature: Option<fn(&str) -> Option<bool>>,
     #[cfg(feature = "rr")]
-    pub(crate) record_support: bool,
-    #[cfg(feature = "rr")]
-    pub(crate) replay_support: bool,
+    pub(crate) rr_config: RRConfig,
 }
 
 /// User-provided configuration for the compiler.
@@ -287,9 +296,7 @@ impl Config {
             #[cfg(not(feature = "std"))]
             detect_host_feature: None,
             #[cfg(feature = "rr")]
-            record_support: false,
-            #[cfg(feature = "rr")]
-            replay_support: false,
+            rr_config: RRConfig::None,
         };
         #[cfg(any(feature = "cranelift", feature = "winch"))]
         {
@@ -2705,11 +2712,9 @@ impl Config {
         self
     }
 
-    /// Enforce deterministic execution configurations. Currently, means the following:
+    /// Enforce deterministic execution configurations. Currently, this means the following:
     /// * Enabling NaN canonicalization with [`Config::cranelift_nan_canonicalization`]
     /// * Enabling deterministic relaxed SIMD with [`Config::relaxed_simd_deterministic`]
-    ///
-    /// Required for faithful record/replay execution.
     #[inline]
     pub fn enforce_determinism(&mut self) -> &mut Self {
         #[cfg(any(feature = "cranelift", feature = "winch"))]
@@ -2718,56 +2723,29 @@ impl Config {
         self
     }
 
-    /// Remove determinstic execution enforcements (if any) applied
-    /// by [`Config::enforce_determinism`].
-    #[inline]
-    pub fn remove_determinism_enforcement(&mut self) -> &mut Self {
-        #[cfg(any(feature = "cranelift", feature = "winch"))]
-        self.cranelift_nan_canonicalization(false);
-        self.relaxed_simd_deterministic(false);
-        self
-    }
-
-    /// Enable execution trace recording with the provided configuration.
+    /// Enable execution trace recording or replaying to the configuration
     ///
-    /// This method implicitly enforces determinism (see [`Config::enforce_determinism`]
-    /// for details).
+    /// When either recording/replaying are enabled, determinism is implicitly
+    /// enforced (see [`Config::enforce_determinism`] for details)
     #[cfg(feature = "rr")]
     #[inline]
-    pub fn recording(&mut self, enable: bool) -> &mut Self {
-        if enable {
-            self.enforce_determinism();
-        } else {
-            self.remove_determinism_enforcement();
-        }
-        self.record_support = enable;
-        self
-    }
-
-    /// Enable execution trace replaying with the provided configuration.
-    ///
-    /// This method implicitly enforces determinism (see [`Config::enforce_determinism`]
-    /// for details).
-    #[cfg(feature = "rr")]
-    #[inline]
-    pub fn replaying(&mut self, enable: bool) -> &mut Self {
-        if enable {
-            self.enforce_determinism();
-        } else {
-            self.remove_determinism_enforcement();
+    pub fn rr(&mut self, cfg: RRConfig) -> &mut Self {
+        self.rr_config = cfg;
+        match self.rr_config {
+            RRConfig::Recording | RRConfig::Replaying => self.enforce_determinism(),
+            _ => self,
         }
-        self.replay_support = enable;
-        self
     }
 
     /// Evaluates to true if current configuration must respect
     /// deterministic execution in its configuration.
-    ///
-    /// Required for faithful record/replay execution.
     #[cfg(feature = "rr")]
     #[inline]
     pub fn is_determinism_enforced(&mut self) -> bool {
-        self.record_support || self.replay_support
+        match self.rr_config {
+            RRConfig::Recording | RRConfig::Replaying => true,
+            RRConfig::None => false,
+        }
     }
 }
 

crates/wasmtime/src/engine.rs

@@ -1,4 +1,6 @@
 use crate::Config;
+#[cfg(feature = "rr")]
+use crate::RRConfig;
 use crate::prelude::*;
 #[cfg(feature = "runtime")]
 pub use crate::runtime::code_memory::CustomCodeMemory;
@@ -258,7 +260,20 @@ impl Engine {
     #[cfg(feature = "rr")]
     #[inline]
     pub fn is_recording(&self) -> bool {
-        self.config().record_support
+        match self.config().rr_config {
+            RRConfig::Recording => true,
+            _ => false,
+        }
+    }
+
+    /// Returns whether the engine is configured to support execution replaying
+    #[cfg(feature = "rr")]
+    #[inline]
+    pub fn is_replaying(&self) -> bool {
+        match self.config().rr_config {
+            RRConfig::Replaying => true,
+            _ => false,
+        }
     }
 
     /// Detects whether the bytes provided are a precompiled object produced by

crates/wasmtime/src/runtime.rs

@@ -88,7 +88,9 @@ pub use memory::*;
 pub use module::{Module, ModuleExport};
 pub use resources::*;
 #[cfg(feature = "rr")]
-pub use rr::{RecordSettings, RecordWriter, ReplayReader, ReplaySettings};
+pub use rr::{
+    RecordSettings, RecordWriter, ReplayEnvironment, ReplayInstance, ReplayReader, ReplaySettings,
+};
 #[cfg(all(feature = "async", feature = "call-hook"))]
 pub use store::CallHookHandler;
 pub use store::{

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

@@ -534,7 +534,7 @@ impl Func {
     /// `LowerParams` and `LowerReturn` type. They must match the type of `self`
     /// for the params/results that are going to be produced. Additionally
     /// these types must be representable with a sequence of `ValRaw` values.
-    unsafe fn call_raw<T, Return, LowerParams, LowerReturn>(
+    pub(crate) unsafe fn call_raw<T, Return, LowerParams, LowerReturn>(
         &self,
         mut store: StoreContextMut<'_, T>,
         lower: impl FnOnce(
@@ -548,6 +548,19 @@ impl Func {
         LowerParams: Copy,
         LowerReturn: Copy,
     {
+        #[cfg(feature = "rr-component")]
+        {
+            use crate::rr::component_events::WasmFuncBeginEvent;
+
+            let component = *self.instance.id().get(store.0).component().checksum();
+            let instance = self.instance.id().instance();
+            let func_idx = self.index;
+            store.0.record_event(|| WasmFuncBeginEvent {
+                component,
+                instance,
+                func_idx,
+            })?;
+        }
         let export = self.lifted_core_func(store.0);
 
         #[repr(C)]
@@ -592,12 +605,11 @@ impl Func {
             ))
             .unwrap();
 
-            component_hooks::record_replay_wasm_func(
+            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.index,
                 self.abi_info(store.0).2,
-                self.instance.id(),
+                self.instance.id().get(store.0).component().types().clone(),
                 &mut store,
             )?;
         }
@@ -810,7 +822,7 @@ impl Func {
         Ok(())
     }
 
-    fn lift_results<'a, 'b>(
+    pub(crate) fn lift_results<'a, 'b>(
         cx: &'a mut LiftContext<'b>,
         results_ty: InterfaceType,
         src: &'a [ValRaw],
@@ -875,7 +887,7 @@ impl Func {
     /// The `lower` closure provided should perform the actual lowering and
     /// return the result of the lowering operation which is then returned from
     /// this function as well.
-    fn with_lower_context<T>(
+    pub(crate) fn with_lower_context<T>(
         self,
         mut store: StoreContextMut<T>,
         may_enter: bool,

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

@@ -257,7 +257,7 @@ where
 
     let types = vminstance.component().types().clone();
 
-    rr::component_hooks::record_replay_host_func_entry(storage, &types, &ty, store.0)?;
+    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);
@@ -813,7 +813,7 @@ where
 
     let types = instance.id().get(store.0).component().types().clone();
 
-    rr::component_hooks::record_replay_host_func_entry(storage, &types, &ty, store.0)?;
+    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];

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

@@ -902,6 +902,13 @@ impl<'a> Instantiator<'a> {
         store.store_data_mut().component_instance_mut(self.id)
     }
 
+    /// Convenience helper to return the instance ID of the `ComponentInstance` that's
+    /// being instantiated
+    #[cfg(feature = "rr-component")]
+    fn id(&self) -> ComponentInstanceId {
+        self.id
+    }
+
     // NB: This method is only intended to be called during the instantiation
     // process because the `Arc::get_mut` here is fallible and won't generally
     // succeed once the instance has been handed to the embedder. Before that
@@ -1016,23 +1023,18 @@ impl<T: 'static> InstancePre<T> {
 
     fn instantiate_impl(&self, mut store: impl AsContextMut<Data = T>) -> Result<Instance> {
         let mut store = store.as_context_mut();
-        #[cfg(feature = "rr-component")]
-        {
-            use crate::rr::{Validate, component_events::InstantiationEvent};
-            store
-                .0
-                .record_event(|| InstantiationEvent::from_component(&self.component))?;
-            // This is a required validation check for functional correctness, so don't use
-            // [`StoreOpaque::next_replay_event_validation`]
-            store.0.next_replay_event_and(|event: InstantiationEvent| {
-                event.validate(&InstantiationEvent::from_component(&self.component))
-            })?;
-        }
         store
             .engine()
             .allocator()
             .increment_component_instance_count()?;
         let mut instantiator = Instantiator::new(&self.component, store.0, &self.imports);
+        #[cfg(feature = "rr-component")]
+        {
+            use crate::rr::component_events::InstantiationEvent;
+            store.0.record_event(|| {
+                InstantiationEvent(*self.component.checksum(), instantiator.id())
+            })?;
+        }
         instantiator.run(&mut store).map_err(|e| {
             store
                 .engine()

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

@@ -112,7 +112,7 @@ mod linker;
 mod matching;
 mod resource_table;
 mod resources;
-mod storage;
+pub(crate) mod storage;
 pub(crate) mod store;
 pub mod types;
 mod values;

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

@@ -3,14 +3,15 @@ use crate::store::{StoreData, StoreId, StoreOpaque};
 #[cfg(feature = "component-model-async")]
 use alloc::vec::Vec;
 use core::pin::Pin;
+use serde::{Deserialize, Serialize};
 use wasmtime_environ::PrimaryMap;
 
 #[derive(Default)]
 pub struct ComponentStoreData {
     instances: PrimaryMap<ComponentInstanceId, Option<OwnedComponentInstance>>,
 }
 
-#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Ord, PartialOrd)]
 pub struct ComponentInstanceId(u32);
 wasmtime_environ::entity_impl!(ComponentInstanceId);
 

crates/wasmtime/src/runtime/func.rs

@@ -2366,7 +2366,7 @@ impl HostContext {
             let func = &state.func;
 
             let func_type_index = state._ty.index();
-            let (params_size, results_size) = {
+            let (flat_size_params, flat_size_results) = {
                 let type_index = state._ty.index();
                 let wasm_func_subtype = caller.engine().signatures().borrow(type_index).unwrap();
                 let wasm_func_type = wasm_func_subtype.unwrap_func();
@@ -2383,13 +2383,13 @@ impl HostContext {
                         .collect::<Vec<_>>(),
                 )
             };
-            let (num_params, num_results) = (params_size.len(), results_size.len());
+            let (num_params, num_results) = (flat_size_params.len(), flat_size_results.len());
 
             // 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_replay_host_func_entry(
                 unsafe { &args.as_ref()[..num_params] },
-                params_size.as_slice(),
+                flat_size_params.as_slice(),
                 &func_type_index,
                 caller.store.0,
             )?;
@@ -2435,7 +2435,7 @@ impl HostContext {
                 // Record the return values
                 rr::core_hooks::record_host_func_return(
                     unsafe { &args.as_ref()[..num_results] },
-                    results_size.as_slice(),
+                    flat_size_results.as_slice(),
                     &func_type_index,
                     caller.store.0,
                 )?;