usdt: permit each probe to have locations from more than one binary path

Fixing issue #1515.

Currently, each usdt probe (provider, probe_name) can only
have locations from the single binary. It is possible that
USDT probes are defined in a header file which eventually causes
the same usdt probe having locations in several different
binary/shared_objects. In such cases, we are not able to attach
the same bpf program to all these locations.

This patch addresses this issue by defining each location to
be `bin_path + addr_offset` vs. previous `addr_offset` only.
This way, each internal Probe class can represent all locations
for the same (provider, probe_name) pair.

The `tplist.py` output is re-organized with the (provider, probe_name)
in the top level like below:
```
...
rtld:lll_futex_wake [sema 0x0]
  location #1 /usr/lib64/ld-2.17.so 0xaac8
    argument #1 8 unsigned bytes @ di
    argument #2 4 signed   bytes @ 1
    argument #3 4 signed   bytes @ 0
  location #2 /usr/lib64/ld-2.17.so 0xe9b9
    argument #1 8 unsigned bytes @ di
    argument #2 4 signed   bytes @ 1
    argument #3 4 signed   bytes @ 0
  location #3 /usr/lib64/ld-2.17.so 0xef3b
    argument #1 8 unsigned bytes @ di
    argument #2 4 signed   bytes @ 1
    argument #3 4 signed   bytes @ 0
...
```

Tested with the following commands
```
  tplist.py
  trace.py -p <pid> 'u::probe "arg1 = %d", arg1'
  trace.py u:<binary_path>:probe "arg1 = %d", arg1'
  argdist.py -p <pid> 'u::probe():s64:arg1'
  argdist.py -C 'u:<binary_path>:probe():s64:arg1'
  funccount.py -p <pid> 'u:<binary_path>:probe'
  funccount.py 'u:<binary_path>:probe'
```

Signed-off-by: Yonghong Song <yhs@fb.com>

src/cc/api/BPF.cc

