Temporary patch: Replay API in-store; separate into independent driver in future

Author: arjunr2

crates/wasmtime/src/config.rs

@@ -24,8 +24,6 @@ use crate::stack::{StackCreator, StackCreatorProxy};
 #[cfg(feature = "async")]
 use wasmtime_fiber::RuntimeFiberStackCreator;
 
-#[cfg(feature = "rr")]
-use crate::rr::ReplayReader;
 #[cfg(feature = "runtime")]
 pub use crate::runtime::code_memory::CustomCodeMemory;
 #[cfg(feature = "cache")]
@@ -178,6 +176,8 @@ pub struct Config {
     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,
 }
 
 /// User-provided configuration for the compiler.
@@ -234,63 +234,6 @@ impl Default for CompilerConfig {
     }
 }
 
-/// Settings for execution replay.
-#[cfg(feature = "rr")]
-#[derive(Debug, Clone)]
-pub struct ReplaySettings {
-    /// Flag to include additional signatures for replay validation.
-    pub validate: bool,
-    /// Static buffer size for deserialization of variable-length types (like [String]).
-    pub deser_buffer_size: usize,
-}
-
-#[cfg(feature = "rr")]
-impl Default for ReplaySettings {
-    fn default() -> Self {
-        Self {
-            validate: false,
-            deser_buffer_size: 64,
-        }
-    }
-}
-
-/// Configuration for replay execution.
-#[cfg(feature = "rr")]
-#[derive(Clone)]
-pub struct ReplayConfig {
-    /// Closure that generates a reader for replaying execution traces.
-    pub reader_initializer: Arc<dyn Fn() -> Box<dyn ReplayReader> + Send + Sync>,
-    /// Flag for dynamic validation checks when replaying events.
-    pub settings: ReplaySettings,
-}
-
-/// Configurations for record/replay (RR) executions.
-#[cfg(feature = "rr")]
-#[derive(Clone)]
-pub enum RRConfig {
-    /// Replay configuration.
-    Replay(ReplayConfig),
-}
-
-#[cfg(feature = "rr")]
-impl From<ReplayConfig> for RRConfig {
-    fn from(value: ReplayConfig) -> Self {
-        Self::Replay(value)
-    }
-}
-
-#[cfg(feature = "rr")]
-impl RRConfig {
-    /// Obtain the replay configuration.
-    ///
-    /// Return [`None`] if it is not configured.
-    pub fn replay(&self) -> Option<&ReplayConfig> {
-        match self {
-            Self::Replay(r) => Some(r),
-        }
-    }
-}
-
 impl Config {
     /// Creates a new configuration object with the default configuration
     /// specified.
@@ -345,6 +288,8 @@ impl Config {
             detect_host_feature: None,
             #[cfg(feature = "rr")]
             record_support: false,
+            #[cfg(feature = "rr")]
+            replay_support: false,
         };
         #[cfg(any(feature = "cranelift", feature = "winch"))]
         {
@@ -2799,14 +2744,30 @@ impl Config {
         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();
+        }
+        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.record_support || self.replay_support
     }
 }
 

crates/wasmtime/src/runtime.rs

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

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

@@ -1,4 +1,4 @@
-use crate::config::{ModuleVersionStrategy, ReplaySettings};
+use crate::config::ModuleVersionStrategy;
 use crate::prelude::*;
 use core::fmt;
 use events::EventActionError;
@@ -36,6 +36,26 @@ impl Default for RecordSettings {
     }
 }
 
+/// Settings for execution replay.
+#[cfg(feature = "rr")]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ReplaySettings {
+    /// Flag to include additional signatures for replay validation.
+    pub validate: bool,
+    /// Static buffer size for deserialization of variable-length types (like [String]).
+    pub deser_buffer_size: usize,
+}
+
+#[cfg(feature = "rr")]
+impl Default for ReplaySettings {
+    fn default() -> Self {
+        Self {
+            validate: false,
+            deser_buffer_size: 64,
+        }
+    }
+}
+
 /// Encapsulation of event types comprising an [`RREvent`] sum type
 mod events;
 /// I/O support for reading and writing traces
