Fix argdist, trace, tplist to use the libbcc USDT support (#698)

* Allow argdist to enable USDT probes without a pid

The current code would only pass the pid to the USDT
class, thereby not allowing USDT probes to be enabled
from the binary path only. If the probe doesn't have
a semaphore, it can actually be enabled for all
processes in a uniform fashion -- which is now
supported.

* Reintroduce USDT support into tplist

To print USDT probe information, tplist needs an API
to return the probe data, including the number of
arguments and locations for each probe. This commit
introduces this API, called bcc_usdt_foreach, and
invokes it from the revised tplist implementation.

Although the result is not 100% identical to the
original tplist, which could also print the probe
argument information, this is not strictly required
for users of the argdist and trace tools, which is
why it was omitted for now.

* Fix trace.py tracepoint support

Somehow, the import of the Perf class was omitted
from tracepoint.py, which would cause failures when
trace enables kernel tracepoints.

* trace: Native bcc USDT support

trace now works again by using the new bcc USDT support
instead of the home-grown Python USDT parser. This
required an additional change in the BPF Python API
to allow multiple USDT context objects to be passed to
the constructor in order to support multiple USDT
probes in a single invocation of trace. Otherwise, the
USDT-related code in trace was greatly simplified, and
uses the `bpf_usdt_readarg` macros to obtain probe
argument values.

One minor inconvenience that was introduced in the bcc
USDT API is that USDT probes with multiple locations
that reside in a shared object *must* have a pid
specified to enable, even if they don't have an
associated semaphore. The reason is that the bcc USDT
code figures out which location invoked the probe by
inspecting `ctx->ip`, which, for shared objects, can
only be determined when the specific process context is
available to figure out where the shared object was
loaded. This limitation did not previously exist,
because instead of looking at `ctx->ip`, the Python
USDT reader generated separate code for each probe
location with an incrementing identifier. It's not a
very big deal because it only means that some probes
can't be enabled without specifying a process id, which
is almost always desired anyway for USDT probes.

argdist has not yet been retrofitted with support for
multiple USDT probes, and needs to be updated in a
separate commit.

* argdist: Support multiple USDT probes

argdist now supports multiple USDT probes, as it did
before the transition to the native bcc USDT support.
This requires aggregating the USDT objects from each
probe and passing them together to the BPF constructor
when the probes are initialized and attached.

Also add a more descriptive exception message to the
USDT class when it fails to enable a probe.

Author: goldshtn

src/cc/bcc_usdt.h

@@ -26,6 +26,18 @@ void *bcc_usdt_new_frompid(int pid);
 void *bcc_usdt_new_frompath(const char *path);
 void bcc_usdt_close(void *usdt);
 
+struct bcc_usdt {
+    const char *provider;
+    const char *name;
+    const char *bin_path;
+    uint64_t semaphore;
+    int num_locations;
+    int num_arguments;
+};
+
+typedef void (*bcc_usdt_cb)(struct bcc_usdt *);
+void bcc_usdt_foreach(void *usdt, bcc_usdt_cb callback);
+
 int bcc_usdt_enable_probe(void *, const char *, const char *);
 const char *bcc_usdt_genargs(void *);
 

src/cc/usdt.cc

@@ -24,6 +24,7 @@
 #include "bcc_proc.h"
 #include "usdt.h"
 #include "vendor/tinyformat.hpp"
+#include "bcc_usdt.h"
 
 namespace USDT {
 
@@ -255,6 +256,19 @@ bool Context::enable_probe(const std::string &probe_name,
   return p && p->enable(fn_name);
 }
 
+void Context::each(each_cb callback) {
+  for (const auto &probe : probes_) {
+    struct bcc_usdt info = {0};
+    info.provider = probe->provider().c_str();
+    info.bin_path = probe->bin_path().c_str();
+    info.name = probe->name().c_str();
+    info.semaphore = probe->semaphore();
+    info.num_locations = probe->num_locations();
+    info.num_arguments = probe->num_arguments();
+    callback(&info);
+  }
+}
+
 void Context::each_uprobe(each_uprobe_cb callback) {
   for (auto &p : probes_) {
     if (!p->enabled())
@@ -288,7 +302,6 @@ Context::~Context() {
 }
 
 extern "C" {
-#include "bcc_usdt.h"
 
 void *bcc_usdt_new_frompid(int pid) {
   USDT::Context *ctx = new USDT::Context(pid);
@@ -331,6 +344,11 @@ const char *bcc_usdt_genargs(void *usdt) {
   return storage_.c_str();
 }
 
+void bcc_usdt_foreach(void *usdt, bcc_usdt_cb callback) {
+  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
+  ctx->each(callback);
+}
+
 void bcc_usdt_foreach_uprobe(void *usdt, bcc_usdt_uprobe_cb callback) {
   USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
   ctx->each_uprobe(callback);

src/cc/usdt.h

@@ -23,6 +23,8 @@
 #include "syms.h"
 #include "vendor/optional.hpp"
 
+struct bcc_usdt;
+
 namespace USDT {
 
 using std::experimental::optional;
@@ -148,6 +150,7 @@ class Probe {
 
   size_t num_locations() const { return locations_.size(); }
   size_t num_arguments() const { return locations_.front().arguments_.size(); }
+  uint64_t semaphore()   const { return semaphore_; }
 
   uint64_t address(size_t n = 0) const { return locations_[n].address_; }
   bool usdt_getarg(std::ostream &stream);
@@ -194,6 +197,9 @@ class Context {
   bool enable_probe(const std::string &probe_name, const std::string &fn_name);
   bool generate_usdt_args(std::ostream &stream);
 
+  typedef void (*each_cb)(struct bcc_usdt *);
+  void each(each_cb callback);
+
   typedef void (*each_uprobe_cb)(const char *, const char *, uint64_t, int);
   void each_uprobe(each_uprobe_cb callback);
 };

src/python/bcc/__init__.py

@@ -149,7 +149,7 @@ def is_exe(fpath):
         return None
 
     def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0,
-            cflags=[], usdt=None):
+            cflags=[], usdt_contexts=[]):
         """Create a a new BPF module with the given source code.
 
         Note:
@@ -179,7 +179,15 @@ def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0,
         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:
+            for usdt_context in usdt_contexts:
+                usdt_text = usdt_context.get_text()
+                if usdt_text is None:
+                    raise Exception("can't generate USDT probe arguments; " +
+                                    "possible cause is missing pid when a " +
+                                    "probe in a shared object has multiple " +
+                                    "locations")
+                text = usdt_context.get_text() + text
 
         if text:
             self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"),
@@ -197,7 +205,8 @@ def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0,
         if not self.module:
             raise Exception("Failed to compile BPF module %s" % src_file)
 
-        if usdt: usdt.attach_uprobes(self)
+        for usdt_context in usdt_contexts:
+            usdt_context.attach_uprobes(self)
 
         # If any "kprobe__" or "tracepoint__" prefixed functions were defined,
         # they will be loaded and attached here.

src/python/bcc/libbcc.py

@@ -157,7 +157,23 @@ class bcc_symbol(ct.Structure):
 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)
+class bcc_usdt(ct.Structure):
+    _fields_ = [
+            ('provider', ct.c_char_p),
+            ('name', ct.c_char_p),
+            ('bin_path', ct.c_char_p),
+            ('semaphore', ct.c_ulonglong),
+            ('num_locations', ct.c_int),
+            ('num_arguments', ct.c_int),
+        ]
+
+_USDT_CB = ct.CFUNCTYPE(None, ct.POINTER(bcc_usdt))
+
+lib.bcc_usdt_foreach.restype = None
+lib.bcc_usdt_foreach.argtypes = [ct.c_void_p, _USDT_CB]
+
+_USDT_PROBE_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]
+lib.bcc_usdt_foreach_uprobe.argtypes = [ct.c_void_p, _USDT_PROBE_CB]

src/python/bcc/tracepoint.py

@@ -16,6 +16,7 @@
 import multiprocessing
 import os
 import re
+from .perf import Perf
 
 class Tracepoint(object):
         enabled_tracepoints = []

src/python/bcc/usdt.py

@@ -12,34 +12,67 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .libbcc import lib, _USDT_CB
+from .libbcc import lib, _USDT_CB, _USDT_PROBE_CB
+
+class USDTProbe(object):
+    def __init__(self, usdt):
+        self.provider = usdt.provider
+        self.name = usdt.name
+        self.bin_path = usdt.bin_path
+        self.semaphore = usdt.semaphore
+        self.num_locations = usdt.num_locations
+        self.num_arguments = usdt.num_arguments
+
+    def __str__(self):
+        return "%s %s:%s [sema 0x%x]\n  %d location(s)\n  %d argument(s)" % \
+               (self.bin_path, self.provider, self.name, self.semaphore,
+                self.num_locations, self.num_arguments)
+
+    def short_name(self):
+        return "%s:%s" % (self.provider, self.name)
 
 class USDT(object):
     def __init__(self, pid=None, path=None):
-        if pid:
+        if pid and pid != -1:
             self.pid = pid
             self.context = lib.bcc_usdt_new_frompid(pid)
             if self.context == None:
-                raise Exception("USDT failed to instrument PID %d" % pid) 
+                raise Exception("USDT failed to instrument PID %d" % pid)
         elif path:
             self.path = path
             self.context = lib.bcc_usdt_new_frompath(path)
             if self.context == None:
-                raise Exception("USDT failed to instrument path %s" % path) 
+                raise Exception("USDT failed to instrument path %s" % path)
+        else:
+            raise Exception("either a pid or a binary path must be specified")
 
     def enable_probe(self, probe, fn_name):
         if lib.bcc_usdt_enable_probe(self.context, probe, fn_name) != 0:
-            raise Exception("failed to enable probe '%s'" % probe)
+            raise Exception(("failed to enable probe '%s'; a possible cause " +
+                            "can be that the probe requires a pid to enable") %
+                            probe)
 
     def get_text(self):
         return lib.bcc_usdt_genargs(self.context)
 
+    def enumerate_probes(self):
+        probes = []
+        def _add_probe(probe):
+            probes.append(USDTProbe(probe.contents))
+
+        lib.bcc_usdt_foreach(self.context, _USDT_CB(_add_probe))
+        return probes
+
+    # This is called by the BPF module's __init__ when it realizes that there
+    # is a USDT context and probes need to be attached.
     def attach_uprobes(self, bpf):
         probes = []
         def _add_probe(binpath, fn_name, addr, pid):
             probes.append((binpath, fn_name, addr, pid))
 
-        lib.bcc_usdt_foreach_uprobe(self.context, _USDT_CB(_add_probe))
+        lib.bcc_usdt_foreach_uprobe(self.context, _USDT_PROBE_CB(_add_probe))
 
         for (binpath, fn_name, addr, pid) in probes:
-            bpf.attach_uprobe(name=binpath, fn_name=fn_name, addr=addr, pid=pid)
+            bpf.attach_uprobe(name=binpath, fn_name=fn_name,
+                              addr=addr, pid=pid)
+

tools/argdist.py

@@ -175,8 +175,9 @@ def _parse_exprs(self, exprs):
                         self._bail("no exprs specified")
                 self.exprs = exprs.split(',')
 
-        def __init__(self, bpf, type, specifier):
-                self.pid = bpf.args.pid
+        def __init__(self, tool, type, specifier):
+                self.usdt_ctx = None
+                self.pid = tool.args.pid
                 self.raw_spec = specifier
                 self._validate_specifier()
 
@@ -200,8 +201,7 @@ def __init__(self, bpf, type, specifier):
                         self.library = parts[1]
                         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)
+                        self._enable_usdt_probe()
                 else:
                         self.library = parts[1]
                 self.is_user = len(self.library) > 0
@@ -242,8 +242,10 @@ def check(expr):
                         (self.function, Probe.next_probe_index)
                 Probe.next_probe_index += 1
 
-        def close(self):
-                pass
+        def _enable_usdt_probe(self):
+                self.usdt_ctx = USDT(path=self.library, pid=self.pid)
+                self.usdt_ctx.enable_probe(
+                        self.function, self.probe_func_name)
 
         def _substitute_exprs(self):
                 def repl(expr):
@@ -262,12 +264,17 @@ def _generate_hash_field(self, i):
                 else:
                         return "%s v%d;\n" % (self.expr_types[i], i)
 
+        def _generate_usdt_arg_assignment(self, i):
+                expr = self.exprs[i]
+                if self.probe_type == "u" and expr[0:3] == "arg":
+                        return ("        u64 %s = 0;\n" +
+                                "        bpf_usdt_readarg(%s, ctx, &%s);\n") % \
+                                (expr, expr[3], expr)
+                else:
+                        return ""
+
         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])
+                text = self._generate_usdt_arg_assignment(i)
                 if self._is_string(self.expr_types[i]):
                         return (text + "        bpf_probe_read(&__key.v%d.s," +
                                 " sizeof(__key.v%d.s), (void *)%s);\n") % \
@@ -291,8 +298,9 @@ def _generate_hash_decl(self):
 
         def _generate_key_assignment(self):
                 if self.type == "hist":
-                        return "%s __key = %s;\n" % \
-                                (self.expr_types[0], self.exprs[0])
+                        return self._generate_usdt_arg_assignment(0) + \
+                               ("%s __key = %s;\n" % \
+                                (self.expr_types[0], self.exprs[0]))
                 else:
                         text = "struct %s_key_t __key = {};\n" % \
                                 self.probe_hash_name
@@ -590,11 +598,6 @@ def _create_probes(self):
                         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]; };
@@ -610,9 +613,13 @@ def _generate_program(self):
                 for probe in self.probes:
                         bpf_source += probe.generate_text()
                 if self.args.verbose:
-                        if self.usdt_ctx: print(self.usdt_ctx.get_text())
+                        for text in [probe.usdt_ctx.get_text() \
+                                     for probe in self.probes if probe.usdt_ctx]:
+                            print(text)
                         print(bpf_source)
-                self.bpf = BPF(text=bpf_source, usdt=self.usdt_ctx)
+                usdt_contexts = [probe.usdt_ctx
+                                 for probe in self.probes if probe.usdt_ctx]
+                self.bpf = BPF(text=bpf_source, usdt_contexts=usdt_contexts)
 
         def _attach(self):
                 Tracepoint.attach(self.bpf)
@@ -637,12 +644,6 @@ def _main_loop(self):
                            count_so_far >= self.args.count:
                                 exit()
 
-        def _close_probes(self):
-                for probe in self.probes:
-                        probe.close()
-                        if self.args.verbose:
-                                print("closed probe: " + str(probe))
-
         def run(self):
                 try:
                         self._create_probes()
@@ -654,7 +655,6 @@ def run(self):
                                 traceback.print_exc()
                         elif sys.exc_info()[0] is not SystemExit:
                                 print(sys.exc_info()[1])
-                self._close_probes()
 
 if __name__ == "__main__":
         Tool().run()