usdt: add a test for one probe in both a.out and shared libraries

Signed-off-by: Yonghong Song <yhs@fb.com>

tests/python/CMakeLists.txt

@@ -72,3 +72,5 @@ add_test(NAME py_test_usdt WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test_usdt sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_usdt.py)
 add_test(NAME py_test_usdt2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test_usdt2 sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_usdt2.py)
+add_test(NAME py_test_usdt3 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test_usdt3 sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_usdt3.py)

tests/python/test_usdt3.py

+#!/usr/bin/python
+#
+# USAGE: test_usdt3.py
+#
+# Copyright 2018 Facebook, Inc
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+from __future__ import print_function
+from bcc import BPF, USDT
+from unittest import main, TestCase
+from subprocess import Popen, PIPE
+import ctypes as ct
+import inspect, os, tempfile
+
+class TestUDST(TestCase):
+    def setUp(self):
+        common_h = b"""
+#include "folly/tracing/StaticTracepoint.h"
+
+static inline void record_val(int val)
+{
+  FOLLY_SDT(test, probe, val);
+}
+
+extern void record_a(int val);
+extern void record_b(int val);
+"""
+
+        a_c = b"""
+#include <stdio.h>
+#include "common.h"
+
+void record_a(int val)
+{
+    record_val(val);
+}
+"""
+
+        b_c = b"""
+#include <stdio.h>
+#include "common.h"
+
+void record_b(int val)
+{
+    record_val(val);
+}
+"""
+
+        m_c = b"""
+#include <stdio.h>
+#include <unistd.h>
+#include "common.h"
+
+int main() {
+   while (1) {
+     record_a(1);
+     record_b(2);
+     record_val(3);
+     sleep(1);
+   }
+   return 0;
+}
+"""
+        # BPF program
+        self.bpf_text = """
+BPF_PERF_OUTPUT(event);
+int do_trace(struct pt_regs *ctx) {
+    int result = 0;
+    bpf_usdt_readarg(1, ctx, &result);
+    event.perf_submit(ctx, &result, sizeof(result));
+    return 0;
+};
+"""
+
+        def _create_file(name, text):
+            text_file = open(name, "w")
+            text_file.write(text)
+            text_file.close()
+
+        # Create source files
+        self.tmp_dir = tempfile.mkdtemp()
+        print("temp directory: " + self.tmp_dir)
+        _create_file(self.tmp_dir + "/common.h", common_h)
+        _create_file(self.tmp_dir + "/a.c", a_c)
+        _create_file(self.tmp_dir + "/b.c", b_c)
+        _create_file(self.tmp_dir + "/m.c", m_c)
+
+        # Compilation
+        # the usdt test:probe exists in liba.so, libb.so and a.out
+        include_path = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + "/include"
+        a_src = self.tmp_dir + "/a.c"
+        a_obj = self.tmp_dir + "/a.o"
+        a_lib = self.tmp_dir + "/liba.so"
+        b_src = self.tmp_dir + "/b.c"
+        b_obj = self.tmp_dir + "/b.o"
+        b_lib = self.tmp_dir + "/libb.so"
+        m_src = self.tmp_dir + "/m.c"
+        m_bin = self.tmp_dir + "/a.out"
+        m_linker_opt = " -L" + self.tmp_dir + " -la -lb"
+        self.assertEqual(os.system("gcc -I" + include_path + " -fpic -c -o " + a_obj + " " + a_src), 0)
+        self.assertEqual(os.system("gcc -I" + include_path + " -fpic -c -o " + b_obj + " " + b_src), 0)
+        self.assertEqual(os.system("gcc -shared -o " + a_lib + " " + a_obj), 0)
+        self.assertEqual(os.system("gcc -shared -o " + b_lib + " " + b_obj), 0)
+        self.assertEqual(os.system("gcc -I" + include_path + " " + m_src + " -o " + m_bin + m_linker_opt), 0)
+
+        # Run the application
+        self.app = Popen([m_bin], env=dict(os.environ, LD_LIBRARY_PATH=self.tmp_dir))
+        # os.system("tplist.py -vvv -p " + str(self.app.pid))
+
+    def test_attach1(self):
+        # enable USDT probe from given PID and verifier generated BPF programs
+        u = USDT(pid=int(self.app.pid))
+        u.enable_probe(probe="probe", fn_name="do_trace")
+        b = BPF(text=self.bpf_text, usdt_contexts=[u])
+
+        # processing events
+        self.probe_value_1 = 0
+        self.probe_value_2 = 0
+        self.probe_value_3 = 0
+        self.probe_value_other = 0
+
+        def print_event(cpu, data, size):
+            result = ct.cast(data, ct.POINTER(ct.c_int)).contents
+            if result.value == 1:
+                self.probe_value_1 = 1
+            elif result.value == 2:
+                self.probe_value_2 = 1
+            elif result.value == 3:
+                self.probe_value_3 = 1
+            else:
+                self.probe_value_other = 1
+
+        b["event"].open_perf_buffer(print_event)
+        for i in range(6):
+            b.kprobe_poll()
+
+        self.assertTrue(self.probe_value_1 != 0)
+        self.assertTrue(self.probe_value_2 != 0)
+        self.assertTrue(self.probe_value_3 != 0)
+        self.assertTrue(self.probe_value_other == 0)
+
+    def tearDown(self):
+        # kill the subprocess, clean the environment
+        self.app.kill()
+        self.app.wait()
+        os.system("rm -rf " + self.tmp_dir)
+
+if __name__ == "__main__":
+    main()