java: Add --java-async-profiler-args, allowing to send arbitrary extra args to AP (#262)

Author: Jongy

gprofiler/profilers/java.py

@@ -134,6 +134,7 @@ def __init__(
         buildids: bool,
         mode: str,
         ap_safemode: int,
+        ap_args: str,
     ):
         self.process = process
         # access the process' root via its topmost parent/ancestor which uses the same mount namespace.
@@ -178,6 +179,7 @@ def __init__(
         assert mode in ("cpu", "itimer"), f"unexpected mode: {mode}"
         self._mode = mode
         self._ap_safemode = ap_safemode
+        self._ap_args = ap_args
 
     def __enter__(self):
         os.makedirs(self._ap_dir_host, 0o755, exist_ok=True)
@@ -286,13 +288,16 @@ def _get_base_cmd(self) -> List[str]:
             "true",
         ]
 
+    def _get_extra_ap_args(self) -> str:
+        return f",{self._ap_args}" if self._ap_args else ""
+
     def _get_start_cmd(self, interval: int) -> List[str]:
         return self._get_base_cmd() + [
             f"start,event={self._mode},file={self._output_path_process},"
             f"{self.OUTPUT_FORMAT},{self.FORMAT_PARAMS},interval={interval},"
             f"log={self._log_path_process}{',buildids' if self._buildids else ''}"
             f"{',fdtransfer' if self._mode == 'cpu' else ''}"
-            f",safemode={self._ap_safemode}"
+            f",safemode={self._ap_safemode}{self._get_extra_ap_args()}"
         ]
 
     def _get_stop_cmd(self, with_output: bool) -> List[str]:
@@ -302,7 +307,7 @@ def _get_stop_cmd(self, with_output: bool) -> List[str]:
             ap_params.append(self.OUTPUT_FORMAT)
             ap_params.append(self.FORMAT_PARAMS)
         ap_params.append(f"log={self._log_path_process}")
-        return self._get_base_cmd() + [",".join(ap_params)]
+        return self._get_base_cmd() + [",".join(ap_params) + self._get_extra_ap_args()]
 
     def _run_async_profiler(self, cmd: List[str]) -> None:
         try:
@@ -464,6 +469,12 @@ def parse_jvm_version(version_string: str) -> JvmVersion:
             " enum in async-profiler's code, in profiler.cpp)."
             " Defaults to '%(default)s').",
         ),
+        ProfilerArgument(
+            "--java-async-profiler-args",
+            dest="java_async_profiler_args",
+            type=str,
+            help="Additional arguments to pass directly to async-profiler (start & stop commands)",
+        ),
         ProfilerArgument(
             "--java-safemode",
             dest="java_safemode",
@@ -499,6 +510,7 @@ def __init__(
         java_version_check: bool,
         java_async_profiler_mode: str,
         java_async_profiler_safemode: int,
+        java_async_profiler_args: str,
         java_safemode: str,
         java_mode: str,
     ):
@@ -514,6 +526,7 @@ def __init__(
             logger.warning("Java version checks are disabled")
         self._mode = java_async_profiler_mode
         self._ap_safemode = java_async_profiler_safemode
+        self._ap_args = java_async_profiler_args
         self._init_java_safemode(java_safemode)
         self._should_profile = True
         # if set, profiling is disabled due to this safemode reason.
@@ -704,7 +717,9 @@ def _profile_process(self, process: Process) -> Optional[StackToSampleCount]:
             self._profiled_pids.add(process.pid)
 
         logger.info(f"Profiling process {process.pid} with async-profiler")
-        with AsyncProfiledProcess(process, self._storage_dir, self._buildids, self._mode, self._ap_safemode) as ap_proc:
+        with AsyncProfiledProcess(
+            process, self._storage_dir, self._buildids, self._mode, self._ap_safemode, self._ap_args
+        ) as ap_proc:
             return self._profile_ap_process(ap_proc, comm)
 
     def _profile_ap_process(self, ap_proc: AsyncProfiledProcess, comm: str) -> Optional[StackToSampleCount]:

tests/test_java.py

@@ -42,11 +42,12 @@ def test_async_profiler_already_running(application_pid, assert_collapsed, tmp_p
         False,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=0,
+        java_async_profiler_args="",
         java_safemode="",
         java_mode="ap",
     ) as profiler:
         process = profiler._select_processes_to_profile()[0]
-        with AsyncProfiledProcess(process, profiler._storage_dir, False, profiler._mode, False) as ap_proc:
+        with AsyncProfiledProcess(process, profiler._storage_dir, False, profiler._mode, False, "") as ap_proc:
             assert ap_proc.start_async_profiler(11)
         assert any("libasyncProfiler.so" in m.path for m in process.memory_maps())
         # run "status"
@@ -56,6 +57,7 @@ def test_async_profiler_already_running(application_pid, assert_collapsed, tmp_p
             False,
             mode="itimer",
             ap_safemode=False,
+            ap_args="",
         ) as ap_proc:
             ap_proc.status_async_profiler()
             # printed the output file, see ACTION_STATUS case in async-profiler/profiler.cpp
@@ -86,6 +88,7 @@ def test_java_async_profiler_cpu_mode(
         True,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=0,
+        java_async_profiler_args="",
         java_safemode="",
         java_mode="ap",
     ) as profiler:
@@ -114,6 +117,7 @@ def test_java_async_profiler_musl_and_cpu(
         True,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=0,
+        java_async_profiler_args="",
         java_safemode="",
         java_mode="ap",
     ) as profiler:
@@ -133,6 +137,7 @@ def test_java_safemode_parameters(tmp_path) -> None:
             True,
             java_async_profiler_mode="cpu",
             java_async_profiler_safemode=0,
+            java_async_profiler_args="",
             java_safemode=JAVA_SAFEMODE_ALL,
             java_mode="ap",
         )
@@ -148,6 +153,7 @@ def test_java_safemode_parameters(tmp_path) -> None:
             False,
             java_async_profiler_mode="cpu",
             java_async_profiler_safemode=127,
+            java_async_profiler_args="",
             java_safemode=JAVA_SAFEMODE_ALL,
             java_mode="ap",
         )
