From 0d928f2b49c7dccf18e3f28e43fd5d4616c7bb99 Mon Sep 17 00:00:00 2001
From: Kevin Turcios <turcioskevinr@gmail.com>
Date: Fri, 10 Apr 2026 09:05:30 -0500
Subject: [PATCH] perf: merge Java tracer into single-pass JVM invocation

Combine JFR profiling and argument capture agent into one
JAVA_TOOL_OPTIONS string, running the target program once instead of
twice. JFR and javaagent are orthogonal JVM features that coexist
without conflict. Keeps build_jfr_env/build_agent_env for standalone
use.
---
 codeflash/languages/java/tracer.py | 38 +++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/codeflash/languages/java/tracer.py b/codeflash/languages/java/tracer.py
index b971e5526..bdeec34e0 100644
--- a/codeflash/languages/java/tracer.py
+++ b/codeflash/languages/java/tracer.py
@@ -61,7 +61,7 @@ ADD_OPENS_FLAGS = (
 
 
 class JavaTracer:
-    """Orchestrates two-stage Java tracing: JFR profiling + argument capture."""
+    """Orchestrates Java tracing: combined JFR profiling + argument capture in a single JVM invocation."""
 
     def trace(
         self,
@@ -72,29 +72,23 @@ class JavaTracer:
         max_function_count: int = 256,
         timeout: int = 0,
     ) -> tuple[Path, Path]:
-        """Run the Java program twice: once for profiling, once for arg capture.
+        """Run the Java program once with both JFR profiling and argument capture.
 
         Returns (trace_db_path, jfr_file_path).
         """
         jfr_file = trace_db_path.with_suffix(".jfr")
         trace_db_path.parent.mkdir(parents=True, exist_ok=True)
 
-        # Stage 1: JFR Profiling
-        logger.info("Stage 1: Running JFR profiling...")
-        jfr_env = self.build_jfr_env(jfr_file)
-        _run_java_with_graceful_timeout(java_command, jfr_env, timeout, "JFR profiling")
-
-        if not jfr_file.exists():
-            logger.warning("JFR file was not created at %s", jfr_file)
-
-        # Stage 2: Argument Capture via Tracing Agent
-        logger.info("Stage 2: Running argument capture...")
         config_path = self.create_tracer_config(
             trace_db_path, packages, project_root=project_root, max_function_count=max_function_count, timeout=timeout
         )
-        agent_env = self.build_agent_env(config_path)
-        _run_java_with_graceful_timeout(java_command, agent_env, timeout, "Argument capture")
+        combined_env = self.build_combined_env(jfr_file, config_path)
 
+        logger.info("Running combined JFR profiling + argument capture...")
+        _run_java_with_graceful_timeout(java_command, combined_env, timeout, "Combined tracing")
+
+        if not jfr_file.exists():
+            logger.warning("JFR file was not created at %s", jfr_file)
         if not trace_db_path.exists():
             logger.error("Trace database was not created at %s", trace_db_path)
 
@@ -141,6 +135,22 @@ class JavaTracer:
         env["JAVA_TOOL_OPTIONS"] = f"{existing} {agent_opts}".strip()
         return env
 
+    def build_combined_env(self, jfr_file: Path, config_path: Path, classpath: str | None = None) -> dict[str, str]:
+        """Build env with both JFR recording and tracing agent in a single JAVA_TOOL_OPTIONS."""
+        env = os.environ.copy()
+        jfr_opts = (
+            f"-XX:StartFlightRecording=filename={jfr_file.resolve()},settings=profile,dumponexit=true"
+            ",jdk.ExecutionSample#period=1ms"
+        )
+        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)
+        agent_opts = f"{ADD_OPENS_FLAGS} -javaagent:{agent_jar}=trace={config_path.resolve()}"
+        existing = env.get("JAVA_TOOL_OPTIONS", "")
+        env["JAVA_TOOL_OPTIONS"] = f"{existing} {jfr_opts} {agent_opts}".strip()
+        return env
+
     @staticmethod
     def detect_packages_from_source(module_root: Path) -> list[str]:
         """Scan Java files for package declarations and return unique package prefixes."""