Merge pull request #862 from palmtenor/cppusdt

Add USDT support to C++ API

examples/cpp/CMakeLists.txt

@@ -26,3 +26,7 @@ install (TARGETS RandomRead DESTINATION share/bcc/examples/cpp)
 add_executable(LLCStat LLCStat.cc)
 target_link_libraries(LLCStat bcc-static)
 install (TARGETS LLCStat DESTINATION share/bcc/examples/cpp)
+
+add_executable(FollyRequestContextSwitch FollyRequestContextSwitch.cc)
+target_link_libraries(FollyRequestContextSwitch bcc-static)
+install (TARGETS FollyRequestContextSwitch DESTINATION share/bcc/examples/cpp)

examples/cpp/FollyRequestContextSwitch.cc

+/*
+ * FollyRequestContextSwitch Monitor RequestContext switch events for any binary
+ *                           uses the class from [folly](http://bit.ly/2h6S1yx).
+ *                           For Linux, uses BCC, eBPF. Embedded C.
+ *
+ * Basic example of using USDT with BCC.
+ *
+ * USAGE: FollyRequestContextSwitch PATH_TO_BINARY
+ *
+ * Copyright (c) Facebook, Inc.
+ * Licensed under the Apache License, Version 2.0 (the "License")
+ */
+
+#include <signal.h>
+#include <iostream>
+#include <vector>
+
+#include "BPF.h"
+
+const std::string BPF_PROGRAM = R"(
+#include <linux/sched.h>
+#include <uapi/linux/ptrace.h>
+
+struct event_t {
+  int pid;
+  char name[16];
+  uint64_t old_addr;
+  uint64_t new_addr;
+};
+
+BPF_PERF_OUTPUT(events);
+
+int on_context_switch(struct pt_regs *ctx) {
+  struct event_t event = {};
+
+  event.pid = bpf_get_current_pid_tgid();
+  bpf_get_current_comm(&event.name, sizeof(event.name));
+  
+  bpf_usdt_readarg(1, ctx, &event.old_addr);
+  bpf_usdt_readarg(2, ctx, &event.new_addr);
+
+  events.perf_submit(ctx, &event, sizeof(event));
+  return 0;
+}
+)";
+
+// Define the same struct to use in user space.
+struct event_t {
+  int pid;
+  char name[16];
+  uint64_t old_addr;
+  uint64_t new_addr;
+};
+
+void handle_output(void* cb_cookie, void* data, int data_size) {
+  auto event = static_cast<event_t*>(data);
+  std::cout << "PID " << event->pid << " (" << event->name << ") ";
+  std::cout << "folly::RequestContext switch from " << event->old_addr << " to "
+            << event->new_addr << std::endl;
+}
+
+ebpf::BPF* bpf;
+
+void signal_handler(int s) {
+  std::cerr << "Terminating..." << std::endl;
+  delete bpf;
+  exit(0);
+}
+
+int main(int argc, char** argv) {
+  if (argc != 2) {
+    std::cout << "USAGE: FollyRequestContextSwitch PATH_TO_BINARY" << std::endl;
+    exit(1);
+  }
+  std::string binary_path(argv[1]);
+
+  bpf = new ebpf::BPF();
+  std::vector<ebpf::USDT> u;
+  u.emplace_back(binary_path, "folly", "request_context_switch_before",
+                 "on_context_switch");
+  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]);
+  if (attach_res.code() != 0) {
+    std::cerr << attach_res.msg() << std::endl;
+    return 1;
+  }
+
+  auto open_res = bpf->open_perf_buffer("events", &handle_output);
+  if (open_res.code() != 0) {
+    std::cerr << open_res.msg() << std::endl;
+    return 1;
+  }
+
+  signal(SIGINT, signal_handler);
+  std::cout << "Started tracing, hit Ctrl-C to terminate." << std::endl;
+  while (true)
+    bpf->poll_perf_buffer("events");
+
+  return 0;
+}

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

src/cc/usdt.cc

@@ -239,7 +239,7 @@ Probe *Context::get(const std::string &probe_name) {
 }
 
 bool Context::generate_usdt_args(std::ostream &stream) {
-  stream << "#include <uapi/linux/ptrace.h>\n";
+  stream << USDT_PROGRAM_HEADER;
   for (auto &p : probes_) {
     if (p->enabled() && !p->usdt_getarg(stream))
       return false;

src/cc/usdt.h

@@ -31,6 +31,9 @@ using std::experimental::optional;
 using std::experimental::nullopt;
 class ArgumentParser;
 
+static const std::string USDT_PROGRAM_HEADER =
+    "#include <uapi/linux/ptrace.h>\n";
+
 class Argument {
 private:
   optional<int> arg_size_;