@@ -61,7 +61,7 @@ def _run_java_with_graceful_timeout(
6161
6262
6363class 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