From 2cfa024855f4b542999c7bf547838e2f85b244a5 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Thu, 15 Dec 2022 16:56:16 -0600
Subject: [PATCH] Support fuel and epoch interruption in the benchmarking API
 (#5449)

When these engine flags are passed be sure to configure the store
appropriately to ensure that the wasm can actually run instead of
crashing immediately, enabling benchmarking comparisons between these
modes.
---
 crates/bench-api/src/lib.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/crates/bench-api/src/lib.rs b/crates/bench-api/src/lib.rs
index 60291d24fe..7d894fe710 100644
--- a/crates/bench-api/src/lib.rs
+++ b/crates/bench-api/src/lib.rs
@@ -406,6 +406,8 @@ struct BenchState {
     make_wasi_cx: Box<dyn FnMut() -> Result<WasiCtx>>,
     module: Option<Module>,
     store_and_instance: Option<(Store<HostState>, Instance)>,
+    epoch_interruption: bool,
+    fuel: Option<u64>,
 }
 
 struct HostState {
@@ -454,6 +456,13 @@ impl BenchState {
             Ok(())
         })?;
 
+        let mut epoch_interruption = false;
+        let mut fuel = None;
+        if let Some(opts) = &options {
+            epoch_interruption = opts.epoch_interruption;
+            fuel = opts.fuel;
+        }
+
         let wasi_modules = options
             .map(|o| o.wasi_modules)
             .flatten()
@@ -484,6 +493,8 @@ impl BenchState {
             make_wasi_cx: Box::new(make_wasi_cx) as _,
             module: None,
             store_and_instance: None,
+            epoch_interruption,
+            fuel,
         })
     }
 
@@ -520,6 +531,13 @@ impl BenchState {
         // stdin/stdout/stderr.
         (self.instantiation_start)(self.instantiation_timer);
         let mut store = Store::new(self.linker.engine(), host);
+        if self.epoch_interruption {
+            store.set_epoch_deadline(1);
+        }
+        if let Some(fuel) = self.fuel {
+            store.add_fuel(fuel).unwrap();
+        }
+
         let instance = self.linker.instantiate(&mut store, &module)?;
         (self.instantiation_end)(self.instantiation_timer);