feat: support heap profiling

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

Signed-off-by: Flavio Castelli <[email protected]>

Cargo.toml

@@ -48,8 +48,11 @@ 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"
+tikv-jemalloc-ctl = "0.5.4"
 
 [dev-dependencies]
 mockall = "0.12"

PROFILING.md

@@ -0,0 +1,56 @@
+# 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
+
+By default memory profiling is not active. This gets automatically activated when the first
+request is sent to the `/debug/pprof/heap` endpoint.
+
+When the memory profiling is active, the CPU usage of the Policy Server will increase.
+Finally, once the memory profiling is activated, there's no way to deactivate it.
+
+Note: the memory profile requires some sampling, since we start with the memory profiling disabled,
+the first memory profile will be empty. The second one will contain the memory usage of the Policy Server.
+
+To obtain a memory profile:
+
+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.

src/api/handlers.rs

@@ -12,7 +12,7 @@ use policy_evaluator::{
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use tokio::task;
-use tracing::{debug, error, Span};
+use tracing::{debug, error, info, Span};
 
 use crate::{
     api::{
@@ -173,6 +173,10 @@ pub(crate) async fn validate_raw_handler(
     Ok(Json(RawReviewResponse::new(response)))
 }
 
+pub(crate) async fn readiness_handler() -> StatusCode {
+    StatusCode::OK
+}
+
 #[derive(Deserialize)]
 pub(crate) struct ProfileParams {
     /// profiling frequency (Hz)
@@ -217,8 +221,43 @@ pub(crate) async fn pprof_get_cpu(
 
     Ok((headers, body))
 }
-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_else(|| handle_pprof_error(ReportGenerationError::CannotGetJemallocControlHandle))?
+        .lock()
+        .await;
+
+    if prof_ctl.activated() {
+        info!("Activating jemalloc profiling");
+        prof_ctl.activate().map_err(|e| {
+            handle_pprof_error(ReportGenerationError::HeapProfilingActivationError(e))
+        })?;
+    }
+
+    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))
 }
 
 async fn acquire_semaphore_and_evaluate(

src/lib.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/main.rs

@@ -17,7 +17,9 @@ 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 but do not activate them by
+/// default. When pprof is activate there's a CPU overhead.
+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:false,lg_prof_sample:19\0";
 
 #[tokio::main]
 async fn main() -> Result<()> {

src/profiling.rs

@@ -29,6 +29,14 @@ pub enum ReportGenerationError {
 
     #[error("cannot encode report to pprof format: {0}")]
     EncodeError(String),
+
+    #[error("failed to dump the profile: {0}")]
+    JemallocError(String),
+
+    #[error("error enabling jemalloc profiling: {0}")]
+    HeapProfilingActivationError(#[from] tikv_jemalloc_ctl::Error),
+    #[error("cannot get jemalloc control handle")]
+    CannotGetJemallocControlHandle,
 }
 
 /// Default frequency of sampling. 99Hz to avoid coincide with special periods