USDT Python API and example (#624)

* Python USDT API

Code from @vmg

* Basic USDT example

* retire procstat.py

* improve/fix USDT exceptions

examples/tracing/nodejs_http_server.py

+#!/usr/bin/python
+#
+# nodejs_http_server    Basic example of node.js USDT tracing.
+#                       For Linux, uses BCC, BPF. Embedded C.
+#
+# USAGE: nodejs_http_server PID
+#
+# Copyright 2016 Netflix, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+from __future__ import print_function
+from bcc import BPF, USDT
+import sys
+
+if len(sys.argv) < 2:
+    print("USAGE: nodejs_http_server PID")
+    exit()
+pid = sys.argv[1]
+debug = 0
+
+# load BPF program
+bpf_text = """
+#include <uapi/linux/ptrace.h>
+int do_trace(struct pt_regs *ctx) {
+    uint64_t addr;
+    char path[128];
+    bpf_usdt_readarg(6, ctx, &addr);
+    bpf_probe_read(&path, sizeof(path), (void *)addr);
+    bpf_trace_printk("path:%s\\n", path);
+    return 0;
+};
+"""
+
+# enable USDT probe from given PID
+u = USDT(pid=int(pid))
+u.enable_probe(probe="http__server__request", fn_name="do_trace")
+if debug:
+    print(u.get_text())
+    print(bpf_text)
+
+# initialize BPF
+b = BPF(text=bpf_text, usdt=u)
+
+# header
+print("%-18s %-16s %-6s %s" % ("TIME(s)", "COMM", "PID", "ARGS"))
+
+# format output
+while 1:
+    try:
+        (task, pid, cpu, flags, ts, msg) = b.trace_fields()
+    except ValueError:
+        print("value error")
+        continue
+    print("%-18.9f %-16s %-6d %s" % (ts, task, pid, msg))

src/cc/bcc_usdt.h

@@ -27,7 +27,7 @@ void *bcc_usdt_new_frompath(const char *path);
 void bcc_usdt_close(void *usdt);
 
 int bcc_usdt_enable_probe(void *, const char *, const char *);
-char *bcc_usdt_genargs(void *);
+const char *bcc_usdt_genargs(void *);
 
 typedef void (*bcc_usdt_uprobe_cb)(const char *, const char *, uint64_t, int);
 void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback);

src/cc/usdt.cc

@@ -319,12 +319,16 @@ int bcc_usdt_enable_probe(void *usdt, const char *probe_name,
   return ctx->enable_probe(probe_name, fn_name) ? 0 : -1;
 }
 
-char *bcc_usdt_genargs(void *usdt) {
+const char *bcc_usdt_genargs(void *usdt) {
+  static std::string storage_;
+
   USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
   std::ostringstream stream;
   if (!ctx->generate_usdt_args(stream))
     return nullptr;
-  return strdup(stream.str().c_str());
+
+  storage_ = stream.str();
+  return storage_.c_str();
 }
 
 void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback) {

src/lua/bcc/usdt.lua

@@ -14,8 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ]]
 local ffi = require("ffi")
-ffi.cdef "void free(void *ptr);"
-
 local libbcc = require("bcc.libbcc")
 local Usdt = class("USDT")
 
@@ -56,10 +54,7 @@ end
 function Usdt:_get_text()
   local argc = libbcc.bcc_usdt_genargs(self.context)
   assert(argc ~= nil)
-
-  local text = ffi.string(argc)
-  ffi.C.free(argc)
-  return text
+  return ffi.string(argc)
 end
 
 function Usdt:_attach_uprobes(bpf)

src/python/bcc/__init__.py

@@ -25,7 +25,6 @@ import sys
 basestring = (unicode if sys.version_info[0] < 3 else str)
 
 from .libbcc import lib, _CB_TYPE, bcc_symbol
-from .procstat import ProcStat, ProcUtils
 from .table import Table
 from .tracepoint import Tracepoint
 from .perf import Perf
@@ -117,7 +116,8 @@ class BPF(object):
                     raise Exception("Could not find file %s" % filename)
         return filename
 
