diff --git a/Cargo.lock b/Cargo.lock index 6440a38a..7f9ddf78 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -235,6 +235,8 @@ dependencies = [ "insta", "itertools", "log", + "mimalloc", + "once_cell", "rand", "rayon", "regex", @@ -749,6 +751,16 @@ version = "0.2.175" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543" +[[package]] +name = "libmimalloc-sys" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "667f4fec20f29dfc6bc7357c582d91796c169ad7e2fce709468aefeb2c099870" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "linux-raw-sys" version = "0.9.4" @@ -782,6 +794,15 @@ version = "2.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" +[[package]] +name = "mimalloc" +version = "0.1.48" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1ee66a4b64c74f4ef288bcbb9192ad9c3feaad75193129ac8509af543894fd8" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "nix" version = "0.30.1" diff --git a/go-runner/Cargo.toml b/go-runner/Cargo.toml index 7c1dfce4..cbc2f91b 100644 --- a/go-runner/Cargo.toml +++ b/go-runner/Cargo.toml @@ -23,6 +23,8 @@ tempfile = "3.14" semver = "1.0.27" dircpy = "0.3.19" rayon = "1" +once_cell = "1.21.3" +mimalloc = "0.1.48" [dev-dependencies] divan = { version = "4.1.0", package = "codspeed-divan-compat" } diff --git a/go-runner/src/builder/patcher.rs b/go-runner/src/builder/patcher.rs index 47a6d346..aedb90ae 100644 --- a/go-runner/src/builder/patcher.rs +++ b/go-runner/src/builder/patcher.rs @@ -4,225 +4,283 @@ use crate::prelude::*; use itertools::Itertools; use rayon::prelude::*; use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::process::Command; -pub fn replace_pkg>(folder: P) -> anyhow::Result<()> { - let codspeed_root = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap(); - let replace_arg = format!( - "github.com/CodSpeedHQ/codspeed-go={}", - codspeed_root.display() - ); - debug!("Replacing codspeed-go with {}", codspeed_root.display()); - - let output = Command::new("go") - .args(["mod", "edit", "-replace", &replace_arg]) - .current_dir(folder.as_ref()) - .output() - .context("Failed to execute 'go mod edit' command")?; +/// Patcher is responsible for patching Go source files to replace imports and package names. +/// +/// It also reverts all changes on drop to avoid breaking tests. This can happen when we +/// rename or move a file, which makes it available in other packages which could lead to duplicate +/// symbols or flag definitions. +pub struct Patcher { + git_repo: PathBuf, +} - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - bail!("Failed to add replace directive: {}", stderr); +impl Patcher { + pub fn new(git_root: &Path) -> Self { + Self { + git_repo: git_root.to_path_buf(), + } } - debug!("Added local replace directive to go.mod"); + /// Replace `package main` with `package main_compat` to allow importing it from other packages. + /// Also replace `package foo_test` with `package main` for external test packages. + /// + /// Returns the previous package name, or None if no replacement was made. + fn patch_package( + source: &mut String, + replacement: Option, + ) -> anyhow::Result> { + let parsed = gosyn::parse_source(&source)?; + let prev_pkg_name = parsed.pkg_name.name; + + let replacement = replacement.or_else(|| { + if prev_pkg_name == "main" { + Some("main_compat".into()) + } else if prev_pkg_name.ends_with("_test") { + // For external test packages (package foo_test), convert to package main + // They will be placed in the codspeed/ subdirectory and built as standalone executables + Some("main".into()) + } else { + None + } + }); - Ok(()) -} + if let Some(new_name) = replacement { + let name_start = parsed.pkg_name.pos; + let name_end = name_start + prev_pkg_name.len(); + source.replace_range(name_start..name_end, &new_name); -pub fn patch_imports>(folder: P) -> anyhow::Result<()> { - let folder = folder.as_ref(); - debug!("Patching imports in folder: {folder:?}"); - - // 1. Find all imports that match "testing" and replace them with codspeed equivalent - let pattern = folder.join("**/*.go"); - 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) + Ok(Some(prev_pkg_name)) + } else { + Ok(None) + } + } + + /// Patches imports and package in specific test files + /// + /// This ensures we only modify the test files that belong to the current test package, + /// avoiding conflicts when multiple test packages exist in the same directory + pub fn patch_packages_for_files(&mut self, files: &[PathBuf]) -> anyhow::Result<()> { + for go_file in files { if !go_file.is_file() { - return None; + continue; } - 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:?}"); + let mut content = fs::read_to_string(go_file) + .context(format!("Failed to read Go file: {go_file:?}"))?; + if Self::patch_package(&mut content, None)?.is_some() { + fs::write(go_file, content) + .context(format!("Failed to write patched Go file: {go_file:?}"))?; } - Some(()) - }) - .count(); - debug!("Patched {} files", patched_files); + } - Ok(()) -} + Ok(()) + } -/// Internal function to apply import patterns to Go source code -pub fn patch_imports_for_source(source: &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; - } + /// Patches all .go files in a directory to rename "package main" to "package main_compat" + /// + /// This is needed when we have a "package main" with benchmarks that need to be imported. + /// By renaming all files in the package to "main_compat", we make it importable. + pub fn patch_all_packages_in_dir>(&mut self, dir: P) -> anyhow::Result<()> { + self.patch_packages_for_files( + &glob::glob(&dir.as_ref().join("*.go").to_string_lossy())? + .filter_map(Result::ok) + .collect::>(), + )?; + + Ok(()) + } - if let Some(import) = parsed - .imports - .iter() - .find(|import| import.path.value == format!("\"{import_path}\"")) - { - let start_pos = import.path.pos; - let end_pos = start_pos + import.path.value.len(); + fn patch_imports_for_source(source: &mut String) -> bool { + let mut modified = false; + + // 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 modified, + }; + + 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; + } - replacements.push((start_pos..end_pos, replacement.to_string())); + if let Some(import) = parsed + .imports + .iter() + .find(|import| import.path.value == format!("\"{import_path}\"")) + { + let start_pos = import.path.pos; + let end_pos = start_pos + import.path.value.len(); + modified = true; + + replacements.push((start_pos..end_pos, replacement.to_string())); + } + }; + + // Then replace sub-packages like "testing/synctest" + for testing_pkg in &["fstest", "iotest", "quick", "slogtest", "synctest"] { + find_replace_range( + &format!("testing/{}", testing_pkg), + &format!( + "{testing_pkg} \"github.com/CodSpeedHQ/codspeed-go/testing/testing/{testing_pkg}\"" + ), + ); } - }; - // Then replace sub-packages like "testing/synctest" - for testing_pkg in &["fstest", "iotest", "quick", "slogtest", "synctest"] { find_replace_range( - &format!("testing/{}", testing_pkg), - &format!( - "{testing_pkg} \"github.com/CodSpeedHQ/codspeed-go/testing/testing/{testing_pkg}\"" - ), + "testing", + "testing \"github.com/CodSpeedHQ/codspeed-go/testing/testing\"", + ); + find_replace_range( + "github.com/thejerf/slogassert", + "\"github.com/CodSpeedHQ/codspeed-go/pkg/slogassert\"", + ); + find_replace_range( + "github.com/frankban/quicktest", + "\"github.com/CodSpeedHQ/codspeed-go/pkg/quicktest\"", ); - } - 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\"", - ); - 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); + } - // 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); + modified } - source -} + pub fn patch_imports>(&mut self, folder: P) -> anyhow::Result<()> { + let folder = folder.as_ref(); + debug!("Patching imports in folder: {folder:?}"); + + // 1. Find all imports that match "testing" and replace them with codspeed equivalent + let pattern = folder.join("**/*.go"); + 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; + } -/// Patches imports and package in specific test files -/// -/// This ensures we only modify the test files that belong to the current test package, -/// avoiding conflicts when multiple test packages exist in the same directory -pub fn patch_packages_for_test_files>(test_files: &[P]) -> anyhow::Result<()> { - debug!("Patching {} test files", test_files.len()); - - let mut patched_files = 0; - for go_file in test_files { - let go_file = go_file.as_ref(); - if !go_file.is_file() { - continue; - } + let Ok(mut content) = fs::read_to_string(&go_file) else { + error!("Failed to read Go file: {go_file:?}"); + return None; + }; - let content = - fs::read_to_string(go_file).context(format!("Failed to read Go file: {go_file:?}"))?; + if Self::patch_imports_for_source(&mut content) { + let Ok(_) = fs::write(&go_file, &content) else { + error!("Failed to write patched Go file: {go_file:?}"); + return None; + }; - let patched_content = patch_package_for_source(content.clone())?; - 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:?}"); + } + Some(()) + }) + .count(); + debug!("Patched {} files", patched_files); - debug!("Patched package in: {go_file:?}"); - patched_files += 1; - } + Ok(()) } - debug!("Patched {patched_files} files"); - Ok(()) -} - -/// Patches all .go files in a directory to rename "package main" to "package main_compat" -/// -/// This is needed when we have a "package main" with benchmarks that need to be imported. -/// By renaming all files in the package to "main_compat", we make it importable. -pub fn patch_all_packages_in_dir>(dir: P) -> anyhow::Result<()> { - let dir = dir.as_ref(); - debug!("Patching all .go files in directory: {dir:?}"); - - let mut patched_files = 0; - let pattern = dir.join("*.go"); - for go_file in glob::glob(pattern.to_str().unwrap())?.filter_map(Result::ok) { - if !go_file.is_file() { - continue; + pub fn rename_test_files>(&mut self, files: &[P]) -> anyhow::Result<()> { + for file in files { + let src_path = file.as_ref(); + let new_path = src_path.with_file_name( + src_path + .file_name() + .unwrap() + .to_string_lossy() + .replace("_test", "_codspeed"), + ); + fs::rename(src_path, new_path.clone())?; } - let content = - fs::read_to_string(&go_file).context(format!("Failed to read Go file: {go_file:?}"))?; - - let patched_content = patch_package_for_source(content.clone())?; - if patched_content != content { - fs::write(&go_file, patched_content) - .context(format!("Failed to write patched Go file: {go_file:?}"))?; + Ok(()) + } - debug!("Patched package in: {go_file:?}"); - patched_files += 1; + pub fn rename_and_move_test_files>( + &mut self, + files: &[P], + dst_dir: &P, + ) -> anyhow::Result<()> { + for src_path in files { + let dst_filename = src_path + .as_ref() + .file_name() + .unwrap() + .to_string_lossy() + .replace("_test.go", "_codspeed.go"); + let dst_path = dst_dir.as_ref().join(&dst_filename); + fs::rename(src_path.as_ref(), &dst_path)?; } + + Ok(()) } - debug!("Patched {patched_files} files in directory"); +} - Ok(()) +impl Drop for Patcher { + fn drop(&mut self) { + let repo_path = &self.git_repo; + + let execute_cmd = |cmd: &str| { + debug!("Executing {cmd:?} in {repo_path:?}"); + + let output = std::process::Command::new("bash") + .args(["-c", cmd]) + .current_dir(repo_path) + .output(); + if let Ok(output) = output { + if !output.status.success() { + error!( + "Failed to execute cmd: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + } else { + panic!("Failed to execute command: {cmd:?}"); + } + }; + + execute_cmd("git reset --hard"); + execute_cmd("git clean -fd"); + execute_cmd("git submodule foreach git reset --hard"); + execute_cmd("git submodule foreach git clean -fd"); + } } -/// Replace `package main` with `package main_compat` to allow importing it from other packages. -/// Also replace `package foo_test` with `package main` for external test packages. -fn patch_package_for_source(source: String) -> anyhow::Result { - let parsed = gosyn::parse_source(&source)?; - let pkg_name = &parsed.pkg_name.name; - - let replacement = if pkg_name == "main" { - Some("main_compat") - } else if pkg_name.ends_with("_test") { - // For external test packages (package foo_test), convert to package main - // They will be placed in the codspeed/ subdirectory and built as standalone executables - Some("main") - } else { - None - }; - - if let Some(new_name) = replacement { - // pkg_name.pos is the position of the identifier in the source - let name_start = parsed.pkg_name.pos; - let name_end = name_start + pkg_name.len(); - - let mut result = source; - result.replace_range(name_start..name_end, new_name); - Ok(result) - } else { - Ok(source) +pub fn replace_pkg>(folder: P) -> anyhow::Result<()> { + let codspeed_root = Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap(); + let replace_arg = format!( + "github.com/CodSpeedHQ/codspeed-go={}", + codspeed_root.display() + ); + debug!("Replacing codspeed-go with {}", codspeed_root.display()); + + let output = Command::new("go") + .args(["mod", "edit", "-replace", &replace_arg]) + .current_dir(folder.as_ref()) + .output() + .context("Failed to execute 'go mod edit' command")?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + bail!("Failed to add replace directive: {}", stderr); } + + debug!("Added local replace directive to go.mod"); + + Ok(()) } /// Installs the codspeed-go dependency in the module @@ -427,8 +485,10 @@ import ( #[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(); + let mut result = source.to_string(); + + Patcher::patch_imports_for_source(&mut result); + Patcher::patch_package(&mut result, None).unwrap(); assert_snapshot!(test_name, result); } } diff --git a/go-runner/src/builder/templater.rs b/go-runner/src/builder/templater.rs index 969edac4..b7e81e25 100644 --- a/go-runner/src/builder/templater.rs +++ b/go-runner/src/builder/templater.rs @@ -2,9 +2,11 @@ use std::fs; use std::path::{Path, PathBuf}; use handlebars::Handlebars; +use once_cell::sync::OnceCell; use serde::{Deserialize, Serialize}; use tempfile::TempDir; +use crate::builder::patcher::Patcher; use crate::builder::{BenchmarkPackage, GoBenchmark}; use crate::utils; use crate::{builder::patcher, prelude::*}; @@ -21,159 +23,269 @@ struct TemplateData { module_name: String, } -/// Runs the templater which sets up a temporary Go project with patched test files and a custom runner. -/// -/// # Returns -/// -/// - `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>( - package: &BenchmarkPackage, - profile_dir: P, -) -> anyhow::Result<(TempDir, PathBuf)> { - // Create a temporary target directory for building the modified Go project. - 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. - // 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); +pub struct CodspeedContext { + package: BenchmarkPackage, + profile_dir: PathBuf, + git_root: PathBuf, + target_dir: PathBuf, + + // Artifacts that have to be cleaned up later + patcher: Patcher, + metadata_path: Option, + runner_path: Option, +} + +impl CodspeedContext { + pub fn for_package>( + package: &BenchmarkPackage, + profile_dir: P, + git_root: PathBuf, + target_dir: PathBuf, + ) -> Self { + Self { + patcher: Patcher::new(&target_dir), + package: package.clone(), + profile_dir: profile_dir.as_ref().to_path_buf(), + git_root, + target_dir, + metadata_path: None, + runner_path: None, + } } - // 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) { - git_dir - } else { - warn!("Could not find git repository root. Falling back to module directory as root"); - PathBuf::from(&package.module.dir) - }; - utils::copy_dir_recursively(&git_root, &target_dir)?; - - // Create a new go-runner.metadata file in the root of the project - // - // The package path will be prepended to the URI. The benchmark will - // find the path relative to the root of the `target_dir`. - // - // This is needed because we could execute a Go project that is a sub-folder - // within a Git repository, then we won't copy the .git folder. Therefore, we - // have to resolve the .git relative path in go-runner and then combine it. - let relative_package_path = utils::get_git_relative_path(&package.dir) - .to_string_lossy() - .into(); - debug!("Relative package path: {relative_package_path}"); - - let metadata = GoRunnerMetadata { - profile_folder: profile_dir.as_ref().to_string_lossy().into(), - relative_package_path, - }; - fs::write( - target_dir.path().join("go-runner.metadata"), - serde_json::to_string_pretty(&metadata)?, - ) - .context("Failed to write go-runner.metadata file")?; - - // Get files that need to be renamed first - let files = package - .test_files() - .with_context(|| anyhow::anyhow!("No test files found for package: {}", package.name))?; - - // Calculate the relative path from git root to package directory - let package_dir = Path::new(&package.dir); - let relative_package_path = package_dir.strip_prefix(&git_root).context(format!( - "Package dir {:?} is not within git root {:?}", - package.dir, git_root - ))?; - debug!("Relative package path: {relative_package_path:?}"); - - // 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.path().join(relative_package_path); - let test_file_paths: Vec = files.iter().map(|f| package_path.join(f)).collect(); - - // If we have external tests (e.g. "package {pkg}_test") they have to be - // changed to "package main" so they can be built within the codspeed/ sub-package. - if package.is_external_test_package() { - info!("Patching external test package files"); - patcher::patch_packages_for_test_files(&test_file_paths)?; - } else if package.name == "main" { - // If this is a "package main" (not external test), we need to patch ALL .go files in the package directory - // so they all become "package main_compat" and can be imported by the runner. - - info!("Package is 'main' - patching all .go files in package directory"); - patcher::patch_all_packages_in_dir(&package_path)?; + fn setup_runner_metadata(&mut self) -> anyhow::Result<()> { + // Create a new go-runner.metadata file in the root of the project + // + // The package path will be prepended to the URI. The benchmark will + // find the path relative to the root of the `target_dir`. + // + // This is needed because we could execute a Go project that is a sub-folder + // within a Git repository, then we won't copy the .git folder. Therefore, we + // have to resolve the .git relative path in go-runner and then combine it. + let relative_package_path = utils::get_git_relative_path(&self.package.dir) + .to_string_lossy() + .into(); + debug!("Relative package path: {relative_package_path}"); + + let metadata = GoRunnerMetadata { + profile_folder: self.profile_dir.to_string_lossy().into(), + relative_package_path, + }; + let metadata_path = self.target_dir.join("go-runner.metadata"); + fs::write(&metadata_path, serde_json::to_string_pretty(&metadata)?) + .context("Failed to write go-runner.metadata file")?; + self.metadata_path = Some(metadata_path); + + Ok(()) } - patcher::patch_imports(&target_dir)?; - - // 3. Install codspeed-go dependency at the package module level - // 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.path().join(relative_module_path)) - .unwrap_or_else(|_| { - // Fall back to target_dir if we can't calculate relative path - 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 - .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() { - // For external test packages: copy test files to codspeed/ subdirectory AND rename them - // (remove _test suffix so Go will compile them with `go build`) - // 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.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); - - fs::rename(&src_path, &dst_path).context(format!( - "Failed to rename external test file from {src_path:?} to {dst_path:?}" - ))?; + + fn patch_files(&mut self) -> anyhow::Result<()> { + let package = &self.package; + let target_dir = &self.target_dir; + + // Get files that need to be renamed first + let files = package.test_files().with_context(|| { + anyhow::anyhow!("No test files found for package: {}", package.name) + })?; + + // 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(self.relative_package_path()?); + let test_file_paths: Vec = files.iter().map(|f| package_path.join(f)).collect(); + + // If we have external tests (e.g. "package {pkg}_test") they have to be + // changed to "package main" so they can be built within the codspeed/ sub-package. + if package.is_external_test_package() { + self.patcher.patch_packages_for_files(&test_file_paths)?; + } else if package.name == "main" { + // If this is a "package main" (not external test), we need to patch ALL .go files in the package directory + // so they all become "package main_compat" and can be imported by the runner. + + info!("Package is 'main' - patching all .go files in package directory"); + self.patcher.patch_all_packages_in_dir(&package_path)?; } - } else { - // 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.path().join(relative_package_path).join(file); - let new_path = old_path.with_file_name( - old_path - .file_name() - .unwrap() - .to_string_lossy() - .replace("_test", "_codspeed"), - ); - - fs::rename(&old_path, &new_path) - .context(format!("Failed to rename {old_path:?} to {new_path:?}"))?; + self.patcher.patch_imports(target_dir)?; + + // Handle test files differently based on whether they're external or internal tests + let codspeed_dir = self.codspeed_dir()?; + fs::create_dir_all(&codspeed_dir).context("Failed to create codspeed directory")?; + + if package.is_external_test_package() { + // For external test packages: copy test files to codspeed/ subdirectory AND rename them + // (remove _test suffix so Go will compile them with `go build`) + // They're now package main and will be built from the subdirectory + self.patcher + .rename_and_move_test_files(&test_file_paths, &codspeed_dir)?; + } else { + // For internal test packages: rename _test.go to _codspeed.go in place + self.patcher.rename_test_files(&test_file_paths)?; + } + + Ok(()) + } + + fn install_codspeed_dependency(&mut self) -> anyhow::Result<()> { + let package = &self.package; + let git_root = &self.git_root; + let target_dir = &self.target_dir; + + // Install codspeed-go dependency at the package module level + // 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)) + .unwrap_or_else(|_| { + // Fall back to target_dir if we can't calculate relative path + target_dir.to_path_buf() + }); + patcher::install_codspeed_dependency(&module_dir)?; + + Ok(()) + } + + fn setup_runner(&mut self) -> anyhow::Result<()> { + let package = &self.package; + + // Generate the codspeed/runner.go file using the template + let mut handlebars = Handlebars::new(); + let template_content = include_str!("template.go"); + handlebars.register_template_string("main", template_content)?; + + // import + // { "", }, + let data = TemplateData { + benchmarks: package.benchmarks.clone(), + module_name: "codspeed_runner".into(), + }; + let rendered = handlebars.render("main", &data)?; + + let runner_path = self.codspeed_dir()?.join("runner.go"); + fs::write(&runner_path, rendered).context("Failed to write runner.go file")?; + self.runner_path = Some(runner_path); + + Ok(()) + } + + pub fn runner_path(&self) -> &PathBuf { + self.runner_path + .as_ref() + .expect("Runner path not set up yet") + } + + fn relative_package_path(&self) -> anyhow::Result { + let package = &self.package; + let git_root = &self.git_root; + + // Calculate the relative path from git root to package directory + let package_dir = Path::new(&package.dir); + let relative_package_path = package_dir.strip_prefix(git_root).context(format!( + "Package dir {:?} is not within git root {:?}", + package.dir, git_root + ))?; + Ok(relative_package_path.to_path_buf()) + } + + fn codspeed_dir(&self) -> anyhow::Result { + let package = &self.package; + let target_dir = &self.target_dir; + let git_root = &self.git_root; + + let package_dir = Path::new(&package.dir); + let relative_package_path = package_dir.strip_prefix(git_root).unwrap_or_else(|_| { + panic!( + "Package dir {:?} is not within git root {:?}", + package.dir, git_root + ) + }); + Ok(target_dir.join(relative_package_path).join("codspeed")) + } +} + +pub struct Templater { + target_dir: OnceCell, +} + +impl Default for Templater { + fn default() -> Self { + Self::new() + } +} + +impl Templater { + pub fn new() -> Self { + Self { + target_dir: OnceCell::new(), } } - // 5. Generate the codspeed/runner.go file using the template - let mut handlebars = Handlebars::new(); - let template_content = include_str!("template.go"); - handlebars.register_template_string("main", template_content)?; + /// Runs the templater which sets up a temporary Go project with patched test files and a custom runner. + /// + /// # 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>( + &self, + package: &BenchmarkPackage, + profile_dir: P, + ) -> anyhow::Result { + // 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) { + git_dir + } else { + warn!("Could not find git repository root. Falling back to module directory as root"); + PathBuf::from(&package.module.dir) + }; + + // Because we added the projects as git submodules, they'll have a symlink to + // the actual git repository. There is a .git file which only contains the path to + // the actual .git folder. + // $ cat .git + // gitdir: ../../../../.git/modules/go-runner/testdata/projects/hugo + // + // In those cases, we'll have to copy the parent directory rather than the submodule. + // NOTE: This is not the case for all tests, so we have to only do it for submodules. + let is_submodule = fs::read_to_string(git_root.join(".git")) + .map(|content| content.starts_with("gitdir:")) + .unwrap_or(false); + let git_root: PathBuf = if cfg!(test) && is_submodule { + utils::get_parent_git_repo_path(git_root.parent().unwrap()).unwrap() + } else { + git_root + }; + info!("Found git root at {git_root:?}"); + + let target_dir = self + .target_dir + .get_or_try_init(|| -> anyhow::Result { + // Create a temporary target directory for building the modified Go project. + let mut target_dir = TempDir::new()?; - // import - // { "", }, - let data = TemplateData { - benchmarks: package.benchmarks.clone(), - module_name: "codspeed_runner".into(), - }; - let rendered = handlebars.render("main", &data)?; + // We don't want to spend time cleanup any temporary files since the code is only + // run on CI servers which clean up themselves. + // 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); + } - let runner_path = codspeed_dir.join("runner.go"); - fs::write(&runner_path, rendered).context("Failed to write runner.go file")?; + utils::copy_dir_recursively(&git_root, target_dir.path())?; - Ok((target_dir, runner_path)) + Ok(target_dir) + })?; + + let mut ctx = CodspeedContext::for_package( + package, + &profile_dir, + git_root, + target_dir.path().to_path_buf(), + ); + ctx.install_codspeed_dependency()?; + ctx.setup_runner_metadata()?; + ctx.patch_files()?; + ctx.setup_runner()?; + + Ok(ctx) + } } diff --git a/go-runner/src/lib.rs b/go-runner/src/lib.rs index 0d051dd0..db278bc9 100644 --- a/go-runner/src/lib.rs +++ b/go-runner/src/lib.rs @@ -1,5 +1,5 @@ use crate::{ - builder::BenchmarkPackage, + builder::{BenchmarkPackage, templater::Templater}, prelude::*, results::{raw_result::RawResult, walltime_results::WalltimeBenchmark}, }; @@ -12,6 +12,9 @@ pub mod results; pub mod runner; pub(crate) mod utils; +#[global_allocator] +static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; + #[cfg(test)] mod integration_tests; @@ -35,13 +38,15 @@ pub fn run_benchmarks>( } // 2. Generate codspeed runners, build binaries, and execute them + let templater = Templater::new(); + let mut threads = Vec::with_capacity(packages.len()); for package in &packages { info!("Generating custom runner for package: {}", package.name); - let (_target_dir, runner_path) = builder::templater::run(package, &profile_dir)?; + let ctx = templater.run(package, &profile_dir)?; info!("Building binary for package: {}", package.name); - let binary_path = match builder::build_binary(&runner_path) { + let binary_path = match builder::build_binary(ctx.runner_path()) { Ok(binary_path) => binary_path, Err(e) => { if cfg!(test) { @@ -61,15 +66,23 @@ pub fn run_benchmarks>( error!("Failed to run benchmarks for {}: {error}", package.name); continue; } + + // Collect the results in a new thread, to avoid blocking the main thread at the end. The + // conversions and computations can take a while when we have a lot of iterations. + let profile_dir = profile_dir.as_ref().to_path_buf(); + threads.push(std::thread::spawn(move || { + collect_walltime_results(profile_dir.as_ref()).unwrap(); + })); } else { info!("Skipping benchmark execution (dry-run mode)"); } } - // 3. Collect the results - if !cli.dry_run { - collect_walltime_results(profile_dir.as_ref())?; - } + // Wait for all result collection threads to finish + threads.into_iter().for_each(|t| { + t.join() + .expect("Failed to join collect_walltime_results thread") + }); Ok(()) } @@ -78,7 +91,8 @@ pub fn run_benchmarks>( pub fn collect_walltime_results(profile_dir: &Path) -> anyhow::Result<()> { let mut benchmarks_by_pid: HashMap> = HashMap::new(); - for (pid, walltime_result) in RawResult::parse_folder(profile_dir)?.into_iter() { + let raw_results_dir = profile_dir.join("raw_results"); + for (pid, walltime_result) in RawResult::parse_folder(&raw_results_dir)?.into_iter() { benchmarks_by_pid .entry(pid) .or_default() diff --git a/go-runner/src/results/raw_result.rs b/go-runner/src/results/raw_result.rs index 785de8d5..ff60ee44 100644 --- a/go-runner/src/results/raw_result.rs +++ b/go-runner/src/results/raw_result.rs @@ -18,24 +18,32 @@ impl RawResult { pub fn parse_folder>( folder: P, ) -> anyhow::Result> { - let glob_pattern = folder.as_ref().join("raw_results").join("*.json"); + let glob_pattern = folder.as_ref().join("*.json"); let result = glob::glob(&glob_pattern.to_string_lossy())? .par_bridge() .filter_map(Result::ok) .filter_map(|path| { - let file = std::fs::File::open(&path).ok()?; - let reader = std::io::BufReader::new(file); - let json: Self = serde_json::from_reader(reader).ok()?; - Some(( - json.pid, - WalltimeBenchmark::from_runtime_data( - json.name, - json.uri, - &json.codspeed_iters_per_round, - &json.codspeed_time_per_round_ns, - None, - ), - )) + let (pid, bench) = { + let file = std::fs::File::open(&path).ok()?; + let reader = std::io::BufReader::new(file); + let json: Self = serde_json::from_reader(reader).ok()?; + + ( + json.pid, + WalltimeBenchmark::from_runtime_data( + json.name, + json.uri, + &json.codspeed_iters_per_round, + &json.codspeed_time_per_round_ns, + None, + ), + ) + }; + + // Remove the file since we processed it + std::fs::remove_file(path).ok(); + + Some((pid, bench)) }) .collect(); Ok(result) diff --git a/go-runner/tests/utils.rs b/go-runner/tests/utils.rs index 4794b96b..e642889b 100644 --- a/go-runner/tests/utils.rs +++ b/go-runner/tests/utils.rs @@ -1,11 +1,16 @@ -use codspeed_go_runner::{builder, builder::BenchmarkPackage, cli::Cli, runner}; +use codspeed_go_runner::{ + builder::{self, BenchmarkPackage, templater::Templater}, + cli::Cli, + runner, +}; use std::path::Path; /// Helper function to run a single package with arguments pub fn run_package_with_args(package: &BenchmarkPackage, args: &[&str]) -> anyhow::Result { let profile_dir = tempfile::TempDir::new()?; - let (_dir, runner_path) = builder::templater::run(package, profile_dir.as_ref())?; - let binary_path = builder::build_binary(&runner_path)?; + let templater = Templater::new(); + let ctx = templater.run(package, profile_dir.as_ref())?; + let binary_path = builder::build_binary(ctx.runner_path())?; runner::run_with_stdout(&binary_path, args) }