gh-98894: Fix dtrace tests in shared builds (#153372)

Generate SystemTap probe definitions targeting libpython for shared
builds and use centralized USDT probe object discovery for readelf
and BPFTrace.

Author: stratakis

Lib/test/dtracedata/call_stack.stp

@@ -10,7 +10,7 @@ function basename:string(path:string)
     return last_token;
 }
 
-probe process.mark("function__entry")
+probe @PYTHON_SYSTEMTAP_PROBE@("function__entry")
 {
     funcname = user_string($arg2);
 
@@ -19,7 +19,8 @@ probe process.mark("function__entry")
     }
 }
 
-probe process.mark("function__entry"), process.mark("function__return")
+probe @PYTHON_SYSTEMTAP_PROBE@("function__entry"),
+      @PYTHON_SYSTEMTAP_PROBE@("function__return")
 {
     filename = user_string($arg1);
     funcname = user_string($arg2);
@@ -31,7 +32,7 @@ probe process.mark("function__entry"), process.mark("function__return")
     }
 }
 
-probe process.mark("function__return")
+probe @PYTHON_SYSTEMTAP_PROBE@("function__return")
 {
     funcname = user_string($arg2);
 

Lib/test/dtracedata/gc.stp

@@ -1,6 +1,6 @@
 global tracing
 
-probe process.mark("function__entry")
+probe @PYTHON_SYSTEMTAP_PROBE@("function__entry")
 {
     funcname = user_string($arg2);
 
@@ -9,14 +9,15 @@ probe process.mark("function__entry")
     }
 }
 
-probe process.mark("gc__start"), process.mark("gc__done")
+probe @PYTHON_SYSTEMTAP_PROBE@("gc__start"),
+      @PYTHON_SYSTEMTAP_PROBE@("gc__done")
 {
     if (tracing) {
         printf("%d\t%s:%ld\n", gettimeofday_us(), $$name, $arg1);
     }
 }
 
-probe process.mark("function__return")
+probe @PYTHON_SYSTEMTAP_PROBE@("function__return")
 {
     funcname = user_string($arg2);
 

Lib/test/test_dtrace.py

@@ -6,11 +6,13 @@
 import subprocess
 import sys
 import sysconfig
+import tempfile
 import types
 import unittest
 
 from test import support
 from test.support import findfile, MS_WINDOWS
+from test.support import os_helper
 
 
 if not support.has_subprocess_support:
@@ -25,6 +27,31 @@ def abspath(filename):
     return os.path.abspath(findfile(filename, subdir="dtracedata"))
 
 
+def get_probe_binary():
+    binary = sys.executable
+    if sysconfig.get_config_var("Py_ENABLE_SHARED"):
+        lib_dir = sysconfig.get_config_var("LIBDIR")
+        if not lib_dir or sysconfig.is_python_build():
+            lib_dir = os.path.abspath(os.path.dirname(sys.executable))
+
+        lib_names = []
+        for name in (
+            sysconfig.get_config_var("INSTSONAME"),
+            sysconfig.get_config_var("LDLIBRARY"),
+        ):
+            if name and name not in lib_names:
+                lib_names.append(name)
+
+        if lib_dir:
+            for name in lib_names:
+                libpython_path = os.path.join(lib_dir, name)
+                if os.path.exists(libpython_path):
+                    binary = libpython_path
+                    break
+
+    return binary
+
+
 def normalize_trace_output(output):
     """Normalize DTrace output for comparison.
 
@@ -180,6 +207,45 @@ class DTraceBackend(TraceBackend):
 class SystemTapBackend(TraceBackend):
     EXTENSION = ".stp"
     COMMAND = ["stap", "-g"]
+    PROBE_PLACEHOLDER = "@PYTHON_SYSTEMTAP_PROBE@"
+
+    @staticmethod
+    def quote_systemtap_string(value):
+        return value.replace("\\", "\\\\").replace('"', '\\"')
+
+    def python_probe(self):
+        executable = self.quote_systemtap_string(sys.executable)
+        probe_binary = get_probe_binary()
+        if probe_binary == sys.executable:
+            return f'process("{executable}").mark'
+
+        # Python built with --enable-shared
+        probe_binary = self.quote_systemtap_string(probe_binary)
+        return f'process("{executable}").library("{probe_binary}").mark'
+
+    def render_script(self, filename):
+        with open(filename) as fp:
+            script = fp.read()
+
+        return script.replace(self.PROBE_PLACEHOLDER, self.python_probe())
+
+    def trace(self, script_file, subcommand=None, *, timeout=None,
+              check_returncode=False):
+        with tempfile.NamedTemporaryFile(
+            mode="w", encoding="utf-8", suffix=self.EXTENSION, delete=False
+        ) as script:
+            script.write(self.render_script(script_file))
+            generated_script_file = script.name
+
+        try:
+            return super().trace(
+                generated_script_file,
+                subcommand,
+                timeout=timeout,
+                check_returncode=check_returncode,
+            )
+        finally:
+            os_helper.unlink(generated_script_file)
 
 
 class BPFTraceBackend(TraceBackend):
@@ -273,7 +339,7 @@ def run_case(self, name, optimize_python=None):
             python_flags.extend(["-O"] * optimize_python)
 
         subcommand = [sys.executable] + python_flags + [python_file]
-        program = self.PROGRAMS[name].format(python=sys.executable)
+        program = self.PROGRAMS[name].format(python=get_probe_binary())
 
         try:
             proc = create_process_group(
@@ -312,7 +378,7 @@ def run_case(self, name, optimize_python=None):
 
     def assert_usable(self):
         # Check if bpftrace is available and can attach to USDT probes
-        program = f'usdt:{sys.executable}:python:function__entry {{ printf("probe: success\\n"); exit(); }}'
+        program = f'usdt:{get_probe_binary()}:python:function__entry {{ printf("probe: success\\n"); exit(); }}'
         try:
             proc = create_process_group(
                 ["bpftrace", "-e", program, "-c",
@@ -455,28 +521,7 @@ def get_readelf_version():
         return int(match.group(1)), int(match.group(2))
 
     def get_readelf_output(self):
-        binary = sys.executable
-        if sysconfig.get_config_var("Py_ENABLE_SHARED"):
-            lib_dir = sysconfig.get_config_var("LIBDIR")
-            if not lib_dir or sysconfig.is_python_build():
-                lib_dir = os.path.abspath(os.path.dirname(sys.executable))
-
-            lib_names = []
-            for name in (
-                sysconfig.get_config_var("INSTSONAME"),
-                sysconfig.get_config_var("LDLIBRARY"),
-            ):
-                if name and name not in lib_names:
-                    lib_names.append(name)
-
-            if lib_dir:
-                for name in lib_names:
-                    libpython_path = os.path.join(lib_dir, name)
-                    if os.path.exists(libpython_path):
-                        binary = libpython_path
-                        break
-
-        return run_readelf(["readelf", "-n", binary])
+        return run_readelf(["readelf", "-n", get_probe_binary()])
 
     def test_check_probes(self):
         readelf_output = self.get_readelf_output()