-    def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0, cflags=[]):
+    def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0,
+            cflags=[], usdt=None):
         """Create a a new BPF module with the given source code.
 
         Note:
@@ -147,6 +147,8 @@ class BPF(object):
         self.tables = {}
         cflags_array = (ct.c_char_p * len(cflags))()
         for i, s in enumerate(cflags): cflags_array[i] = s.encode("ascii")
+        if usdt and text: text = usdt.get_text() + text
+
         if text:
             self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"),
                     self.debug, cflags_array, len(cflags_array))
@@ -163,6 +165,8 @@ class BPF(object):
         if not self.module:
             raise Exception("Failed to compile BPF module %s" % src_file)
 
+        if usdt: usdt.attach_uprobes(self)
+
         # If any "kprobe__" or "tracepoint__" prefixed functions were defined,
         # they will be loaded and attached here.
         self._trace_autoload()
@@ -785,4 +789,4 @@ class BPF(object):
             self.tracefile.close()
 
 
-from .usdt import USDTReader
+from .usdt import USDT

src/python/bcc/libbcc.py

@@ -135,3 +135,23 @@ lib.bcc_symcache_resolve_name.argtypes = [
 
 lib.bcc_symcache_refresh.restype = None
 lib.bcc_symcache_refresh.argtypes = [ct.c_void_p]
+
+lib.bcc_usdt_new_frompid.restype = ct.c_void_p
+lib.bcc_usdt_new_frompid.argtypes = [ct.c_int]
+
+lib.bcc_usdt_new_frompath.restype = ct.c_void_p
+lib.bcc_usdt_new_frompath.argtypes = [ct.c_char_p]
+
+lib.bcc_usdt_close.restype = None
+lib.bcc_usdt_close.argtypes = [ct.c_void_p]
+
+lib.bcc_usdt_enable_probe.restype = ct.c_int
+lib.bcc_usdt_enable_probe.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_char_p]
+
+lib.bcc_usdt_genargs.restype = ct.c_char_p
+lib.bcc_usdt_genargs.argtypes = [ct.c_void_p]
+
+_USDT_CB = ct.CFUNCTYPE(None, ct.c_char_p, ct.c_char_p, ct.c_ulonglong, ct.c_int)
+
+lib.bcc_usdt_foreach_uprobe.restype = None
+lib.bcc_usdt_foreach_uprobe.argtypes = [ct.c_void_p, _USDT_CB]

src/python/bcc/procstat.py

-# Copyright 2016 Sasha Goldshtein
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import os
-
-class ProcStat(object):
-        def __init__(self, pid):
-                self.pid = pid
-                self.exe = self._get_exe()
-                self.start_time = self._get_start_time()
-
-        def is_stale(self):
-                return self.exe != self._get_exe() or \
-                       self.start_time != self._get_start_time()
-
-        def _get_exe(self):
-                return os.popen("readlink -f /proc/%d/exe" % self.pid).read()
-
-        def _get_start_time(self):
-                return os.popen("cut -d' ' -f 22 /proc/%d/stat" %
-                                self.pid).read()
-
-class ProcUtils(object):
-        @staticmethod
-        def get_load_address(pid, bin_path):
-                """
-                get_load_address(pid, bin_path)
-
-                Returns the address at which the specified module is loaded
-                in the specified process. The module path must match exactly
-                the file system path, not a symbolic link.
-                """
-                with open("/proc/%d/maps" % pid) as m:
-                        maps = m.readlines()
-                addrs = map(lambda l: l.split('-')[0],
-                            filter(lambda l: bin_path in l, maps)
-                            )
-                if len(addrs) == 0:
-                        raise ValueError("lib %s not loaded in pid %d"
-                                        % (bin_path, pid))
-                return int(addrs[0], 16)
-
-        @staticmethod
-        def get_modules(pid):
-                """
-                get_modules(pid)
-
-                Returns a list of all the modules loaded into the specified
-                process. Modules are enumerated by looking at /proc/$PID/maps
-                and returning the module name for regions that contain
-                executable code.
-                """
-                with open("/proc/%d/maps" % pid) as f:
-                        maps = f.readlines()
-                modules = []
-                for line in maps:
-                        parts = line.strip().split()
-                        if len(parts) < 6:
-                                continue
-                        if parts[5][0] == '[' or not 'x' in parts[1]:
-                                continue
-                        modules.append(parts[5])
-                return modules
-
-        @staticmethod
-        def is_shared_object(bin_path):
-                """
-                is_shared_object(bin_path)
-
-                Returns whether the specified binary is a shared object, rather
-                than an executable. If it is neither, an error is raised.
-                """
-                mime_type = os.popen("file --mime-type -b %s" % bin_path
-                                ).read().strip()
-                if mime_type == "application/x-sharedlib":
-                        return True
-                if mime_type == "application/x-executable":
-                        return False
-                raise ValueError("invalid mime type %s for binary %s" %
-                                (mime_type, bin_path))
-
-        @staticmethod
-        def traverse_symlink(path):
-                """Returns the actual path behind the specified symlink."""
-                return os.popen("readlink -f %s" % path).read().strip()
-
-        @staticmethod
-        def which(bin_path):
-                """
-                which(bin_path)
-
-                Traverses the PATH environment variable, looking for the first
-                directory that contains an executable file named bin_path, and
-                returns the full path to that file, or None if no such file
-                can be found. This is meant to replace invocations of the
-                "which" shell utility, which doesn't have portable semantics
-                for skipping aliases.
-                """
-                # Source: http://stackoverflow.com/a/377028
-                def is_exe(fpath):
-                        return os.path.isfile(fpath) and \
-                               os.access(fpath, os.X_OK)
-
-                fpath, fname = os.path.split(bin_path)
-                if fpath:
-                        if is_exe(bin_path):
-                                return bin_path
-                else:
-                        for path in os.environ["PATH"].split(os.pathsep):
-                                path = path.strip('"')
-                                exe_file = os.path.join(path, bin_path)
-                                if is_exe(exe_file):
-                                        return exe_file
-                return None

tools/argdist.py

@@ -12,7 +12,7 @@
 # Licensed under the Apache License, Version 2.0 (the "License")
 # Copyright (C) 2016 Sasha Goldshtein.
 
-from bcc import BPF, Tracepoint, Perf, ProcUtils, USDTReader
+from bcc import BPF, Tracepoint, Perf, USDT
 from time import sleep, strftime
 import argparse
 import re
@@ -175,8 +175,8 @@ u64 __time = bpf_ktime_get_ns();
                         self._bail("no exprs specified")
                 self.exprs = exprs.split(',')
 
-        def __init__(self, type, specifier, pid):
-                self.pid = pid
+        def __init__(self, bpf, type, specifier):
+                self.pid = bpf.args.pid
                 self.raw_spec = specifier
                 self._validate_specifier()
 
@@ -198,8 +198,10 @@ u64 __time = bpf_ktime_get_ns();
                         self.function = "perf_trace_" + self.function
                 elif self.probe_type == "u":
                         self.library = parts[1]
-                        self._find_usdt_probe()
-                        self._enable_usdt_probe()
+                        self.probe_func_name = "%s_probe%d" % \
+                                (self.function, Probe.next_probe_index)
+                        bpf.enable_usdt_probe(self.function,
+                                        fn_name=self.probe_func_name)
                 else:
                         self.library = parts[1]
                 self.is_user = len(self.library) > 0
@@ -240,26 +242,8 @@ u64 __time = bpf_ktime_get_ns();
                         (self.function, Probe.next_probe_index)
                 Probe.next_probe_index += 1
 
-        def _enable_usdt_probe(self):
-                if self.usdt.need_enable():
-                        if self.pid is None:
-                                self._bail("probe needs pid to enable")
-                        self.usdt.enable(self.pid)
-
-        def _disable_usdt_probe(self):
-                if self.probe_type == "u" and self.usdt.need_enable():
-                        self.usdt.disable(self.pid)
-
         def close(self):
-                self._disable_usdt_probe()
-
-        def _find_usdt_probe(self):
-                reader = USDTReader(bin_path=self.library)
-                for probe in reader.probes:
-                        if probe.name == self.function:
-                                self.usdt = probe
-                                return
-                self._bail("unrecognized USDT probe %s" % self.function)
+                pass
 
         def _substitute_exprs(self):
                 def repl(expr):
@@ -279,12 +263,18 @@ u64 __time = bpf_ktime_get_ns();
                         return "%s v%d;\n" % (self.expr_types[i], i)
 
         def _generate_field_assignment(self, i):
+                text = ""
+                if self.probe_type == "u" and self.exprs[i][0:3] == "arg":
+                    text = ("        u64 %s;\n" + 
+                           "        bpf_usdt_readarg(%s, ctx, &%s);\n") % \
+                           (self.exprs[i], self.exprs[i][3], self.exprs[i])
                 if self._is_string(self.expr_types[i]):
-                        return ("        bpf_probe_read(&__key.v%d.s," +
+                        return (text + "        bpf_probe_read(&__key.v%d.s," +
                                 " sizeof(__key.v%d.s), (void *)%s);\n") % \
                                 (i, i, self.exprs[i])
                 else:
-                        return "        __key.v%d = %s;\n" % (i, self.exprs[i])
+                        return text + "        __key.v%d = %s;\n" % \
+                               (i, self.exprs[i])
 
         def _generate_hash_decl(self):
                 if self.type == "hist":
@@ -331,7 +321,7 @@ u64 __time = bpf_ktime_get_ns();
                 probe_text = """
 DATA_DECL
 
