Merge pull request iovisor#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

@@ -0,0 +1,105 @@
+/*
+ * FollyRequestContextSwitch Monitor RequestContext switch events for any binary
+ * uses the class from [folly](http:https://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 @@ class BPF {
  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,11 +157,49 @@ class BPF {
 
  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_;
  std::map<std::string, BPFPerfBuffer*> perf_buffers_;
  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_;