diff --git a/codeflash/languages/java/function_optimizer.py b/codeflash/languages/java/function_optimizer.py index 5700a907a..08980032e 100644 --- a/codeflash/languages/java/function_optimizer.py +++ b/codeflash/languages/java/function_optimizer.py @@ -404,7 +404,9 @@ class JavaFunctionOptimizer(FunctionOptimizer): line_profiler_output_path = get_run_tmp_file(Path("line_profiler_output.json")) success = self.language_support.instrument_source_for_line_profiler( - func_info=self.function_to_optimize, line_profiler_output_file=line_profiler_output_path + func_info=self.function_to_optimize, + line_profiler_output_file=line_profiler_output_path, + project_classpath=self._get_project_classpath(), ) if not success: return {"timings": {}, "unit": 0, "str_out": ""} diff --git a/codeflash/languages/java/line_profiler.py b/codeflash/languages/java/line_profiler.py index 854a8549d..706b543a8 100644 --- a/codeflash/languages/java/line_profiler.py +++ b/codeflash/languages/java/line_profiler.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import logging +import os import re from pathlib import Path from typing import TYPE_CHECKING, Any @@ -130,9 +131,9 @@ class JavaLineProfiler: config_output_path.write_text(json.dumps(config, indent=2), encoding="utf-8") return config_output_path - def build_javaagent_arg(self, config_path: Path) -> str: + def build_javaagent_arg(self, config_path: Path, classpath: str | None = None) -> str: """Return the -javaagent JVM argument string.""" - agent_jar = find_agent_jar() + agent_jar = find_agent_jar(classpath=classpath) if agent_jar is None: msg = f"{AGENT_JAR_NAME} not found in resources or dev build directory" raise FileNotFoundError(msg) @@ -565,12 +566,20 @@ def find_method_for_line( return Path(file_path).name, line_num -def find_agent_jar() -> Path | None: +def find_agent_jar(classpath: str | None = None) -> Path | None: """Locate the profiler agent JAR file (now bundled in codeflash-runtime). - Checks local Maven repo, package resources, and development build directory. + Checks the resolved classpath (if provided), local Maven repo, package resources, + and development build directory. """ - # Check local Maven repository first (fastest) + # Check resolved classpath first (Gradle projects resolve here, not ~/.m2) + if classpath: + for entry in classpath.split(os.pathsep): + jar_path = Path(entry) + if "codeflash-runtime" in jar_path.name and jar_path.suffix == ".jar" and jar_path.exists(): + return jar_path + + # Check local Maven repository (Maven projects resolve here) m2_jar = ( Path.home() / ".m2" diff --git a/codeflash/languages/java/support.py b/codeflash/languages/java/support.py index ab3818348..7115d2225 100644 --- a/codeflash/languages/java/support.py +++ b/codeflash/languages/java/support.py @@ -590,7 +590,7 @@ class JavaSupport(LanguageSupport): ) def instrument_source_for_line_profiler( - self, func_info: FunctionToOptimize, line_profiler_output_file: Path + self, func_info: FunctionToOptimize, line_profiler_output_file: Path, project_classpath: str | None = None ) -> bool: """Prepare line profiling via the bytecode-instrumentation agent. @@ -602,6 +602,7 @@ class JavaSupport(LanguageSupport): Args: func_info: Function to profile. line_profiler_output_file: Path where profiling results will be written by the agent. + project_classpath: Resolved classpath from the build tool, used to locate the agent JAR. Returns: True if preparation succeeded, False otherwise. @@ -619,7 +620,7 @@ class JavaSupport(LanguageSupport): source=source, file_path=func_info.file_path, functions=[func_info], config_output_path=config_path ) - self.line_profiler_agent_arg = profiler.build_javaagent_arg(config_path) + self.line_profiler_agent_arg = profiler.build_javaagent_arg(config_path, classpath=project_classpath) self.line_profiler_warmup_iterations = profiler.warmup_iterations return True except Exception: diff --git a/codeflash/languages/java/tracer.py b/codeflash/languages/java/tracer.py index ab8f19514..50506797e 100644 --- a/codeflash/languages/java/tracer.py +++ b/codeflash/languages/java/tracer.py @@ -132,9 +132,9 @@ class JavaTracer: env["JAVA_TOOL_OPTIONS"] = f"{existing} {jfr_opts}".strip() return env - def build_agent_env(self, config_path: Path) -> dict[str, str]: + def build_agent_env(self, config_path: Path, classpath: str | None = None) -> dict[str, str]: env = os.environ.copy() - agent_jar = find_agent_jar() + agent_jar = find_agent_jar(classpath=classpath) if agent_jar is None: msg = "codeflash-runtime JAR not found, cannot run tracing agent" raise FileNotFoundError(msg)