feat(go-runner): optimize walltime result parsing

go-runner/Cargo.toml

@@ -22,6 +22,7 @@ thiserror = "2.0"
 tempfile = "3.14"
 semver = "1.0.27"
 dircpy = "0.3.19"
+rayon = "1"
 
 [dev-dependencies]
 divan = { version = "4.1.0", package = "codspeed-divan-compat" }

go-runner/benches/go_runner.rs

@@ -1,4 +1,5 @@
 use codspeed_go_runner::results::raw_result::RawResult;
+use std::time::Duration;
 use tempfile::TempDir;
 
 #[divan::bench(max_time = std::time::Duration::from_secs(5))]
@@ -30,12 +31,15 @@ fn bench_go_runner(bencher: divan::Bencher) {
 }
 
 const TIME_ENTRIES: [usize; 5] = [100_000, 500_000, 1_000_000, 5_000_000, 10_000_000];
+const FILE_COUNT: [usize; 3] = [5, 10, 25];
 
-#[divan::bench(consts = TIME_ENTRIES, max_time = std::time::Duration::from_secs(5))]
-fn bench_collect_results<const N: usize>(bencher: divan::Bencher) {
-    fn random_raw_result<const N: usize>() -> RawResult {
-        let times_per_round = (0..N).map(|_| rand::random::<u64>() % 1_000_000).collect();
-        let iters_per_round = (0..N).map(|_| rand::random::<u64>() % 1_000 + 1).collect();
+#[divan::bench(args = FILE_COUNT, consts = TIME_ENTRIES, max_time = Duration::from_secs(5))]
+fn bench_collect_results<const N: usize>(bencher: divan::Bencher, file_count: usize) {
+    use rand::prelude::*;
+
+    fn random_raw_result<const N: usize>(rng: &mut StdRng) -> RawResult {
+        let times_per_round = (0..N).map(|_| rng.random::<u64>() % 1_000_000).collect();
+        let iters_per_round = (0..N).map(|_| rng.random::<u64>() % 1_000 + 1).collect();
         RawResult {
             name: "foo".into(),
             uri: "foo".into(),
@@ -47,11 +51,16 @@ fn bench_collect_results<const N: usize>(bencher: divan::Bencher) {
 
     bencher
         .with_inputs(|| {
+            let mut rng = StdRng::seed_from_u64(42);
+
             let profile_dir = TempDir::new().unwrap();
             let raw_results = profile_dir.path().join("raw_results");
             std::fs::create_dir(&raw_results).unwrap();
 
-            for (i, raw_result) in (0..10).map(|_| random_raw_result::<N>()).enumerate() {
+            for (i, raw_result) in (0..file_count)
+                .map(|_| random_raw_result::<N>(&mut rng))
+                .enumerate()
+            {
                 let json = serde_json::to_string(&raw_result).unwrap();
                 std::fs::write(raw_results.join(format!("{i}.json")), json).unwrap();
             }

go-runner/src/builder/patcher.rs

@@ -1,6 +1,8 @@
 //! Patches the imports to use codspeed rather than the official "testing" package.
 
 use crate::prelude::*;
+use itertools::Itertools;
+use rayon::prelude::*;
 use std::fs;
 use std::path::Path;
 use std::process::Command;
@@ -34,47 +36,56 @@ pub fn patch_imports<P: AsRef<Path>>(folder: P) -> anyhow::Result<()> {
     debug!("Patching imports in folder: {folder:?}");
 
     // 1. Find all imports that match "testing" and replace them with codspeed equivalent
-    let mut patched_files = 0;
-
     let pattern = folder.join("**/*.go");
-    for go_file in glob::glob(pattern.to_str().unwrap())?.filter_map(Result::ok) {
-        // Skip directories - glob can match directories ending in .go (e.g., vendor/github.com/nats-io/nats.go)
-        if !go_file.is_file() {
-            continue;
-        }
-
-        let content =
-            fs::read_to_string(&go_file).context(format!("Failed to read Go file: {go_file:?}"))?;
-
-        let patched_content = patch_imports_for_source(&content);
-        if patched_content != content {
-            fs::write(&go_file, patched_content)
-                .context(format!("Failed to write patched Go file: {go_file:?}"))?;
-
-            debug!("Patched imports in: {go_file:?}");
-            patched_files += 1;
-        }
-    }
-    debug!("Patched {patched_files} files");
+    let patched_files = glob::glob(pattern.to_str().unwrap())?
+        .par_bridge()
+        .filter_map(Result::ok)
+        .filter_map(|go_file| {
+            // Skip directories - glob can match directories ending in .go (e.g., vendor/github.com/nats-io/nats.go)
+            if !go_file.is_file() {
+                return None;
+            }
+
+            let Ok(content) = fs::read_to_string(&go_file) else {
+                error!("Failed to read Go file: {go_file:?}");
+                return None;
+            };
+
+            let patched_content = patch_imports_for_source(&content);
+            if patched_content != content {
+                let Ok(_) = fs::write(&go_file, patched_content) else {
+                    error!("Failed to write patched Go file: {go_file:?}");
+                    return None;
+                };
+
+                debug!("Patched imports in: {go_file:?}");
+            }
+            Some(())
+        })
+        .count();
+    debug!("Patched {} files", patched_files);
 
     Ok(())
 }
 
 /// Internal function to apply import patterns to Go source code
 pub fn patch_imports_for_source(source: &str) -> String {
-    let replace_import = |mut source: String, import_path: &str, replacement: &str| -> String {
+    let mut source = source.to_string();
+
+    // If we can't parse the source, skip this replacement
+    // This can happen with template files or malformed Go code
+    let parsed = match gosyn::parse_source(&source) {
+        Ok(p) => p,
+        Err(_) => return source,
+    };
+
+    let mut replacements = vec![];
+    let mut find_replace_range = |import_path: &str, replacement: &str| {
         // Optimization: check if the import path exists in the source before parsing
         if !source.contains(import_path) {
-            return source;
+            return;
         }
 
-        // If we can't parse the source, skip this replacement
-        // This can happen with template files or malformed Go code
-        let parsed = match gosyn::parse_source(&source) {
-            Ok(p) => p,
-            Err(_) => return source,
-        };
-
         if let Some(import) = parsed
             .imports
             .iter()
@@ -83,39 +94,43 @@ pub fn patch_imports_for_source(source: &str) -> String {
             let start_pos = import.path.pos;
             let end_pos = start_pos + import.path.value.len();
 
-            source.replace_range(start_pos..end_pos, replacement);
+            replacements.push((start_pos..end_pos, replacement.to_string()));
         }
-
-        source
     };
 
-    let mut source = replace_import(
-        source.to_string(),
-        "testing",
-        "testing \"github.com/CodSpeedHQ/codspeed-go/testing/testing\"",
-    );
-
     // Then replace sub-packages like "testing/synctest"
     for testing_pkg in &["fstest", "iotest", "quick", "slogtest", "synctest"] {
-        source = replace_import(
-            source.to_string(),
+        find_replace_range(
             &format!("testing/{}", testing_pkg),
             &format!(
                 "{testing_pkg} \"github.com/CodSpeedHQ/codspeed-go/testing/testing/{testing_pkg}\""
             ),
         );
     }
 
-    let source = replace_import(
-        source,
+    find_replace_range(
+        "testing",
+        "testing \"github.com/CodSpeedHQ/codspeed-go/testing/testing\"",
+    );
+    find_replace_range(
         "github.com/thejerf/slogassert",
         "\"github.com/CodSpeedHQ/codspeed-go/pkg/slogassert\"",
     );
-    replace_import(
-        source,
+    find_replace_range(
         "github.com/frankban/quicktest",
         "\"github.com/CodSpeedHQ/codspeed-go/pkg/quicktest\"",
-    )
+    );
+
+    // Apply replacements in reverse order to avoid shifting positions
+    for (range, replacement) in replacements
+        .into_iter()
+        .sorted_by_key(|(range, _)| range.start)
+        .rev()
+    {
+        source.replace_range(range, &replacement);
+    }
+
+    source
 }
 
 /// Patches imports and package in specific test files
@@ -381,6 +396,17 @@ func BenchmarkExample(b *testing.B) {
 func TestExample(t *testing.T) {
     s := "package main"
 }
+"#;
+
+    const MANY_TESTING_IMPORTS: &str = r#"package subpackages
+import (
+	"bytes"
+	"io"
+	"testing"
+	"testing/fstest"
+	"testing/iotest"
+	"testing/synctest"
+)
 "#;
 
     #[rstest]
@@ -399,6 +425,7 @@ func TestExample(t *testing.T) {
         MULTILINE_IMPORT_WITH_TESTING_STRING
     )]
     #[case("package_main", PACKAGE_MAIN)]
+    #[case("many_testing_imports", MANY_TESTING_IMPORTS)]
     fn test_patch_go_source(#[case] test_name: &str, #[case] source: &str) {
         let result = patch_imports_for_source(source);
         let result = patch_package_for_source(result).unwrap();

go-runner/src/builder/snapshots/codspeed_go_runner__builder__patcher__tests__many_testing_imports.snap

@@ -0,0 +1,13 @@
+---
+source: go-runner/src/builder/patcher.rs
+expression: result
+---
+package subpackages
+import (
+	"bytes"
+	"io"
+	testing "github.com/CodSpeedHQ/codspeed-go/testing/testing"
+	fstest "github.com/CodSpeedHQ/codspeed-go/testing/testing/fstest"
+	iotest "github.com/CodSpeedHQ/codspeed-go/testing/testing/iotest"
+	synctest "github.com/CodSpeedHQ/codspeed-go/testing/testing/synctest"
+)

go-runner/src/builder/templater.rs

@@ -25,13 +25,23 @@ struct TemplateData {
 ///
 /// # Returns
 ///
-/// The path to the generated runner.go file. This should be passed to the `build_binary` function to build
-/// the binary that will execute the benchmarks.
-pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow::Result<PathBuf> {
+/// - `TempDir`: The temporary directory containing the modified Go project. This directory will be automatically deleted when dropped (only in tests).
+/// - `PathBuf`: The path to the generated runner.go file. This should be passed to the `build_binary` function to build
+///   the binary that will execute the benchmarks.
+pub fn run<P: AsRef<Path>>(
+    package: &BenchmarkPackage,
+    profile_dir: P,
+) -> anyhow::Result<(TempDir, PathBuf)> {
     // Create a temporary target directory for building the modified Go project.
-    // NOTE: We don't want to spend time cleanup any temporary files since the code is only
+    let mut target_dir = TempDir::new()?;
+
+    // We don't want to spend time cleanup any temporary files since the code is only
     // run on CI servers which clean up themselves.
-    let target_dir = TempDir::new()?.keep();
+    // However, when running tests we don't want to fill the disk with temporary files, which
+    // can cause the tests to fail due to lack of disk space.
+    if cfg!(not(test)) {
+        target_dir.disable_cleanup(true);
+    }
 
     // 1. Copy the whole git repository to a build directory
     let git_root = if let Ok(git_dir) = utils::get_parent_git_repo_path(&package.module.dir) {
@@ -60,7 +70,7 @@ pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow
         relative_package_path,
     };
     fs::write(
-        target_dir.join("go-runner.metadata"),
+        target_dir.path().join("go-runner.metadata"),
         serde_json::to_string_pretty(&metadata)?,
     )
     .context("Failed to write go-runner.metadata file")?;
@@ -81,7 +91,7 @@ pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow
     // 2. Patch the imports and package of the test files
     // - Renames package declarations (to support main package tests and external tests)
     // - Fixes imports to use our compat packages (e.g., testing/quicktest/testify)
-    let package_path = target_dir.join(relative_package_path);
+    let package_path = target_dir.path().join(relative_package_path);
     let test_file_paths: Vec<PathBuf> = files.iter().map(|f| package_path.join(f)).collect();
 
     // If we have external tests (e.g. "package {pkg}_test") they have to be
@@ -102,15 +112,18 @@ pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow
     // Find the module directory by getting the relative path from git root
     let module_dir = Path::new(&package.module.dir)
         .strip_prefix(&git_root)
-        .map(|relative_module_path| target_dir.join(relative_module_path))
+        .map(|relative_module_path| target_dir.path().join(relative_module_path))
         .unwrap_or_else(|_| {
             // Fall back to target_dir if we can't calculate relative path
-            target_dir.to_path_buf()
+            target_dir.path().to_path_buf()
         });
     patcher::install_codspeed_dependency(&module_dir)?;
 
     // 4. Handle test files differently based on whether they're external or internal tests
-    let codspeed_dir = target_dir.join(relative_package_path).join("codspeed");
+    let codspeed_dir = target_dir
+        .path()
+        .join(relative_package_path)
+        .join("codspeed");
     fs::create_dir_all(&codspeed_dir).context("Failed to create codspeed directory")?;
 
     if package.is_external_test_package() {
@@ -119,7 +132,7 @@ pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow
         // They're now package main and will be built from the subdirectory
         debug!("Handling external test package - moving files to codspeed/ subdirectory");
         for file in files {
-            let src_path = target_dir.join(relative_package_path).join(file);
+            let src_path = target_dir.path().join(relative_package_path).join(file);
             // Rename _test.go to _codspeed.go so it's not treated as a test file
             let dst_filename = file.replace("_test.go", "_codspeed.go");
             let dst_path = codspeed_dir.join(&dst_filename);
@@ -132,7 +145,7 @@ pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow
         // For internal test packages: rename _test.go to _codspeed.go in place
         debug!("Handling internal test package - renaming files in place");
         for file in files {
-            let old_path = target_dir.join(relative_package_path).join(file);
+            let old_path = target_dir.path().join(relative_package_path).join(file);
             let new_path = old_path.with_file_name(
                 old_path
                     .file_name()
@@ -162,5 +175,5 @@ pub fn run<P: AsRef<Path>>(package: &BenchmarkPackage, profile_dir: P) -> anyhow
     let runner_path = codspeed_dir.join("runner.go");
     fs::write(&runner_path, rendered).context("Failed to write runner.go file")?;
 
-    Ok(runner_path)
+    Ok((target_dir, runner_path))
 }

go-runner/src/lib.rs

@@ -1,4 +1,8 @@
-use crate::{builder::BenchmarkPackage, prelude::*};
+use crate::{
+    builder::BenchmarkPackage,
+    prelude::*,
+    results::{raw_result::RawResult, walltime_results::WalltimeBenchmark},
+};
 use std::{collections::HashMap, path::Path};
 
 pub mod builder;
@@ -35,7 +39,7 @@ pub fn run_benchmarks<P: AsRef<Path>>(
     // 2. Generate codspeed runners, build binaries, and execute them
     for package in &packages {
         info!("Generating custom runner for package: {}", package.name);
-        let runner_path = builder::templater::run(package, &profile_dir)?;
+        let (_target_dir, runner_path) = builder::templater::run(package, &profile_dir)?;
 
         info!("Building binary for package: {}", package.name);
 
@@ -74,16 +78,13 @@ pub fn run_benchmarks<P: AsRef<Path>>(
 
 // TODO: This should be merged with codspeed-rust/codspeed/walltime_results.rs
 pub fn collect_walltime_results(profile_dir: &Path) -> anyhow::Result<()> {
-    let raw_results = results::raw_result::RawResult::parse_folder(profile_dir)?;
-    info!("Parsed {} raw results", raw_results.len());
+    let mut benchmarks_by_pid: HashMap<u32, Vec<WalltimeBenchmark>> = HashMap::new();
 
-    let mut benchmarks_by_pid: HashMap<u32, Vec<results::walltime_results::WalltimeBenchmark>> =
-        HashMap::new();
-    for raw in raw_results {
+    for (pid, walltime_result) in RawResult::parse_folder(profile_dir)?.into_iter() {
         benchmarks_by_pid
-            .entry(raw.pid)
+            .entry(pid)
             .or_default()
-            .push(raw.into_walltime_benchmark());
+            .push(walltime_result);
     }
 
     for (pid, walltime_benchmarks) in benchmarks_by_pid {