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