Merge pull request #693 from flavio/add-pprof-support

add pprof support

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

@@ -1,5 +1,5 @@
 use clap::builder::PossibleValue;
-use clap::{crate_authors, crate_description, crate_name, crate_version, Arg, Command};
+use clap::{crate_authors, crate_description, crate_name, crate_version, Arg, ArgAction, Command};
 use itertools::Itertools;
 use lazy_static::lazy_static;
 use policy_evaluator::burrego;
@@ -21,11 +21,7 @@ lazy_static! {
 }
 
 pub(crate) fn build_cli() -> Command {
-    Command::new(crate_name!())
-        .author(crate_authors!())
-        .version(crate_version!())
-        .about(crate_description!())
-        .arg(
+    let mut args = vec![
             Arg::new("log-level")
                 .long("log-level")
                 .value_name("LOG_LEVEL")
@@ -39,8 +35,6 @@ pub(crate) fn build_cli() -> Command {
                     PossibleValue::new("error"),
                 ])
                 .help("Log level"),
-        )
-        .arg(
             Arg::new("log-fmt")
                 .long("log-fmt")
                 .value_name("LOG_FMT")
@@ -52,169 +46,132 @@ pub(crate) fn build_cli() -> Command {
                     PossibleValue::new("otlp"),
                 ])
                 .help("Log output format"),
-        )
-        .arg(
             Arg::new("log-no-color")
                 .long("log-no-color")
                 .env("NO_COLOR")
-                .required(false)
+                .action(ArgAction::SetTrue)
                 .help("Disable colored output for logs"),
-        )
-        .arg(
             Arg::new("address")
                 .long("addr")
                 .value_name("BIND_ADDRESS")
                 .default_value("0.0.0.0")
                 .env("KUBEWARDEN_BIND_ADDRESS")
                 .help("Bind against ADDRESS"),
-        )
-        .arg(
             Arg::new("port")
                 .long("port")
                 .value_name("PORT")
                 .default_value("3000")
                 .env("KUBEWARDEN_PORT")
                 .help("Listen on PORT"),
-        )
-        .arg(
             Arg::new("workers")
                 .long("workers")
                 .value_name("WORKERS_NUMBER")
                 .env("KUBEWARDEN_WORKERS")
                 .help("Number of workers thread to create"),
-        )
-        .arg(
             Arg::new("cert-file")
                 .long("cert-file")
                 .value_name("CERT_FILE")
                 .default_value("")
                 .env("KUBEWARDEN_CERT_FILE")
                 .help("Path to an X.509 certificate file for HTTPS"),
-        )
-        .arg(
             Arg::new("key-file")
                 .long("key-file")
                 .value_name("KEY_FILE")
                 .default_value("")
                 .env("KUBEWARDEN_KEY_FILE")
                 .help("Path to an X.509 private key file for HTTPS"),
-        )
-        .arg(
             Arg::new("policies")
                 .long("policies")
                 .value_name("POLICIES_FILE")
                 .env("KUBEWARDEN_POLICIES")
                 .default_value("policies.yml")
                 .help("YAML file holding the policies to be loaded and their settings"),
-        )
-        .arg(
             Arg::new("policies-download-dir")
                 .long("policies-download-dir")
                 .value_name("POLICIES_DOWNLOAD_DIR")
                 .default_value(".")
                 .env("KUBEWARDEN_POLICIES_DOWNLOAD_DIR")
                 .help("Download path for the policies"),
-        )
-        .arg(
             Arg::new("sigstore-cache-dir")
                 .long("sigstore-cache-dir")
                 .value_name("SIGSTORE_CACHE_DIR")
                 .default_value("sigstore-data")
                 .env("KUBEWARDEN_SIGSTORE_CACHE_DIR")
                 .help("Directory used to cache sigstore data"),
-        )
-        .arg(
             Arg::new("sources-path")
                 .long("sources-path")
                 .value_name("SOURCES_PATH")
                 .env("KUBEWARDEN_SOURCES_PATH")
                 .help("YAML file holding source information (https, registry insecure hosts, custom CA's...)"),
-        )
-        .arg(
             Arg::new("verification-path")
                 .long("verification-path")
                 .value_name("VERIFICATION_CONFIG_PATH")
                 .env("KUBEWARDEN_VERIFICATION_CONFIG_PATH")
                 .help("YAML file holding verification information (URIs, keys, annotations...)"),
-        )
-        .arg(
             Arg::new("docker-config-json-path")
                 .long("docker-config-json-path")
                 .value_name("DOCKER_CONFIG")
                 .env("KUBEWARDEN_DOCKER_CONFIG_JSON_PATH")
                 .help("Path to a Docker config.json-like path. Can be used to indicate registry authentication details"),
-        )
-        .arg(
             Arg::new("enable-metrics")
                 .long("enable-metrics")
                 .env("KUBEWARDEN_ENABLE_METRICS")
-                .required(false)
+                .action(ArgAction::SetTrue)
                 .help("Enable metrics"),
-        )
-        .arg(
-            Arg::new("enable-verification")
-                .long("enable-verification")
-                .env("KUBEWARDEN_ENABLE_VERIFICATION")
-                .required(false)
-                .help("Enable Sigstore verification"),
-        )
-        .arg(
             Arg::new("always-accept-admission-reviews-on-namespace")
                 .long("always-accept-admission-reviews-on-namespace")
                 .value_name("NAMESPACE")
                 .env("KUBEWARDEN_ALWAYS_ACCEPT_ADMISSION_REVIEWS_ON_NAMESPACE")
                 .required(false)
                 .help("Always accept AdmissionReviews that target the given namespace"),
-        )
-        .arg(
             Arg::new("disable-timeout-protection")
                 .long("disable-timeout-protection")
+                .action(ArgAction::SetTrue)
                 .env("KUBEWARDEN_DISABLE_TIMEOUT_PROTECTION")
-                .required(false)
                 .help("Disable policy timeout protection"),
-        )
-        .arg(
             Arg::new("policy-timeout")
                 .long("policy-timeout")
                 .env("KUBEWARDEN_POLICY_TIMEOUT")
                 .value_name("MAXIMUM_EXECUTION_TIME_SECONDS")
                 .default_value("2")
                 .help("Interrupt policy evaluation after the given time"),
-        )
-        .arg(
             Arg::new("daemon")
                 .long("daemon")
                 .env("KUBEWARDEN_DAEMON")
-                .required(false)
+                .action(ArgAction::SetTrue)
                 .help("If set, runs policy-server in detached mode as a daemon"),
-        )
-        .arg(
             Arg::new("daemon-pid-file")
                 .long("daemon-pid-file")
                 .env("KUBEWARDEN_DAEMON_PID_FILE")
                 .default_value("policy-server.pid")
                 .help("Path to pid file, used only when running in daemon mode"),
-        )
-        .arg(
             Arg::new("daemon-stdout-file")
                 .long("daemon-stdout-file")
                 .env("KUBEWARDEN_DAEMON_STDOUT_FILE")
                 .required(false)
                 .help("Path to file holding stdout, used only when running in daemon mode"),
-        )
-        .arg(
             Arg::new("daemon-stderr-file")
                 .long("daemon-stderr-file")
                 .env("KUBEWARDEN_DAEMON_STDERR_FILE")
                 .required(false)
                 .help("Path to file holding stderr, used only when running in daemon mode"),
-        )
-        .arg(
             Arg::new("ignore-kubernetes-connection-failure")
                 .long("ignore-kubernetes-connection-failure")
                 .env("KUBEWARDEN_IGNORE_KUBERNETES_CONNECTION_FAILURE")
-                .required(false)
+                .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()));
+
+    Command::new(crate_name!())
+        .author(crate_authors!())
+        .version(crate_version!())
+        .about(crate_description!())
         .long_version(VERSION_AND_BUILTINS.as_str())
+        .args(args)
 }