feat: introduce pprof debugging endpoint

Allow CPU profiling using pprof. There's a new boolean flag
`--enable-pprof` that exposes a new HTTP endpoint named
`/debug/pprof/cpu`.

By doing a GET request against the endpoint, a pprof profile is
generated and downloaded.

The endpoint takes two optional query parameters:

- interval: the length in seconds of the profile (30 secs by default)
- frequency: profiling frequency (99 Hz by default)

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

Author: flavio

Cargo.lock

@@ -14,12 +14,14 @@ edition = "2021"
 anyhow = "1.0"
 clap = { version = "4.5", features = ["cargo", "env"] }
 daemonize = "0.5"
+futures = "0.3"
 humansize = "2.1"
 itertools = "0.12.1"
 k8s-openapi = { version = "0.21.1", default-features = false, features = [
   "v1_29",
 ] }
 lazy_static = "1.4.0"
+mime = "0.3"
 num_cpus = "1.16.0"
 opentelemetry-otlp = { version = "0.14.0", features = ["metrics", "tonic"] }
 opentelemetry = { version = "0.21", default-features = false, features = [
@@ -28,8 +30,10 @@ opentelemetry = { version = "0.21", default-features = false, features = [
 ] }
 opentelemetry_sdk = { version = "0.21", features = ["rt-tokio"] }
 procfs = "0.16"
+pprof = { version = "0.13", features = ["prost-codec"] }
 policy-evaluator = { git = "https://github.com/kubewarden/policy-evaluator", tag = "v0.16.1" }
 rayon = "1.9"
+regex = "1.10"
 serde_json = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_yaml = "0.9.32"
@@ -42,7 +46,7 @@ tracing-opentelemetry = "0.22.0"
 tracing-subscriber = { version = "0.3", features = ["ansi", "fmt", "json"] }
 semver = { version = "1.0.22", features = ["serde"] }
 mockall_double = "0.3"
-axum = { version = "0.7.4", features = ["macros"] }
+axum = { version = "0.7.4", features = ["macros", "query"] }
 axum-server = { version = "0.6", features = ["tls-rustls"] }
 tower-http = { version = "0.5.2", features = ["trace"] }
 

Cargo.toml

@@ -1,26 +1,30 @@
 use axum::{
-    extract::{self, FromRequest},
-    http::StatusCode,
+    extract::{self, FromRequest, Query},
+    http::{header, StatusCode},
     response::IntoResponse,
     Json,
 };
 use policy_evaluator::{
     admission_request::AdmissionRequest, admission_response::AdmissionResponse,
     policy_evaluator::ValidateRequest,
 };
-use serde::Serialize;
+
+use serde::{Deserialize, Serialize};
 use std::sync::Arc;
 use tokio::task;
 use tracing::{debug, error, Span};
 
-use crate::api::{
-    admission_review::{AdmissionReviewRequest, AdmissionReviewResponse},
-    api_error::ApiError,
-    raw_review::{RawReviewRequest, RawReviewResponse},
-    service::{evaluate, RequestOrigin},
-    state::ApiServerState,
+use crate::{
+    api::{
+        admission_review::{AdmissionReviewRequest, AdmissionReviewResponse},
+        api_error::ApiError,
+        raw_review::{RawReviewRequest, RawReviewResponse},
+        service::{evaluate, RequestOrigin},
+        state::ApiServerState,
+    },
+    profiling,
 };
-use crate::evaluation::errors::EvaluationError;
+use crate::{evaluation::errors::EvaluationError, profiling::ReportGenerationError};
 
 // create an extractor that internally uses `axum::Json` but has a custom rejection
 #[derive(FromRequest)]
@@ -169,6 +173,50 @@ pub(crate) async fn validate_raw_handler(
     Ok(Json(RawReviewResponse::new(response)))
 }
 
+#[derive(Deserialize)]
+pub(crate) struct ProfileParams {
+    /// profiling frequency (Hz)
+    #[serde(default = "profiling::default_profiling_frequency")]
+    pub frequency: i32,
+
+    /// profiling time interval (seconds)
+    #[serde(default = "profiling::default_profiling_interval")]
+    pub interval: u64,
+}
+
+// Generate a pprof CPU profile using google's pprof format
+// The report is generated and sent to the user as binary data
+pub(crate) async fn pprof_get_cpu(
+    profiling_params: Query<ProfileParams>,
+) -> Result<impl axum::response::IntoResponse, (StatusCode, ApiError)> {
+    let frequency = profiling_params.frequency;
+    let interval = profiling_params.interval;
+
+    let end = async move {
+        tokio::time::sleep(tokio::time::Duration::from_secs(interval)).await;
+        Ok(())
+    };
+
+    let body = profiling::start_one_cpu_profile(end, frequency)
+        .await
+        .map_err(handle_pprof_error)?;
+
+    let mut headers = header::HeaderMap::new();
+    headers.insert(
+        header::CONTENT_DISPOSITION,
+        r#"attachment; filename="cpu_profile"#.parse().unwrap(),
+    );
+    headers.insert(
+        header::CONTENT_LENGTH,
+        body.len().to_string().parse().unwrap(),
+    );
+    headers.insert(
+        header::CONTENT_TYPE,
+        mime::APPLICATION_OCTET_STREAM.to_string().parse().unwrap(),
+    );
+
+    Ok((headers, body))
+}
 pub(crate) async fn readiness_handler() -> StatusCode {
     StatusCode::OK
 }
@@ -261,3 +309,15 @@ fn handle_evaluation_error(error: EvaluationError) -> (StatusCode, ApiError) {
         }
     }
 }
+
+fn handle_pprof_error(error: ReportGenerationError) -> (StatusCode, ApiError) {
+    error!("pprof error: {}", error);
+
+    (
+        StatusCode::INTERNAL_SERVER_ERROR,
+        ApiError {
+            status: StatusCode::INTERNAL_SERVER_ERROR,
+            message: "Something went wrong".to_owned(),
+        },
+    )
+}

src/api/handlers.rs

@@ -160,6 +160,11 @@ pub(crate) fn build_cli() -> Command {
                 .env("KUBEWARDEN_IGNORE_KUBERNETES_CONNECTION_FAILURE")
                 .action(ArgAction::SetTrue)
                 .help("Do not exit with an error if the Kubernetes connection fails. This will cause context aware policies to break when there's no connection with Kubernetes."),
+            Arg::new("enable-pprof")
+                .long("enable-pprof")
+                .env("KUBEWARDEN_ENABLE_PPROF")
+                .action(ArgAction::SetTrue)
+                .help("Enable pprof profiling"),
     ];
     args.sort_by(|a, b| a.get_id().cmp(b.get_id()));
 

src/cli.rs

@@ -41,6 +41,7 @@ pub struct Config {
     pub log_fmt: String,
     pub log_no_color: bool,
     pub daemon: bool,
+    pub enable_pprof: bool,
     pub daemon_pid_file: String,
     pub daemon_stdout_file: Option<String>,
     pub daemon_stderr_file: Option<String>,
@@ -131,6 +132,10 @@ impl Config {
                 key_file,
             })
         };