-QUALIFIER int PROBENAME(struct pt_regs *ctx SIGNATURE)
+int PROBENAME(struct pt_regs *ctx SIGNATURE)
 {
         PID_FILTER
         PREFIX
@@ -342,7 +332,6 @@ QUALIFIER int PROBENAME(struct pt_regs *ctx SIGNATURE)
 }
 """
                 prefix = ""
-                qualifier = ""
                 signature = ""
 
                 # If any entry arguments are probed in a ret probe, we need
@@ -357,10 +346,6 @@ QUALIFIER int PROBENAME(struct pt_regs *ctx SIGNATURE)
                 if self.probe_type == "t":
                         program += self.tp.generate_struct()
                         prefix += self.tp.generate_get_struct()
-                elif self.probe_type == "u":
-                        qualifier = "static inline"
-                        signature = ", int __loc_id"
-                        prefix += self.usdt.generate_usdt_cases()
                 elif self.probe_type == "p" and len(self.signature) > 0:
                         # Only entry uprobes/kprobes can have user-specified
                         # signatures. Other probes force it to ().
@@ -380,12 +365,6 @@ QUALIFIER int PROBENAME(struct pt_regs *ctx SIGNATURE)
                         "1" if len(self.filter) == 0 else self.filter)
                 program = program.replace("COLLECT", collect)
                 program = program.replace("PREFIX", prefix)
-                program = program.replace("QUALIFIER", qualifier)
-
-                if self.probe_type == "u":
-                        self.usdt_thunk_names = []
-                        program += self.usdt.generate_usdt_thunks(
-                                self.probe_func_name, self.usdt_thunk_names)
 
                 return program
 
@@ -396,13 +375,7 @@ QUALIFIER int PROBENAME(struct pt_regs *ctx SIGNATURE)
                 if libpath is None or len(libpath) == 0:
                         self._bail("unable to find library %s" % self.library)
 
-                if self.probe_type == "u":
-                        for i, location in enumerate(self.usdt.locations):
-                                self.bpf.attach_uprobe(name=libpath,
-                                        addr=location.address,
-                                        fn_name=self.usdt_thunk_names[i],
-                                        pid=self.pid or -1)
-                elif self.probe_type == "r":
+                if self.probe_type == "r":
                         self.bpf.attach_uretprobe(name=libpath,
                                                   sym=self.function,
                                                   fn_name=self.probe_func_name,
@@ -423,6 +396,7 @@ QUALIFIER int PROBENAME(struct pt_regs *ctx SIGNATURE)
 
         def attach(self, bpf):
                 self.bpf = bpf
+                if self.probe_type == "u": return;
                 if self.is_user:
                         self._attach_u()
                 else:
@@ -604,19 +578,23 @@ argdist -p 2780 -z 120 \\
                   metavar="header",
                   help="additional header files to include in the BPF program")
                 self.args = parser.parse_args()
+                self.usdt_ctx = None
 
         def _create_probes(self):
                 self.probes = []
                 for specifier in (self.args.countspecifier or []):
-                        self.probes.append(Probe(
-                                "freq", specifier, self.args.pid))
+                        self.probes.append(Probe(self, "freq", specifier))
                 for histspecifier in (self.args.histspecifier or []):
-                        self.probes.append(
-                                Probe("hist", histspecifier, self.args.pid))
+                        self.probes.append(Probe(self, "hist", histspecifier))
                 if len(self.probes) == 0:
                         print("at least one specifier is required")
                         exit()
 
+        def enable_usdt_probe(self, probe_name, fn_name):
+                if not self.usdt_ctx:
+                        self.usdt_ctx = USDT(pid=self.args.pid)
+                self.usdt_ctx.enable_probe(probe_name, fn_name)
+
         def _generate_program(self):
                 bpf_source = """
 struct __string_t { char s[%d]; };
@@ -632,8 +610,9 @@ struct __string_t { char s[%d]; };
                 for probe in self.probes:
                         bpf_source += probe.generate_text()
                 if self.args.verbose:
+                        if self.usdt_ctx: print(self.usdt_ctx.get_text())
                         print(bpf_source)
-                self.bpf = BPF(text=bpf_source)
+                self.bpf = BPF(text=bpf_source, usdt=self.usdt_ctx)
 
         def _attach(self):
                 Tracepoint.attach(self.bpf)

tools/trace.py

@@ -9,7 +9,7 @@
 # Licensed under the Apache License, Version 2.0 (the "License")
 # Copyright (C) 2016 Sasha Goldshtein.
 
-from bcc import BPF, Tracepoint, Perf, ProcUtils, USDTReader
+from bcc import BPF, Tracepoint, Perf, ProcUtils, USDT
 from time import sleep, strftime
 import argparse
 import re