Unify and improve C++'s USDT implementation (#1841)

* Add interface to Probe's getargs call

This commit allows the Probe instance to generate argument for arbitary
probe function

* Refactor C++ USDT implementation

This commit makes C++ USDT implementation uses the common USDT::Context
and USDT::Probe logic

* Add test case for C++ USDT API

* Improve FollyRequestContextSwitch example

examples/cpp/FollyRequestContextSwitch.cc

@@ -12,6 +12,7 @@
  */
 
 #include <signal.h>
+#include <functional>
 #include <iostream>
 #include <vector>
 
@@ -59,43 +60,50 @@ void handle_output(void* cb_cookie, void* data, int data_size) {
             << event->new_addr << std::endl;
 }
 
-ebpf::BPF* bpf;
+std::function<void(int)> shutdown_handler;
 
-void signal_handler(int s) {
-  std::cerr << "Terminating..." << std::endl;
-  delete bpf;
-  exit(0);
-}
+void signal_handler(int s) { shutdown_handler(s); }
 
 int main(int argc, char** argv) {
-  if (argc != 2) {
-    std::cout << "USAGE: FollyRequestContextSwitch PATH_TO_BINARY" << std::endl;
-    exit(1);
+  std::string binary;
+  pid_t pid = -1;
+  for (int i = 0; i < argc; i++) {
+    if (strncmp(argv[i], "--pid", 5) == 0) {
+      pid = std::stoi(argv[i + 1]);
+      i++;
+      continue;
+    }
+    if (strncmp(argv[i], "--binary", 8) == 0) {
+      binary = argv[i + 1];
+      i++;
+      continue;
+    }
   }
-  std::string binary_path(argv[1]);
-
-  std::vector<ebpf::USDT> u;
-  u.emplace_back(binary_path, "folly", "request_context_switch_before",
-                 "on_context_switch");
-  auto usdt_init_res = u[0].init();
-  if (usdt_init_res.code() != 0) {
-    std::cerr << usdt_init_res.msg() << std::endl;
-    return 1;
+
+  if (pid <= 0 && binary.empty()) {
+    std::cout << "Must specify at least one of binary or PID:" << std::endl
+              << "FollyRequestContextSwitch [--pid PID] [--binary BINARY]"
+              << std::endl;
+    exit(1);
   }
 
-  bpf = new ebpf::BPF();
-  auto init_res = bpf->init(BPF_PROGRAM, {}, u);
+  ebpf::USDT u(binary, pid, "folly", "request_context_switch_before",
+               "on_context_switch");
+
+  ebpf::BPF* bpf = new ebpf::BPF();
+
+  auto init_res = bpf->init(BPF_PROGRAM, {}, {u});
   if (init_res.code() != 0) {
     std::cerr << init_res.msg() << std::endl;
     return 1;
   }
 
-  auto attach_res = bpf->attach_usdt(u[0]);
+  auto attach_res = bpf->attach_usdt(u);
   if (attach_res.code() != 0) {
     std::cerr << attach_res.msg() << std::endl;
     return 1;
   } else {
-    std::cout << "Attached to USDT " << u[0];
+    std::cout << "Attached to USDT " << u;
   }
 
   auto open_res = bpf->open_perf_buffer("events", &handle_output);
@@ -104,7 +112,14 @@ int main(int argc, char** argv) {
     return 1;
   }
 
+  shutdown_handler = [&](int s) {
+    std::cerr << "Terminating..." << std::endl;
+    bpf->detach_usdt(u);
+    delete bpf;
+    exit(0);
+  };
   signal(SIGINT, signal_handler);
+
   std::cout << "Started tracing, hit Ctrl-C to terminate." << std::endl;
   auto perf_buffer = bpf->get_perf_buffer("events");
   if (perf_buffer)

src/cc/api/BPF.cc

@@ -59,11 +59,10 @@ StatusTuple BPF::init(const std::string& bpf_program,
                       const std::vector<USDT>& usdt) {
   std::string all_bpf_program;
 
-  for (auto u : usdt) {
-    if (!u.initialized_)
-      TRY2(u.init());
-    all_bpf_program += u.program_text_;
-    usdt_.push_back(std::move(u));
+  for (const auto& u : usdt) {
+    usdt_.emplace_back(u);
+    TRY2(usdt_.back().init());
+    all_bpf_program += usdt_.back().program_text_;
   }
 
   auto flags_len = cflags.size();
@@ -223,17 +222,22 @@ StatusTuple BPF::attach_uprobe(const std::string& binary_path,
 }
 
 StatusTuple BPF::attach_usdt(const USDT& usdt, pid_t pid) {
-  for (const auto& u : usdt_)
+  for (const auto& u : usdt_) {
     if (u == usdt) {
+      auto& probe = *static_cast<::USDT::Probe*>(u.probe_.get());
+      if (!probe.enable(u.probe_func_))
+        return StatusTuple(-1, "Unable to enable USDT " + u.print_name());
+
       bool failed = false;
       std::string err_msg;
       int cnt = 0;
-      for (auto addr : u.addresses_) {
-        auto res =
-            attach_uprobe(u.binary_path_, std::string(), u.probe_func_, addr);
+      for (const auto& loc : probe.locations_) {
+        auto res = attach_uprobe(loc.bin_path_, std::string(), u.probe_func_,
+                                 loc.address_, BPF_PROBE_ENTRY, pid);
         if (res.code() != 0) {
           failed = true;
-          err_msg += "USDT " + u.print_name() + " at " + std::to_string(addr);
+          err_msg += "USDT " + u.print_name() + " at " + loc.bin_path_ +
+                     " address " + std::to_string(loc.address_);
           err_msg += ": " + res.msg() + "\n";
           break;
         }
@@ -242,13 +246,18 @@ StatusTuple BPF::attach_usdt(const USDT& usdt, pid_t pid) {
       if (failed) {
         for (int i = 0; i < cnt; i++) {
           auto res =
-              detach_uprobe(u.binary_path_, std::string(), u.addresses_[i]);
-          err_msg += "During clean up: " + res.msg() + "\n";
+              detach_uprobe(probe.locations_[i].bin_path_, std::string(),
+                            probe.locations_[i].address_, BPF_PROBE_ENTRY, pid);
+          if (res.code() != 0)
+            err_msg += "During clean up: " + res.msg() + "\n";
         }
         return StatusTuple(-1, err_msg);
-      } else
+      } else {
         return StatusTuple(0);
+      }
     }
+  }
+
   return StatusTuple(-1, "USDT %s not found", usdt.print_name().c_str());
 }
 
@@ -398,24 +407,35 @@ StatusTuple BPF::detach_uprobe(const std::string& binary_path,
   return StatusTuple(0);
 }
 
-StatusTuple BPF::detach_usdt(const USDT& usdt) {
-  for (const auto& u : usdt_)
+StatusTuple BPF::detach_usdt(const USDT& usdt, pid_t pid) {
+  for (const auto& u : usdt_) {
     if (u == usdt) {
+      auto& probe = *static_cast<::USDT::Probe*>(u.probe_.get());
       bool failed = false;
       std::string err_msg;
-      for (auto addr : u.addresses_) {
-        auto res = detach_uprobe(u.binary_path_, std::string(), addr);
+      for (const auto& loc : probe.locations_) {
+        auto res = detach_uprobe(loc.bin_path_, std::string(), loc.address_,
+                                 BPF_PROBE_ENTRY, pid);
         if (res.code() != 0) {
           failed = true;
-          err_msg += "USDT " + u.print_name() + " at " + std::to_string(addr);
+          err_msg += "USDT " + u.print_name() + " at " + loc.bin_path_ +
+                     " address " + std::to_string(loc.address_);
           err_msg += ": " + res.msg() + "\n";
         }
       }
+
+      if (!probe.disable()) {
+        failed = true;
+        err_msg += "Unable to disable USDT " + u.print_name();
+      }
+
       if (failed)
         return StatusTuple(-1, err_msg);
       else
         return StatusTuple(0);
     }
+  }
+
   return StatusTuple(-1, "USDT %s not found", usdt.print_name().c_str());
 }
 
@@ -677,26 +697,85 @@ StatusTuple BPF::detach_perf_event_all_cpu(open_probe_t& attr) {
   return StatusTuple(0);
 }
 
+USDT::USDT(const std::string& binary_path, const std::string& provider,
+           const std::string& name, const std::string& probe_func)
+    : initialized_(false),
+      binary_path_(binary_path),
+      pid_(-1),
+      provider_(provider),
+      name_(name),
+      probe_func_(probe_func) {}
+
+USDT::USDT(pid_t pid, const std::string& provider, const std::string& name,
+           const std::string& probe_func)
+    : initialized_(false),
+      binary_path_(),
+      pid_(pid),
+      provider_(provider),
+      name_(name),
+      probe_func_(probe_func) {}
+
+USDT::USDT(const std::string& binary_path, pid_t pid,
+           const std::string& provider, const std::string& name,
+           const std::string& probe_func)
+    : initialized_(false),
+      binary_path_(binary_path),
+      pid_(pid),
+      provider_(provider),
+      name_(name),
+      probe_func_(probe_func) {}
+
+USDT::USDT(const USDT& usdt)
+    : initialized_(false),
+      binary_path_(usdt.binary_path_),
+      pid_(usdt.pid_),
+      provider_(usdt.provider_),
+      name_(usdt.name_),
+      probe_func_(usdt.probe_func_) {}
+
+bool USDT::operator==(const USDT& other) const {
+  return (provider_ == other.provider_) && (name_ == other.name_) &&
+         (binary_path_ == other.binary_path_) && (pid_ == other.pid_) &&
+         (probe_func_ == other.probe_func_);
+}
+
 StatusTuple USDT::init() {
-  ::USDT::Context ctx(binary_path_);
-  if (!ctx.loaded())
+  std::unique_ptr<::USDT::Context> ctx;
+  if (!binary_path_.empty() && pid_ > 0)
+    ctx.reset(new ::USDT::Context(pid_, binary_path_));
+  else if (!binary_path_.empty())
+    ctx.reset(new ::USDT::Context(binary_path_));
+  else if (pid_ > 0)
+    ctx.reset(new ::USDT::Context(pid_));
+  else
+    return StatusTuple(-1, "No valid Binary Path or PID provided");
+
+  if (!ctx->loaded())
     return StatusTuple(-1, "Unable to load USDT " + print_name());
-  auto probe = ctx.get(provider_, name_);
-  if (probe == nullptr)
+
+  auto deleter = [](void* probe) { delete static_cast<::USDT::Probe*>(probe); };
+  for (auto& p : ctx->probes_) {
+    if (p->provider_ == provider_ && p->name_ == name_) {
+      // Take ownership of the probe that we are interested in, and avoid it
+      // being destrcuted when we destruct the USDT::Context instance
+      probe_ = std::unique_ptr<void, std::function<void(void*)>>(p.release(),
+                                                                 deleter);
+      p.swap(ctx->probes_.back());
+      ctx->probes_.pop_back();
+      break;
+    }
+  }
+  if (!probe_)
     return StatusTuple(-1, "Unable to find USDT " + print_name());
+  ctx.reset(nullptr);
+  auto& probe = *static_cast<::USDT::Probe*>(probe_.get());
 
-  if (!probe->enable(probe_func_))
-    return StatusTuple(-1, "Failed to enable USDT " + print_name());
   std::ostringstream stream;
-  if (!probe->usdt_getarg(stream))
+  if (!probe.usdt_getarg(stream, probe_func_))
     return StatusTuple(
         -1, "Unable to generate program text for USDT " + print_name());
   program_text_ = ::USDT::USDT_PROGRAM_HEADER + stream.str();
 
-  addresses_.reserve(probe->num_locations());
-  for (size_t i = 0; i < probe->num_locations(); i++)
-    addresses_.emplace_back(probe->address(i));
-
   initialized_ = true;
   return StatusTuple(0);
 }

src/cc/api/BPF.h

@@ -75,7 +75,7 @@ class BPF {
                             bpf_probe_attach_type attach_type = BPF_PROBE_ENTRY,
                             pid_t pid = -1);
   StatusTuple attach_usdt(const USDT& usdt, pid_t pid = -1);
-  StatusTuple detach_usdt(const USDT& usdt);
+  StatusTuple detach_usdt(const USDT& usdt, pid_t pid = -1);
 
   StatusTuple attach_tracepoint(const std::string& tracepoint,
                                 const std::string& probe_func);
@@ -239,41 +239,39 @@ class BPF {
 class USDT {
  public:
   USDT(const std::string& binary_path, const std::string& provider,
-       const std::string& name, const std::string& probe_func)
-      : initialized_(false),
-        binary_path_(binary_path),
-        provider_(provider),
-        name_(name),
-        probe_func_(probe_func) {}
+       const std::string& name, const std::string& probe_func);
+  USDT(pid_t pid, const std::string& provider, const std::string& name,
+       const std::string& probe_func);
+  USDT(const std::string& binary_path, pid_t pid, const std::string& provider,
+       const std::string& name, const std::string& probe_func);
+  USDT(const USDT& usdt);
 
   StatusTuple init();
 
-  bool operator==(const USDT& other) const {
-    return (provider_ == other.provider_) && (name_ == other.name_) &&
-           (binary_path_ == other.binary_path_) &&
-           (probe_func_ == other.probe_func_);
-  }
+  bool operator==(const USDT& other) const;
 
   std::string print_name() const {
-    return provider_ + ":" + name_ + " from " + binary_path_ + " for probe " +
-           "probe_func_";
+    return provider_ + ":" + name_ + " from binary " + binary_path_ + " PID " +
+           std::to_string(pid_) + " for probe " + "probe_func_";
   }
 
   friend std::ostream& operator<<(std::ostream& out, const USDT& usdt) {
-    return out << usdt.provider_ << ":" << usdt.name_ << " from "
-               << usdt.binary_path_ << " for probe " << usdt.probe_func_;
+    return out << usdt.provider_ << ":" << usdt.name_ << " from binary "
+               << usdt.binary_path_ << " PID " << usdt.pid_ << " for probe "
+               << usdt.probe_func_;
   }
 
  private:
   bool initialized_;
 
   std::string binary_path_;
+  pid_t pid_;
+
   std::string provider_;
   std::string name_;
   std::string probe_func_;
 
-  std::vector<uintptr_t> addresses_;
-
+  std::unique_ptr<void, std::function<void(void*)>> probe_;
   std::string program_text_;
 
   friend class BPF;

src/cc/usdt.h

@@ -26,6 +26,11 @@
 
 struct bcc_usdt;
 
+namespace ebpf {
+  class BPF;
+  class USDT;
+}
+
 namespace USDT {
 
 using std::experimental::optional;
@@ -209,7 +214,9 @@ public:
   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);
+  bool usdt_getarg(std::ostream &stream, const std::string& probe_func);
   std::string get_arg_ctype(int arg_index) {
     return largest_arg_type(arg_index);
   }
@@ -226,6 +233,9 @@ public:
   const std::string &provider() { return provider_; }
 
   friend class Context;
+
+  friend class ::ebpf::BPF;
+  friend class ::ebpf::USDT;
 };
 
 class Context {
@@ -269,5 +279,8 @@ public:
 
   typedef void (*each_uprobe_cb)(const char *, const char *, uint64_t, int);
   void each_uprobe(each_uprobe_cb callback);
+
+  friend class ::ebpf::BPF;
+  friend class ::ebpf::USDT;
 };
 }

src/cc/usdt/usdt.cc

@@ -159,11 +159,15 @@ std::string Probe::largest_arg_type(size_t arg_n) {
 }
 
 bool Probe::usdt_getarg(std::ostream &stream) {
-  const size_t arg_count = locations_[0].arguments_.size();
-
-  if (!attached_to_)
+  if (!attached_to_ || attached_to_->empty())
     return false;
 
+  return usdt_getarg(stream, attached_to_.value());
+}
+
+bool Probe::usdt_getarg(std::ostream &stream, const std::string& probe_func) {
+  const size_t arg_count = locations_[0].arguments_.size();
+
   if (arg_count == 0)
     return true;
 
@@ -175,7 +179,7 @@ bool Probe::usdt_getarg(std::ostream &stream) {
                 "static __always_inline int _bpf_readarg_%s_%d("
                 "struct pt_regs *ctx, void *dest, size_t len) {\n"
                 "  if (len != sizeof(%s)) return -1;\n",
-                attached_to_.value(), arg_n + 1, ctype);
+                probe_func, arg_n + 1, ctype);
 
     if (locations_.size() == 1) {
       Location &location = locations_.front();

tests/cc/test_usdt_probes.cc

@@ -20,6 +20,7 @@
 
 #include "catch.hpp"
 #include "usdt.h"
+#include "api/BPF.h"
 
 #ifdef HAVE_SDT_HEADER
 /* required to insert USDT probes on this very executable --
@@ -54,6 +55,34 @@ TEST_CASE("test finding a probe in our own process", "[usdt]") {
     REQUIRE(a_probed_function() != 0);
   }
 }
+
+TEST_CASE("test fine a probe in our own binary with C++ API", "[usdt]") {
+    ebpf::BPF bpf;
+    ebpf::USDT u("/proc/self/exe", "libbcc_test", "sample_probe_1", "on_event");
+
+    auto res = bpf.init("int on_event() { return 0; }", {}, {u});
+    REQUIRE(res.code() == 0);
+
+    res = bpf.attach_usdt(u);
+    REQUIRE(res.code() == 0);
+
+    res = bpf.detach_usdt(u);
+    REQUIRE(res.code() == 0);
+}
+
+TEST_CASE("test fine a probe in our Process with C++ API", "[usdt]") {
+    ebpf::BPF bpf;
+    ebpf::USDT u(::getpid(), "libbcc_test", "sample_probe_1", "on_event");
+
+    auto res = bpf.init("int on_event() { return 0; }", {}, {u});
+    REQUIRE(res.code() == 0);
+
+    res = bpf.attach_usdt(u);
+    REQUIRE(res.code() == 0);
+
+    res = bpf.detach_usdt(u);
+    REQUIRE(res.code() == 0);
+}
 #endif  // HAVE_SDT_HEADER
 
 class ChildProcess {