servers/http/
pprof.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#[cfg(feature = "pprof")]
16pub mod handler {
17    use std::num::NonZeroI32;
18    use std::time::Duration;
19
20    use axum::extract::Query;
21    use axum::http::StatusCode;
22    use axum::response::IntoResponse;
23    use common_pprof::Profiling;
24    use common_telemetry::info;
25    use serde::{Deserialize, Serialize};
26    use snafu::ResultExt;
27
28    use crate::error::{DumpPprofSnafu, Result};
29
30    /// Output format.
31    #[derive(Debug, Serialize, Deserialize)]
32    #[serde(rename_all = "snake_case")]
33    pub enum Output {
34        /// google’s pprof format report in protobuf.
35        Proto,
36        /// Simple text format.
37        Text,
38        /// svg flamegraph.
39        Flamegraph,
40    }
41
42    #[derive(Serialize, Deserialize, Debug)]
43    #[serde(default)]
44    pub struct PprofQuery {
45        seconds: u64,
46        frequency: NonZeroI32,
47        output: Output,
48    }
49
50    impl Default for PprofQuery {
51        fn default() -> PprofQuery {
52            PprofQuery {
53                seconds: 5,
54                // Safety: 99 is non zero.
55                frequency: NonZeroI32::new(99).unwrap(),
56                output: Output::Proto,
57            }
58        }
59    }
60
61    #[axum_macros::debug_handler]
62    pub async fn pprof_handler(Query(req): Query<PprofQuery>) -> Result<impl IntoResponse> {
63        info!("start pprof, request: {:?}", req);
64
65        let profiling = Profiling::new(Duration::from_secs(req.seconds), req.frequency.into());
66        let body = match req.output {
67            Output::Proto => profiling.dump_proto().await.context(DumpPprofSnafu)?,
68            Output::Text => {
69                let report = profiling.dump_text().await.context(DumpPprofSnafu)?;
70                report.into_bytes()
71            }
72            Output::Flamegraph => profiling.dump_flamegraph().await.context(DumpPprofSnafu)?,
73        };
74
75        info!("finish pprof");
76
77        info!("Dump data success, size: {}", body.len());
78
79        Ok((StatusCode::OK, body))
80    }
81}
82
83#[cfg(not(feature = "pprof"))]
84pub mod handler {
85    use axum::http::StatusCode;
86    use axum::response::IntoResponse;
87
88    use crate::error::Result;
89
90    #[axum_macros::debug_handler]
91    pub async fn pprof_handler() -> Result<impl IntoResponse> {
92        Ok((
93            StatusCode::NOT_IMPLEMENTED,
94            "The 'pprof' feature is disabled",
95        ))
96    }
97}
98
99pub use handler::pprof_handler;
servers/http/pprof.rs

servers/http/
pprof.rs
