Skip to content

Commit f347084

Browse files
authored
Merge pull request #693 from flavio/add-pprof-support
add pprof support
2 parents a5d53ca + 6f41df7 commit f347084

File tree

8 files changed

+600
-152
lines changed

8 files changed

+600
-152
lines changed

Cargo.lock

+279-70
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

+5-1
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,14 @@ edition = "2021"
1414
anyhow = "1.0"
1515
clap = { version = "4.5", features = ["cargo", "env"] }
1616
daemonize = "0.5"
17+
futures = "0.3"
1718
humansize = "2.1"
1819
itertools = "0.12.1"
1920
k8s-openapi = { version = "0.21.1", default-features = false, features = [
2021
"v1_29",
2122
] }
2223
lazy_static = "1.4.0"
24+
mime = "0.3"
2325
num_cpus = "1.16.0"
2426
opentelemetry-otlp = { version = "0.14.0", features = ["metrics", "tonic"] }
2527
opentelemetry = { version = "0.21", default-features = false, features = [
@@ -28,8 +30,10 @@ opentelemetry = { version = "0.21", default-features = false, features = [
2830
] }
2931
opentelemetry_sdk = { version = "0.21", features = ["rt-tokio"] }
3032
procfs = "0.16"
33+
pprof = { version = "0.13", features = ["prost-codec"] }
3134
policy-evaluator = { git = "https://github.com/kubewarden/policy-evaluator", tag = "v0.16.1" }
3235
rayon = "1.9"
36+
regex = "1.10"
3337
serde_json = "1.0"
3438
serde = { version = "1.0", features = ["derive"] }
3539
serde_yaml = "0.9.32"
@@ -42,7 +46,7 @@ tracing-opentelemetry = "0.22.0"
4246
tracing-subscriber = { version = "0.3", features = ["ansi", "fmt", "json"] }
4347
semver = { version = "1.0.22", features = ["serde"] }
4448
mockall_double = "0.3"
45-
axum = { version = "0.7.4", features = ["macros"] }
49+
axum = { version = "0.7.4", features = ["macros", "query"] }
4650
axum-server = { version = "0.6", features = ["tls-rustls"] }
4751
tower-http = { version = "0.5.2", features = ["trace"] }
4852

src/api/handlers.rs

+70-10
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,30 @@
11
use axum::{
2-
extract::{self, FromRequest},
3-
http::StatusCode,
2+
extract::{self, FromRequest, Query},
3+
http::{header, StatusCode},
44
response::IntoResponse,
55
Json,
66
};
77
use policy_evaluator::{
88
admission_request::AdmissionRequest, admission_response::AdmissionResponse,
99
policy_evaluator::ValidateRequest,
1010
};
11-
use serde::Serialize;
11+
12+
use serde::{Deserialize, Serialize};
1213
use std::sync::Arc;
1314
use tokio::task;
1415
use tracing::{debug, error, Span};
1516

16-
use crate::api::{
17-
admission_review::{AdmissionReviewRequest, AdmissionReviewResponse},
18-
api_error::ApiError,
19-
raw_review::{RawReviewRequest, RawReviewResponse},
20-
service::{evaluate, RequestOrigin},
21-
state::ApiServerState,
17+
use crate::{
18+
api::{
19+
admission_review::{AdmissionReviewRequest, AdmissionReviewResponse},
20+
api_error::ApiError,
21+
raw_review::{RawReviewRequest, RawReviewResponse},
22+
service::{evaluate, RequestOrigin},
23+
state::ApiServerState,
24+
},
25+
profiling,
2226
};
23-
use crate::evaluation::errors::EvaluationError;
27+
use crate::{evaluation::errors::EvaluationError, profiling::ReportGenerationError};
2428

2529
// create an extractor that internally uses `axum::Json` but has a custom rejection
2630
#[derive(FromRequest)]
@@ -169,6 +173,50 @@ pub(crate) async fn validate_raw_handler(
169173
Ok(Json(RawReviewResponse::new(response)))
170174
}
171175

176+
#[derive(Deserialize)]
177+
pub(crate) struct ProfileParams {
178+
/// profiling frequency (Hz)
179+
#[serde(default = "profiling::default_profiling_frequency")]
180+
pub frequency: i32,
181+
182+
/// profiling time interval (seconds)
183+
#[serde(default = "profiling::default_profiling_interval")]
184+
pub interval: u64,
185+
}
186+
187+
// Generate a pprof CPU profile using google's pprof format
188+
// The report is generated and sent to the user as binary data
189+
pub(crate) async fn pprof_get_cpu(
190+
profiling_params: Query<ProfileParams>,
191+
) -> Result<impl axum::response::IntoResponse, (StatusCode, ApiError)> {
192+
let frequency = profiling_params.frequency;
193+
let interval = profiling_params.interval;
194+
195+
let end = async move {
196+
tokio::time::sleep(tokio::time::Duration::from_secs(interval)).await;
197+
Ok(())
198+
};
199+
200+
let body = profiling::start_one_cpu_profile(end, frequency)
201+
.await
202+
.map_err(handle_pprof_error)?;
203+
204+
let mut headers = header::HeaderMap::new();
205+
headers.insert(
206+
header::CONTENT_DISPOSITION,
207+
r#"attachment; filename="cpu_profile"#.parse().unwrap(),
208+
);
209+
headers.insert(
210+
header::CONTENT_LENGTH,
211+
body.len().to_string().parse().unwrap(),
212+
);
213+
headers.insert(
214+
header::CONTENT_TYPE,
215+
mime::APPLICATION_OCTET_STREAM.to_string().parse().unwrap(),
216+
);
217+
218+
Ok((headers, body))
219+
}
172220
pub(crate) async fn readiness_handler() -> StatusCode {
173221
StatusCode::OK
174222
}
@@ -261,3 +309,15 @@ fn handle_evaluation_error(error: EvaluationError) -> (StatusCode, ApiError) {
261309
}
262310
}
263311
}
312+
313+
fn handle_pprof_error(error: ReportGenerationError) -> (StatusCode, ApiError) {
314+
error!("pprof error: {}", error);
315+
316+
(
317+
StatusCode::INTERNAL_SERVER_ERROR,
318+
ApiError {
319+
status: StatusCode::INTERNAL_SERVER_ERROR,
320+
message: "Something went wrong".to_owned(),
321+
},
322+
)
323+
}

src/cli.rs

+20-63
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use clap::builder::PossibleValue;
2-
use clap::{crate_authors, crate_description, crate_name, crate_version, Arg, Command};
2+
use clap::{crate_authors, crate_description, crate_name, crate_version, Arg, ArgAction, Command};
33
use itertools::Itertools;
44
use lazy_static::lazy_static;
55
use policy_evaluator::burrego;
@@ -21,11 +21,7 @@ lazy_static! {
2121
}
2222

2323
pub(crate) fn build_cli() -> Command {
24-
Command::new(crate_name!())
25-
.author(crate_authors!())
26-
.version(crate_version!())
27-
.about(crate_description!())
28-
.arg(
24+
let mut args = vec![
2925
Arg::new("log-level")
3026
.long("log-level")
3127
.value_name("LOG_LEVEL")
@@ -39,8 +35,6 @@ pub(crate) fn build_cli() -> Command {
3935
PossibleValue::new("error"),
4036
])
4137
.help("Log level"),
42-
)
43-
.arg(
4438
Arg::new("log-fmt")
4539
.long("log-fmt")
4640
.value_name("LOG_FMT")
@@ -52,169 +46,132 @@ pub(crate) fn build_cli() -> Command {
5246
PossibleValue::new("otlp"),
5347
])
5448
.help("Log output format"),
55-
)
56-
.arg(
5749
Arg::new("log-no-color")
5850
.long("log-no-color")
5951
.env("NO_COLOR")
60-
.required(false)
52+
.action(ArgAction::SetTrue)
6153
.help("Disable colored output for logs"),
62-
)
63-
.arg(
6454
Arg::new("address")
6555
.long("addr")
6656
.value_name("BIND_ADDRESS")
6757
.default_value("0.0.0.0")
6858
.env("KUBEWARDEN_BIND_ADDRESS")
6959
.help("Bind against ADDRESS"),
70-
)
71-
.arg(
7260
Arg::new("port")
7361
.long("port")
7462
.value_name("PORT")
7563
.default_value("3000")
7664
.env("KUBEWARDEN_PORT")
7765
.help("Listen on PORT"),
78-
)
79-
.arg(
8066
Arg::new("workers")
8167
.long("workers")
8268
.value_name("WORKERS_NUMBER")
8369
.env("KUBEWARDEN_WORKERS")
8470
.help("Number of workers thread to create"),
85-
)
86-
.arg(
8771
Arg::new("cert-file")
8872
.long("cert-file")
8973
.value_name("CERT_FILE")
9074
.default_value("")
9175
.env("KUBEWARDEN_CERT_FILE")
9276
.help("Path to an X.509 certificate file for HTTPS"),
93-
)
94-
.arg(
9577
Arg::new("key-file")
9678
.long("key-file")
9779
.value_name("KEY_FILE")
9880
.default_value("")
9981
.env("KUBEWARDEN_KEY_FILE")
10082
.help("Path to an X.509 private key file for HTTPS"),
101-
)
102-
.arg(
10383
Arg::new("policies")
10484
.long("policies")
10585
.value_name("POLICIES_FILE")
10686
.env("KUBEWARDEN_POLICIES")
10787
.default_value("policies.yml")
10888
.help("YAML file holding the policies to be loaded and their settings"),
109-
)
110-
.arg(
11189
Arg::new("policies-download-dir")
11290
.long("policies-download-dir")
11391
.value_name("POLICIES_DOWNLOAD_DIR")
11492
.default_value(".")
11593
.env("KUBEWARDEN_POLICIES_DOWNLOAD_DIR")
11694
.help("Download path for the policies"),
117-
)
118-
.arg(
11995
Arg::new("sigstore-cache-dir")
12096
.long("sigstore-cache-dir")
12197
.value_name("SIGSTORE_CACHE_DIR")
12298
.default_value("sigstore-data")
12399
.env("KUBEWARDEN_SIGSTORE_CACHE_DIR")
124100
.help("Directory used to cache sigstore data"),
125-
)
126-
.arg(
127101
Arg::new("sources-path")
128102
.long("sources-path")
129103
.value_name("SOURCES_PATH")
130104
.env("KUBEWARDEN_SOURCES_PATH")
131105
.help("YAML file holding source information (https, registry insecure hosts, custom CA's...)"),
132-
)
133-
.arg(
134106
Arg::new("verification-path")
135107
.long("verification-path")
136108
.value_name("VERIFICATION_CONFIG_PATH")
137109
.env("KUBEWARDEN_VERIFICATION_CONFIG_PATH")
138110
.help("YAML file holding verification information (URIs, keys, annotations...)"),
139-
)
140-
.arg(
141111
Arg::new("docker-config-json-path")
142112
.long("docker-config-json-path")
143113
.value_name("DOCKER_CONFIG")
144114
.env("KUBEWARDEN_DOCKER_CONFIG_JSON_PATH")
145115
.help("Path to a Docker config.json-like path. Can be used to indicate registry authentication details"),
146-
)
147-
.arg(
148116
Arg::new("enable-metrics")
149117
.long("enable-metrics")
150118
.env("KUBEWARDEN_ENABLE_METRICS")
151-
.required(false)
119+
.action(ArgAction::SetTrue)
152120
.help("Enable metrics"),
153-
)
154-
.arg(
155-
Arg::new("enable-verification")
156-
.long("enable-verification")
157-
.env("KUBEWARDEN_ENABLE_VERIFICATION")
158-
.required(false)
159-
.help("Enable Sigstore verification"),
160-
)
161-
.arg(
162121
Arg::new("always-accept-admission-reviews-on-namespace")
163122
.long("always-accept-admission-reviews-on-namespace")
164123
.value_name("NAMESPACE")
165124
.env("KUBEWARDEN_ALWAYS_ACCEPT_ADMISSION_REVIEWS_ON_NAMESPACE")
166125
.required(false)
167126
.help("Always accept AdmissionReviews that target the given namespace"),
168-
)
169-
.arg(
170127
Arg::new("disable-timeout-protection")
171128
.long("disable-timeout-protection")
129+
.action(ArgAction::SetTrue)
172130
.env("KUBEWARDEN_DISABLE_TIMEOUT_PROTECTION")
173-
.required(false)
174131
.help("Disable policy timeout protection"),
175-
)
176-
.arg(
177132
Arg::new("policy-timeout")
178133
.long("policy-timeout")
179134
.env("KUBEWARDEN_POLICY_TIMEOUT")
180135
.value_name("MAXIMUM_EXECUTION_TIME_SECONDS")
181136
.default_value("2")
182137
.help("Interrupt policy evaluation after the given time"),
183-
)
184-
.arg(
185138
Arg::new("daemon")
186139
.long("daemon")
187140
.env("KUBEWARDEN_DAEMON")
188-
.required(false)
141+
.action(ArgAction::SetTrue)
189142
.help("If set, runs policy-server in detached mode as a daemon"),
190-
)
191-
.arg(
192143
Arg::new("daemon-pid-file")
193144
.long("daemon-pid-file")
194145
.env("KUBEWARDEN_DAEMON_PID_FILE")
195146
.default_value("policy-server.pid")
196147
.help("Path to pid file, used only when running in daemon mode"),
197-
)
198-
.arg(
199148
Arg::new("daemon-stdout-file")
200149
.long("daemon-stdout-file")
201150
.env("KUBEWARDEN_DAEMON_STDOUT_FILE")
202151
.required(false)
203152
.help("Path to file holding stdout, used only when running in daemon mode"),
204-
)
205-
.arg(
206153
Arg::new("daemon-stderr-file")
207154
.long("daemon-stderr-file")
208155
.env("KUBEWARDEN_DAEMON_STDERR_FILE")
209156
.required(false)
210157
.help("Path to file holding stderr, used only when running in daemon mode"),
211-
)
212-
.arg(
213158
Arg::new("ignore-kubernetes-connection-failure")
214159
.long("ignore-kubernetes-connection-failure")
215160
.env("KUBEWARDEN_IGNORE_KUBERNETES_CONNECTION_FAILURE")
216-
.required(false)
161+
.action(ArgAction::SetTrue)
217162
.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."),
218-
)
163+
Arg::new("enable-pprof")
164+
.long("enable-pprof")
165+
.env("KUBEWARDEN_ENABLE_PPROF")
166+
.action(ArgAction::SetTrue)
167+
.help("Enable pprof profiling"),
168+
];
169+
args.sort_by(|a, b| a.get_id().cmp(b.get_id()));
170+
171+
Command::new(crate_name!())
172+
.author(crate_authors!())
173+
.version(crate_version!())
174+
.about(crate_description!())
219175
.long_version(VERSION_AND_BUILTINS.as_str())
176+
.args(args)
220177
}

0 commit comments

Comments
 (0)