55import subprocess
66import sys
77import sysconfig
8+ import tempfile
89import types
910import unittest
1011
@@ -24,6 +25,31 @@ def abspath(filename):
2425 return os .path .abspath (findfile (filename , subdir = "dtracedata" ))
2526
2627
28+ def get_probe_binary ():
29+ binary = sys .executable
30+ if sysconfig .get_config_var ("Py_ENABLE_SHARED" ):
31+ lib_dir = sysconfig .get_config_var ("LIBDIR" )
32+ if not lib_dir or sysconfig .is_python_build ():
33+ lib_dir = os .path .abspath (os .path .dirname (sys .executable ))
34+
35+ lib_names = []
36+ for name in (
37+ sysconfig .get_config_var ("INSTSONAME" ),
38+ sysconfig .get_config_var ("LDLIBRARY" ),
39+ ):
40+ if name and name not in lib_names :
41+ lib_names .append (name )
42+
43+ if lib_dir :
44+ for name in lib_names :
45+ libpython_path = os .path .join (lib_dir , name )
46+ if os .path .exists (libpython_path ):
47+ binary = libpython_path
48+ break
49+
50+ return binary
51+
52+
2753def normalize_trace_output (output ):
2854 """Normalize DTrace output for comparison.
2955
@@ -179,6 +205,43 @@ class DTraceBackend(TraceBackend):
179205class SystemTapBackend (TraceBackend ):
180206 EXTENSION = ".stp"
181207 COMMAND = ["stap" , "-g" ]
208+ PROBE_PLACEHOLDER = "@PYTHON_SYSTEMTAP_PROBE@"
209+
210+ @staticmethod
211+ def _quote_systemtap_string (value ):
212+ return value .replace ("\\ " , "\\ \\ " ).replace ('"' , '\\ "' )
213+
214+ def _python_probe (self ):
215+ executable = self ._quote_systemtap_string (sys .executable )
216+ probe_binary = get_probe_binary ()
217+ if probe_binary != sys .executable :
218+ probe_binary = self ._quote_systemtap_string (probe_binary )
219+ return f'process("{ executable } ").library("{ probe_binary } ").mark'
220+ return f'process("{ executable } ").mark'
221+
222+ def _render_script (self , script_file ):
223+ with open (script_file ) as script :
224+ return script .read ().replace (
225+ self .PROBE_PLACEHOLDER , self ._python_probe ()
226+ )
227+
228+ def trace (self , script_file , subcommand = None , * , timeout = None ,
229+ check_returncode = False ):
230+ with tempfile .NamedTemporaryFile (
231+ mode = "w" , encoding = "utf-8" , suffix = self .EXTENSION , delete = False
232+ ) as script :
233+ script .write (self ._render_script (script_file ))
234+ generated_script_file = script .name
235+
236+ try :
237+ return super ().trace (
238+ generated_script_file ,
239+ subcommand ,
240+ timeout = timeout ,
241+ check_returncode = check_returncode ,
242+ )
243+ finally :
244+ os .unlink (generated_script_file )
182245
183246
184247class BPFTraceBackend (TraceBackend ):
@@ -272,7 +335,7 @@ def run_case(self, name, optimize_python=None):
272335 python_flags .extend (["-O" ] * optimize_python )
273336
274337 subcommand = [sys .executable ] + python_flags + [python_file ]
275- program = self .PROGRAMS [name ].format (python = sys . executable )
338+ program = self .PROGRAMS [name ].format (python = get_probe_binary () )
276339
277340 try :
278341 proc = create_process_group (
@@ -311,7 +374,7 @@ def run_case(self, name, optimize_python=None):
311374
312375 def assert_usable (self ):
313376 # Check if bpftrace is available and can attach to USDT probes
314- program = f'usdt:{ sys . executable } :python:function__entry {{ printf("probe: success\\ n"); exit(); }}'
377+ program = f'usdt:{ get_probe_binary () } :python:function__entry {{ printf("probe: success\\ n"); exit(); }}'
315378 try :
316379 proc = create_process_group (
317380 ["bpftrace" , "-e" , program , "-c" , f"{ sys .executable } -c pass" ],
@@ -453,28 +516,7 @@ def get_readelf_version():
453516 return int (match .group (1 )), int (match .group (2 ))
454517
455518 def get_readelf_output (self ):
456- binary = sys .executable
457- if sysconfig .get_config_var ("Py_ENABLE_SHARED" ):
458- lib_dir = sysconfig .get_config_var ("LIBDIR" )
459- if not lib_dir or sysconfig .is_python_build ():
460- lib_dir = os .path .abspath (os .path .dirname (sys .executable ))
461-
462- lib_names = []
463- for name in (
464- sysconfig .get_config_var ("INSTSONAME" ),
465- sysconfig .get_config_var ("LDLIBRARY" ),
466- ):
467- if name and name not in lib_names :
468- lib_names .append (name )
469-
470- if lib_dir :
471- for name in lib_names :
472- libpython_path = os .path .join (lib_dir , name )
473- if os .path .exists (libpython_path ):
474- binary = libpython_path
475- break
476-
477- return run_readelf (["readelf" , "-n" , binary ])
519+ return run_readelf (["readelf" , "-n" , get_probe_binary ()])
478520
479521 def test_check_probes (self ):
480522 readelf_output = self .get_readelf_output ()
0 commit comments