Allow to enable wasmtime's profiling WebAssembly

This enables wasmtime's builtin profiling support.
https://docs.wasmtime.dev/examples-profiling.html

cli/src/main.rs

@@ -35,7 +35,7 @@ use {
 /// Create a new server, bind it to an address, and serve responses until an error occurs.
 pub async fn serve(opts: Opts) -> Result<(), Error> {
     // Load the wasm module into an execution context
-    let mut ctx = ExecuteCtx::new(opts.input())?
+    let mut ctx = ExecuteCtx::new(opts.input(), opts.profiling_strategy())?
         .with_log_stderr(opts.log_stderr())
         .with_log_stdout(opts.log_stdout());
 

cli/src/opts.rs

@@ -7,7 +7,7 @@ use {
         path::{Path, PathBuf},
     },
     structopt::StructOpt,
-    viceroy_lib::Error,
+    viceroy_lib::{Error, ProfilingStrategy},
 };
 
 // Command-line arguments for the Viceroy CLI.
@@ -42,6 +42,9 @@ pub struct Opts {
     /// to a value before starting Viceroy
     #[structopt(short = "v", parse(from_occurrences))]
     verbosity: usize,
+    // Whether to enable wasmtime's builtin profiler.
+    #[structopt(long = "profiler", parse(try_from_str = check_wasmtime_profiler_mode))]
+    profiler: Option<ProfilingStrategy>,
 }
 
 impl Opts {
@@ -77,6 +80,11 @@ impl Opts {
     pub fn verbosity(&self) -> usize {
         self.verbosity
     }
+
+    // Whether to enable wasmtime's builtin profiler.
+    pub fn profiling_strategy(&self) -> ProfilingStrategy {
+        self.profiler.unwrap_or(ProfilingStrategy::None)
+    }
 }
 
 /// A parsing function used by [`Opts`][opts] to check that the input is a valid Wasm module in
@@ -92,6 +100,17 @@ fn check_module(s: &str) -> Result<PathBuf, Error> {
     }
 }
 
+/// A parsing function used by [`Opts`][opts] to check that the input is valid wasmtime's profiling strategy.
+///
+/// [opts]: struct.Opts.html
+fn check_wasmtime_profiler_mode(s: &str) -> Result<ProfilingStrategy, Error> {
+    match s {
+        "jitdump" => Ok(ProfilingStrategy::JitDump),
+        "vtune" => Ok(ProfilingStrategy::VTune),
+        _ => Err(Error::ProfilingStrategy),
+    }
+}
+
 /// A collection of unit tests for our CLI argument parsing.
 ///
 /// Note: When using [`StructOpt::from_iter_safe`][from] to test how command line arguments are
@@ -224,4 +243,49 @@ mod opts_tests {
             res => panic!("unexpected result: {:?}", res),
         }
     }
+
+    /// Test that wasmtime's jitdump profiling strategy is accepted.
+    #[test]
+    fn wasmtime_profiling_strategy_jitdump_is_accepted() -> TestResult {
+        let args = &[
+            "dummy-program-name",
+            "--profiler",
+            "jitdump",
+            &test_file("minimal.wat"),
+        ];
+        match Opts::from_iter_safe(args) {
+            Ok(_) => Ok(()),
+            res => panic!("unexpected result: {:?}", res),
+        }
+    }
+
+    /// Test that wasmtime's VTune profiling strategy is accepted.
+    #[test]
+    fn wasmtime_profiling_strategy_vtune_is_accepted() -> TestResult {
+        let args = &[
+            "dummy-program-name",
+            "--profiler",
+            "vtune",
+            &test_file("minimal.wat"),
+        ];
+        match Opts::from_iter_safe(args) {
+            Ok(_) => Ok(()),
+            res => panic!("unexpected result: {:?}", res),
+        }
+    }
+
+    /// Test that an invalid wasmtime's profiling strategy rejected.
+    #[test]
+    fn invalid_wasmtime_profiling_strategy_is_rejected() -> TestResult {
+        let args = &[
+            "dummy-program-name",
+            "--profiler",
+            "invalid_profiling_strategy",
+            &test_file("minimal.wat"),
+        ];
+        match Opts::from_iter_safe(args) {
+            Ok(_) => panic!("unexpected result"),
+            Err(_) => Ok(()),
+        }
+    }
 }

