From 560029410209f9b6ff3f7328120f3faab9f7028e Mon Sep 17 00:00:00 2001 From: Paul Osborne Date: Mon, 23 Mar 2026 14:23:15 -0500 Subject: [PATCH] Support guest profiling for components This change follows with the work I did in upstream wasmtime to support guest profiling for components: https://github.com/bytecodealliance/wasmtime/pull/10507 Integration tests for both components and core modules were added to ensure we have some coverage of the guest profiling (just ensuring it generates profiler output). --- CHANGELOG.md | 1 + Cargo.lock | 1 + cli/Cargo.toml | 2 + cli/tests/integration/common.rs | 16 ++++- cli/tests/integration/main.rs | 1 + cli/tests/integration/profiling.rs | 52 ++++++++++++++ src/execute.rs | 88 +++++++++++++++-------- wasm_abi/adapter/src/fastly/http_cache.rs | 5 +- 8 files changed, 131 insertions(+), 35 deletions(-) create mode 100644 cli/tests/integration/profiling.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 5462960c..571fae00 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ - Update to Rust 2024 Edition. ([#588](https://github.com/fastly/Viceroy/pull/588)) - Add stub implementations for bot detection hostcalls. - Add no-op implementations for stale-if-error hostcalls ([#591](https://github.com/fastly/Viceroy/pull/591)) +- Guest profiling support added for components. ([#593](https://github.com/fastly/Viceroy/pull/593)) ## 0.16.4 (2026-01-26) diff --git a/Cargo.lock b/Cargo.lock index 69da00e7..fe9c4bfe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2849,6 +2849,7 @@ dependencies = [ "rustls", "rustls-pemfile", "serde_json", + "tempfile", "tls-listener", "tokio", "tokio-rustls", diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 952b372f..83c3650c 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -54,3 +54,5 @@ anyhow = { workspace = true } futures = { workspace = true } url = { workspace = true } tls-listener = { version = "^0.7.0", features = ["rustls", "hyper-h1", "tokio-net", "rt"] } +tempfile = "3" +serde_json = { workspace = true } diff --git a/cli/tests/integration/common.rs b/cli/tests/integration/common.rs index b827aba7..9a52b692 100644 --- a/cli/tests/integration/common.rs +++ b/cli/tests/integration/common.rs @@ -92,6 +92,8 @@ pub struct Test { via_hyper: bool, unknown_import_behavior: UnknownImportBehavior, adapt_component: bool, + profiling_strategy: ProfilingStrategy, + guest_profile_config: Option, } impl Test { @@ -116,6 +118,8 @@ impl Test { via_hyper: false, unknown_import_behavior: Default::default(), adapt_component: false, + profiling_strategy: ProfilingStrategy::None, + guest_profile_config: None, } } @@ -140,6 +144,8 @@ impl Test { via_hyper: false, unknown_import_behavior: Default::default(), adapt_component: false, + profiling_strategy: ProfilingStrategy::None, + guest_profile_config: None, } } @@ -287,6 +293,12 @@ impl Test { self } + /// Enable guest profiling with the specified configuration. + pub fn with_guest_profiling(mut self, config: viceroy_lib::GuestProfileConfig) -> Self { + self.guest_profile_config = Some(config); + self + } + /// Pass the given requests through this test, returning the associated responses. /// /// A `Test` can be used repeatedly against different requests, either individually (as with @@ -331,9 +343,9 @@ impl Test { let ctx = ExecuteCtx::build( &self.module_path, - ProfilingStrategy::None, + self.profiling_strategy.clone(), HashSet::new(), - None, + self.guest_profile_config.clone(), self.unknown_import_behavior, self.adapt_component, )? diff --git a/cli/tests/integration/main.rs b/cli/tests/integration/main.rs index 168c66e6..80f57c4f 100644 --- a/cli/tests/integration/main.rs +++ b/cli/tests/integration/main.rs @@ -20,6 +20,7 @@ mod invalid_status_code; mod kv_store; mod logging; mod memory; +mod profiling; mod request; mod response; mod reusable_sessions; diff --git a/cli/tests/integration/profiling.rs b/cli/tests/integration/profiling.rs new file mode 100644 index 00000000..c4fdf436 --- /dev/null +++ b/cli/tests/integration/profiling.rs @@ -0,0 +1,52 @@ +//! Tests for guest profiling functionality. + +use { + crate::{ + common::{Test, TestResult}, + viceroy_test, + }, + hyper::StatusCode, + std::time::Duration, + viceroy_lib::GuestProfileConfig, +}; + +viceroy_test!(guest_profiling_works, |is_component| { + // Create a temporary directory for the profile output + let temp_dir = tempfile::tempdir()?; + let profile_dir = temp_dir.path().join("profiles"); + std::fs::create_dir(&profile_dir)?; + + let resp = Test::using_fixture("noop.wasm") + .adapt_component(is_component) + .with_guest_profiling(GuestProfileConfig { + path: profile_dir.clone(), + sample_period: Duration::from_micros(50), + }) + .against_empty() + .await?; + + // Verify the request succeeded + assert_eq!(resp.status(), StatusCode::OK); + + // Verify that a profile file was created in the directory + let profile_files: Vec<_> = std::fs::read_dir(&profile_dir)? + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().and_then(|s| s.to_str()) == Some("json")) + .collect(); + + assert!( + !profile_files.is_empty(), + "At least one profile JSON file should be created in {:?}", + profile_dir + ); + + // Verify the first profile file has content and is valid JSON + let first_profile = &profile_files[0]; + let metadata = std::fs::metadata(first_profile.path())?; + assert!(metadata.len() > 0, "Profile file should not be empty"); + + let profile_content = std::fs::read_to_string(first_profile.path())?; + let _: serde_json::Value = serde_json::from_str(&profile_content)?; + + Ok(()) +}); diff --git a/src/execute.rs b/src/execute.rs index 9167fd6c..d239d9d4 100644 --- a/src/execute.rs +++ b/src/execute.rs @@ -61,14 +61,17 @@ const REGION_NONE: &str = "none"; enum Instance { Module(Module, InstancePre), - Component(compute::bindings::AdapterServicePre), + Component( + Component, + compute::bindings::AdapterServicePre, + ), } impl Instance { fn unwrap_module(&self) -> (&Module, &InstancePre) { match self { Instance::Module(m, i) => (m, i), - Instance::Component(_) => panic!("unwrap_module called on a component"), + Instance::Component(_, _) => panic!("unwrap_module called on a component"), } } } @@ -245,7 +248,10 @@ impl ExecuteCtx { } let instance_pre = linker.instantiate_pre(&component)?; - Instance::Component(compute::bindings::AdapterServicePre::new(instance_pre)?) + Instance::Component( + component, + compute::bindings::AdapterServicePre::new(instance_pre)?, + ) } else { let mut linker = Linker::new(&engine); link_host_functions(&mut linker, &wasi_modules)?; @@ -625,15 +631,21 @@ impl ExecuteCtx { }); match self.instance_pre.as_ref() { - Instance::Component(instance_pre) => { - if self.guest_profile_config.is_some() { - warn!("Components do not currently support the guest profiler"); - } + Instance::Component(component, instance_pre) => { + let profiler = self.guest_profile_config.as_deref().map(|pcfg| { + let program_name = "main"; + GuestProfiler::new_component( + program_name, + pcfg.sample_period, + component.clone(), + std::iter::empty(), + ) + }); let req = session.downstream_request(); let body = session.downstream_request_body(); - let mut store = ComponentCtx::create_store(&self, session, None, |ctx| { + let mut store = ComponentCtx::create_store(&self, session, profiler, |ctx| { ctx.arg("compute-app"); }) .map_err(ExecutionError::Context)?; @@ -671,6 +683,9 @@ impl ExecuteCtx { } }; + // If we collected a profile, write it to the file + write_profile_component(&mut store, guest_profile_path.as_ref()); + // Ensure the downstream response channel is closed, whether or not a response was // sent during execution. let resp = outcome @@ -793,7 +808,7 @@ impl ExecuteCtx { let session = Session::new(downstream, active_cpu_time_us.clone(), self.clone()); - if let Instance::Component(_) = self.instance_pre.as_ref() { + if let Instance::Component(_, _) = self.instance_pre.as_ref() { panic!("components not currently supported with `run`"); } @@ -884,7 +899,7 @@ impl ExecuteCtx { } pub fn is_component(&self) -> bool { - matches!(self.instance_pre.as_ref(), Instance::Component(_)) + matches!(self.instance_pre.as_ref(), Instance::Component(_, _)) } } @@ -976,29 +991,44 @@ impl ExecuteCtxBuilder { } } +fn write_profile_to_file(profile: Box, path: &PathBuf) { + match std::fs::File::create(path) + .map_err(anyhow::Error::new) + .and_then(|output| profile.finish(std::io::BufWriter::new(output))) + { + Err(e) => { + event!( + Level::ERROR, + "failed writing profile at {}: {e:#}", + path.display() + ); + } + _ => { + event!( + Level::INFO, + "\nProfile written to: {}\nView this profile at https://profiler.firefox.com/.", + path.display() + ); + } + } +} + fn write_profile(store: &mut wasmtime::Store, guest_profile_path: Option<&PathBuf>) { if let (Some(profile), Some(path)) = (store.data_mut().take_guest_profiler(), guest_profile_path) { - match std::fs::File::create(path) - .map_err(anyhow::Error::new) - .and_then(|output| profile.finish(std::io::BufWriter::new(output))) - { - Err(e) => { - event!( - Level::ERROR, - "failed writing profile at {}: {e:#}", - path.display() - ); - } - _ => { - event!( - Level::INFO, - "\nProfile written to: {}\nView this profile at https://profiler.firefox.com/.", - path.display() - ); - } - } + write_profile_to_file(profile, path); + } +} + +fn write_profile_component( + store: &mut wasmtime::Store, + guest_profile_path: Option<&PathBuf>, +) { + if let (Some(profile), Some(path)) = + (store.data_mut().take_guest_profiler(), guest_profile_path) + { + write_profile_to_file(profile, path); } } diff --git a/wasm_abi/adapter/src/fastly/http_cache.rs b/wasm_abi/adapter/src/fastly/http_cache.rs index 201f64f9..c839735d 100644 --- a/wasm_abi/adapter/src/fastly/http_cache.rs +++ b/wasm_abi/adapter/src/fastly/http_cache.rs @@ -147,10 +147,7 @@ mod http_cache { STALE_WHILE_REVALIDATE_NS, (*options).stale_while_revalidate_ns ), - stale_if_error_ns: when_enabled!( - STALE_IF_ERROR_NS, - (*options).stale_if_error_ns - ), + stale_if_error_ns: when_enabled!(STALE_IF_ERROR_NS, (*options).stale_if_error_ns), surrogate_keys, length: when_enabled!(LENGTH, (*options).length), sensitive_data: mask.contains(HttpCacheWriteOptionsMask::SENSITIVE_DATA),