@@ -249,7 +269,7 @@ pub trait Recorder {
 /// essentially operates as an iterator over the recorded events
 pub trait Replayer: Iterator<Item = RREvent> {
     /// Constructs a reader on buffer
-    fn new_replayer(reader: Box<dyn ReplayReader>, settings: ReplaySettings) -> Result<Self>
+    fn new_replayer(reader: impl ReplayReader + 'static, settings: ReplaySettings) -> Result<Self>
     where
         Self: Sized;
 
@@ -459,7 +479,10 @@ impl Drop for ReplayBuffer {
 }
 
 impl Replayer for ReplayBuffer {
-    fn new_replayer(mut reader: Box<dyn ReplayReader>, settings: ReplaySettings) -> Result<Self> {
+    fn new_replayer(
+        mut reader: impl ReplayReader + 'static,
+        settings: ReplaySettings,
+    ) -> Result<Self> {
         let mut scratch = [0u8; 12];
         // Ensure module versions match
         let version = io::from_replay_reader::<&str, _>(&mut reader, &mut scratch)?;
@@ -479,6 +502,7 @@ impl Replayer for ReplayBuffer {
         }
 
         let deser_buffer = vec![0; settings.deser_buffer_size];
+        let reader = Box::new(reader);
 
         Ok(ReplayBuffer {
             reader,

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

@@ -132,6 +132,7 @@ where
         if self == expect {
             Ok(())
         } else {
+            log::error!("Validation against {:?} failed!", expect);
             Err(ReplayError::FailedValidation)
         }
     }

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

@@ -47,9 +47,11 @@ where
         store
             .0
             .record_event_validation(|| WasmFuncReturnEvent::from_anyhow_result(&result))?;
-        store
-            .0
-            .next_replay_event_validation::<WasmFuncReturnEvent, Result<RRFuncArgVals>>(&result)?;
+        // TODO: After adding validation support, replay with `next_replay_event_validation`
+        store.0.next_replay_event_and(|_r: WasmFuncReturnEvent| {
+            log::warn!("Yet to implement validation for WasmFuncReturnEvent; skipping for now");
+            Ok(())
+        })?;
         result?;
         return Ok(());
     }

crates/wasmtime/src/runtime/store.rs

@@ -90,7 +90,7 @@ use crate::rr::Validate;
 #[cfg(feature = "rr")]
 use crate::rr::{
     RREvent, RecordBuffer, RecordSettings, RecordWriter, Recorder, ReplayBuffer, ReplayError,
-    Replayer,
+    ReplayReader, ReplaySettings, Replayer,
 };
 #[cfg(feature = "gc")]
 use crate::runtime::vm::GcRootsList;
@@ -1025,6 +1025,20 @@ impl<T> Store<T> {
     ) -> Result<()> {
         self.inner.init_recording(recorder, settings)
     }
+
+    /// Configure a [`Store`] to enable execution replaying
+    ///
+    /// This feature must be initialized before instantiating any module within
+    /// the Store. Replay of events is performed according to provided settings, and
+    /// read from the provided reader.
+    #[cfg(feature = "rr")]
+    pub fn init_replaying(
+        &mut self,
+        replayer: impl ReplayReader + 'static,
+        settings: ReplaySettings,
+    ) -> Result<()> {
+        self.inner.init_replaying(replayer, settings)
+    }
 }
 
 impl<'a, T> StoreContext<'a, T> {
@@ -1473,6 +1487,24 @@ impl StoreOpaque {
         Ok(())
     }
 
+    #[cfg(feature = "rr")]
+    pub fn init_replaying(
+        &mut self,
+        replayer: impl ReplayReader + 'static,
+        settings: ReplaySettings,
+    ) -> Result<()> {
+        ensure!(
+            self.instance_count == 0,
+            "replaying store must not initialize any modules"
+        );
+        ensure!(
+            !self.engine().is_recording(),
+            "store recording cannot be enabled when initializing replay"
+        );
+        self.replay_buffer = Some(ReplayBuffer::new_replayer(replayer, settings)?);
+        Ok(())
+    }
+
     #[cfg(feature = "rr")]
     #[inline(always)]
     pub fn record_buffer_mut(&mut self) -> Option<&mut RecordBuffer> {

src/commands/replay.rs

@@ -3,8 +3,7 @@
 use crate::commands::run::RunCommand;
 use anyhow::Result;
 use clap::Parser;
-use std::{fs, io::BufReader, path::PathBuf, sync::Arc};
-use wasmtime::{ReplayConfig, ReplaySettings};
+use std::path::PathBuf;
 
 #[derive(Parser)]
 /// Replay-specific options for CLI
@@ -16,19 +15,19 @@ pub struct ReplayOptions {
     ///
     /// Note: The module used for replay must exactly match that used during recording
     #[arg(short, long, required = true, value_name = "RECORDED TRACE")]
-    trace: PathBuf,
+    pub trace: PathBuf,
 
     /// Dynamic checks of record signatures to validate replay consistency.
     ///
     /// Requires record traces to be generated with `validation_metadata` enabled.
     #[arg(short, long, default_value_t = false)]
-    validate: bool,
+    pub validate: bool,
 
     /// Size of static buffer needed to deserialized variable-length types like String. This is not
     /// not relevant for basic functional recording/replaying, but may be required to replay traces where
     /// `validation-metadata` was enabled for recording
     #[arg(short, long, default_value_t = 64)]
-    deser_buffer_size: usize,
+    pub deser_buffer_size: usize,
 }
 
 /// Execute a deterministic, embedding-agnostic replay of a Wasm modules given its associated recorded trace
@@ -48,19 +47,7 @@ impl ReplayCommand {
         if self.replay_opts.validate {
             anyhow::bail!("Cannot use `validate` when `rr-validate` feature is disabled");
         }
-        let replay_cfg = ReplayConfig {
-            reader_initializer: Arc::new(move || {
-                Box::new(BufReader::new(
-                    fs::File::open(&self.replay_opts.trace).unwrap(),
-                ))
-            }),
-            settings: ReplaySettings {
-                validate: self.replay_opts.validate,
-                deser_buffer_size: self.replay_opts.deser_buffer_size,
-                ..Default::default()
-            },
-        };
         // Replay uses the `run` command harness
-        self.run_cmd.execute(Some(replay_cfg))
+        self.run_cmd.execute(Some(self.replay_opts))
     }
 }

src/commands/run.rs

@@ -5,23 +5,25 @@
     allow(irrefutable_let_patterns, unreachable_patterns)
 )]
 
+#[cfg(feature = "rr")]
+use crate::commands::ReplayOptions;
 use crate::common::{Profile, RunCommon, RunTarget};
 use anyhow::{Context as _, Error, Result, anyhow, bail};
 use clap::Parser;
 use std::ffi::OsString;
+#[cfg(feature = "rr")]
+use std::io::BufReader;
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 use std::thread;
 use wasi_common::sync::{Dir, TcpListener, WasiCtxBuilder, ambient_authority};
-#[cfg(feature = "rr")]
-use wasmtime::ReplayConfig;
 use wasmtime::{Engine, Func, Module, Store, StoreLimits, Val, ValType};
 use wasmtime_wasi::{WasiCtxView, WasiView};
 
 #[cfg(feature = "rr")]
 use std::{fs, io};
 #[cfg(feature = "rr")]
-use wasmtime::RecordSettings;
+use wasmtime::{RecordSettings, ReplaySettings};
 #[cfg(feature = "wasi-config")]
 use wasmtime_wasi_config::{WasiConfig, WasiConfigVariables};
 #[cfg(feature = "wasi-http")]
@@ -93,7 +95,7 @@ impl RunCommand {
     /// Executes the command.
     pub fn execute(
         mut self,
-        #[cfg(feature = "rr")] replay_cfg: Option<ReplayConfig>,
+        #[cfg(feature = "rr")] replay_opts: Option<ReplayOptions>,
     ) -> Result<()> {
         self.run.common.init_logging()?;
 
@@ -115,8 +117,8 @@ impl RunCommand {
         }
 
         #[cfg(feature = "rr")]
-        if let Some(cfg) = replay_cfg {
-            //config.enable_replay(cfg)?;
+        if replay_opts.is_some() {
+            config.replaying(true);
         }
 
         let engine = Engine::new(&config)?;
@@ -195,6 +197,18 @@ impl RunCommand {
                     store.init_recording(fs::File::create(&path)?, settings)?;
                 }
             }
+
+            if let Some(opts) = replay_opts {
+                let settings = ReplaySettings {
+                    validate: opts.validate,
+                    deser_buffer_size: opts.deser_buffer_size,
+                    ..Default::default()
+                };
+                store.init_replaying(
+                    BufReader::new(fs::File::open(opts.trace).unwrap()),
+                    settings,
+                )?;
+            }
         }
 
         // Always run the module asynchronously to ensure that the module can be