Merge pull request #1263 from iovisor/yhs_dev

permit multiple pids attaching to the same probe

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()