Add profiling using flamescope/speedscope

coolreader18 · coolreader18 · commit ecdbe0408b1f · 2019-07-10T13:14:01.000-05:00
diff --git a/.gitignore b/.gitignore
@@ -10,4 +10,7 @@ __pycache__
 wasm-pack.log
 .idea/
 tests/snippets/resources
+
 flame-graph.html
+flame.txt
+flamescope.json
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -16,7 +16,7 @@ path = "./benchmarks/bench.rs"
 
 [features]
 default = []
-flame-it = ["rustpython-vm/flame-it", "flame"]
+flame-it = ["rustpython-vm/flame-it", "flame", "flamescope"]
 
 [dependencies]
 log="0.4.1"
@@ -29,6 +29,7 @@ rustyline = "4.1.0"
 xdg = "2.2.0"
 
 flame = { version = "0.2", optional = true }
+flamescope = { version = "0.1", optional = true }
 
 [dev-dependencies.cpython]
 version = "0.2"
diff --git a/src/main.rs b/src/main.rs
@@ -29,6 +29,8 @@ fn main() {
 }
 
 fn run() -> Result<(), Box<dyn std::error::Error>> {
+    #[cfg(feature = "flame-it")]
+    let main_guard = flame::start_guard("RustPython main");
     env_logger::init();
     let app = App::new("RustPython")
         .version(crate_version!())
@@ -55,15 +57,25 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
         )
         .arg(Arg::from_usage("[pyargs] 'args for python'").multiple(true));
     #[cfg(feature = "flame-it")]
-    let app = app.arg(
-        Arg::with_name("profile_output")
-            .long("profile-output")
-            .takes_value(true)
-            .help(
-                "the file to output the profile graph to. present due to being \
-                 built with feature 'flame-it'",
-            ),
-    );
+    let app = app
+        .arg(
+            Arg::with_name("profile_output")
+                .long("profile-output")
+                .takes_value(true)
+                .help(
+                    "the file to output the profiling information to. present due to being \
+                     built with feature 'flame-it'",
+                ),
+        )
+        .arg(
+            Arg::with_name("profile_format")
+                .long("profile-format")
+                .takes_value(true)
+                .help(
+                    "the profile format to output the profiling information in. present due to \
+                     being built with feature 'flame-it'",
+                ),
+        );
     let matches = app.get_matches();
 
     // Construct vm:
@@ -92,13 +104,43 @@ fn run() -> Result<(), Box<dyn std::error::Error>> {
     {
         use std::fs::File;
 
-        let profile_output = matches
-            .value_of_os("profile_output")
-            .unwrap_or_else(|| "flame-graph.html".as_ref());
-        if profile_output == "-" {
-            flame::dump_stdout();
+        main_guard.end();
+
+        enum ProfileFormat {
+            Html,
+            Text,
+            Speedscope,
+        }
+
+        let profile_output = matches.value_of_os("profile_output");
+
+        let profile_format = match matches.value_of("profile_format") {
+            Some("html") => ProfileFormat::Html,
+            Some("text") => ProfileFormat::Text,
+            None if profile_output == Some("-".as_ref()) => ProfileFormat::Text,
+            Some("speedscope") | None => ProfileFormat::Speedscope,
+            Some(other) => {
+                error!("Unknown profile format {}", other);
+                process::exit(1);
+            }
+        };
+
+        let profile_output = profile_output.unwrap_or_else(|| match profile_format {
+            ProfileFormat::Html => "flame-graph.html".as_ref(),
+            ProfileFormat::Text => "flame.txt".as_ref(),
+            ProfileFormat::Speedscope => "flamescope.json".as_ref(),
+        });
+
+        let profile_output: Box<dyn std::io::Write> = if profile_output == "-" {
+            Box::new(std::io::stdout())
         } else {
-            flame::dump_html(&mut File::create(profile_output)?)?;
+            Box::new(File::create(profile_output)?)
+        };
+
+        match profile_format {
+            ProfileFormat::Html => flame::dump_html(profile_output)?,
+            ProfileFormat::Text => flame::dump_text_to_writer(profile_output)?,
+            ProfileFormat::Speedscope => flamescope::dump(profile_output)?,
         }
     }
 
diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -271,11 +271,6 @@ impl Frame {
     // #[cfg_attr(feature = "flame-it", flame("Frame"))]
     pub fn run(&self, vm: &VirtualMachine) -> Result<ExecutionResult, PyObjectRef> {
         flame_guard!(format!("Frame::run({})", self.code.obj_name));
-        // flame doesn't include notes in the html graph :(
-        flame_note!(
-            flame::StrCow::from("CodeObj name"),
-            self.code.obj_name.clone().into()
-        );
 
         let filename = &self.code.source_path.to_string();
 
@@ -342,10 +337,11 @@ impl Frame {
 
     /// Execute a single instruction.
     #[allow(clippy::cognitive_complexity)]
-    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn execute_instruction(&self, vm: &VirtualMachine) -> FrameResult {
         let instruction = self.fetch_instruction();
 
+        flame_guard!(format!("Frame::execute_instruction({:?})", instruction));
+
         #[cfg(feature = "vm-tracing-logging")]
         {
             trace!("=======");