+        let enable_pprof = matches
+            .get_one::<bool>("enable-pprof")
+            .expect("clap should have assigned a default value")
+            .to_owned();
 
         Ok(Self {
             addr,
@@ -152,6 +157,7 @@ impl Config {
             daemon_pid_file,
             daemon_stdout_file,
             daemon_stderr_file,
+            enable_pprof,
         })
     }
 }

src/config.rs

@@ -7,6 +7,7 @@ mod test_utils;
 pub mod api;
 pub mod config;
 pub mod metrics;
+pub mod profiling;
 pub mod tracing;
 
 use ::tracing::{debug, error, info, Level};
@@ -29,7 +30,7 @@ use tokio::time;
 use tower_http::trace::{self, TraceLayer};
 
 use crate::api::handlers::{
-    audit_handler, readiness_handler, validate_handler, validate_raw_handler,
+    audit_handler, pprof_get_cpu, readiness_handler, validate_handler, validate_raw_handler,
 };
 use crate::api::state::ApiServerState;
 use crate::evaluation::{
@@ -176,7 +177,7 @@ impl PolicyServer {
             None
         };
 
-        let router = Router::new()
+        let mut router = Router::new()
             .route("/audit/:policy_id", post(audit_handler))
             .route("/validate/:policy_id", post(validate_handler))
             .route("/validate_raw/:policy_id", post(validate_raw_handler))
@@ -187,6 +188,10 @@ impl PolicyServer {
                     .make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
                     .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));
