Add runtime benchmark profiling command

Author: Kobzol

collector/src/bin/collector.rs

@@ -31,8 +31,9 @@ use tokio::runtime::Runtime;
 use collector::compile::execute::bencher::BenchProcessor;
 use collector::compile::execute::profiler::{ProfileProcessor, Profiler};
 use collector::runtime::{
-    bench_runtime, runtime_benchmark_dir, BenchmarkFilter, BenchmarkSuite,
-    BenchmarkSuiteCompilation, CargoIsolationMode, DEFAULT_RUNTIME_ITERATIONS,
+    bench_runtime, prepare_runtime_benchmark_suite, profile_runtime, runtime_benchmark_dir,
+    BenchmarkFilter, BenchmarkSuite, BenchmarkSuiteCompilation, CargoIsolationMode,
+    DEFAULT_RUNTIME_ITERATIONS,
 };
 use collector::toolchain::{
     create_toolchain_from_published_version, get_local_toolchain, Sysroot, Toolchain,
@@ -245,7 +246,7 @@ fn cg_annotate(cgout: &Path, path: &Path) -> anyhow::Result<()> {
 }
 
 #[allow(clippy::too_many_arguments)]
-fn profile(
+fn profile_compile(
     toolchain: &Toolchain,
     profiler: Profiler,
     out_dir: &Path,
@@ -492,6 +493,16 @@ enum Commands {
         #[arg(long = "no-isolate")]
         no_isolate: bool,
     },
+
+    /// Profiles a runtime benchmark.
+    ProfileRuntime {
+        #[command(flatten)]
+        local: LocalOptions,
+
+        /// Name of the benchmark that should be profiled
+        benchmark: String,
+    },
+
     /// Benchmarks a local rustc
     BenchLocal {
         #[command(flatten)]
@@ -640,15 +651,7 @@ fn main_result() -> anyhow::Result<i32> {
             no_isolate,
         } => {
             log_db(&db);
-            let toolchain = get_local_toolchain(
-                &[Profile::Opt],
-                &local.rustc,
-                None,
-                local.cargo.as_deref(),
-                local.id.as_deref(),
-                "",
-                target_triple,
-            )?;
+            let toolchain = get_local_toolchain_for_runtime_benchmarks(&local, &target_triple)?;
             let pool = Pool::open(&db.db);
 
             let isolation_mode = if no_isolate {
@@ -679,6 +682,17 @@ fn main_result() -> anyhow::Result<i32> {
             run_benchmarks(&mut rt, conn, shared, None, Some(config))?;
             Ok(0)
         }
+        Commands::ProfileRuntime { local, benchmark } => {
+            let toolchain = get_local_toolchain_for_runtime_benchmarks(&local, &target_triple)?;
+            let suite = prepare_runtime_benchmark_suite(
+                &toolchain,
+                &runtime_benchmark_dir,
+                CargoIsolationMode::Cached,
+            )?
+            .suite;
+            profile_runtime(suite, &benchmark)?;
+            Ok(0)
+        }
         Commands::BenchLocal {
             local,
             opts,
@@ -894,7 +908,7 @@ fn main_result() -> anyhow::Result<i32> {
                         target_triple.clone(),
                     )?;
                     let id = toolchain.id.clone();
-                    profile(
+                    profile_compile(
                         &toolchain,
                         profiler,
                         &out_dir,
@@ -995,6 +1009,21 @@ Make sure to modify `{dir}/perf-config.json` if the category/artifact don't matc
     }
 }
 
+fn get_local_toolchain_for_runtime_benchmarks(
+    local: &LocalOptions,
+    target_triple: &str,
+) -> anyhow::Result<Toolchain> {
+    get_local_toolchain(
+        &[Profile::Opt],
+        &local.rustc,
+        None,
+        local.cargo.as_deref(),
+        local.id.as_deref(),
+        "",
+        target_triple.to_string(),
+    )
+}
+
 async fn load_runtime_benchmarks(
     conn: &mut dyn Connection,
     benchmark_dir: &Path,

collector/src/runtime/benchmark.rs

@@ -59,6 +59,15 @@ impl BenchmarkSuite {
             .iter()
             .flat_map(|suite| suite.benchmark_names.iter().map(|n| n.as_ref()))
     }
+
+    pub fn get_group_by_benchmark(&self, benchmark: &str) -> Option<&BenchmarkGroup> {
+        self.groups.iter().find(|group| {
+            group
+                .benchmark_names
+                .iter()
+                .any(|b| b.as_str() == benchmark)
+        })
+    }
 }
 
 pub struct BenchmarkFilter {

collector/src/runtime/mod.rs

@@ -1,3 +1,4 @@
+use std::ffi::OsStr;
 use std::io::{BufRead, BufReader, Cursor};
 use std::path::Path;
 use std::process::{Command, Stdio};
@@ -13,7 +14,7 @@ pub use benchmark::{
 use database::{ArtifactIdNumber, CollectionId, Connection};
 
 use crate::utils::git::get_rustc_perf_commit;
-use crate::{run_command_with_output, CollectorCtx};
+use crate::{run_command, run_command_with_output, CollectorCtx};
 
 mod benchmark;
 
@@ -105,6 +106,30 @@ pub async fn bench_runtime(
     Ok(())
 }
 
+pub fn profile_runtime(suite: BenchmarkSuite, benchmark: &str) -> anyhow::Result<()> {
+    let Some(group) = suite.get_group_by_benchmark(benchmark) else {
+        return Err(anyhow::anyhow!("Benchmark `{benchmark}` not found"));
+    };
+    let mut cmd = prepare_command(&group.binary);
+    cmd.arg("profile").arg(benchmark);
+    run_command(&mut cmd)?;
+    Ok(())
+}
+
+/// Prepares a command for execution, adding some shared flags.
+fn prepare_command<S: AsRef<OsStr>>(binary: S) -> Command {
+    // Turn off ASLR
+    let mut command = Command::new("setarch");
+    command.arg(std::env::consts::ARCH).arg("-R").arg(binary);
+
+    // We want to see a backtrace if the program panics
+    command.env("RUST_BACKTRACE", "1");
+
+    command.stdout(Stdio::piped());
+    command.stderr(Stdio::piped());
+    command
+}
+
 /// Records the results (stats) of a benchmark into the database.
 async fn record_stats(
     conn: &dyn Connection,
@@ -191,28 +216,18 @@ fn execute_runtime_benchmark_binary(
     filter: &BenchmarkFilter,
     iterations: u32,
 ) -> anyhow::Result<impl Iterator<Item = anyhow::Result<BenchmarkMessage>>> {
-    // Turn off ASLR
-    let mut command = Command::new("setarch");
-    command.arg(std::env::consts::ARCH);
-    command.arg("-R");
-    command.arg(binary);
+    let mut command = prepare_command(binary);
     command.arg("run");
     command.arg("--iterations");
     command.arg(&iterations.to_string());
 
-    // We want to see a backtrace if the benchmark panics
-    command.env("RUST_BACKTRACE", "1");
-
     if let Some(ref exclude) = filter.exclude {
         command.args(["--exclude", exclude]);
     }
     if let Some(ref include) = filter.include {
         command.args(["--include", include]);
     }
 
-    command.stdout(Stdio::piped());
-    command.stderr(Stdio::piped());
-
     let output = run_command_with_output(&mut command)?;
     if !output.status.success() {
         return Err(anyhow::anyhow!(