cli/tests/integration/common.rs

@@ -9,7 +9,7 @@ use tracing_subscriber::filter::EnvFilter;
 use viceroy_lib::{
     body::Body,
     config::{Backend, Backends, Dictionaries, FastlyConfig, ObjectStore},
-    ExecuteCtx, ViceroyService,
+    ExecuteCtx, ProfilingStrategy, ViceroyService,
 };
 
 /// A shorthand for the path to our test fixtures' build artifacts for Rust tests.
@@ -183,7 +183,7 @@ impl Test {
             .try_init()
             .ok();
 
-        let ctx = ExecuteCtx::new(&self.module_path)
+        let ctx = ExecuteCtx::new(&self.module_path, ProfilingStrategy::None)
             .expect("failed to set up execution context")
             .with_backends(self.backends.clone())
             .with_dictionaries(self.dictionaries.clone())

lib/src/error.rs

@@ -20,6 +20,9 @@ pub enum Error {
     #[error("Expected a valid Wasm file")]
     FileFormat,
 
+    #[error("Expected a valid wastime's profiling strategy")]
+    ProfilingStrategy,
+
     #[error(transparent)]
     FastlyConfig(#[from] FastlyConfigError),
 
@@ -155,6 +158,7 @@ impl Error {
             | Error::IoError(_)
             | Error::NotAvailable(_)
             | Error::Other(_)
+            | Error::ProfilingStrategy
             | Error::StreamingChunkSend
             | Error::UnknownBackend(_)
             | Error::Utf8Expected(_)

lib/src/execute.rs

@@ -23,7 +23,7 @@ use {
     },
     tokio::sync::oneshot::{self, Sender},
     tracing::{event, info, info_span, Instrument, Level},
-    wasmtime::{Engine, InstancePre, Linker, Module},
+    wasmtime::{Engine, InstancePre, Linker, Module, ProfilingStrategy},
 };
 
 /// Execution context used by a [`ViceroyService`](struct.ViceroyService.html).
@@ -57,8 +57,12 @@ pub struct ExecuteCtx {
 
 impl ExecuteCtx {
     /// Create a new execution context, given the path to a module.
-    pub fn new(module_path: impl AsRef<Path>) -> Result<Self, Error> {
-        let engine = Engine::new(&configure_wasmtime())?;
+    pub fn new(
+        module_path: impl AsRef<Path>,
+        profiling_strategy: ProfilingStrategy,
+    ) -> Result<Self, Error> {
+        let config = &configure_wasmtime(profiling_strategy)?;
+        let engine = Engine::new(config)?;
         let mut linker = Linker::new(&engine);
         link_host_functions(&mut linker)?;
         let module = Module::from_file(&engine, module_path)?;
@@ -307,7 +311,7 @@ impl ExecuteCtx {
     }
 }
 
-fn configure_wasmtime() -> wasmtime::Config {
+fn configure_wasmtime(profiling_strategy: ProfilingStrategy) -> Result<wasmtime::Config, Error> {
     use wasmtime::{
         Config, InstanceAllocationStrategy, InstanceLimits, ModuleLimits,
         PoolingAllocationStrategy, WasmBacktraceDetails,
@@ -318,6 +322,7 @@ fn configure_wasmtime() -> wasmtime::Config {
     config.wasm_backtrace_details(WasmBacktraceDetails::Enable);
     config.async_support(true);
     config.consume_fuel(true);
+    config.profiler(profiling_strategy)?;
 
     let module_limits = ModuleLimits {
         // allow for up to 128MiB of linear memory
@@ -340,5 +345,5 @@ fn configure_wasmtime() -> wasmtime::Config {
         instance_limits: InstanceLimits::default(),
     });
 
-    config
+    Ok(config)
 }

lib/src/lib.rs

@@ -29,4 +29,5 @@ mod streaming_body;
 mod upstream;
 mod wiggle_abi;
 
+pub use wasmtime::ProfilingStrategy;
 pub use {error::Error, execute::ExecuteCtx, service::ViceroyService, upstream::BackendConnector};