Add flamegraph profiling

Author: coolreader18

.gitignore

@@ -9,4 +9,5 @@ __pycache__
 .vscode
 wasm-pack.log
 .idea/
-tests/snippets/resources
+tests/snippets/resources
+flame-graph.html

Cargo.lock

@@ -14,6 +14,9 @@ members = [".", "derive", "vm", "wasm/lib", "parser", "compiler", "bytecode"]
 name = "bench"
 path = "./benchmarks/bench.rs"
 
+[features]
+default = []
+flame-it = ["rustpython-vm/flame-it", "flame"]
 
 [dependencies]
 log="0.4.1"
@@ -25,5 +28,7 @@ rustpython-vm = {path = "vm", version = "0.1.0"}
 rustyline = "4.1.0"
 xdg = "2.2.0"
 
+flame = { version = "0.2", optional = true }
+
 [dev-dependencies.cpython]
 version = "0.2"

Cargo.toml

@@ -19,10 +19,18 @@ use rustpython_vm::{
 
 use rustyline::{error::ReadlineError, Editor};
 use std::path::PathBuf;
+use std::process;
 
 fn main() {
+    if let Err(err) = run() {
+        error!("Error: {}", err);
+        process::exit(1);
+    }
+}
+
+fn run() -> Result<(), Box<std::error::Error>> {
     env_logger::init();
-    let matches = App::new("RustPython")
+    let app = App::new("RustPython")
         .version(crate_version!())
         .author(crate_authors!())
         .about("Rust implementation of the Python language")
@@ -45,8 +53,18 @@ fn main() {
                 .takes_value(true)
                 .help("run library module as script"),
         )
-        .arg(Arg::from_usage("[pyargs] 'args for python'").multiple(true))
-        .get_matches();
+        .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 matches = app.get_matches();
 
     // Construct vm:
     let vm = VirtualMachine::new();
@@ -69,6 +87,22 @@ fn main() {
 
     // See if any exception leaked out:
     handle_exception(&vm, result);
+
+    #[cfg(feature = "flame-it")]
+    {
+        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();
+        } else {
+            flame::dump_html(&mut File::create(profile_output)?)?;
+        }
+    }
+
+    Ok(())
 }
 
 fn _run_string(vm: &VirtualMachine, source: &str, source_path: String) -> PyResult {
@@ -84,7 +118,7 @@ fn _run_string(vm: &VirtualMachine, source: &str, source_path: String) -> PyResu
 fn handle_exception(vm: &VirtualMachine, result: PyResult) {
     if let Err(err) = result {
         print_exception(vm, &err);
-        std::process::exit(1);
+        process::exit(1);
     }
 }
 
@@ -116,14 +150,14 @@ fn run_script(vm: &VirtualMachine, script_file: &str) -> PyResult {
                 "can't find '__main__' module in '{}'",
                 file_path.to_str().unwrap()
             );
-            std::process::exit(1);
+            process::exit(1);
         }
     } else {
         error!(
             "can't open file '{}': No such file or directory",
             file_path.to_str().unwrap()
         );
-        std::process::exit(1);
+        process::exit(1);
     };
 
     let dir = file_path.parent().unwrap().to_str().unwrap().to_string();
@@ -138,7 +172,7 @@ fn run_script(vm: &VirtualMachine, script_file: &str) -> PyResult {
                 file_path.to_str().unwrap(),
                 err.kind()
             );
-            std::process::exit(1);
+            process::exit(1);
         }
     }
 }

src/main.rs

@@ -11,6 +11,7 @@ include = ["src/**/*.rs", "Cargo.toml", "build.rs", "Lib/**/*.py"]
 [features]
 default = ["rustpython-parser", "rustpython-compiler"]
 vm-tracing-logging = []
+flame-it = ["flame", "flamer"]
 
 [dependencies]
 # Crypto:
@@ -57,5 +58,8 @@ maplit = "1.0"
 proc-macro-hack = "0.5"
 bitflags = "1.1"
 
+flame = { version = "0.2", optional = true }
+flamer = { version = "0.3", optional = true }
+
 [target.'cfg(all(unix, not(target_os = "android")))'.dependencies]
 pwd = "1"

vm/Cargo.toml

