gh-98894: Fix dtrace tests in shared builds

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()