Fully encapsulate trace as RREvents for new record/replay embedders

Author: arjunr2

crates/rr/src/events.rs

@@ -172,30 +172,6 @@ macro_rules! event_error_types {
     );
 }
 
-/// Events used as markers for debugging/testing in traces
-///
-/// Marker events should be injectable at any point in a record
-/// trace without impacting functional correctness of replay
-pub mod marker_events {
-    use serde::{Deserialize, Serialize};
-
-    /// A Nop event
-    #[derive(Debug, Clone, Serialize, Deserialize)]
-    pub struct NopEvent;
-
-    /// An event for custom String messages
-    #[derive(Debug, Clone, Serialize, Deserialize)]
-    pub struct CustomMessageEvent(pub String);
-    impl<T> From<T> for CustomMessageEvent
-    where
-        T: Into<String>,
-    {
-        fn from(v: T) -> Self {
-            Self(v.into())
-        }
-    }
-}
-
 pub mod common_events;
 pub mod component_events;
 pub mod core_events;

crates/rr/src/events/common_events.rs

@@ -4,6 +4,7 @@
 //! or [`core_events`]
 
 use super::*;
+use crate::RecordSettings;
 use serde::{Deserialize, Serialize};
 
 /// A call event from Wasm (core or component) into the host
@@ -40,3 +41,37 @@ impl Validate<&Result<RRFuncArgVals>> for WasmFuncReturnEvent {
 event_error_types! {
     pub struct WasmFuncReturnError(..)
 }
+
+/// Signature of recorded trace.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TraceSignatureEvent {
+    /// Checksum of the trace contents.
+    ///
+    /// This can be used to verify integrity of the trace during replay.
+    pub checksum: String,
+    /// Settings used during trace recording.
+    pub settings: RecordSettings,
+}
+
+impl Validate<str> for TraceSignatureEvent {
+    fn validate(&self, expect: &str) -> Result<(), ReplayError> {
+        self.log();
+        if self.checksum == expect {
+            Ok(())
+        } else {
+            Err(ReplayError::FailedValidation)
+        }
+    }
+}
+
+/// A diagnostic event for custom String messages.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CustomMessageEvent(pub String);
+impl<T> From<T> for CustomMessageEvent
+where
+    T: Into<String>,
+{
+    fn from(v: T) -> Self {
+        Self(v.into())
+    }
+}

crates/rr/src/io.rs

@@ -1,7 +1,7 @@
+use super::RREvent;
 use anyhow::Result;
 use core::any::Any;
 use postcard;
-use serde::{Deserialize, Serialize};
 
 cfg_if::cfg_if! {
     if #[cfg(feature = "std")] {
@@ -130,12 +130,12 @@ cfg_if::cfg_if! {
     }
 }
 
-/// Serialize and write `value` to a `RecordWriter`
+/// Serialize and write an `RREvent` to a `RecordWriter`
 ///
-/// Currently uses `postcard` serializer
-pub fn to_record_writer<T, W>(value: &T, writer: &mut W) -> Result<()>
+/// This is the lowest-level underlying writer function for RR events,
+/// helpful for implementing `Recorder`s. Currently uses [`postcard`] serializer.
+pub fn to_record_writer<W>(value: &RREvent, writer: &mut W) -> Result<()>
 where
-    T: Serialize + ?Sized,
     W: RecordWriter + ?Sized,
 {
     #[cfg(feature = "std")]
@@ -150,13 +150,12 @@ where
     Ok(())
 }
 
-/// Read and deserialize a `value` from a `ReplayReader`.
+/// Read and deserialize an `RREvent` from a `ReplayReader`.
 ///
-/// Currently uses `postcard` deserializer, with optional scratch
-/// buffer to deserialize into
-pub fn from_replay_reader<'a, T, R>(reader: &'a mut R, scratch: &'a mut [u8]) -> Result<T>
+/// This is the lowest-level underlying reader function for RR events,
+/// helpful for implementing `Replayer`s. Currently uses [`postcard`] deserializer.
+pub fn from_replay_reader<'a, R>(reader: &'a mut R, scratch: &'a mut [u8]) -> Result<RREvent>
 where
