Refactor buffering and add CLI option

Author: arjunr2

crates/cli-flags/src/lib.rs

@@ -506,9 +506,8 @@ wasmtime_option_group! {
         /// Include (optional) signatures to facilitate validation checks during replay
         /// (see `wasmtime replay` for details).
         pub validation_metadata: Option<bool>,
-        /// Window size of internal buffering for record events (large windows offer more opportunities
-        /// for coalescing events at the cost of memory usage).
-        pub event_window_size: Option<usize>,
+        /// Writer buffer size (in bytes) for internal buffering during recording. Defaults to 8 KiB.
+        pub buffer_size: Option<usize>,
     }
 
     enum Record {

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

@@ -110,8 +110,6 @@ impl From<RRModuleInstanceId> for InstanceId {
 ///
 /// This type can be optimized for [`RREvent`] data configurations.
 pub struct RecordBuffer {
-    /// In-memory event buffer to enable windows for coalescing
-    buf: Vec<RREvent>,
     /// Writer to store data into
     writer: Box<dyn RecordWriter>,
     /// Settings in record configuration
@@ -121,10 +119,7 @@ pub struct RecordBuffer {
 impl RecordBuffer {
     /// Push a new record event [`RREvent`] to the buffer
     fn push_event(&mut self, event: RREvent) -> Result<()> {
-        self.buf.push(event);
-        if self.buf.len() >= self.settings().event_window_size {
-            self.flush()?;
-        }
+        to_record_writer(&event, &mut *self.writer)?;
         Ok(())
     }
 
@@ -140,7 +135,6 @@ impl Recorder for 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,
         };
@@ -170,10 +164,8 @@ impl Recorder for RecordBuffer {
 
     fn flush(&mut self) -> Result<()> {
         log::debug!("Flushing record buffer...");
-        for e in self.buf.drain(..) {
-            to_record_writer(&e, &mut *self.writer)?;
-        }
-        return Ok(());
+        self.writer.flush()?;
+        Ok(())
     }
 
     #[inline]
@@ -273,19 +265,11 @@ impl Replayer for ReplayBuffer {
     }
 
     #[inline]
-    #[allow(
-        unused,
-        reason = "method only used for gated validation, but will be extended in the future"
-    )]
     fn settings(&self) -> &ReplaySettings {
         &self.settings
     }
 
     #[inline]
-    #[allow(
-        unused,
-        reason = "method only used for gated validation, but will be extended in the future"
-    )]
     fn trace_settings(&self) -> &RecordSettings {
         &self.trace_settings
     }

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

@@ -21,15 +21,12 @@ mod io;
 pub struct RecordSettings {
     /// Flag to include additional signatures for replay validation.
     pub add_validation: bool,
-    /// Maximum window size of internal event buffer.
-    pub event_window_size: usize,
 }
 
 impl Default for RecordSettings {
     fn default() -> Self {
         Self {
             add_validation: false,
-            event_window_size: 16,
         }
     }
 }

src/commands/replay.rs

@@ -36,6 +36,10 @@ pub struct ReplayOptions {
     /// `validate` was enabled for recording.
     #[arg(short, long, default_value_t = 64)]
     pub deserialize_buffer_size: usize,
+
+    /// Size of the read buffer (in bytes) for internal buffering during replay.
+    #[arg(short, long, default_value_t = 8 * 1024)]
+    pub read_buffer_size: usize,
 }
 
 /// Execute a deterministic, embedding-agnostic replay of a Wasm modules given its associated recorded trace.
@@ -104,7 +108,7 @@ impl ReplayCommand {
 
         let allow_unknown_exports = self.run_cmd.run.common.wasm.unknown_exports_allow;
         let mut replay_instance = renv.instantiate_with(
-            io::BufReader::new(fs::File::open(opts.trace)?),
+            io::BufReader::with_capacity(opts.read_buffer_size, fs::File::open(opts.trace)?),
             |store| {
                 // If fuel has been configured, we want to add the configured
                 // fuel amount to this store.

src/commands/run.rs

@@ -135,14 +135,15 @@ pub mod rr_impl {
                 add_validation: record
                     .validation_metadata
                     .unwrap_or(default_settings.add_validation),
-                event_window_size: record
-                    .event_window_size
-                    .unwrap_or(default_settings.event_window_size),
             };
             if path.trim().is_empty() {
                 store.record(io::sink(), settings)?;
             } else {
-                store.record(fs::File::create(&path)?, settings)?;
+                let file = io::BufWriter::with_capacity(
+                    record.buffer_size.unwrap_or(8 * 1024), // Default to 8 KiB buffer
+                    fs::File::create(&path)?,
+                );
+                store.record(file, settings)?;
             }
         }
         Ok(())