Add USDT support to C++ API

src/cc/BPF.cc

@@ -30,6 +30,7 @@
 #include "bpf_module.h"
 #include "libbpf.h"
 #include "perf_reader.h"
+#include "usdt.h"
 
 #include "BPF.h"
 
@@ -50,13 +51,25 @@ std::string sanitize_str(std::string str, bool (*validator)(char),
 }
 
 StatusTuple BPF::init(const std::string& bpf_program,
-                      std::vector<std::string> cflags) {
+                      std::vector<std::string> cflags, 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));
+  }
+
   auto flags_len = cflags.size();
   const char* flags[flags_len];
   for (size_t i = 0; i < flags_len; i++)
     flags[i] = cflags[i].c_str();
-  if (bpf_module_->load_string(bpf_program, flags, flags_len) != 0)
+
+  all_bpf_program += bpf_program;
+  if (bpf_module_->load_string(all_bpf_program, flags, flags_len) != 0)
     return StatusTuple(-1, "Unable to initialize BPF program");
+
   return StatusTuple(0);
 };
 
@@ -206,6 +219,37 @@ StatusTuple BPF::attach_uprobe(const std::string& binary_path,
   return StatusTuple(0);
 }
 
+StatusTuple BPF::attach_usdt(const USDT& usdt, pid_t pid, int cpu,
+                             int group_fd) {
+  for (auto& u : usdt_)
+    if (u == usdt) {
+      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);
+        if (res.code() != 0) {
+          failed = true;
+          err_msg += "USDT " + u.print_name() + " at " + std::to_string(addr);
+          err_msg += ": " + res.msg() + "\n";
+          break;
+        }
+        cnt++;
+      }
+      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";
+        }
+        return StatusTuple(-1, err_msg);
+      } else
+        return StatusTuple(0);
+    }
+  return StatusTuple(-1, "USDT %s not found", usdt.print_name().c_str());
+}
+
 StatusTuple BPF::attach_tracepoint(const std::string& tracepoint,
                                    const std::string& probe_func,
                                    pid_t pid, int cpu, int group_fd,
@@ -311,6 +355,27 @@ StatusTuple BPF::detach_uprobe(const std::string& binary_path,
   return StatusTuple(0);
 }
 
+StatusTuple BPF::detach_usdt(const USDT& usdt) {
+  for (auto& u : usdt_)
+    if (u == usdt) {
+      bool failed = false;
+      std::string err_msg;
+      for (auto addr : u.addresses_) {
+        auto res = detach_uprobe(u.binary_path_, std::string(), addr);
+        if (res.code() != 0) {
+          failed = true;
+          err_msg += "USDT " + u.print_name() + " at " + std::to_string(addr);
+          err_msg += ": " + res.msg() + "\n";
+        }
+      }
+      if (failed)
+        return StatusTuple(-1, err_msg);
+      else
+        return StatusTuple(0);
+    }
+  return StatusTuple(-1, "USDT %s not found", usdt.print_name().c_str());
+}
+
 StatusTuple BPF::detach_tracepoint(const std::string& tracepoint) {
   auto it = tracepoints_.find(tracepoint);
   if (it == tracepoints_.end())
@@ -383,7 +448,7 @@ StatusTuple BPF::load_func(const std::string& func_name,
 StatusTuple BPF::unload_func(const std::string& func_name) {
   auto it = funcs_.find(func_name);
   if (it == funcs_.end())
-    return StatusTuple(-1, "Probe function %s not loaded", func_name.c_str());
+    return StatusTuple(0);
 
   int res = close(it->second);
   if (res != 0)
@@ -478,4 +543,28 @@ StatusTuple BPF::detach_perf_event_all_cpu(open_probe_t& attr) {
   return StatusTuple(0);
 }
 
+StatusTuple USDT::init() {
+  auto ctx =
+      std::unique_ptr<::USDT::Context>(new ::USDT::Context(binary_path_));
+  if (!ctx->loaded())
+    return StatusTuple(-1, "Unable to load USDT " + print_name());
+  auto probe = ctx->get(name_);
+  if (probe == nullptr)
+    return StatusTuple(-1, "Unable to find USDT " + print_name());
+
+  if (!probe->enable(probe_func_))
+    return StatusTuple(-1, "Failed to enable USDT " + print_name());
+  std::ostringstream stream;
+  if (!probe->usdt_getarg(stream))
+    return StatusTuple(
+        -1, "Unable to generate program text for USDT " + print_name());
+  program_text_ = ::USDT::USDT_PROGRAM_HEADER + stream.str();
+
+  for (size_t i = 0; i < probe->num_locations(); i++)
+    addresses_.push_back(probe->address(i));
+
+  initialized_ = true;
+  return StatusTuple(0);
+}
+
 }  // namespace ebpf

src/cc/BPF.h

@@ -40,13 +40,16 @@ struct open_probe_t {
   std::map<int, int>* per_cpu_fd;
 };
 
+class USDT;
+
 class BPF {
 public:
   static const int BPF_MAX_STACK_DEPTH = 127;
 
   explicit BPF(unsigned int flag = 0) : bpf_module_(new BPFModule(flag)) {}
   StatusTuple init(const std::string& bpf_program,
-                   std::vector<std::string> cflags = {});
+                   std::vector<std::string> cflags = {},
+                   std::vector<USDT> usdt = {});
 
   ~BPF();
   StatusTuple detach_all();
@@ -70,6 +73,9 @@ public:
       const std::string& binary_path, const std::string& symbol,
       uint64_t symbol_addr = 0,
       bpf_attach_type attach_type = bpf_attach_type::probe_entry);
+  StatusTuple attach_usdt(const USDT& usdt, pid_t pid = -1, int cpu = 0,
+                          int group_fd = -1);
+  StatusTuple detach_usdt(const USDT& usdt);
 
   StatusTuple attach_tracepoint(const std::string& tracepoint,
                                 const std::string& probe_func,
@@ -151,6 +157,8 @@ private:
 
   std::map<std::string, int> funcs_;
 
+  std::vector<USDT> usdt_;
+
   std::map<std::string, open_probe_t> kprobes_;
   std::map<std::string, open_probe_t> uprobes_;
   std::map<std::string, open_probe_t> tracepoints_;
@@ -158,4 +166,40 @@ private:
   std::map<std::pair<uint32_t, uint32_t>, open_probe_t> perf_events_;
 };
 
+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) {}
+
+  bool operator==(const USDT& other) const {
+    return (provider_ == other.provider_) && (name_ == other.name_) &&
+           (binary_path_ == other.binary_path_) &&
+           (probe_func_ == other.probe_func_);
+  }
+
+  std::string print_name() const {
+    return provider_ + ":" + name_ + " from " + binary_path_;
+  }
+
+private:
+  StatusTuple init();
+  bool initialized_;
+
+  std::string binary_path_;
+  std::string provider_;
+  std::string name_;
+  std::string probe_func_;
+
+  std::vector<intptr_t> addresses_;
+
+  std::string program_text_;
+
+  friend class BPF;
+};
+
 }  // namespace ebpf