Skip to content

Commit 0d928f2

Browse files
committed
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.
1 parent ecf4e63 commit 0d928f2

1 file changed

Lines changed: 24 additions & 14 deletions

File tree

codeflash/languages/java/tracer.py

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def _run_java_with_graceful_timeout(
6161

6262

6363
class JavaTracer:
64-
"""Orchestrates two-stage Java tracing: JFR profiling + argument capture."""
64+
"""Orchestrates Java tracing: combined JFR profiling + argument capture in a single JVM invocation."""
6565

6666
def trace(
6767
self,
@@ -72,29 +72,23 @@ def trace(
7272
max_function_count: int = 256,
7373
timeout: int = 0,
7474
) -> tuple[Path, Path]:
75-
"""Run the Java program twice: once for profiling, once for arg capture.
75+
"""Run the Java program once with both JFR profiling and argument capture.
7676
7777
Returns (trace_db_path, jfr_file_path).
7878
"""
7979
jfr_file = trace_db_path.with_suffix(".jfr")
8080
trace_db_path.parent.mkdir(parents=True, exist_ok=True)
8181

82-
# Stage 1: JFR Profiling
83-
logger.info("Stage 1: Running JFR profiling...")
84-
jfr_env = self.build_jfr_env(jfr_file)
85-
_run_java_with_graceful_timeout(java_command, jfr_env, timeout, "JFR profiling")
86-
87-
if not jfr_file.exists():
88-
logger.warning("JFR file was not created at %s", jfr_file)
89-
90-
# Stage 2: Argument Capture via Tracing Agent
91-
logger.info("Stage 2: Running argument capture...")
9282
config_path = self.create_tracer_config(
9383
trace_db_path, packages, project_root=project_root, max_function_count=max_function_count, timeout=timeout
9484
)
95-
agent_env = self.build_agent_env(config_path)
96-
_run_java_with_graceful_timeout(java_command, agent_env, timeout, "Argument capture")
85+
combined_env = self.build_combined_env(jfr_file, config_path)
86+
87+
logger.info("Running combined JFR profiling + argument capture...")
88+
_run_java_with_graceful_timeout(java_command, combined_env, timeout, "Combined tracing")
9789

90+
if not jfr_file.exists():
91+
logger.warning("JFR file was not created at %s", jfr_file)
9892
if not trace_db_path.exists():
9993
logger.error("Trace database was not created at %s", trace_db_path)
10094

@@ -141,6 +135,22 @@ def build_agent_env(self, config_path: Path, classpath: str | None = None) -> di
141135
env["JAVA_TOOL_OPTIONS"] = f"{existing} {agent_opts}".strip()
142136
return env
143137

138+
def build_combined_env(self, jfr_file: Path, config_path: Path, classpath: str | None = None) -> dict[str, str]:
139+
"""Build env with both JFR recording and tracing agent in a single JAVA_TOOL_OPTIONS."""
140+
env = os.environ.copy()
141+
jfr_opts = (
142+
f"-XX:StartFlightRecording=filename={jfr_file.resolve()},settings=profile,dumponexit=true"
143+
",jdk.ExecutionSample#period=1ms"
144+
)
145+
agent_jar = find_agent_jar(classpath=classpath)
146+
if agent_jar is None:
147+
msg = "codeflash-runtime JAR not found, cannot run tracing agent"
148+
raise FileNotFoundError(msg)
149+
agent_opts = f"{ADD_OPENS_FLAGS} -javaagent:{agent_jar}=trace={config_path.resolve()}"
150+
existing = env.get("JAVA_TOOL_OPTIONS", "")
151+
env["JAVA_TOOL_OPTIONS"] = f"{existing} {jfr_opts} {agent_opts}".strip()
152+
return env
153+
144154
@staticmethod
145155
def detect_packages_from_source(module_root: Path) -> list[str]:
146156
"""Scan Java files for package declarations and return unique package prefixes."""

0 commit comments

Comments
 (0)