Allow to enable wasmtime's profiling support

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())

cli/tests/trap-test/src/main.rs

@@ -2,7 +2,7 @@ use {
     crate::common::{TestResult, RUST_FIXTURE_PATH},
     http::header::WARNING,
     hyper::{Body, Request, StatusCode},
-    viceroy_lib::ExecuteCtx,
+    viceroy_lib::{ExecuteCtx, ProfilingStrategy},
 };
 
 #[path = "../../integration/common.rs"]
@@ -11,7 +11,7 @@ mod common;
 #[tokio::test(flavor = "multi_thread")]
 async fn fatal_error_traps() -> TestResult {
     let module_path = format!("../../{}/response.wasm", RUST_FIXTURE_PATH);
-    let ctx = ExecuteCtx::new(module_path)?;
+    let ctx = ExecuteCtx::new(module_path, ProfilingStrategy::None)?;
     let req = Request::get("http://127.0.0.1:7878/").body(Body::from(""))?;
     let resp = ctx
         .handle_request(req, "127.0.0.1".parse().unwrap())

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)?;
@@ -167,10 +171,10 @@ impl ExecuteCtx {
     ///
     /// ```no_run
     /// # use hyper::{Body, http::Request};
-    /// # use viceroy_lib::{Error, ExecuteCtx, ViceroyService};
+    /// # use viceroy_lib::{Error, ExecuteCtx, ProfilingStrategy, ViceroyService};
     /// # async fn f() -> Result<(), Error> {
     /// # let req = Request::new(Body::from(""));
-    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm")?;
+    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None)?;
     /// let resp = ctx.handle_request(req, "127.0.0.1".parse().unwrap()).await?;
     /// # Ok(())
     /// # }
@@ -307,7 +311,7 @@ impl ExecuteCtx {
     }
 }
 
-fn configure_wasmtime() -> wasmtime::Config {
+fn configure_wasmtime(profiling_strategy: ProfilingStrategy) -> wasmtime::Config {
     use wasmtime::{
         Config, InstanceAllocationStrategy, InstanceLimits, 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);
 
     const MB: usize = 1 << 20;
 

lib/src/lib.rs

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

lib/src/service.rs

@@ -40,9 +40,9 @@ impl ViceroyService {
     /// # Example
     ///
     /// ```no_run
-    /// # use viceroy_lib::{Error, ExecuteCtx, ViceroyService};
+    /// # use viceroy_lib::{Error, ExecuteCtx, ProfilingStrategy, ViceroyService};
     /// # fn f() -> Result<(), Error> {
-    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm")?;
+    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None)?;
     /// let svc = ViceroyService::new(ctx);
     /// # Ok(())
     /// # }