Enable custom profiling interval via CLI flag (#357)

cli/Cargo.toml

@@ -50,6 +50,7 @@ wat = "^1.0.38"
 wasmtime = { workspace = true }
 wasmtime-wasi = { workspace = true }
 libc = "^0.2.139"
+humantime = "2.1.0"
 
 [dev-dependencies]
 anyhow = { workspace = true }

cli/src/main.rs

@@ -40,8 +40,13 @@ use {
 /// Create a new server, bind it to an address, and serve responses until an error occurs.
 pub async fn serve(serve_args: ServeArgs) -> Result<(), Error> {
     // Load the wasm module into an execution context
-    let ctx =
-        create_execution_context(serve_args.shared(), true, serve_args.profile_guest()).await?;
+    let ctx = create_execution_context(
+        serve_args.shared(),
+        true,
+        serve_args.profile_guest(),
+        serve_args.profile_guest_interval(),
+    )
+    .await?;
 
     if let Some(guest_profile_path) = serve_args.profile_guest() {
         std::fs::create_dir_all(guest_profile_path)?;
@@ -161,7 +166,13 @@ pub async fn main() -> ExitCode {
 /// Execute a Wasm program in the Viceroy environment.
 pub async fn run_wasm_main(run_args: RunArgs) -> Result<(), anyhow::Error> {
     // Load the wasm module into an execution context
-    let ctx = create_execution_context(run_args.shared(), false, run_args.profile_guest()).await?;
+    let ctx = create_execution_context(
+        run_args.shared(),
+        false,
+        run_args.profile_guest(),
+        run_args.profile_guest_interval(),
+    )
+    .await?;
     let input = run_args.shared().input();
     let program_name = match input.file_stem() {
         Some(stem) => stem.to_string_lossy(),
@@ -297,13 +308,15 @@ async fn create_execution_context(
     args: &SharedArgs,
     check_backends: bool,
     guest_profile_path: Option<PathBuf>,
+    guest_profile_interval: Option<Duration>,
 ) -> Result<ExecuteCtx, anyhow::Error> {
     let input = args.input();
     let mut ctx = ExecuteCtx::new(
         input,
         args.profiling_strategy(),
         args.wasi_modules(),
         guest_profile_path,
+        guest_profile_interval,
         args.unknown_import_behavior(),
         args.adapt(),
     )?

cli/src/opts.rs

@@ -1,5 +1,7 @@
 //! Command line arguments.
 
+use std::{str::FromStr, time::Duration};
+
 use viceroy_lib::config::UnknownImportBehavior;
 
 use {
@@ -92,9 +94,10 @@ pub struct SharedArgs {
     ///
     /// The `guest` option can be additionally configured as:
     ///
-    ///     --profile=guest[,path]
+    ///     --profile=guest[,path[,interval]]
     ///
-    /// where `path` is the directory or filename to write the profile(s) to.
+    /// where `path` is the directory or filename to write the profile(s) to
+    /// and `interval` is the duration in milliseconds between samples.
     #[arg(long = "profile", value_name = "STRATEGY", value_parser = check_wasmtime_profiler_mode)]
     profile: Option<Profile>,
     /// Set of experimental WASI modules to link against.
@@ -120,7 +123,10 @@ pub struct SharedArgs {
 #[derive(Debug, Clone)]
 enum Profile {
     Native(ProfilingStrategy),
-    Guest { path: Option<String> },
+    Guest {
+        path: Option<String>,
+        interval: Option<Duration>,
+    },
 }
 
 impl ServeArgs {
@@ -132,7 +138,7 @@ impl ServeArgs {
 
     /// The path to write guest profiles to
     pub fn profile_guest(&self) -> Option<PathBuf> {
-        if let Some(Profile::Guest { path }) = &self.shared.profile {
+        if let Some(Profile::Guest { path, .. }) = &self.shared.profile {
             Some(
                 path.clone()
                     .unwrap_or_else(|| "guest-profiles".to_string())
@@ -143,6 +149,15 @@ impl ServeArgs {
         }
     }
 
+    /// The interval for guest profiling
+    pub fn profile_guest_interval(&self) -> Option<Duration> {
+        if let Some(Profile::Guest { interval, .. }) = &self.shared.profile {
+            *interval
+        } else {
+            None
+        }
+    }
+
     pub fn shared(&self) -> &SharedArgs {
         &self.shared
     }
@@ -160,7 +175,7 @@ impl RunArgs {
 
     /// The path to write a guest profile to
     pub fn profile_guest(&self) -> Option<PathBuf> {
-        if let Some(Profile::Guest { path }) = &self.shared.profile {
+        if let Some(Profile::Guest { path, .. }) = &self.shared.profile {
             Some(
                 path.clone()
                     .unwrap_or_else(|| "guest-profile.json".to_string())
@@ -170,6 +185,15 @@ impl RunArgs {
             None
         }
     }
+
+    /// The interval for guest profiling
+    pub fn profile_guest_interval(&self) -> Option<Duration> {
+        if let Some(Profile::Guest { interval, .. }) = &self.shared.profile {
+            *interval
+        } else {
+            None
+        }
+    }
 }
 
 impl SharedArgs {
@@ -323,14 +347,37 @@ fn check_wasmtime_profiler_mode(s: &str) -> Result<Profile, Error> {
         ["jitdump"] => Ok(Profile::Native(ProfilingStrategy::JitDump)),
         ["perfmap"] => Ok(Profile::Native(ProfilingStrategy::PerfMap)),
         ["vtune"] => Ok(Profile::Native(ProfilingStrategy::VTune)),
-        ["guest"] => Ok(Profile::Guest { path: None }),
+        ["guest"] => Ok(Profile::Guest {
+            path: None,
+            interval: None,
+        }),
         ["guest", path] => Ok(Profile::Guest {
             path: Some(path.to_string()),
+            interval: None,
         }),
+        ["guest", path, interval] => {
+            let interval_duration = parse_duration(interval)?;
+            Ok(Profile::Guest {
+                path: Some(path.to_string()),
+                interval: Some(interval_duration),
+            })
+        }
         _ => Err(Error::ProfilingStrategy),
     }
 }
 
+/// Parse duration from either a string slice as with a duration unit (e.g. 100ms)
+/// or without (e.g. 100) and default to milliseconds.
+fn parse_duration(dur: &str) -> Result<Duration, Error> {
+    if let Ok(duration) = humantime::Duration::from_str(dur) {
+        return Ok(duration.into());
+    }
+    if let Ok(millis) = dur.parse::<u64>() {
+        return Ok(std::time::Duration::from_millis(millis));
+    }
+    Err(Error::ProfilingStrategy)
+}
+
 /// A collection of unit tests for our CLI argument parsing.
 ///
 /// Note: When using [`Clap::try_parse_from`][from] to test how command line arguments are

cli/tests/integration/common.rs

@@ -325,6 +325,7 @@ impl Test {
             ProfilingStrategy::None,
             HashSet::new(),
             None,
+            None,
             self.unknown_import_behavior,
             self.adapt_component,
         )?

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

@@ -20,6 +20,7 @@ async fn fatal_error_traps_impl(adapt_core_wasm: bool) -> TestResult {
         ProfilingStrategy::None,
         HashSet::new(),
         None,
+        None,
         viceroy_lib::config::UnknownImportBehavior::LinkError,
         adapt_core_wasm,
     )?;

lib/src/execute.rs

@@ -103,6 +103,8 @@ pub struct ExecuteCtx {
     /// this must refer to a directory, while in run mode it names
     /// a file.
     guest_profile_path: Arc<Option<PathBuf>>,
+    /// Duration for the interval between samples.
+    guest_profile_interval: Arc<Option<Duration>>,
 }
 
 impl ExecuteCtx {
@@ -112,6 +114,7 @@ impl ExecuteCtx {
         profiling_strategy: ProfilingStrategy,
         wasi_modules: HashSet<ExperimentalModule>,
         guest_profile_path: Option<PathBuf>,
+        guest_profile_interval: Option<Duration>,
         unknown_import_behavior: UnknownImportBehavior,
         adapt_components: bool,
     ) -> Result<Self, Error> {
@@ -223,6 +226,7 @@ impl ExecuteCtx {
             epoch_increment_thread,
             epoch_increment_stop,
             guest_profile_path: Arc::new(guest_profile_path),
+            guest_profile_interval: Arc::new(guest_profile_interval),
         })
     }
 
@@ -352,7 +356,7 @@ impl ExecuteCtx {
     /// # async fn f() -> Result<(), Error> {
     /// # let req = Request::new(Body::from(""));
     /// let adapt_core_wasm = false;
-    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None, HashSet::new(), None, Default::default(), adapt_core_wasm)?;
+    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None, HashSet::new(), None, None, Default::default(), adapt_core_wasm)?;
     /// let local = "127.0.0.1:80".parse().unwrap();
     /// let remote = "127.0.0.1:0".parse().unwrap();
     /// let resp = ctx.handle_request(req, local, remote).await?;
@@ -536,9 +540,13 @@ impl ExecuteCtx {
             Instance::Module(module, instance_pre) => {
                 let profiler = self.guest_profile_path.is_some().then(|| {
                     let program_name = "main";
+                    let interval = self
+                        .guest_profile_interval
+                        .as_ref()
+                        .unwrap_or(EPOCH_INTERRUPTION_PERIOD);
                     GuestProfiler::new(
                         program_name,
-                        EPOCH_INTERRUPTION_PERIOD,
+                        interval,
                         vec![(program_name.to_string(), module.clone())],
                     )
                 });
@@ -636,9 +644,13 @@ impl ExecuteCtx {
         let (module, instance_pre) = self.instance_pre.unwrap_module();
 
         let profiler = self.guest_profile_path.is_some().then(|| {
+            let interval = self
+                .guest_profile_interval
+                .as_ref()
+                .unwrap_or(EPOCH_INTERRUPTION_PERIOD);
             GuestProfiler::new(
                 program_name,
-                EPOCH_INTERRUPTION_PERIOD,
+                interval,
                 vec![(program_name.to_string(), module.clone())],
             )
         });

lib/src/service.rs

@@ -44,7 +44,7 @@ impl ViceroyService {
     /// use viceroy_lib::{Error, ExecuteCtx, ProfilingStrategy, ViceroyService};
     /// # fn f() -> Result<(), Error> {
     /// let adapt_core_wasm = false;
-    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None, HashSet::new(), None, Default::default(), adapt_core_wasm)?;
+    /// let ctx = ExecuteCtx::new("path/to/a/file.wasm", ProfilingStrategy::None, HashSet::new(), None, None, Default::default(), adapt_core_wasm)?;
     /// let svc = ViceroyService::new(ctx);
     /// # Ok(())
     /// # }