permit multiple pids attaching to the same probe

Currently, if more than one pid-associated USDT attaching to
the same probe, usdt readarg code will be generated twice and
the compiler will complain.

This patch solves issue by preventing code duplication if
a previous context with the same mnt point and exec binary
has generated the code for the same probe. The event name is
also changed to have pid embedded so different pid-associated
uprobe event will have different names.

This patch introduces an internal uprobe event name
discrepency. It is a good idea to have event name
generation in libbpf so that both C++ API and Python API
will have consistent name conventions. This will be
addressed in a subsequent commit as it is largely
a different issue.
Signed-off-by: Yonghong Song <yhs@fb.com>

src/cc/bcc_usdt.h

@@ -67,7 +67,7 @@ int bcc_usdt_get_argument(void *usdt, const char *probe_name,
                           struct bcc_usdt_argument *argument);
 
 int bcc_usdt_enable_probe(void *, const char *, const char *);
-const char *bcc_usdt_genargs(void *);
+const char *bcc_usdt_genargs(void **ctx_array, int len);
 const char *bcc_usdt_get_probe_argctype(
   void *ctx, const char* probe_name, const int arg_index
 );

src/cc/ns_guard.cc

@@ -21,7 +21,7 @@
 
 #include "ns_guard.h"
 
-ProcMountNS::ProcMountNS(int pid) {
+ProcMountNS::ProcMountNS(int pid) : target_ino_(0) {
   if (pid < 0)
     return;
 
@@ -38,6 +38,7 @@ ProcMountNS::ProcMountNS(int pid) {
   if (fstat(target_fd, &target_stat) != 0)
     return;
 
+  target_ino_ = target_stat.st_ino;
   if (self_stat.st_ino == target_stat.st_ino)
     // Both current and target Process are in same mount namespace
     return;

src/cc/ns_guard.h

@@ -32,10 +32,12 @@ class ProcMountNS {
   explicit ProcMountNS(int pid);
   int self() const { return self_fd_; }
   int target() const { return target_fd_; }
+  ino_t target_ino() const { return target_ino_; }
 
  private:
   ebpf::FileDesc self_fd_;
   ebpf::FileDesc target_fd_;
+  ino_t target_ino_;
 };
 
 // ProcMountNSGuard switches to the target mount namespace and restores the

src/cc/usdt.cc

@@ -15,6 +15,7 @@
  */
 #include <cstring>
 #include <sstream>
+#include <unordered_set>
 
 #include <fcntl.h>
 #include <sys/types.h>
@@ -250,15 +251,6 @@ Probe *Context::get(const std::string &probe_name) {
   return nullptr;
 }
 
-bool Context::generate_usdt_args(std::ostream &stream) {
-  stream << USDT_PROGRAM_HEADER;
-  for (auto &p : probes_) {
-    if (p->enabled() && !p->usdt_getarg(stream))
-      return false;
-  }
-  return true;
-}
-
 bool Context::enable_probe(const std::string &probe_name,
                            const std::string &fn_name) {
   if (pid_stat_ && pid_stat_->is_stale())
@@ -296,15 +288,30 @@ void Context::each_uprobe(each_uprobe_cb callback) {
 Context::Context(const std::string &bin_path) : loaded_(false) {
   std::string full_path = resolve_bin_path(bin_path);
   if (!full_path.empty()) {
-    if (bcc_elf_foreach_usdt(full_path.c_str(), _each_probe, this) == 0)
+    if (bcc_elf_foreach_usdt(full_path.c_str(), _each_probe, this) == 0) {
+      cmd_bin_path_ = full_path;
       loaded_ = true;
+    }
   }
 }
 
 Context::Context(int pid) : pid_(pid), pid_stat_(pid),
   mount_ns_instance_(new ProcMountNS(pid)), loaded_(false) {
-  if (bcc_procutils_each_module(pid, _each_module, this) == 0)
+  if (bcc_procutils_each_module(pid, _each_module, this) == 0) {
+    // get exe command from /proc/<pid>/exe
+    // assume the maximum path length 4096, which should be
+    // sufficiently large to cover all use cases
+    char source[64];
+    char cmd_buf[4096];
+    snprintf(source, sizeof(source), "/proc/%d/exe", pid);
+    ssize_t cmd_len = readlink(source, cmd_buf, sizeof(cmd_buf) - 1);
+    if (cmd_len == -1)
+      return;
+    cmd_buf[cmd_len] = '\0';
+    cmd_bin_path_.assign(cmd_buf, cmd_len + 1);
+
     loaded_ = true;
+  }
 }
 
 Context::~Context() {
@@ -345,13 +352,32 @@ int bcc_usdt_enable_probe(void *usdt, const char *probe_name,
   return ctx->enable_probe(probe_name, fn_name) ? 0 : -1;
 }
 
-const char *bcc_usdt_genargs(void *usdt) {
+const char *bcc_usdt_genargs(void **usdt_array, int len) {
   static std::string storage_;
-
-  USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
   std::ostringstream stream;
-  if (!ctx->generate_usdt_args(stream))
-    return nullptr;
+
+  stream << USDT::USDT_PROGRAM_HEADER;
+  // Generate genargs codes for an array of USDT Contexts.
+  //
+  // Each mnt_point + cmd_bin_path + probe_provider + probe_name
+  // uniquely identifies a probe.
+  std::unordered_set<std::string> generated_probes;
+  for (int i = 0; i < len; i++) {
+    USDT::Context *ctx = static_cast<USDT::Context *>(usdt_array[i]);
+
+    for (size_t j = 0; j < ctx->num_probes(); j++) {
+      USDT::Probe *p = ctx->get(j);
+      if (p->enabled()) {
+        std::string key = std::to_string(ctx->inode()) + "*"
+          + ctx->cmd_bin_path() + "*" + p->provider() + "*" + p->name();
+        if (generated_probes.find(key) != generated_probes.end())
+          continue;
+        if (!p->usdt_getarg(stream))
+          return nullptr;
+        generated_probes.insert(key);
+      }
+    }
+  }
 
   storage_ = stream.str();
   return storage_.c_str();

src/cc/usdt.h

@@ -201,6 +201,7 @@ class Context {
   optional<int> pid_;
   optional<ProcStat> pid_stat_;
   std::unique_ptr<ProcMountNS> mount_ns_instance_;
+  std::string cmd_bin_path_;
   bool loaded_;
 
   static void _each_probe(const char *binpath, const struct bcc_elf_usdt *probe,
@@ -218,12 +219,13 @@ public:
   optional<int> pid() const { return pid_; }
   bool loaded() const { return loaded_; }
   size_t num_probes() const { return probes_.size(); }
+  const std::string & cmd_bin_path() const { return cmd_bin_path_; }
+  ino_t inode() const { return mount_ns_instance_->target_ino(); }
 
   Probe *get(const std::string &probe_name);
   Probe *get(int pos) { return probes_[pos].get(); }
 
   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);

src/python/bcc/__init__.py

@@ -268,14 +268,15 @@ class BPF(object):
         cflags_array = (ct.c_char_p * len(cflags))()
         for i, s in enumerate(cflags): cflags_array[i] = s.encode("ascii")
         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_text + text
+            ctx_array = (ct.c_void_p * len(usdt_contexts))()
+            for i, usdt in enumerate(usdt_contexts): ctx_array[i] = ct.c_void_p(usdt.get_context())
+            usdt_text = lib.bcc_usdt_genargs(ctx_array, len(usdt_contexts)).decode()
+            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_text + text
 
         if text:
             self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"),
@@ -781,6 +782,14 @@ class BPF(object):
             raise Exception("Error %d enumerating symbols in %s" % (res, name))
         return addresses
 
+    def _get_uprobe_evname(self, prefix, path, addr, pid):
+        if pid == -1:
+            return "%s_%s_0x%x" % (prefix, self._probe_repl.sub("_", path), addr)
+        else:
+            # if pid is valid, put pid in the name, so different pid
+            # can have different event names
+            return "%s_%s_0x%x_%d" % (prefix, self._probe_repl.sub("_", path), addr, pid)
+
     def attach_uprobe(self, name="", sym="", sym_re="", addr=None,
             fn_name="", pid=-1, cpu=0, group_fd=-1):
         """attach_uprobe(name="", sym="", sym_re="", addr=None, fn_name=""
@@ -819,7 +828,7 @@ class BPF(object):
 
         self._check_probe_quota(1)
         fn = self.load_func(fn_name, BPF.KPROBE)
-        ev_name = "p_%s_0x%x" % (self._probe_repl.sub("_", path), addr)
+        ev_name = self._get_uprobe_evname("p", path, addr, pid)
         res = lib.bpf_attach_uprobe(fn.fd, 0, ev_name.encode("ascii"),
                 path.encode("ascii"), addr, pid, cpu, group_fd,
                 self._reader_cb_impl, ct.cast(id(self), ct.py_object))
@@ -838,7 +847,7 @@ class BPF(object):
 
         name = str(name)
         (path, addr) = BPF._check_path_symbol(name, sym, addr, pid)
-        ev_name = "p_%s_0x%x" % (self._probe_repl.sub("_", path), addr)
+        ev_name = self._get_uprobe_evname("p", path, addr, pid)
         if ev_name not in self.open_uprobes:
             raise Exception("Uprobe %s is not attached" % ev_name)
         lib.perf_reader_free(self.open_uprobes[ev_name])

src/python/bcc/libbcc.py

@@ -191,7 +191,7 @@ 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]
+lib.bcc_usdt_genargs.argtypes = [ct.POINTER(ct.c_void_p), ct.c_int]
 
 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]

src/python/bcc/usdt.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import ctypes as ct
-import sys
+import os, sys
 from .libbcc import lib, _USDT_CB, _USDT_PROBE_CB, \
                     bcc_usdt_location, bcc_usdt_argument, \
                     BCC_USDT_ARGUMENT_FLAGS
@@ -157,8 +157,8 @@ USDT probes. Look for a configure flag similar to --with-dtrace or
 tplist tool.""")
             sys.exit(1)
 
-    def get_text(self):
-        return lib.bcc_usdt_genargs(self.context).decode()
+    def get_context(self):
+        return self.context
 
     def get_probe_arg_ctype(self, probe_name, arg_index):
         return lib.bcc_usdt_get_probe_argctype(

tests/python/CMakeLists.txt

@@ -70,3 +70,5 @@ add_test(NAME py_test_tools_memleak WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR
   COMMAND ${TEST_WRAPPER} py_test_tools_memleak sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_tools_memleak.py)
 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)

tests/python/test_usdt2.py

+#!/usr/bin/python
+#
+# USAGE: test_usdt2.py
+#
+# Copyright 2017 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
+from tempfile import NamedTemporaryFile
+import ctypes as ct
+import inspect
+import os
+import signal
+
+class TestUDST(TestCase):
+    def setUp(self):
+        # Application, minimum, to define three trace points
+        app_text = b"""
+#include <stdlib.h>
+#include <unistd.h>
+#include "folly/tracing/StaticTracepoint.h"
+
+int main(int argc, char **argv) {
+  int t = atoi(argv[1]);
+  while (1) {
+    FOLLY_SDT(test, probe_point_1, t);
+    FOLLY_SDT(test, probe_point_2, t + 1);
+    FOLLY_SDT(test, probe_point_3, t + 2);
+    sleep(1);
+  }
+  return 1;
+}
+"""
+        # BPF program
+        self.bpf_text = """
+#include <uapi/linux/ptrace.h>
+
+BPF_PERF_OUTPUT(event1);
+BPF_PERF_OUTPUT(event2);
+BPF_PERF_OUTPUT(event3);
+BPF_PERF_OUTPUT(event4);
+BPF_PERF_OUTPUT(event5);
+BPF_PERF_OUTPUT(event6);
+
+int do_trace1(struct pt_regs *ctx) {
+    u32 pid = bpf_get_current_pid_tgid();
+    int result = 0;
+    bpf_usdt_readarg(1, ctx, &result);
+    if (FILTER)
+      event1.perf_submit(ctx, &result, sizeof(result));
+    else
+      event4.perf_submit(ctx, &result, sizeof(result));
+    return 0;
+};
+int do_trace2(struct pt_regs *ctx) {
+    u32 pid = bpf_get_current_pid_tgid();
+    int result = 0;
+    bpf_usdt_readarg(1, ctx, &result);
+    if (FILTER)
+      event2.perf_submit(ctx, &result, sizeof(result));
+    else
+      event5.perf_submit(ctx, &result, sizeof(result));
+    return 0;
+}
+int do_trace3(struct pt_regs *ctx) {
+    u32 pid = bpf_get_current_pid_tgid();
+    int result = 0;
+    bpf_usdt_readarg(1, ctx, &result);
+    if (FILTER)
+      event3.perf_submit(ctx, &result, sizeof(result));
+    else
+      event6.perf_submit(ctx, &result, sizeof(result));
+    return 0;
+}
+"""
+
+        # Compile and run the application
+        self.ftemp = NamedTemporaryFile(delete=False)
+        self.ftemp.close()
+        comp = Popen(["gcc", "-I", "%s/include" % os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))),
+                      "-x", "c", "-o", self.ftemp.name, "-"],
+                     stdin=PIPE)
+        comp.stdin.write(app_text)
+        comp.stdin.close()
+        self.assertEqual(comp.wait(), 0)
+
+        # create 3 applications, 2 applications will have usdt attached and
+        # the third one does not, and the third one should not call into
+        # bpf program.
+        self.app = Popen([self.ftemp.name, "1"])
+        self.app2 = Popen([self.ftemp.name, "11"])
+        self.app3 = Popen([self.ftemp.name, "21"])
+
+    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_point_1", fn_name="do_trace1")
+        u.enable_probe(probe="probe_point_2", fn_name="do_trace2")
+        u2 = USDT(pid=int(self.app2.pid))
+        u2.enable_probe(probe="probe_point_2", fn_name="do_trace2")
+        u2.enable_probe(probe="probe_point_3", fn_name="do_trace3")
+        self.bpf_text = self.bpf_text.replace("FILTER", "pid == %d" % self.app.pid)
+        b = BPF(text=self.bpf_text, usdt_contexts=[u, u2])
+
+        # Event states for each event:
+        # 0 - probe not caught, 1 - probe caught with correct value,
+        # 2 - probe caught with incorrect value
+        self.evt_st_1 = 0
+        self.evt_st_2 = 0
+        self.evt_st_3 = 0
+        self.evt_st_4 = 0
+        self.evt_st_5 = 0
+        self.evt_st_6 = 0
+
+        def check_event_val(data, event_state, expected_val):
+            result = ct.cast(data, ct.POINTER(ct.c_int)).contents
+            if result.value == expected_val:
+                if (event_state == 0 or event_state == 1):
+                    return 1
+            return 2
+
+        def print_event1(cpu, data, size):
+            self.evt_st_1 = check_event_val(data, self.evt_st_1, 1)
+
+        def print_event2(cpu, data, size):
+            self.evt_st_2 = check_event_val(data, self.evt_st_2, 2)
+
+        def print_event3(cpu, data, size):
+            self.evt_st_3 = check_event_val(data, self.evt_st_3, 3)
+
+        def print_event4(cpu, data, size):
+            self.evt_st_4 = check_event_val(data, self.evt_st_4, 11)
+
+        def print_event5(cpu, data, size):
+            self.evt_st_5 = check_event_val(data, self.evt_st_5, 12)
+
+        def print_event6(cpu, data, size):
+            self.evt_st_6 = check_event_val(data, self.evt_st_6, 13)
+
+        # loop with callback to print_event
+        b["event1"].open_perf_buffer(print_event1)
+        b["event2"].open_perf_buffer(print_event2)
+        b["event3"].open_perf_buffer(print_event3)
+        b["event4"].open_perf_buffer(print_event4)
+        b["event5"].open_perf_buffer(print_event5)
+        b["event6"].open_perf_buffer(print_event6)
+
+        # three iterations to make sure we get some probes and have time to process them
+        for i in range(3):
+            b.kprobe_poll()
+
+        # note that event1 and event4 do not really fire, so their state should be 0
+        # use separate asserts so that if test fails we know which one is the culprit
+        self.assertTrue(self.evt_st_1 == 1)
+        self.assertTrue(self.evt_st_2 == 1)
+        self.assertTrue(self.evt_st_3 == 0)
+        self.assertTrue(self.evt_st_4 == 0)
+        self.assertTrue(self.evt_st_5 == 1)
+        self.assertTrue(self.evt_st_6 == 1)
+
+    def tearDown(self):
+        # kill the subprocess, clean the environment
+        self.app.kill()
+        self.app.wait()
+        self.app2.kill()
+        self.app2.wait()
+        self.app3.kill()
+        self.app3.wait()
+        os.unlink(self.ftemp.name)
+
+if __name__ == "__main__":
+    main()