1#!/usr/bin/env python 2# 3# USAGE: test_usdt2.py 4# 5# Copyright 2017 Facebook, Inc 6# Licensed under the Apache License, Version 2.0 (the "License") 7 8from __future__ import print_function 9from bcc import BPF, USDT 10from unittest import main, TestCase 11from subprocess import Popen, PIPE 12from tempfile import NamedTemporaryFile 13import ctypes as ct 14import inspect 15import os 16import signal 17 18class TestUDST(TestCase): 19 def setUp(self): 20 # Application, minimum, to define three trace points 21 app_text = b""" 22#include <stdlib.h> 23#include <unistd.h> 24#include "folly/tracing/StaticTracepoint.h" 25 26int main(int argc, char **argv) { 27 int t = atoi(argv[1]); 28 while (1) { 29 FOLLY_SDT(test, probe_point_1, t); 30 FOLLY_SDT(test, probe_point_2, t + 1); 31 FOLLY_SDT(test, probe_point_3, t + 2); 32 sleep(1); 33 } 34 return 1; 35} 36""" 37 # BPF program 38 self.bpf_text = """ 39#include <uapi/linux/ptrace.h> 40 41BPF_PERF_OUTPUT(event1); 42BPF_PERF_OUTPUT(event2); 43BPF_PERF_OUTPUT(event3); 44BPF_PERF_OUTPUT(event4); 45BPF_PERF_OUTPUT(event5); 46BPF_PERF_OUTPUT(event6); 47 48int do_trace1(struct pt_regs *ctx) { 49 u32 pid = bpf_get_current_pid_tgid(); 50 int result = 0; 51 bpf_usdt_readarg(1, ctx, &result); 52 if (FILTER) 53 event1.perf_submit(ctx, &result, sizeof(result)); 54 else 55 event4.perf_submit(ctx, &result, sizeof(result)); 56 return 0; 57}; 58int do_trace2(struct pt_regs *ctx) { 59 u32 pid = bpf_get_current_pid_tgid(); 60 int result = 0; 61 bpf_usdt_readarg(1, ctx, &result); 62 if (FILTER) 63 event2.perf_submit(ctx, &result, sizeof(result)); 64 else 65 event5.perf_submit(ctx, &result, sizeof(result)); 66 return 0; 67} 68int do_trace3(struct pt_regs *ctx) { 69 u32 pid = bpf_get_current_pid_tgid(); 70 int result = 0; 71 bpf_usdt_readarg(1, ctx, &result); 72 if (FILTER) 73 event3.perf_submit(ctx, &result, sizeof(result)); 74 else 75 event6.perf_submit(ctx, &result, sizeof(result)); 76 return 0; 77} 78""" 79 80 # Compile and run the application 81 self.ftemp = NamedTemporaryFile(delete=False) 82 self.ftemp.close() 83 comp = Popen(["gcc", "-I", "%s/include" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))), 84 "-x", "c", "-o", self.ftemp.name, "-"], 85 stdin=PIPE) 86 comp.stdin.write(app_text) 87 comp.stdin.close() 88 self.assertEqual(comp.wait(), 0) 89 90 # create 3 applications, 2 applications will have usdt attached and 91 # the third one does not, and the third one should not call into 92 # bpf program. 93 self.app = Popen([self.ftemp.name, "1"]) 94 self.app2 = Popen([self.ftemp.name, "11"]) 95 self.app3 = Popen([self.ftemp.name, "21"]) 96 97 def test_attach1(self): 98 # Enable USDT probe from given PID and verifier generated BPF programs. 99 u = USDT(pid=int(self.app.pid)) 100 u.enable_probe(probe="probe_point_1", fn_name="do_trace1") 101 u.enable_probe(probe="probe_point_2", fn_name="do_trace2") 102 u2 = USDT(pid=int(self.app2.pid)) 103 u2.enable_probe(probe="probe_point_2", fn_name="do_trace2") 104 u2.enable_probe(probe="probe_point_3", fn_name="do_trace3") 105 self.bpf_text = self.bpf_text.replace("FILTER", "pid == %d" % self.app.pid) 106 b = BPF(text=self.bpf_text, usdt_contexts=[u, u2]) 107 108 # Event states for each event: 109 # 0 - probe not caught, 1 - probe caught with correct value, 110 # 2 - probe caught with incorrect value 111 self.evt_st_1 = 0 112 self.evt_st_2 = 0 113 self.evt_st_3 = 0 114 self.evt_st_4 = 0 115 self.evt_st_5 = 0 116 self.evt_st_6 = 0 117 118 def check_event_val(data, event_state, expected_val): 119 result = ct.cast(data, ct.POINTER(ct.c_int)).contents 120 if result.value == expected_val: 121 if (event_state == 0 or event_state == 1): 122 return 1 123 return 2 124 125 def print_event1(cpu, data, size): 126 self.evt_st_1 = check_event_val(data, self.evt_st_1, 1) 127 128 def print_event2(cpu, data, size): 129 self.evt_st_2 = check_event_val(data, self.evt_st_2, 2) 130 131 def print_event3(cpu, data, size): 132 self.evt_st_3 = check_event_val(data, self.evt_st_3, 3) 133 134 def print_event4(cpu, data, size): 135 self.evt_st_4 = check_event_val(data, self.evt_st_4, 11) 136 137 def print_event5(cpu, data, size): 138 self.evt_st_5 = check_event_val(data, self.evt_st_5, 12) 139 140 def print_event6(cpu, data, size): 141 self.evt_st_6 = check_event_val(data, self.evt_st_6, 13) 142 143 # loop with callback to print_event 144 b["event1"].open_perf_buffer(print_event1) 145 b["event2"].open_perf_buffer(print_event2) 146 b["event3"].open_perf_buffer(print_event3) 147 b["event4"].open_perf_buffer(print_event4) 148 b["event5"].open_perf_buffer(print_event5) 149 b["event6"].open_perf_buffer(print_event6) 150 151 # three iterations to make sure we get some probes and have time to process them 152 for i in range(3): 153 b.perf_buffer_poll() 154 155 # note that event1 and event4 do not really fire, so their state should be 0 156 # use separate asserts so that if test fails we know which one is the culprit 157 self.assertTrue(self.evt_st_1 == 1) 158 self.assertTrue(self.evt_st_2 == 1) 159 self.assertTrue(self.evt_st_3 == 0) 160 self.assertTrue(self.evt_st_4 == 0) 161 self.assertTrue(self.evt_st_5 == 1) 162 self.assertTrue(self.evt_st_6 == 1) 163 164 def tearDown(self): 165 # kill the subprocess, clean the environment 166 self.app.kill() 167 self.app.wait() 168 self.app2.kill() 169 self.app2.wait() 170 self.app3.kill() 171 self.app3.wait() 172 os.unlink(self.ftemp.name) 173 174if __name__ == "__main__": 175 main() 176