enable pprof cpu profiling

This adds two routes to control pprof
Enable CPU profiling (100hz)
http://localhost:7280/pprof/start
Get the profile as flamegraph
http://localhost:7280/pprof/stop

The routes are behind the `pprof` feature flag

quickwit/Cargo.toml

@@ -171,6 +171,7 @@ postcard = { version = "1.0.4", features = [
 ], default-features = false }
 predicates = "3"
 prettyplease = "0.2.0"
+pprof = { version = "0.13", features = ["flamegraph"] }
 proc-macro2 = "1.0.50"
 prometheus = { version = "0.13", features = ["process"] }
 proptest = "1"

quickwit/quickwit-cli/Cargo.toml

@@ -81,6 +81,7 @@ quickwit-storage = { workspace = true, features = ["testsuite"] }
 [features]
 jemalloc = ["dep:tikv-jemalloc-ctl", "dep:tikv-jemallocator"]
 ci-test = []
+pprof = ["quickwit-serve/pprof"]
 openssl-support = ["openssl-probe"]
 # Requires to enable tokio unstable via RUSTFLAGS="--cfg tokio_unstable"
 tokio-console = ["console-subscriber", "quickwit-common/named_tasks"]

quickwit/quickwit-serve/Cargo.toml

@@ -30,6 +30,7 @@ mime_guess = { workspace = true }
 once_cell = { workspace = true }
 opentelemetry = { workspace = true }
 percent-encoding = { workspace = true }
+pprof = { workspace = true, optional = true }
 prost = { workspace = true }
 prost-types = { workspace = true }
 regex = { workspace = true }
@@ -91,3 +92,9 @@ quickwit-opentelemetry = { workspace = true, features = ["testsuite"] }
 quickwit-proto = { workspace = true, features = ["testsuite"] }
 quickwit-search = { workspace = true, features = ["testsuite"] }
 quickwit-storage = { workspace = true, features = ["testsuite"] }
+
+[features]
+pprof = [
+  "dep:pprof"
+]
+

quickwit/quickwit-serve/src/lib.rs

@@ -38,6 +38,7 @@ mod metrics_api;
 mod node_info_handler;
 mod openapi;
 mod otlp_api;
+mod pprof;
 mod rate_modulator;
 mod rest;
 mod rest_api_response;

quickwit/quickwit-serve/src/pprof.rs

@@ -0,0 +1,106 @@
+// Copyright (C) 2024 Quickwit, Inc.
+//
+// Quickwit is offered under the AGPL v3.0 and as commercial software.
+// For commercial licensing, contact us at [email protected].
+//
+// AGPL:
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as
+// published by the Free Software Foundation, either version 3 of the
+// License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+use warp::Filter;
+
+/// pprof/start to start cpu profiling
+/// pprof/stop to stop cpu profiling and return a flamegraph
+#[cfg(not(feature = "pprof"))]
+pub fn pprof_routes() -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
+    let start_profiler = {
+        warp::path!("pprof" / "start").map(move || {
+            warp::reply::with_status(
+                "not compiled with pprof feature",
+                warp::http::StatusCode::BAD_REQUEST,
+            )
+        })
+    };
+
+    let stop_profiler = {
+        warp::path!("pprof" / "stop").map(move || {
+            warp::reply::with_status(
+                "not compiled with pprof feature",
+                warp::http::StatusCode::BAD_REQUEST,
+            )
+        })
+    };
+
+    start_profiler.or(stop_profiler)
+}
+
+/// pprof/start to start cpu profiling
+/// pprof/stop to stop cpu profiling and return a flamegraph
+#[cfg(feature = "pprof")]
+pub fn pprof_routes() -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection> + Clone {
+    use std::sync::{Arc, Mutex};
+
+    use pprof::ProfilerGuard;
+    use warp::reply::Reply;
+    let profiler_guard: Arc<Mutex<Option<ProfilerGuard<'_>>>> = Arc::new(Mutex::new(None));
+
+    let start_profiler = {
+        let profiler_guard = Arc::clone(&profiler_guard);
+        warp::path!("pprof" / "start").map(move || {
+            let mut guard = profiler_guard.lock().unwrap();
+            if guard.is_none() {
+                *guard = Some(pprof::ProfilerGuard::new(100).unwrap());
+                warp::reply::with_status("CPU profiling started", warp::http::StatusCode::OK)
+            } else {
+                warp::reply::with_status(
+                    "CPU profiling is already running",
+                    warp::http::StatusCode::BAD_REQUEST,
+                )
+            }
+        })
+    };
+
+    let stop_profiler = {
+        let profiler_guard = Arc::clone(&profiler_guard);
+        warp::path!("pprof" / "stop").map(move || {
+            let profiler_guard = Arc::clone(&profiler_guard);
+            get_flamegraph(profiler_guard)
+        })
+    };
+
+    fn get_flamegraph(profiler_guard: Arc<Mutex<Option<ProfilerGuard>>>) -> impl warp::Reply {
+        let mut guard = profiler_guard.lock().unwrap();
+        if let Some(profiler) = guard.take() {
+            if let Ok(report) = profiler.report().build() {
+                let mut buffer = Vec::new();
+                if report.flamegraph(&mut buffer).is_ok() {
+                    return warp::reply::with_header(buffer, "Content-Type", "image/svg+xml")
+                        .into_response();
+                }
+            }
+            warp::reply::with_status(
+                "Failed to generate flamegraph",
+                warp::http::StatusCode::INTERNAL_SERVER_ERROR,
+            )
+            .into_response()
+        } else {
+            warp::reply::with_status(
+                "CPU profiling is not running",
+                warp::http::StatusCode::BAD_REQUEST,
+            )
+            .into_response()
+        }
+    }
+
+    start_profiler.or(stop_profiler)
+}