Fix trace.py USDT argument filtering (#710)

src/cc/bcc_usdt.h

@@ -40,6 +40,9 @@ 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 *);
+const char *bcc_usdt_get_probe_argctype(
+  void *ctx, const char* probe_name, const int arg_index
+);
 
 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

@@ -344,6 +344,14 @@ const char *bcc_usdt_genargs(void *usdt) {
   return storage_.c_str();
 }
 
+const char *bcc_usdt_get_probe_argctype(
+  void *ctx, const char* probe_name, const int arg_index
+) {
+  USDT::Probe *p = static_cast<USDT::Context *>(ctx)->get(probe_name);
+  std::string res = p ? p->get_arg_ctype(arg_index) : "";
+  return res.c_str();
+}
+
 void bcc_usdt_foreach(void *usdt, bcc_usdt_cb callback) {
   USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
   ctx->each(callback);

src/cc/usdt.h

@@ -154,6 +154,9 @@ public:
 
   uint64_t address(size_t n = 0) const { return locations_[n].address_; }
   bool usdt_getarg(std::ostream &stream);
+  std::string get_arg_ctype(int arg_index) {
+    return largest_arg_type(arg_index);
+  }
 
   bool need_enable() const { return semaphore_ != 0x0; }
   bool enable(const std::string &fn_name);

src/python/bcc/libbcc.py

@@ -157,6 +157,9 @@ 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]
 
+lib.bcc_usdt_get_probe_argctype.restype = ct.c_char_p
+lib.bcc_usdt_get_probe_argctype.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_int]
+
 class bcc_usdt(ct.Structure):
     _fields_ = [
             ('provider', ct.c_char_p),

src/python/bcc/usdt.py

@@ -55,6 +55,10 @@ class USDT(object):
     def get_text(self):
         return lib.bcc_usdt_genargs(self.context)
 
+    def get_probe_arg_ctype(self, probe_name, arg_index):
+        return lib.bcc_usdt_get_probe_argctype(
+            self.context, probe_name, arg_index)
+
     def enumerate_probes(self):
         probes = []
         def _add_probe(probe):

tools/trace.py

@@ -299,6 +299,26 @@ BPF_PERF_OUTPUT(%s);
                                         (idx, Probe.c_type[field_type], expr)
                 self._bail("unrecognized field type %s" % field_type)
 
+        def _generate_usdt_filter_read(self):
+            text = ""
+            if self.probe_type == "u":
+                    for arg, _ in Probe.aliases.items():
+                        if not (arg.startswith("arg") and (arg in self.filter)):
+                                continue
+                        arg_index = int(arg.replace("arg", ""))
+                        arg_ctype = self.usdt.get_probe_arg_ctype(
+                                self.usdt_name, arg_index)
+                        if not arg_ctype:
+                                self._bail("Unable to determine type of {} "
+                                           "in the filter".format(arg))
+                        text += """
+        {} {}_filter;
+        bpf_usdt_readarg({}, ctx, &{}_filter);
+                        """.format(arg_ctype, arg, arg_index, arg)
+                        self.filter = self.filter.replace(
+                                arg, "{}_filter".format(arg))
+            return text
+
         def generate_program(self, include_self):
                 data_decl = self._generate_data_decl()
                 # kprobes don't have built-in pid filters, so we have to add
@@ -329,6 +349,7 @@ BPF_PERF_OUTPUT(%s);
                 text = """
 int %s(%s)
 {
+        %s
         %s
         %s
         if (!(%s)) return 0;
@@ -343,7 +364,8 @@ int %s(%s)
 }
 """
                 text = text % (self.probe_name, signature,
-                               pid_filter, prefix, self.filter,
+                               pid_filter, prefix,
+                               self._generate_usdt_filter_read(), self.filter,
                                self.struct_name, data_fields, self.events_name)
 
                 return data_decl + "\n" + text