Merge pull request #590 from goldshtn/bcc-tp-support

bcc: Tracepoint support in libbpf and BPF

src/cc/libbpf.c

@@ -342,6 +342,34 @@ int bpf_detach_uprobe(const char *event_desc) {
   return bpf_detach_probe(event_desc, "uprobe");
 }
 
+void * bpf_attach_tracepoint(int progfd, const char *tp_category,
+                             const char *tp_name, int pid, int cpu,
+                             int group_fd, perf_reader_cb cb, void *cb_cookie) {
+  char buf[256];
+  struct perf_reader *reader = NULL;
+
+  reader = perf_reader_new(cb, NULL, cb_cookie);
+  if (!reader)
+    goto error;
+
+  snprintf(buf, sizeof(buf), "/sys/kernel/debug/tracing/events/%s/%s",
+           tp_category, tp_name);
+  if (bpf_attach_tracing_event(progfd, buf, reader, pid, cpu, group_fd) < 0)
+    goto error;
+
+  return reader;
+
+error:
+  perf_reader_free(reader);
+  return NULL;
+}
+
+int bpf_detach_tracepoint(const char *tp_category, const char *tp_name) {
+  // Right now, there is nothing to do, but it's a good idea to encourage
+  // callers to detach anything they attach.
+  return 0;
+}
+
 void * bpf_open_perf_buffer(perf_reader_raw_cb raw_cb, void *cb_cookie, int pid, int cpu) {
   int pfd;
   struct perf_event_attr attr = {};

src/libbpf.h

@@ -54,6 +54,11 @@ void * bpf_attach_uprobe(int progfd, const char *event, const char *event_desc,
                          void *cb_cookie);
 int bpf_detach_uprobe(const char *event_desc);
 
+void * bpf_attach_tracepoint(int progfd, const char *tp_category,
+                             const char *tp_name, int pid, int cpu,
+                             int group_fd, perf_reader_cb cb, void *cb_cookie);
+int bpf_detach_tracepoint(const char *tp_category, const char *tp_name);
+
 void * bpf_open_perf_buffer(perf_reader_raw_cb raw_cb, void *cb_cookie, int pid, int cpu);
 
 #define LOG_BUF_SIZE 65536

src/python/bcc/__init__.py

@@ -33,6 +33,7 @@ from .usyms import ProcessSymbols
 
 open_kprobes = {}
 open_uprobes = {}
+open_tracepoints = {}
 tracefile = None
 TRACEFS = "/sys/kernel/debug/tracing"
 _kprobe_limit = 1000
@@ -53,8 +54,14 @@ def cleanup_kprobes():
         if isinstance(k, str):
             desc = "-:uprobes/%s" % k
             lib.bpf_detach_uprobe(desc.encode("ascii"))
+    for k, v in open_tracepoints.items():
+        lib.perf_reader_free(v)
+        if isinstance(k, str):
+            (tp_category, tp_name) = k.split(':')
+            lib.bpf_detach_tracepoint(tp_category, tp_name)
     open_kprobes.clear()
     open_uprobes.clear()
+    open_tracepoints.clear()
     if tracefile:
         tracefile.close()
 
@@ -85,6 +92,7 @@ class BPF(object):
     KPROBE = 2
     SCHED_CLS = 3
     SCHED_ACT = 4
+    TRACEPOINT = 5
 
     _probe_repl = re.compile("[^a-zA-Z0-9_]")
     _sym_caches = {}
@@ -385,8 +393,13 @@ class BPF(object):
 
     @staticmethod
     def open_uprobes():
-            global open_uprobes
-            return open_uprobes
+        global open_uprobes
+        return open_uprobes
+
+    @staticmethod
+    def open_tracepoints():
+        global open_tracepoints
+        return open_tracepoints
 
     @staticmethod
     def detach_kprobe(event):
@@ -453,6 +466,51 @@ class BPF(object):
     def find_library(libname):
         return lib.bcc_procutils_which_so(libname.encode("ascii")).decode()
 
+    def attach_tracepoint(self, tp="", fn_name="", pid=-1, cpu=0, group_fd=-1):
+        """attach_tracepoint(tp="", fn_name="", pid=-1, cpu=0, group_fd=-1)
+
+        Run the bpf function denoted by fn_name every time the kernel tracepoint
+        specified by 'tp' is hit. The optional parameters pid, cpu, and group_fd
+        can be used to filter the probe. The tracepoint specification is simply
+        the tracepoint category and the tracepoint name, separated by a colon.
+        For example: sched:sched_switch, syscalls:sys_enter_bind, etc.
+
+        To obtain a list of kernel tracepoints, use the tplist tool or cat the
+        file /sys/kernel/debug/tracing/available_events.
+
+        Example: BPF(text).attach_tracepoint("sched:sched_switch", "on_switch")
+        """
+
+        fn = self.load_func(fn_name, BPF.TRACEPOINT)
+        (tp_category, tp_name) = tp.split(':')
+        res = lib.bpf_attach_tracepoint(fn.fd, tp_category.encode("ascii"),
+                tp_name.encode("ascii"), pid, cpu, group_fd,
+                self._reader_cb_impl, ct.cast(id(self), ct.py_object))
+        res = ct.cast(res, ct.c_void_p)
+        if not res:
+            raise Exception("Failed to attach BPF to tracepoint")
+        open_tracepoints[tp] = res
+        return self
+
+    def detach_tracepoint(self, tp=""):
+        """detach_tracepoint(tp="")
+
+        Stop running a bpf function that is attached to the kernel tracepoint
+        specified by 'tp'.
+
+        Example: bpf.detach_tracepoint("sched:sched_switch")
+        """
+
+        if tp not in open_tracepoints:
+            raise Exception("Tracepoint %s is not attached" % tp)
+        lib.perf_reader_free(open_tracepoints[tp])
+        (tp_category, tp_name) = tp.split(':')
+        res = lib.bpf_detach_tracepoint(tp_category.encode("ascii"),
+                                        tp_name.encode("ascii"))
+        if res < 0:
+            raise Exception("Failed to detach BPF from tracepoint")
+        del open_tracepoints[tp]
+
     def attach_uprobe(self, name="", sym="", addr=None,
             fn_name="", pid=-1, cpu=0, group_fd=-1):
         """attach_uprobe(name="", sym="", addr=None, fn_name=""

tests/python/CMakeLists.txt

@@ -50,6 +50,8 @@ add_test(NAME py_uprobes WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_uprobes sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_uprobes.py)
 add_test(NAME py_test_stackid WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_stackid sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_stackid.py)
+add_test(NAME py_test_tracepoint WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test_tracepoint sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_tracepoint.py)
 
 add_test(NAME py_test_dump_func WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_dump_func simple ${CMAKE_CURRENT_SOURCE_DIR}/test_dump_func.py)

tests/python/test_tracepoint.py

+#!/usr/bin/env python
+# Copyright (c) Sasha Goldshtein
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+import bcc
+import unittest
+from time import sleep
+import distutils.version
+import os
+
+def kernel_version_ge(major, minor):
+    # True if running kernel is >= X.Y
+    version = distutils.version.LooseVersion(os.uname()[2]).version
+    if version[0] > major:
+        return True
+    if version[0] < major:
+        return False
+    if minor and version[1] < minor:
+        return False
+    return True
+
+@unittest.skipUnless(kernel_version_ge(4,7), "requires kernel >= 4.7")
+class TestTracepoint(unittest.TestCase):
+    def test_tracepoint(self):
+        text = """#include <linux/ptrace.h>
+        struct tp_args {
+            unsigned long long __unused__;
+            char prev_comm[16];
+            pid_t prev_pid;
+            int prev_prio;
+            long prev_state;
+            char next_comm[16];
+            pid_t next_pid;
+            int next_prio;
+        };
+        BPF_HASH(switches, u32, u64);
+        int probe_switch(struct tp_args *args) {
+            if (args == 0)
+                return 0;
+            u64 val = 0;
+            u32 pid = args->next_pid;
+            u64 *existing = switches.lookup_or_init(&pid, &val);
+            (*existing)++;
+            return 0;
+        }
+        """
+        b = bcc.BPF(text=text)
+        b.attach_tracepoint("sched:sched_switch", "probe_switch")
+        sleep(1)
+        total_switches = 0
+        for k, v in b["switches"].items():
+            total_switches += v.value
+        self.assertNotEqual(0, total_switches)
+        b.detach_tracepoint("sched:sched_switch")
+
+if __name__ == "__main__":
+    unittest.main()