-    T: Deserialize<'a>,
     R: ReplayReader + ?Sized,
 {
     #[cfg(feature = "std")]
@@ -167,7 +166,7 @@ where
     {
         let flavor = ReplayReaderFlavor { reader, scratch };
         let mut deserializer = postcard::Deserializer::from_flavor(flavor);
-        let t = T::deserialize(&mut deserializer)?;
+        let t = serde::Deserialize::deserialize(&mut deserializer)?;
         deserializer.finalize()?;
         Ok(t)
     }

crates/rr/src/lib.rs

@@ -2,7 +2,6 @@ use anyhow::Result;
 use core::fmt;
 pub use events::{
     EventError, RRFuncArgVals, ResultEvent, Validate, common_events, component_events, core_events,
-    marker_events,
 };
 pub use io::{RecordWriter, ReplayReader, from_replay_reader, to_record_writer};
 use serde::{Deserialize, Serialize};
@@ -54,21 +53,28 @@ mod io;
 /// Macro template for [`RREvent`] and its conversion to/from specific
 /// event types
 macro_rules! rr_event {
-        (
-            $(
-                $(#[doc = $doc:literal])*
-                $variant:ident($event:ty)
-            ),*
-        ) => (
+    (
+        $(
+            $(#[doc = $doc_np:literal])*
+            $variant_no_payload:ident
+        ),*
+        ;
+        $(
+            $(#[doc = $doc:literal])*
+            $variant:ident($event:ty)
+        ),*
+    ) => (
         /// A single, unified, low-level recording/replay event
         ///
         /// This type is the narrow waist for serialization/deserialization.
         /// Higher-level events (e.g. import calls consisting of lifts and lowers
         /// of parameter/return types) may drop down to one or more [`RREvent`]s
         #[derive(Debug, Clone, Serialize, Deserialize)]
         pub enum RREvent {
-            /// Event signalling the end of a trace
-            Eof,
+            $(
+                $(#[doc = $doc_np])*
+                $variant_no_payload,
+            )*
             $(
                 $(#[doc = $doc])*
                 $variant($event),
@@ -78,9 +84,11 @@ macro_rules! rr_event {
         impl fmt::Display for RREvent {
             fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                 match self {
-                    Self::Eof => write!(f, "Eof event"),
                     $(
-                    Self::$variant(e) => write!(f, "{:?}", e),
+                        Self::$variant_no_payload => write!(f, "{}", stringify!($variant_no_payload)),
+                    )*
+                    $(
+                        Self::$variant(payload) => write!(f, "{:?}", payload),
                     )*
                 }
             }
@@ -104,16 +112,25 @@ macro_rules! rr_event {
                 }
             }
         )*
-   );
+    );
+
 }
 
 // Set of supported record/replay events
 rr_event! {
-    // Marker events
     /// Nop Event
-    Nop(marker_events::NopEvent),
-    /// A custom message
-    CustomMessage(marker_events::CustomMessageEvent),
+    Nop,
+    /// Event signalling the end of a trace
+    Eof
+    ;
+    /// The signature of the trace, enabling trace integrity during replay.
+    ///
+    /// This is always at the start of any valid trace.
+    TraceSignature(common_events::TraceSignatureEvent),
+    /// A custom message in the trace, useful for diagnostics.
+    ///
+    /// Does not affect trace replay functionality
+    CustomMessage(common_events::CustomMessageEvent),
 
     // Common events for both core or component wasm
     // REQUIRED events
@@ -172,11 +189,11 @@ rr_event! {
 }
 
 impl RREvent {
-    /// Indicates whether current event is a marker event
+    /// Indicates whether current event is a diagnostic event
     #[inline]
-    pub fn is_marker(&self) -> bool {
+    pub fn is_diagnostic(&self) -> bool {
         match self {
-            Self::Nop(_) | Self::CustomMessage(_) => true,
+            Self::Nop | Self::CustomMessage(_) => true,
             _ => false,
         }
     }

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

@@ -138,15 +138,18 @@ impl RecordBuffer {
 }
 
 impl Recorder for RecordBuffer {
-    fn new_recorder(mut writer: impl RecordWriter, settings: RecordSettings) -> Result<Self> {
-        // Replay requires the Module version and record settings
-        to_record_writer(ModuleVersionStrategy::WasmtimeVersion.as_str(), &mut writer)?;
-        to_record_writer(&settings, &mut writer)?;
-        Ok(RecordBuffer {
+    fn new_recorder(writer: impl RecordWriter, settings: RecordSettings) -> Result<Self> {
+        let settings_local = settings.clone();
+        let mut buf = RecordBuffer {
             buf: Vec::new(),
             writer: Box::new(writer),
             settings,
-        })
+        };
+        buf.record_event(|| common_events::TraceSignatureEvent {
+            checksum: ModuleVersionStrategy::WasmtimeVersion.as_str().to_string(),
+            settings: settings_local,
+        })?;
+        Ok(buf)
     }
 
     #[inline]
@@ -211,7 +214,7 @@ impl Iterator for ReplayBuffer {
                     if let RREvent::Eof = &event {
                         self.eof_encountered = true;
                         break 'event_loop None;
-                    } else if event.is_marker() {
+                    } else if event.is_diagnostic() {
                         continue 'event_loop;
                     } else {
                         log::debug!("Read replay event => {event}");
@@ -245,38 +248,29 @@ impl Drop for ReplayBuffer {
 }
 
 impl Replayer for ReplayBuffer {
-    fn new_replayer(
-        mut reader: impl ReplayReader + 'static,
-        settings: ReplaySettings,
-    ) -> Result<Self> {
-        let mut scratch = [0u8; 12];
-        // Ensure module versions match
-        let version = from_replay_reader::<&str, _>(&mut reader, &mut scratch)?;
-        assert_eq!(
-            version,
-            ModuleVersionStrategy::WasmtimeVersion.as_str(),
-            "Wasmtime version mismatch between engine used for record and replay"
-        );
-
-        // Read the recording settings
-        let trace_settings: RecordSettings = from_replay_reader(&mut reader, &mut scratch)?;
-
-        if settings.validate && !trace_settings.add_validation {
+    fn new_replayer(reader: impl ReplayReader + 'static, settings: ReplaySettings) -> Result<Self> {
+        let mut buf = ReplayBuffer {
+            reader: Box::new(reader),
+            deser_buffer: vec![0; settings.deserialize_buffer_size],
+            settings,
+            // This doesn't matter now; will override after reading header
+            trace_settings: RecordSettings::default(),
+            eof_encountered: false,
+        };
+
+        let signature: common_events::TraceSignatureEvent = buf.next_event_typed()?;
+        // Ensure the trace integrity
+        signature.validate(ModuleVersionStrategy::WasmtimeVersion.as_str())?;
+        // Update the trace settings
+        buf.trace_settings = signature.settings;
+
+        if buf.settings.validate && !buf.trace_settings.add_validation {
             log::warn!(
                 "Replay validation will be omitted since the recorded trace has no validation metadata..."
             );
         }
 
-        let deser_buffer = vec![0; settings.deserialize_buffer_size];
-        let reader = Box::new(reader);
-
-        Ok(ReplayBuffer {
-            reader,
-            settings,
-            trace_settings,
-            deser_buffer,
-            eof_encountered: false,
-        })
+        Ok(buf)
     }
 
     #[inline]