@@ -616,7 +616,7 @@ StatusTuple USDT::init() {
   ::USDT::Context ctx(binary_path_);
   if (!ctx.loaded())
     return StatusTuple(-1, "Unable to load USDT " + print_name());
-  auto probe = ctx.get(name_);
+  auto probe = ctx.get(provider_, name_);
   if (probe == nullptr)
     return StatusTuple(-1, "Unable to find USDT " + print_name());
 

src/cc/bcc_usdt.h

@@ -37,6 +37,7 @@ struct bcc_usdt {
 
 struct bcc_usdt_location {
     uint64_t address;
+    const char *bin_path;
 };
 
 #define BCC_USDT_ARGUMENT_NONE                0x0
@@ -60,9 +61,11 @@ struct bcc_usdt_argument {
 
 typedef void (*bcc_usdt_cb)(struct bcc_usdt *);
 void bcc_usdt_foreach(void *usdt, bcc_usdt_cb callback);
-int bcc_usdt_get_location(void *usdt, const char *probe_name,
+int bcc_usdt_get_location(void *usdt, const char *provider_name,
+                          const char *probe_name,
                           int index, struct bcc_usdt_location *location);
-int bcc_usdt_get_argument(void *usdt, const char *probe_name,
+int bcc_usdt_get_argument(void *usdt, const char *provider_name,
+                          const char *probe_name,
                           int location_index, int argument_index,
                           struct bcc_usdt_argument *argument);
 

src/cc/usdt.h

@@ -170,12 +170,13 @@ class ArgumentParser_x64 : public ArgumentParser {
 
 struct Location {
   uint64_t address_;
+  std::string bin_path_;
   std::vector<Argument> arguments_;
-  Location(uint64_t addr, const char *arg_fmt);
+  Location(uint64_t addr, const std::string &bin_path, const char *arg_fmt);
 };
 
 class Probe {
-  std::string bin_path_;
+  std::string bin_path_; // initial bin_path when Probe is created
   std::string provider_;
   std::string name_;
   uint64_t semaphore_;
@@ -184,17 +185,18 @@ class Probe {
 
   optional<int> pid_;
   ProcMountNS *mount_ns_;
-  optional<bool> in_shared_object_;
+  std::unordered_map<std::string, bool> object_type_map_; // bin_path => is shared lib?
 
   optional<std::string> attached_to_;
   optional<uint64_t> attached_semaphore_;
 
   std::string largest_arg_type(size_t arg_n);
 
   bool add_to_semaphore(int16_t val);
-  bool resolve_global_address(uint64_t *global, const uint64_t addr);
+  bool resolve_global_address(uint64_t *global, const std::string &bin_path,
+                              const uint64_t addr);
   bool lookup_semaphore_addr(uint64_t *address);
-  void add_location(uint64_t addr, const char *fmt);
+  void add_location(uint64_t addr, const std::string &bin_path, const char *fmt);
 
 public:
   Probe(const char *bin_path, const char *provider, const char *name,
@@ -205,6 +207,7 @@ class Probe {
   uint64_t semaphore()   const { return semaphore_; }
 
   uint64_t address(size_t n = 0) const { return locations_[n].address_; }
+  const char *location_bin_path(size_t n = 0) const { return locations_[n].bin_path_.c_str(); }
   const Location &location(size_t n) const { return locations_[n]; }
   bool usdt_getarg(std::ostream &stream);
   std::string get_arg_ctype(int arg_index) {
@@ -217,7 +220,7 @@ class Probe {
   bool disable();
   bool enabled() const { return !!attached_to_; }
 
-  bool in_shared_object();
+  bool in_shared_object(const std::string &bin_path);
   const std::string &name() { return name_; }
   const std::string &bin_path() { return bin_path_; }
   const std::string &provider() { return provider_; }
@@ -255,6 +258,7 @@ class Context {
   ino_t inode() const { return mount_ns_instance_->target_ino(); }
 
   Probe *get(const std::string &probe_name);
+  Probe *get(const std::string &provider_name, 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);

src/cc/usdt/usdt.cc

@@ -31,7 +31,10 @@
 
 namespace USDT {
 
-Location::Location(uint64_t addr, const char *arg_fmt) : address_(addr) {
+Location::Location(uint64_t addr, const std::string &bin_path, const char *arg_fmt)
+    : address_(addr),
+      bin_path_(bin_path) {
+
 #ifdef __aarch64__
   ArgumentParser_aarch64 parser(arg_fmt);
 #elif __powerpc64__
@@ -56,18 +59,19 @@ Probe::Probe(const char *bin_path, const char *provider, const char *name,
       pid_(pid),
       mount_ns_(ns) {}
 
-bool Probe::in_shared_object() {
-  if (!in_shared_object_) {
-    ProcMountNSGuard g(mount_ns_);
-    in_shared_object_ = bcc_elf_is_shared_obj(bin_path_.c_str());
-  }
-  return in_shared_object_.value();
+bool Probe::in_shared_object(const std::string &bin_path) {
+    if (object_type_map_.find(bin_path) == object_type_map_.end()) {
+      ProcMountNSGuard g(mount_ns_);
+      return (object_type_map_[bin_path] = bcc_elf_is_shared_obj(bin_path.c_str()));
+    }
+    return object_type_map_[bin_path];
 }
 
-bool Probe::resolve_global_address(uint64_t *global, const uint64_t addr) {
-  if (in_shared_object()) {
+bool Probe::resolve_global_address(uint64_t *global, const std::string &bin_path,
+                                   const uint64_t addr) {
+  if (in_shared_object(bin_path)) {
     return (pid_ &&
-            !bcc_resolve_global_addr(*pid_, bin_path_.c_str(), addr, global));
+            !bcc_resolve_global_addr(*pid_, bin_path.c_str(), addr, global));
   }
 
   *global = addr;
@@ -79,7 +83,7 @@ bool Probe::add_to_semaphore(int16_t val) {
 
   if (!attached_semaphore_) {
     uint64_t addr;
-    if (!resolve_global_address(&addr, semaphore_))
+    if (!resolve_global_address(&addr, bin_path_, semaphore_))
       return false;
     attached_semaphore_ = addr;
   }
@@ -175,7 +179,7 @@ bool Probe::usdt_getarg(std::ostream &stream) {
     if (locations_.size() == 1) {
       Location &location = locations_.front();
       stream << "  ";
-      if (!location.arguments_[arg_n].assign_to_local(stream, cptr, bin_path_,
+      if (!location.arguments_[arg_n].assign_to_local(stream, cptr, location.bin_path_,
                                                       pid_))
         return false;
       stream << "\n  return 0;\n}\n";
@@ -184,11 +188,12 @@ bool Probe::usdt_getarg(std::ostream &stream) {
       for (Location &location : locations_) {
         uint64_t global_address;
 
-        if (!resolve_global_address(&global_address, location.address_))
+        if (!resolve_global_address(&global_address, location.bin_path_,
+                                    location.address_))
           return false;
 
         tfm::format(stream, "  case 0x%xULL: ", global_address);
-        if (!location.arguments_[arg_n].assign_to_local(stream, cptr, bin_path_,
+        if (!location.arguments_[arg_n].assign_to_local(stream, cptr, location.bin_path_,
                                                         pid_))
           return false;
 
@@ -201,18 +206,18 @@ bool Probe::usdt_getarg(std::ostream &stream) {
   return true;
 }
 
-void Probe::add_location(uint64_t addr, const char *fmt) {
-  locations_.emplace_back(addr, fmt);
+void Probe::add_location(uint64_t addr, const std::string &bin_path, const char *fmt) {
+  locations_.emplace_back(addr, bin_path, fmt);
 }
 
 void Probe::finalize_locations() {
   std::sort(locations_.begin(), locations_.end(),
             [](const Location &a, const Location &b) {
-              return a.address_ < b.address_;
+              return a.bin_path_ < b.bin_path_ || a.address_ < b.address_;
             });
   auto last = std::unique(locations_.begin(), locations_.end(),
                           [](const Location &a, const Location &b) {
-                            return a.address_ == b.address_;
+                            return a.bin_path_ == b.bin_path_ && a.address_ == b.address_;
                           });
   locations_.erase(last, locations_.end());
 }
@@ -239,15 +244,15 @@ int Context::_each_module(const char *modpath, uint64_t, uint64_t, uint64_t,
 void Context::add_probe(const char *binpath, const struct bcc_elf_usdt *probe) {
   for (auto &p : probes_) {
     if (p->provider_ == probe->provider && p->name_ == probe->name) {
-      p->add_location(probe->pc, probe->arg_fmt);
+      p->add_location(probe->pc, binpath, probe->arg_fmt);
       return;
     }
   }
 
   probes_.emplace_back(
       new Probe(binpath, probe->provider, probe->name, probe->semaphore, pid_,
 	mount_ns_instance_.get()));
-  probes_.back()->add_location(probe->pc, probe->arg_fmt);
+  probes_.back()->add_location(probe->pc, binpath, probe->arg_fmt);
 }
 
 std::string Context::resolve_bin_path(const std::string &bin_path) {
@@ -272,13 +277,41 @@ Probe *Context::get(const std::string &probe_name) {
   return nullptr;
 }
 
+Probe *Context::get(const std::string &provider_name,
+                    const std::string &probe_name) {
+  for (auto &p : probes_) {
+    if (p->provider_ == provider_name && p->name_ == probe_name)
+      return p.get();
+  }
+  return nullptr;
+}
+
 bool Context::enable_probe(const std::string &probe_name,
                            const std::string &fn_name) {
   if (pid_stat_ && pid_stat_->is_stale())
     return false;
 
-  auto p = get(probe_name);
-  return p && p->enable(fn_name);
+  // FIXME: we may have issues here if the context has two same probes's
+  // but different providers. For example, libc:setjmp and rtld:setjmp,
+  // libc:lll_futex_wait and rtld:lll_futex_wait.
+  Probe *found_probe = nullptr;
+  for (auto &p : probes_) {
+    if (p->name_ == probe_name) {
+      if (found_probe != nullptr) {
+         fprintf(stderr, "Two same-name probes (%s) but different providers\n",
+                 probe_name.c_str());
+         return false;
+      }
+      found_probe = p.get();
+    }
+  }
+
+  if (found_probe != nullptr) {
+    found_probe->enable(fn_name);
+    return true;
+  }
+
+  return false;
 }
 
 void Context::each(each_cb callback) {
@@ -300,7 +333,7 @@ void Context::each_uprobe(each_uprobe_cb callback) {
       continue;
 
     for (Location &loc : p->locations_) {
-      callback(p->bin_path_.c_str(), p->attached_to_->c_str(), loc.address_,
+      callback(loc.bin_path_.c_str(), p->attached_to_->c_str(), loc.address_,
                pid_.value_or(-1));
     }
   }
@@ -417,23 +450,26 @@ void bcc_usdt_foreach(void *usdt, bcc_usdt_cb callback) {
   ctx->each(callback);
 }
 
-int bcc_usdt_get_location(void *usdt, const char *probe_name,
+int bcc_usdt_get_location(void *usdt, const char *provider_name,
+                          const char *probe_name,
                           int index, struct bcc_usdt_location *location) {
   USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
-  USDT::Probe *probe = ctx->get(probe_name);
+  USDT::Probe *probe = ctx->get(provider_name, probe_name);
   if (!probe)
     return -1;
   if (index < 0 || (size_t)index >= probe->num_locations())
     return -1;
   location->address = probe->address(index);
+  location->bin_path = probe->location_bin_path(index);
   return 0;
 }
 
-int bcc_usdt_get_argument(void *usdt, const char *probe_name,
+int bcc_usdt_get_argument(void *usdt, const char *provider_name,
+                          const char *probe_name,
                           int location_index, int argument_index,
                           struct bcc_usdt_argument *argument) {
   USDT::Context *ctx = static_cast<USDT::Context *>(usdt);
-  USDT::Probe *probe = ctx->get(probe_name);
+  USDT::Probe *probe = ctx->get(provider_name, probe_name);
   if (!probe)
     return -1;
   if (argument_index < 0 || (size_t)argument_index >= probe->num_arguments())

src/python/bcc/libbcc.py

@@ -208,7 +208,8 @@ class bcc_usdt(ct.Structure):
 
 class bcc_usdt_location(ct.Structure):
     _fields_ = [
-            ('address', ct.c_ulonglong)
+            ('address', ct.c_ulonglong),
+            ('bin_path', ct.c_char_p),
         ]
 
 class BCC_USDT_ARGUMENT_FLAGS(object):
@@ -238,11 +239,11 @@ class bcc_usdt_argument(ct.Structure):
 lib.bcc_usdt_foreach.argtypes = [ct.c_void_p, _USDT_CB]
 
 lib.bcc_usdt_get_location.restype = ct.c_int
-lib.bcc_usdt_get_location.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_int,
+lib.bcc_usdt_get_location.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_char_p, ct.c_int,
                                       ct.POINTER(bcc_usdt_location)]
 
 lib.bcc_usdt_get_argument.restype = ct.c_int
-lib.bcc_usdt_get_argument.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_int,
+lib.bcc_usdt_get_argument.argtypes = [ct.c_void_p, ct.c_char_p, ct.c_char_p, ct.c_int,
                                       ct.c_int, ct.POINTER(bcc_usdt_argument)]
 
 _USDT_PROBE_CB = ct.CFUNCTYPE(None, ct.c_char_p, ct.c_char_p,

src/python/bcc/usdt.py

@@ -82,13 +82,15 @@ def __init__(self, probe, index, location):
         self.index = index
         self.num_arguments = probe.num_arguments
         self.address = location.address
+        self.bin_path = location.bin_path
 
     def __str__(self):
-        return "0x%x" % self.address
+        return "%s 0x%x" % (self.bin_path, self.address)
 
     def get_argument(self, index):
         arg = bcc_usdt_argument()
-        res = lib.bcc_usdt_get_argument(self.probe.context, self.probe.name,
+        res = lib.bcc_usdt_get_argument(self.probe.context, self.probe.provider,
+                                        self.probe.name,
                                         self.index, index, ct.byref(arg))
         if res != 0:
             raise USDTException(
@@ -107,15 +109,15 @@ def __init__(self, context, probe):
         self.num_arguments = probe.num_arguments
 
     def __str__(self):
-        return "%s %s:%s [sema 0x%x]" % \
-               (self.bin_path, self.provider, self.name, self.semaphore)
+        return "%s:%s [sema 0x%x]" % \
+               (self.provider, self.name, self.semaphore)
 
     def short_name(self):
         return "%s:%s" % (self.provider, self.name)
 
     def get_location(self, index):
         loc = bcc_usdt_location()
-        res = lib.bcc_usdt_get_location(self.context, self.name,
+        res = lib.bcc_usdt_get_location(self.context, self.provider, self.name,
                                         index, ct.byref(loc))
         if res != 0:
             raise USDTException("error retrieving probe location %d" % index)

tests/cc/test_usdt_probes.cc

@@ -42,7 +42,7 @@ TEST_CASE("test finding a probe in our own process", "[usdt]") {
     auto probe = ctx.get("sample_probe_1");
     REQUIRE(probe);
 
-    REQUIRE(probe->in_shared_object() == false);
+    REQUIRE(probe->in_shared_object(probe->bin_path()) == false);
     REQUIRE(probe->name() == "sample_probe_1");
     REQUIRE(probe->provider() == "libbcc_test");
     REQUIRE(probe->bin_path().find("/test_libbcc") != std::string::npos);
@@ -132,7 +132,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
       auto probe = ctx.get(name);
       REQUIRE(probe);
 
-      REQUIRE(probe->in_shared_object() == true);
+      REQUIRE(probe->in_shared_object(probe->bin_path()) == true);
       REQUIRE(probe->name() == name);
       REQUIRE(probe->provider() == "ruby");
 
@@ -155,7 +155,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
       auto probe = ctx.get(name);
       REQUIRE(probe);
 
-      REQUIRE(probe->in_shared_object() == true);
+      REQUIRE(probe->in_shared_object(probe->bin_path()) == true);
       REQUIRE(probe->name() == name);
       REQUIRE(probe->provider() == "ruby");
 
@@ -205,7 +205,7 @@ TEST_CASE("test listing all USDT probes in Ruby/MRI", "[usdt]") {
       auto probe = ctx.get(name);
       REQUIRE(probe);
 
-      REQUIRE(probe->in_shared_object() == true);
+      REQUIRE(probe->in_shared_object(probe->bin_path()) == true);
       REQUIRE(probe->name() == name);
       REQUIRE(probe->provider() == "ruby");