feat: add feature pprof

GreptimeTeam · Jun 5, 2023 · df393fb · df393fb
1 parent e7f1260
commit df393fb
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 66 deletions.
diff --git a/src/common/pprof/README.md b/src/common/pprof/README.md
@@ -1,5 +1,10 @@
 # Profiling CPU
 
+## Build GreptimeDB with `pprof` feature
+
+```bash
+cargo build --features=pprof
+```
 
 ## HTTP API
 Sample at 99 Hertz, for 5 seconds, output report in protobuf format.

diff --git a/src/servers/Cargo.toml b/src/servers/Cargo.toml
@@ -5,6 +5,7 @@ edition.workspace = true
 license.workspace = true
 
 [features]
+pprof = ["dep:common-pprof"]
 mem-prof = ["dep:common-mem-prof"]
 dashboard = []
 
@@ -25,7 +26,7 @@ common-error = { path = "../common/error" }
 common-grpc = { path = "../common/grpc" }
 common-grpc-expr = { path = "../common/grpc-expr" }
 common-mem-prof = { path = "../common/mem-prof", optional = true }
-common-pprof = { path = "../common/pprof" }
+common-pprof = { path = "../common/pprof", optional = true }
 common-query = { path = "../common/query" }
 common-recordbatch = { path = "../common/recordbatch" }
 common-runtime = { path = "../common/runtime" }

diff --git a/src/servers/src/error.rs b/src/servers/src/error.rs
@@ -267,6 +267,7 @@ pub enum Error {
         location: Location,
     },
 
+    #[cfg(feature = "pprof")]
     #[snafu(display("Failed to dump pprof data, source: {}", source))]
     DumpPprof {
         #[snafu(backtrace)]
@@ -348,6 +349,7 @@ impl ErrorExt for Error {
                 }
             }
 
+            #[cfg(feature = "pprof")]
             DumpPprof { source, .. } => source.status_code(),
         }
     }

diff --git a/src/servers/src/http.rs b/src/servers/src/http.rs
@@ -503,20 +503,6 @@ impl HttpServer {
             );
         }
 
-        // prof routers
-        // router = router.nest(
-        //     &format!("/{HTTP_API_VERSION}/prof"),
-        //     Router::new()
-        //         .route(
-        //             "/cpu",
-        //             routing::get(pprof::pprof_handler).post(pprof::pprof_handler),
-        //         )
-        //         .route(
-        //             "/mem",
-        //             routing::get(mem_prof::mem_prof_handler).post(mem_prof::mem_prof_handler),
-        //         ),
-        // );
-
         if let Some(metrics_handler) = self.metrics_handler {
             router = router.nest("", self.route_metrics(metrics_handler));
         }

diff --git a/src/servers/src/http/pprof.rs b/src/servers/src/http/pprof.rs
@@ -12,66 +12,87 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::num::NonZeroI32;
-use std::time::Duration;
+#[cfg(feature = "pprof")]
+pub mod handler {
+    use std::num::NonZeroI32;
+    use std::time::Duration;
 
-use axum::extract::Query;
-use axum::http::StatusCode;
-use axum::response::IntoResponse;
-use common_pprof::Profiling;
-use common_telemetry::logging;
-use schemars::JsonSchema;
-use serde::{Deserialize, Serialize};
-use snafu::ResultExt;
+    use axum::extract::Query;
+    use axum::http::StatusCode;
+    use axum::response::IntoResponse;
+    use common_pprof::Profiling;
+    use common_telemetry::logging;
+    use schemars::JsonSchema;
+    use serde::{Deserialize, Serialize};
+    use snafu::ResultExt;
 
-use crate::error::{DumpPprofSnafu, Result};
+    use crate::error::{DumpPprofSnafu, Result};
 
-/// Output format.
-#[derive(Debug, Serialize, Deserialize, JsonSchema)]
-#[serde(rename_all = "snake_case")]
-pub enum Output {
-    /// google’s pprof format report in protobuf.
-    Proto,
-    /// Simple text format.
-    Text,
-    /// svg flamegraph.
-    Flamegraph,
-}
+    /// Output format.
+    #[derive(Debug, Serialize, Deserialize, JsonSchema)]
+    #[serde(rename_all = "snake_case")]
+    pub enum Output {
+        /// google’s pprof format report in protobuf.
+        Proto,
+        /// Simple text format.
+        Text,
+        /// svg flamegraph.
+        Flamegraph,
+    }
 
-#[derive(Serialize, Deserialize, Debug, JsonSchema)]
-#[serde(default)]
-pub struct PprofQuery {
-    seconds: u64,
-    frequency: NonZeroI32,
-    output: Output,
-}
+    #[derive(Serialize, Deserialize, Debug, JsonSchema)]
+    #[serde(default)]
+    pub struct PprofQuery {
+        seconds: u64,
+        frequency: NonZeroI32,
+        output: Output,
+    }
 
-impl Default for PprofQuery {
-    fn default() -> PprofQuery {
-        PprofQuery {
-            seconds: 5,
-            // Safety: 99 is non zero.
-            frequency: NonZeroI32::new(99).unwrap(),
-            output: Output::Proto,
+    impl Default for PprofQuery {
+        fn default() -> PprofQuery {
+            PprofQuery {
+                seconds: 5,
+                // Safety: 99 is non zero.
+                frequency: NonZeroI32::new(99).unwrap(),
+                output: Output::Proto,
+            }
         }
     }
-}
 
-#[axum_macros::debug_handler]
-pub async fn pprof_handler(Query(req): Query<PprofQuery>) -> Result<impl IntoResponse> {
-    logging::info!("start pprof, request: {:?}", req);
+    #[axum_macros::debug_handler]
+    pub async fn pprof_handler(Query(req): Query<PprofQuery>) -> Result<impl IntoResponse> {
+        logging::info!("start pprof, request: {:?}", req);
 
-    let profiling = Profiling::new(Duration::from_secs(req.seconds), req.frequency.into());
-    let body = match req.output {
-        Output::Proto => profiling.dump_proto().await.context(DumpPprofSnafu)?,
-        Output::Text => {
-            let report = profiling.report().await.context(DumpPprofSnafu)?;
-            format!("{:?}", report).into_bytes()
-        }
-        Output::Flamegraph => profiling.dump_flamegraph().await.context(DumpPprofSnafu)?,
-    };
+        let profiling = Profiling::new(Duration::from_secs(req.seconds), req.frequency.into());
+        let body = match req.output {
+            Output::Proto => profiling.dump_proto().await.context(DumpPprofSnafu)?,
+            Output::Text => {
+                let report = profiling.report().await.context(DumpPprofSnafu)?;
+                format!("{:?}", report).into_bytes()
+            }
+            Output::Flamegraph => profiling.dump_flamegraph().await.context(DumpPprofSnafu)?,
+        };
 
-    logging::info!("finish pprof");
+        logging::info!("finish pprof");
 
-    Ok((StatusCode::OK, body))
+        Ok((StatusCode::OK, body))
+    }
 }
+
+#[cfg(not(feature = "pprof"))]
+pub mod handler {
+    use axum::http::StatusCode;
+    use axum::response::IntoResponse;
+
+    use crate::error::Result;
+
+    #[axum_macros::debug_handler]
+    pub async fn pprof_handler() -> Result<impl IntoResponse> {
+        Ok((
+            StatusCode::NOT_IMPLEMENTED,
+            "The 'pprof' feature is disabled",
+        ))
+    }
+}
+
+pub use handler::pprof_handler;