Merge pull request RustPython#4290 from DimitrisJim/flame_it

youknowone · web-flow · commit 3660ac20cff4 · 2022-11-24T20:09:31.000+09:00
Fix issues with flame-it.
diff --git 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<dyn std::error::Error>> {
-    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<dyn io::Write> = 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<Scope> {
     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<dyn std::error::Error>> {
+    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<dyn io::Write> = 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
@@ -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<OsString>,
+    #[cfg(feature = "flame-it")]
+    pub profile_format: Option<String>,
 }
 
 /// 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,
         }
     }
 }
