feat: support heap profiling

Support the generation of pprof profiles of heap usage by using jemalloc
facilities.

Signed-off-by: Flavio Castelli <fcastelli@suse.com>

Author: flavio

Cargo.lock

@@ -48,8 +48,10 @@ axum = { version = "0.7.5", features = ["macros", "query"] }
 axum-server = { version = "0.6", features = ["tls-rustls"] }
 tower-http = { version = "0.5.2", features = ["trace"] }
 tikv-jemallocator = { version = "0.5.4", features = [
+  "profiling",
   "unprefixed_malloc_on_supported_platforms",
 ] }
+jemalloc_pprof = "0.1.0"
 
 [dev-dependencies]
 mockall = "0.12"

Cargo.toml

@@ -0,0 +1,45 @@
+# Profiling
+
+Ensure the `policy-server` is running with the `--enable-pprof` flag. This can also be set with the `KUBEWARDEN_ENABLE_PPROF=1` environment variable.
+
+## CPU
+
+1. Port-forward policy-server port (8443):
+
+```sh
+kubectl port-forward -n kubewarden service/policy-server-default 8443:8443
+```
+
+1. Download the CPU profile using `curl`:
+
+```sh
+curl --insecure https://localhost:8443/debug/pprof/cpu -o cpu_profile.prof
+```
+
+1. Use [`pprof`](https://github.com/google/pprof) to analyze the profile:
+
+```sh
+pprof -http=:8080 cpu_profile.prof
+```
+
+## Memory
+
+1. Port-forward policy-server port (8443):
+
+```sh
+kubectl port-forward -n kubewarden service/policy-server-default 8443:8443
+```
+
+1. Download the CPU profile using `curl`:
+
+```sh
+curl --insecure https://localhost:8443/debug/pprof/heap -o heap_profile.prof
+```
+
+1. Use [`pprof`](https://github.com/google/pprof) to analyze the profile:
+
+```sh
+pprof -http=:8080 heap_profile.prof
+```
+
+**Warning:** do not use the `pprof` tool from the `go tool` binary, it cannot handle the memory profile generated by Policy Server.

PROFILING.md

@@ -221,6 +221,50 @@ pub(crate) async fn readiness_handler() -> StatusCode {
     StatusCode::OK
 }
 
+// Generate a pprof heap profile using google's pprof format
+// The report is generated and sent to the user as binary data
+pub(crate) async fn pprof_get_heap(
+) -> Result<impl axum::response::IntoResponse, (StatusCode, ApiError)> {
+    let mut prof_ctl = jemalloc_pprof::PROF_CTL
+        .as_ref()
+        .ok_or(handle_pprof_error(
+            ReportGenerationError::HeapProfilingNotEnabled,
+        ))?
+        .lock()
+        .await;
+    require_profiling_activated(&prof_ctl).map_err(handle_pprof_error)?;
+    let pprof = prof_ctl
+        .dump_pprof()
+        .map_err(|e| handle_pprof_error(ReportGenerationError::JemallocError(e.to_string())))?;
+
+    let mut headers = header::HeaderMap::new();
+    headers.insert(
+        header::CONTENT_DISPOSITION,
+        r#"attachment; filename="heap_profile"#.parse().unwrap(),
+    );
+    headers.insert(
+        header::CONTENT_LENGTH,
+        pprof.len().to_string().parse().unwrap(),
+    );
+    headers.insert(
+        header::CONTENT_TYPE,
+        mime::APPLICATION_OCTET_STREAM.to_string().parse().unwrap(),
+    );
+
+    Ok((headers, pprof))
+}
+
+/// Checks whether jemalloc profiling is activated an returns an error response if not.
+fn require_profiling_activated(
+    prof_ctl: &jemalloc_pprof::JemallocProfCtl,
+) -> Result<(), ReportGenerationError> {
+    if prof_ctl.activated() {
+        Ok(())
+    } else {
+        Err(ReportGenerationError::HeapProfilingNotEnabled)
+    }
+}
+
 async fn acquire_semaphore_and_evaluate(
     state: Arc<ApiServerState>,
     policy_id: String,

src/api/handlers.rs

@@ -38,7 +38,8 @@ use tokio::{
 use tower_http::trace::{self, TraceLayer};
 
 use crate::api::handlers::{
-    audit_handler, pprof_get_cpu, readiness_handler, validate_handler, validate_raw_handler,
+    audit_handler, pprof_get_cpu, pprof_get_heap, readiness_handler, validate_handler,
+    validate_raw_handler,
 };
 use crate::api::state::ApiServerState;
 use crate::evaluation::{
@@ -211,7 +212,9 @@ impl PolicyServer {
                     .on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
             );
         if config.enable_pprof {
-            let pprof_router = Router::new().route("/debug/pprof/cpu", get(pprof_get_cpu));
+            let pprof_router = Router::new()
+                .route("/debug/pprof/cpu", get(pprof_get_cpu))
+                .route("/debug/pprof/heap", get(pprof_get_heap));
             router = Router::new().merge(router).merge(pprof_router);
         }
 

src/lib.rs

@@ -17,7 +17,8 @@ static GLOBAL: Jemalloc = Jemalloc;
 
 #[allow(non_upper_case_globals)]
 #[export_name = "malloc_conf"]
-pub static malloc_conf: &[u8] = b"background_thread:true,tcache_max:4096,dirty_decay_ms:5000,muzzy_decay_ms:5000,abort_conf:true\0";
+/// Prioritize memory usage, then enable features request by pprof
+pub static malloc_conf: &[u8] = b"background_thread:true,tcache_max:4096,dirty_decay_ms:5000,muzzy_decay_ms:5000,abort_conf:true,prof:true,prof_active:true,lg_prof_sample:19\0";
 
 #[tokio::main]
 async fn main() -> Result<()> {

src/main.rs

@@ -29,6 +29,12 @@ pub enum ReportGenerationError {
 
     #[error("cannot encode report to pprof format: {0}")]
     EncodeError(String),
+
+    #[error("failed to dump the profile: {0}")]
+    JemallocError(String),
+
+    #[error("heap profiling is not activated")]
+    HeapProfilingNotEnabled,
 }
 
 /// Default frequency of sampling. 99Hz to avoid coincide with special periods