+            router = Router::new().merge(router).merge(pprof_router);
+        }
 
         Ok(Self {
             router,

src/lib.rs

@@ -0,0 +1,149 @@
+use futures::{
+    future::BoxFuture,
+    task::{Context, Poll},
+    Future, FutureExt,
+};
+use lazy_static::lazy_static;
+use pprof::protos::Message;
+use regex::Regex;
+use std::{pin::Pin, sync::Mutex};
+use thiserror::Error;
+
+lazy_static! {
+    // If it's some it means there are already a CPU profiling.
+    static ref CPU_PROFILE_ACTIVE: Mutex<Option<()>> = Mutex::new(None);
+
+    // To normalize thread names.
+    static ref THREAD_NAME_RE: Regex =
+        Regex::new(r"^(?P<thread_name>[a-z-_ :]+?)(-?\d)*$").unwrap();
+    static ref THREAD_NAME_REPLACE_SEPERATOR_RE: Regex = Regex::new(r"[_ ]").unwrap();
+}
+
+#[derive(Debug, Error)]
+pub enum ReportGenerationError {
+    #[error("CPU profile already running")]
+    CPUAlreadyProfiling,
+
+    #[error("pprof error: {0}")]
+    PprofError(#[from] pprof::Error),
+
+    #[error("cannot encode report to pprof format: {0}")]
+    EncodeError(String),
+}
+
+/// Default frequency of sampling. 99Hz to avoid coincide with special periods
+pub fn default_profiling_frequency() -> i32 {
+    99
+}
+
+/// Default profiling interval time - 30 seconds
+pub fn default_profiling_interval() -> u64 {
+    30
+}
+
+/// Trigger one cpu profile.
+pub async fn start_one_cpu_profile<F>(
+    end: F,
+    frequency: i32,
+) -> Result<Vec<u8>, ReportGenerationError>
+where
+    F: Future<Output = Result<(), ReportGenerationError>> + Send + 'static,
+{
+    if CPU_PROFILE_ACTIVE.lock().unwrap().is_some() {
+        return Err(ReportGenerationError::CPUAlreadyProfiling);
+    }
+
+    let on_start = || {
+        let mut activate = CPU_PROFILE_ACTIVE.lock().unwrap();
+        assert!(activate.is_none());
+        *activate = Some(());
+        let guard = pprof::ProfilerGuardBuilder::default()
+            .frequency(frequency)
+            .blocklist(&["libc", "libgcc", "pthread", "vdso"])
+            .build()?;
+        Ok(guard)
+    };
+
+    let on_end = move |guard: pprof::ProfilerGuard<'static>| {
+        let report = guard
+            .report()
+            .frames_post_processor(move |frames| {
+                let name = extract_thread_name(&frames.thread_name);
+                frames.thread_name = name;
+            })
+            .build()?;
+        let mut body = Vec::new();
+        let profile = report.pprof()?;
+
+        profile
+            .encode(&mut body)
+            .map_err(|e| ReportGenerationError::EncodeError(e.to_string()))?;
+
+        drop(guard);
+        *CPU_PROFILE_ACTIVE.lock().unwrap() = None;
+
+        Ok(body)
+    };
+
+    ProfileRunner::new(on_start, on_end, end.boxed())?.await
+}
+
+fn extract_thread_name(thread_name: &str) -> String {
+    THREAD_NAME_RE
+        .captures(thread_name)
+        .and_then(|cap| {
+            cap.name("thread_name").map(|thread_name| {
+                THREAD_NAME_REPLACE_SEPERATOR_RE
+                    .replace_all(thread_name.as_str(), "-")
+                    .into_owned()
+            })
+        })
+        .unwrap_or_else(|| thread_name.to_owned())
+}
+
+type OnEndFn<I, T> = Box<dyn FnOnce(I) -> Result<T, ReportGenerationError> + Send + 'static>;
+
+struct ProfileRunner<I, T> {
+    item: Option<I>,
+    on_end: Option<OnEndFn<I, T>>,
+    end: BoxFuture<'static, Result<(), ReportGenerationError>>,
+}
+
+impl<I, T> Unpin for ProfileRunner<I, T> {}
+
+impl<I, T> ProfileRunner<I, T> {
+    fn new<F1, F2>(
+        on_start: F1,
+        on_end: F2,
+        end: BoxFuture<'static, Result<(), ReportGenerationError>>,
+    ) -> Result<Self, ReportGenerationError>
+    where
+        F1: FnOnce() -> Result<I, ReportGenerationError>,
+        F2: FnOnce(I) -> Result<T, ReportGenerationError> + Send + 'static,
+    {
+        let item = on_start()?;
+        Ok(ProfileRunner {
+            item: Some(item),
+            on_end: Some(Box::new(on_end) as OnEndFn<I, T>),
+            end,
+        })
+    }
+}
+
+impl<I, T> Future for ProfileRunner<I, T> {
+    type Output = Result<T, ReportGenerationError>;
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.end.as_mut().poll(cx) {
+            Poll::Ready(res) => {
+                let item = self.item.take().unwrap();
+                let on_end = self.on_end.take().unwrap();
+                let r = match (res, on_end(item)) {
+                    (Ok(_), r) => r,
+                    (Err(errmsg), _) => Err(errmsg),
+                };
+                Poll::Ready(r)
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}

src/profiling.rs

@@ -69,6 +69,7 @@ pub(crate) async fn app() -> Router {
         daemon_pid_file: "policy_server.pid".to_owned(),
         daemon_stdout_file: None,
         daemon_stderr_file: None,
+        enable_pprof: true,
     };
 
     let server = PolicyServer::new_from_config(config).await.unwrap();