Skip to content

Commit

Permalink
Merge pull request #693 from flavio/add-pprof-support
Browse files Browse the repository at this point in the history 
add pprof support
  • Loading branch information
@flavio
flavio authored Mar 7, 2024
2 parents a5d53ca + 6f41df7 commit f347084
Show file tree
Hide file tree
Showing 8 changed files with 600 additions and 152 deletions.
349 changes: 279 additions & 70 deletions Cargo.lock
View file

Large diffs are not rendered by default.

6 changes: 5 additions & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand All @@ -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"
Expand All @@ -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"] }

Expand Down
80 changes: 70 additions & 10 deletions src/api/handlers.rs
View file
Original file line number Diff line number Diff line change
@@ -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)]
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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(),
},
)
}
83 changes: 20 additions & 63 deletions src/cli.rs
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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")
Expand All @@ -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")
Expand All @@ -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)
}
Loading

0 comments on commit f347084

Please sign in to comment.