Draft support for pprof profile.proto format

Cargo.toml

@@ -1,5 +1,6 @@
 [features]
 unwind = ["remoteprocess/unwind"]
+pprof = ["prost-build", "protobuf-src", "prost", "prost-types", "flate2"]
 
 [package]
 name = "py-spy"
@@ -14,6 +15,10 @@ license = "MIT"
 build="build.rs"
 edition="2021"
 
+[build-dependencies]
+prost-build = { version = "0.14", optional = true }
+protobuf-src = { version = "2.0", optional = true }
+
 [dependencies]
 anyhow = "1"
 clap = {version="3.2", features=["wrap_help", "cargo", "derive"]}
@@ -42,6 +47,10 @@ rand = "0.8"
 rand_distr = "0.4"
 remoteprocess = "0.5.1"
 chrono = "0.4.26"
+prost = { version = "0.14", optional = true }
+prost-build = { version = "0.14", optional = true }
+prost-types = { version = "0.14", optional = true }
+flate2 = { version = "1.1.5", features = ["rust_backend"], optional = true }
 
 [dev-dependencies]
 py-spy-testdata = "0.1.0"

build.rs

@@ -9,4 +9,15 @@ fn main() {
         }
         _ => {}
     }
+
+    // Only compile protobuf if pprof feature is enabled
+    #[cfg(feature = "pprof")]
+    {
+        // Use bundled protoc from protobuf-src
+        std::env::set_var("PROTOC", protobuf_src::protoc());
+
+        // Add prost-build configuration
+        prost_build::compile_protos(&["src/pprof/profile.proto"], &["src/pprof/"])
+            .expect("Failed to compile protos");
+    }
 }

src/config.rs

@@ -68,6 +68,10 @@ pub enum FileFormat {
     raw,
     speedscope,
     chrometrace,
+    #[cfg(feature = "pprof")]
+    pprof,
+    #[cfg(feature = "pprof")]
+    pprof_gzip,
 }
 
 impl FileFormat {
@@ -369,6 +373,18 @@ impl Config {
             std::process::exit(1);
         }
 
+        // Check if pprof format was requested but pprof feature is not enabled
+        if !cfg!(feature = "pprof") {
+            if let Some(format_str) = matches.value_of("format") {
+                if format_str == "pprof" || format_str == "pprof_gzip" {
+                    eprintln!(
+                        "pprof output format is not supported. Please rebuild with --features pprof"
+                    );
+                    std::process::exit(1);
+                }
+            }
+        }
+
         match subcommand {
             "record" => {
                 config.sampling_rate = matches.value_of_t("rate")?;

src/main.rs

@@ -15,6 +15,8 @@ mod dump;
 mod flamegraph;
 #[cfg(feature = "unwind")]
 mod native_stack_trace;
+#[cfg(feature = "pprof")]
+mod pprof;
 mod python_bindings;
 mod python_data_access;
 mod python_interpreters;
@@ -118,6 +120,17 @@ impl Recorder for chrometrace::Chrometrace {
     }
 }
 
+#[cfg(feature = "pprof")]
+impl Recorder for pprof::Pprof {
+    fn increment(&mut self, trace: &StackTrace) -> Result<(), Error> {
+        Ok(self.increment(trace)?)
+    }
+
+    fn write(&self, w: &mut dyn Write) -> Result<(), Error> {
+        self.write(w)
+    }
+}
+
 pub struct RawFlamegraph(flamegraph::Flamegraph);
 
 impl Recorder for RawFlamegraph {
@@ -142,6 +155,10 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
         Some(FileFormat::chrometrace) => {
             Box::new(chrometrace::Chrometrace::new(config.show_line_numbers))
         }
+        #[cfg(feature = "pprof")]
+        Some(FileFormat::pprof) => Box::new(pprof::Pprof::new(false)),
+        #[cfg(feature = "pprof")]
+        Some(FileFormat::pprof_gzip) => Box::new(pprof::Pprof::new(true)),
         None => return Err(format_err!("A file format is required to record samples")),
     };
 
@@ -153,6 +170,10 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
                 Some(FileFormat::speedscope) => "json",
                 Some(FileFormat::raw) => "txt",
                 Some(FileFormat::chrometrace) => "json",
+                #[cfg(feature = "pprof")]
+                Some(FileFormat::pprof) => "pb",
+                #[cfg(feature = "pprof")]
+                Some(FileFormat::pprof_gzip) => "pb.gz",
                 None => return Err(format_err!("A file format is required to record samples")),
             };
             let local_time = Local::now().to_rfc3339_opts(SecondsFormat::Secs, true);
@@ -361,6 +382,14 @@ fn record_samples(pid: remoteprocess::Pid, config: &Config) -> Result<(), Error>
             );
             println!("{lede}Visit chrome://tracing or https://ui.perfetto.dev/ to view");
         }
+        #[cfg(feature = "pprof")]
+        FileFormat::pprof | FileFormat::pprof_gzip => {
+            println!(
+                "{lede}Wrote pprof profile to '{filename}'. Samples: {samples} Errors: {errors}"
+            );
+            println!("{lede}Use 'pprof -http=:8080 {filename}' to view");
+            println!("{lede}Or visit https://pprof.me/");
+        }
     };
 
     Ok(())