@@ -168,6 +174,7 @@ def test_java_safemode_version_check(
         True,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=127,
+        java_async_profiler_args="",
         java_safemode=JAVA_SAFEMODE_ALL,
         java_mode="ap",
     ) as profiler:
@@ -193,6 +200,7 @@ def test_java_safemode_build_number_check(
         True,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=127,
+        java_async_profiler_args="",
         java_safemode=JAVA_SAFEMODE_ALL,
         java_mode="ap",
     ) as profiler:
@@ -235,6 +243,7 @@ def sap_and_crash(self, *args, **kwargs):
         True,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=127,
+        java_async_profiler_args="",
         java_safemode=JAVA_SAFEMODE_ALL,
         java_mode="ap",
     )
@@ -253,7 +262,7 @@ def sap_and_crash(self, *args, **kwargs):
 def test_disable_java_profiling(application_pid, tmp_path, monkeypatch, caplog):
     caplog.set_level(logging.DEBUG)
 
-    profiler = JavaProfiler(1, 5, Event(), str(tmp_path), False, False, "cpu", 0, False, "ap")
+    profiler = JavaProfiler(1, 5, Event(), str(tmp_path), False, False, "cpu", 0, "", False, "ap")
     dummy_reason = "dummy reason"
     monkeypatch.setattr(profiler, "_safemode_disable_reason", dummy_reason)
     with profiler:
@@ -275,6 +284,7 @@ def test_already_loaded_async_profiler_profiling_failure(tmp_path, monkeypatch,
             True,
             java_async_profiler_mode="cpu",
             java_async_profiler_safemode=127,
+            java_async_profiler_args="",
             java_safemode=JAVA_SAFEMODE_ALL,
             java_mode="ap",
         ) as profiler:
@@ -289,6 +299,7 @@ def test_already_loaded_async_profiler_profiling_failure(tmp_path, monkeypatch,
         True,
         java_async_profiler_mode="cpu",
         java_async_profiler_safemode=127,
+        java_async_profiler_args="",
         java_safemode=JAVA_SAFEMODE_ALL,
         java_mode="ap",
     ) as profiler:

tests/test_sanity.py

@@ -41,6 +41,7 @@ def test_java_from_host(
         True,
         java_async_profiler_mode="itimer",
         java_async_profiler_safemode=0,
+        java_async_profiler_args="",
         java_safemode="",
         java_mode="ap",
     ) as profiler: