From d07547860b271093d7004c5ee8bfc6d25b10c065 Mon Sep 17 00:00:00 2001 From: Dimitris Fasarakis Hilliard Date: Wed, 23 Nov 2022 18:31:13 +0200 Subject: [PATCH] Fix issues with flame-it. --- src/lib.rs | 116 +++++++++++++++++++++---------------------- vm/src/vm/setting.rs | 12 +++++ 2 files changed, 70 insertions(+), 58 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 436a8f70b2..4036f0922c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,9 @@ extern crate env_logger; #[macro_use] extern crate log; +#[cfg(feature = "flame-it")] +use vm::Settings; + mod interpreter; mod settings; mod shell; @@ -80,70 +83,13 @@ pub fn run(init: impl FnOnce(&mut VirtualMachine) + 'static) -> ExitCode { config = config.init_stdlib(); } config = config.init_hook(Box::new(init)); - let interp = config.interpreter(); - - #[cfg(feature = "flame-it")] - let main_guard = flame::start_guard("RustPython main"); + let interp = config.interpreter(); let exitcode = interp.run(move |vm| run_rustpython(vm, run_mode)); - #[cfg(feature = "flame-it")] - { - main_guard.end(); - if let Err(e) = write_profile(&matches) { - error!("Error writing profile information: {}", e); - } - } ExitCode::from(exitcode) } -#[cfg(feature = "flame-it")] -fn write_profile(matches: &ArgMatches) -> Result<(), Box> { - use std::{fs, io}; - - 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); - // TODO: Need to change to ExitCode or Termination - std::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 = if profile_output == "-" { - Box::new(io::stdout()) - } else { - Box::new(fs::File::create(profile_output)?) - }; - - let profile_output = io::BufWriter::new(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)?, - } - - Ok(()) -} - fn setup_main_module(vm: &VirtualMachine) -> PyResult { let scope = vm.new_scope_with_builtins(); let main_module = vm.new_module("__main__", scope.globals.clone(), None); @@ -206,6 +152,9 @@ fn install_pip(_installer: &str, _scope: Scope, vm: &VirtualMachine) -> PyResult } fn run_rustpython(vm: &VirtualMachine, run_mode: RunMode) -> PyResult<()> { + #[cfg(feature = "flame-it")] + let main_guard = flame::start_guard("RustPython main"); + let scope = setup_main_module(vm)?; let site_result = vm.import("site", None, 0); @@ -244,6 +193,57 @@ fn run_rustpython(vm: &VirtualMachine, run_mode: RunMode) -> PyResult<()> { } } } + #[cfg(feature = "flame-it")] + { + main_guard.end(); + if let Err(e) = write_profile(&vm.state.as_ref().settings) { + error!("Error writing profile information: {}", e); + } + } + Ok(()) +} + +#[cfg(feature = "flame-it")] +fn write_profile(settings: &Settings) -> Result<(), Box> { + use std::{fs, io}; + + enum ProfileFormat { + Html, + Text, + Speedscope, + } + let profile_output = settings.profile_output.as_deref(); + let profile_format = match settings.profile_format.as_deref() { + 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); + // TODO: Need to change to ExitCode or Termination + std::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 = if profile_output == "-" { + Box::new(io::stdout()) + } else { + Box::new(fs::File::create(profile_output)?) + }; + + let profile_output = io::BufWriter::new(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)?, + } Ok(()) } diff --git a/vm/src/vm/setting.rs b/vm/src/vm/setting.rs index 62b0d6848c..53a4c59a38 100644 --- a/vm/src/vm/setting.rs +++ b/vm/src/vm/setting.rs @@ -1,3 +1,6 @@ +#[cfg(feature = "flame-it")] +use std::ffi::OsString; + /// Struct containing all kind of settings for the python vm. #[non_exhaustive] pub struct Settings { @@ -67,6 +70,11 @@ pub struct Settings { /// false for wasm. Not a command-line option pub allow_external_library: bool, + + #[cfg(feature = "flame-it")] + pub profile_output: Option, + #[cfg(feature = "flame-it")] + pub profile_format: Option, } /// Sensible default settings. @@ -95,6 +103,10 @@ impl Default for Settings { stdio_unbuffered: false, check_hash_based_pycs: "default".to_owned(), allow_external_library: cfg!(feature = "importlib"), + #[cfg(feature = "flame-it")] + profile_output: None, + #[cfg(feature = "flame-it")] + profile_format: None, } } }