@@ -33,11 +33,17 @@ def normalize_trace_output(output):
3333 result = [
3434 row .split ("\t " )
3535 for row in output .splitlines ()
36- if row and not row .startswith ('#' )
36+ if row and not row .startswith ('#' ) and not row . startswith ( '@' )
3737 ]
3838 result .sort (key = lambda row : int (row [0 ]))
3939 result = [row [1 ] for row in result ]
40- return "\n " .join (result )
40+ # Normalize paths to basenames (bpftrace outputs full paths)
41+ normalized = []
42+ for line in result :
43+ # Replace full paths with just the filename
44+ line = re .sub (r'/[^:]+/([^/:]+\.py)' , r'\1' , line )
45+ normalized .append (line )
46+ return "\n " .join (normalized )
4147 except (IndexError , ValueError ):
4248 raise AssertionError (
4349 "tracer produced unparsable output:\n {}" .format (output )
@@ -95,14 +101,163 @@ def assert_usable(self):
95101
96102class DTraceBackend (TraceBackend ):
97103 EXTENSION = ".d"
98- COMMAND = ["dtrace" , "-q" , "-s" ]
104+ COMMAND = ["dtrace" , "-q" , "-Z" , "- s" ]
99105
100106
101107class SystemTapBackend (TraceBackend ):
102108 EXTENSION = ".stp"
103109 COMMAND = ["stap" , "-g" ]
104110
105111
112+ class BPFTraceBackend (TraceBackend ):
113+ EXTENSION = ".bt"
114+ COMMAND = ["bpftrace" ]
115+
116+ # Inline bpftrace programs for each test case
117+ PROGRAMS = {
118+ "call_stack" : """
119+ usdt:{python}:python:function__entry {{
120+ printf("%lld\\ tfunction__entry:%s:%s:%d\\ n",
121+ nsecs, str(arg0), str(arg1), arg2);
122+ }}
123+ usdt:{python}:python:function__return {{
124+ printf("%lld\\ tfunction__return:%s:%s:%d\\ n",
125+ nsecs, str(arg0), str(arg1), arg2);
126+ }}
127+ """ ,
128+ "gc" : """
129+ usdt:{python}:python:function__entry {{
130+ if (str(arg1) == "start") {{ @tracing = 1; }}
131+ }}
132+ usdt:{python}:python:function__return {{
133+ if (str(arg1) == "start") {{ @tracing = 0; }}
134+ }}
135+ usdt:{python}:python:gc__start {{
136+ if (@tracing) {{
137+ printf("%lld\\ tgc__start:%d\\ n", nsecs, arg0);
138+ }}
139+ }}
140+ usdt:{python}:python:gc__done {{
141+ if (@tracing) {{
142+ printf("%lld\\ tgc__done:%lld\\ n", nsecs, arg0);
143+ }}
144+ }}
145+ END {{ clear(@tracing); }}
146+ """ ,
147+ }
148+
149+ # Which test scripts to filter by filename (None = use @tracing flag)
150+ FILTER_BY_FILENAME = {"call_stack" : "call_stack.py" }
151+
152+ # Expected outputs for each test case
153+ # Note: bpftrace captures <module> entry/return and may have slight timing
154+ # differences compared to SystemTap due to probe firing order
155+ EXPECTED = {
156+ "call_stack" : """function__entry:call_stack.py:<module>:0
157+ function__entry:call_stack.py:start:23
158+ function__entry:call_stack.py:function_1:1
159+ function__entry:call_stack.py:function_3:9
160+ function__return:call_stack.py:function_3:10
161+ function__return:call_stack.py:function_1:2
162+ function__entry:call_stack.py:function_2:5
163+ function__entry:call_stack.py:function_1:1
164+ function__entry:call_stack.py:function_3:9
165+ function__return:call_stack.py:function_3:10
166+ function__return:call_stack.py:function_1:2
167+ function__return:call_stack.py:function_2:6
168+ function__entry:call_stack.py:function_3:9
169+ function__return:call_stack.py:function_3:10
170+ function__entry:call_stack.py:function_4:13
171+ function__return:call_stack.py:function_4:14
172+ function__entry:call_stack.py:function_5:18
173+ function__return:call_stack.py:function_5:21
174+ function__return:call_stack.py:start:28
175+ function__return:call_stack.py:<module>:30""" ,
176+ "gc" : """gc__start:0
177+ gc__done:0
178+ gc__start:1
179+ gc__done:0
180+ gc__start:2
181+ gc__done:0
182+ gc__start:2
183+ gc__done:1""" ,
184+ }
185+
186+ def run_case (self , name , optimize_python = None ):
187+ if name not in self .PROGRAMS :
188+ raise unittest .SkipTest (f"No bpftrace program for { name } " )
189+
190+ python_file = abspath (name + ".py" )
191+ python_flags = []
192+ if optimize_python :
193+ python_flags .extend (["-O" ] * optimize_python )
194+
195+ subcommand = [sys .executable ] + python_flags + [python_file ]
196+ program = self .PROGRAMS [name ].format (python = sys .executable )
197+
198+ try :
199+ proc = subprocess .Popen (
200+ ["bpftrace" , "-e" , program , "-c" , " " .join (subcommand )],
201+ stdout = subprocess .PIPE ,
202+ stderr = subprocess .PIPE ,
203+ universal_newlines = True ,
204+ )
205+ stdout , stderr = proc .communicate (timeout = 60 )
206+ except subprocess .TimeoutExpired :
207+ proc .kill ()
208+ raise AssertionError ("bpftrace timed out" )
209+ except (FileNotFoundError , PermissionError ) as e :
210+ raise unittest .SkipTest (f"bpftrace not available: { e } " )
211+
212+ if proc .returncode != 0 :
213+ raise AssertionError (
214+ f"bpftrace failed with code { proc .returncode } :\n { stderr } "
215+ )
216+
217+ # Filter output by filename if specified (bpftrace captures everything)
218+ if name in self .FILTER_BY_FILENAME :
219+ filter_filename = self .FILTER_BY_FILENAME [name ]
220+ filtered_lines = [
221+ line for line in stdout .splitlines ()
222+ if filter_filename in line
223+ ]
224+ stdout = "\n " .join (filtered_lines )
225+
226+ actual_output = normalize_trace_output (stdout )
227+ expected_output = self .EXPECTED [name ].strip ()
228+
229+ return (expected_output , actual_output )
230+
231+ def assert_usable (self ):
232+ # Check if bpftrace is available and can attach to USDT probes
233+ program = f'usdt:{ sys .executable } :python:function__entry {{ printf("probe: success\\ n"); exit(); }}'
234+ try :
235+ proc = subprocess .Popen (
236+ ["bpftrace" , "-e" , program , "-c" , f"{ sys .executable } -c pass" ],
237+ stdout = subprocess .PIPE ,
238+ stderr = subprocess .PIPE ,
239+ universal_newlines = True ,
240+ )
241+ stdout , stderr = proc .communicate (timeout = 10 )
242+ except subprocess .TimeoutExpired :
243+ proc .kill ()
244+ proc .communicate () # Clean up
245+ raise unittest .SkipTest ("bpftrace timed out during usability check" )
246+ except OSError as e :
247+ raise unittest .SkipTest (f"bpftrace not available: { e } " )
248+
249+ # Check for permission errors (bpftrace usually requires root)
250+ if proc .returncode != 0 :
251+ raise unittest .SkipTest (
252+ f"bpftrace(1) failed with code { proc .returncode } : { stderr } "
253+ )
254+
255+ if "probe: success" not in stdout :
256+ raise unittest .SkipTest (
257+ f"bpftrace(1) failed: stdout={ stdout !r} stderr={ stderr !r} "
258+ )
259+
260+
106261@unittest .skipIf (MS_WINDOWS , "Tests not compliant with trace on Windows." )
107262class TraceTests :
108263 # unittest.TestCase options
@@ -127,7 +282,8 @@ def test_function_entry_return(self):
127282 def test_verify_call_opcodes (self ):
128283 """Ensure our call stack test hits all function call opcodes"""
129284
130- opcodes = set (["CALL_FUNCTION" , "CALL_FUNCTION_EX" , "CALL_FUNCTION_KW" ])
285+ # Modern Python uses CALL, CALL_KW, and CALL_FUNCTION_EX
286+ opcodes = set (["CALL" , "CALL_FUNCTION_EX" , "CALL_KW" ])
131287
132288 with open (abspath ("call_stack.py" )) as f :
133289 code_string = f .read ()
@@ -152,9 +308,6 @@ def get_function_instructions(funcname):
152308 def test_gc (self ):
153309 self .run_case ("gc" )
154310
155- def test_line (self ):
156- self .run_case ("line" )
157-
158311
159312class DTraceNormalTests (TraceTests , unittest .TestCase ):
160313 backend = DTraceBackend ()
@@ -175,6 +328,17 @@ class SystemTapOptimizedTests(TraceTests, unittest.TestCase):
175328 backend = SystemTapBackend ()
176329 optimize_python = 2
177330
331+
332+ class BPFTraceNormalTests (TraceTests , unittest .TestCase ):
333+ backend = BPFTraceBackend ()
334+ optimize_python = 0
335+
336+
337+ class BPFTraceOptimizedTests (TraceTests , unittest .TestCase ):
338+ backend = BPFTraceBackend ()
339+ optimize_python = 2
340+
341+
178342class CheckDtraceProbes (unittest .TestCase ):
179343 @classmethod
180344 def setUpClass (cls ):
@@ -235,6 +399,8 @@ def test_check_probes(self):
235399 "Name: audit" ,
236400 "Name: gc__start" ,
237401 "Name: gc__done" ,
402+ "Name: function__entry" ,
403+ "Name: function__return" ,
238404 ]
239405
240406 for probe_name in available_probe_names :
@@ -247,8 +413,6 @@ def test_missing_probes(self):
247413
248414 # Missing probes will be added in the future.
249415 missing_probe_names = [
250- "Name: function__entry" ,
251- "Name: function__return" ,
252416 "Name: line" ,
253417 ]
254418
0 commit comments