@@ -268,6 +268,7 @@ impl Frame {
         }
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     pub fn run(&self, vm: &VirtualMachine) -> Result<ExecutionResult, PyObjectRef> {
         let filename = &self.code.source_path.to_string();
 
@@ -334,6 +335,7 @@ 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();
 
@@ -891,6 +893,7 @@ impl Frame {
         }
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn get_elements(
         &self,
         vm: &VirtualMachine,
@@ -912,6 +915,7 @@ impl Frame {
         }
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn import(
         &self,
         vm: &VirtualMachine,
@@ -938,6 +942,7 @@ impl Frame {
         Ok(None)
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn import_star(&self, vm: &VirtualMachine, module: &str, level: usize) -> FrameResult {
         let module = vm.import(module, &vm.ctx.new_tuple(vec![]), level)?;
 
@@ -951,6 +956,7 @@ impl Frame {
     }
 
     // Unwind all blocks:
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn unwind_blocks(&self, vm: &VirtualMachine) -> Option<PyObjectRef> {
         while let Some(block) = self.pop_block() {
             match block.typ {
@@ -978,6 +984,7 @@ impl Frame {
         None
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn unwind_loop(&self, vm: &VirtualMachine) -> Block {
         loop {
             let block = self.current_block().expect("not in a loop");
@@ -1003,6 +1010,7 @@ impl Frame {
         }
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn unwind_exception(&self, vm: &VirtualMachine, exc: PyObjectRef) -> Option<PyObjectRef> {
         // unwind block stack on exception and find any handlers:
         while let Some(block) = self.pop_block() {
@@ -1105,6 +1113,7 @@ impl Frame {
         }
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn load_name(
         &self,
         vm: &VirtualMachine,
@@ -1150,6 +1159,7 @@ impl Frame {
         *self.lasti.borrow_mut() = target_pc;
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn execute_binop(
         &self,
         vm: &VirtualMachine,
@@ -1194,6 +1204,7 @@ impl Frame {
         Ok(None)
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn execute_unop(&self, vm: &VirtualMachine, op: &bytecode::UnaryOperator) -> FrameResult {
         let a = self.pop_value();
         let value = match *op {
@@ -1234,6 +1245,7 @@ impl Frame {
         Ok(result)
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn execute_compare(
         &self,
         vm: &VirtualMachine,
@@ -1326,6 +1338,7 @@ impl Frame {
         stack[stack.len() - depth - 1].clone()
     }
 
+    #[cfg_attr(feature = "flame-it", flame("Frame"))]
     fn get_exception(&self, vm: &VirtualMachine, none_allowed: bool) -> PyResult {
         let exception = self.pop_value();
         if none_allowed && vm.get_none().is(&exception)

vm/src/frame.rs

@@ -11,13 +11,15 @@ use crate::vm::VirtualMachine;
 use rustpython_compiler::compile;
 
 pub fn init_importlib(vm: &VirtualMachine, external: bool) -> PyResult {
+    flame_guard!("init importlib");
     let importlib = import_frozen(vm, "_frozen_importlib")?;
     let impmod = import_builtin(vm, "_imp")?;
     let install = vm.get_attribute(importlib.clone(), "_install")?;
     vm.invoke(install, vec![vm.sys_module.clone(), impmod])?;
     vm.import_func
         .replace(vm.get_attribute(importlib.clone(), "__import__")?);
     if external && cfg!(feature = "rustpython-compiler") {
+        flame_guard!("install_external");
         let install_external =
             vm.get_attribute(importlib.clone(), "_install_external_importers")?;
         vm.invoke(install_external, vec![])?;

vm/src/import.rs

@@ -12,6 +12,10 @@
     clippy::implicit_hasher
 )]
 
+#[cfg(feature = "flame-it")]
+#[macro_use]
+extern crate flamer;
+
 #[macro_use]
 extern crate bitflags;
 #[macro_use]
@@ -30,8 +34,7 @@ extern crate self as rustpython_vm;
 
 pub use rustpython_derive::*;
 
-use proc_macro_hack::proc_macro_hack;
-#[proc_macro_hack]
+#[proc_macro_hack::proc_macro_hack]
 pub use rustpython_derive::py_compile_bytecode;
 
 //extern crate eval; use eval::eval::*;

vm/src/lib.rs

@@ -266,3 +266,10 @@ macro_rules! vm_trace {
         trace!($($arg)+);
     }
 }
+
+macro_rules! flame_guard {
+    ($name:expr) => {
+        #[cfg(feature = "flame-it")]
+        let _guard = ::flame::start_guard($name);
+    };
+}