common_pprof/
nix.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
15use std::time::Duration;
16
17use pprof::protos::Message;
18use snafu::ResultExt;
19
20use crate::error::{PprofSnafu, Result};
21
22/// CPU profiler utility.
23// Inspired by https://github.com/datafuselabs/databend/blob/67f445e83cd4eceda98f6c1c114858929d564029/src/common/base/src/base/profiling.rs
24#[derive(Debug)]
25pub struct Profiling {
26    /// Sample duration.
27    duration: Duration,
28    /// Sample frequency.
29    frequency: i32,
30}
31
32impl Profiling {
33    /// Creates a new profiler.
34    pub fn new(duration: Duration, frequency: i32) -> Profiling {
35        Profiling {
36            duration,
37            frequency,
38        }
39    }
40
41    /// Profiles and returns a generated pprof report.
42    pub async fn report(&self) -> Result<pprof::Report> {
43        let guard = pprof::ProfilerGuardBuilder::default()
44            .frequency(self.frequency)
45            .blocklist(&["libc", "libgcc", "pthread", "vdso"])
46            .build()
47            .context(PprofSnafu)?;
48        tokio::time::sleep(self.duration).await;
49        guard.report().build().context(PprofSnafu)
50    }
51
52    /// Profiles and returns a generated text.
53    pub async fn dump_text(&self) -> Result<String> {
54        let report = self.report().await?;
55        let text = format!("{report:?}");
56        Ok(text)
57    }
58
59    /// Profiles and returns a generated flamegraph.
60    pub async fn dump_flamegraph(&self) -> Result<Vec<u8>> {
61        let mut body: Vec<u8> = Vec::new();
62
63        let report = self.report().await?;
64        report.flamegraph(&mut body).context(PprofSnafu)?;
65
66        Ok(body)
67    }
68
69    /// Profiles and returns a generated proto.
70    pub async fn dump_proto(&self) -> Result<Vec<u8>> {
71        let report = self.report().await?;
72        // Generate google’s pprof format report.
73        let profile = report.pprof().context(PprofSnafu)?;
74        let body = profile.encode_to_vec();
75
76        Ok(body)
77    }
78}