From dfd9c63be2f079455382e80ccf36fe9fc7d6e02a Mon Sep 17 00:00:00 2001
From: Teng Qin <qinteng@fb.com>
Date: Thu, 8 Dec 2016 17:46:13 -0800
Subject: [PATCH] Expose PMU event support in C++ API

---
 src/cc/BPF.cc | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 src/cc/BPF.h  | 10 ++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/src/cc/BPF.cc b/src/cc/BPF.cc
index cc15f29..265b4df 100644
--- a/src/cc/BPF.cc
+++ b/src/cc/BPF.cc
@@ -22,6 +22,7 @@
 #include <iostream>
 #include <memory>
 #include <sstream>
+#include <utility>
 #include <vector>
 
 #include "bcc_exception.h"
@@ -107,11 +108,19 @@ StatusTuple BPF::detach_all() {
     delete it.second;
   }
 
+  for (auto it : perf_events_) {
+    auto res = detach_perf_event_all_cpu(it.second);
+    if (res.code() != 0) {
+      error_msg += res.msg() + "\n";
+      has_error = true;
+    }
+  }
+
   for (auto it : funcs_) {
     int res = close(it.second);
     if (res != 0) {
       error_msg += "Failed to unload BPF program for " + it.first + ": ";
-      error_msg += std::string(std::strerror(res)) + "\n";
+      error_msg += std::string(std::strerror(errno)) + "\n";
       has_error = true;
     }
   }
@@ -231,6 +240,44 @@ StatusTuple BPF::attach_tracepoint(const std::string& tracepoint,
   return StatusTuple(0);
 }
 
+StatusTuple BPF::attach_perf_event(uint32_t ev_type, uint32_t ev_config,
+                                   const std::string& probe_func,
+                                   uint64_t sample_period, uint64_t sample_freq,
+                                   pid_t pid, int cpu, int group_fd) {
+  auto ev_pair = std::make_pair(ev_type, ev_config);
+  if (perf_events_.find(ev_pair) != perf_events_.end())
+    return StatusTuple(-1, "Perf event type %d config %d already attached",
+                       ev_type, ev_config);
+
+  int probe_fd;
+  TRY2(load_func(probe_func, BPF_PROG_TYPE_PERF_EVENT, probe_fd));
+
+  auto fds = new std::map<int, int>();
+  int cpu_st = 0;
+  int cpu_en = sysconf(_SC_NPROCESSORS_ONLN) - 1;
+  if (cpu >= 0)
+    cpu_st = cpu_en = cpu;
+  for (int i = cpu_st; i <= cpu_en; i++) {
+    int fd = bpf_attach_perf_event(probe_fd, ev_type, ev_config, sample_period,
+                                   sample_freq, pid, i, group_fd);
+    if (fd < 0) {
+      for (auto it : *fds)
+        close(it.second);
+      delete fds;
+      TRY2(unload_func(probe_func));
+      return StatusTuple(-1, "Failed to attach perf event type %d config %d",
+                         ev_type, ev_config);
+    }
+    fds->emplace(i, fd);
+  }
+
+  open_probe_t p = {};
+  p.func = probe_func;
+  p.per_cpu_fd = fds;
+  perf_events_[ev_pair] = std::move(p);
+  return StatusTuple(0);
+}
+
 StatusTuple BPF::detach_kprobe(const std::string& kernel_func,
                                bpf_attach_type attach_type) {
   std::string event = get_kprobe_event(kernel_func, attach_type);
@@ -274,6 +321,16 @@ StatusTuple BPF::detach_tracepoint(const std::string& tracepoint) {
   return StatusTuple(0);
 }
 
+StatusTuple BPF::detach_perf_event(uint32_t ev_type, uint32_t ev_config) {
+  auto it = perf_events_.find(std::make_pair(ev_type, ev_config));
+  if (it == perf_events_.end())
+    return StatusTuple(-1, "Perf Event type %d config %d not attached",
+                       ev_type, ev_config);
+  TRY2(detach_perf_event_all_cpu(it->second));
+  perf_events_.erase(it);
+  return StatusTuple(0);
+}
+
 StatusTuple BPF::open_perf_buffer(const std::string& name,
                                   perf_reader_raw_cb cb, void* cb_cookie) {
   if (perf_buffers_.find(name) == perf_buffers_.end())
@@ -401,4 +458,24 @@ StatusTuple BPF::detach_tracepoint_event(const std::string& tracepoint,
   return StatusTuple(0);
 }
 
+StatusTuple BPF::detach_perf_event_all_cpu(open_probe_t& attr) {
+  bool has_error = false;
+  std::string err_msg;
+  for (auto it : *attr.per_cpu_fd) {
+    int res = close(it.second);
+    if (res < 0) {
+      has_error = true;
+      err_msg += "Failed to close perf event FD " + std::to_string(it.second) +
+                 " For CPU " + std::to_string(it.first) + ": ";
+      err_msg += std::string(std::strerror(errno)) + "\n";
+    }
+  }
+  delete attr.per_cpu_fd;
+  TRY2(unload_func(attr.func));
+
+  if (has_error)
+    return StatusTuple(-1, err_msg);
+  return StatusTuple(0);
+}
+
 }  // namespace ebpf
diff --git a/src/cc/BPF.h b/src/cc/BPF.h
index b5f1ae3..420dd6b 100644
--- a/src/cc/BPF.h
+++ b/src/cc/BPF.h
@@ -37,6 +37,7 @@ enum class bpf_attach_type {
 struct open_probe_t {
   void* reader_ptr;
   std::string func;
+  std::map<int, int>* per_cpu_fd;
 };
 
 class BPF {
@@ -77,6 +78,13 @@ public:
                                 void* cb_cookie = nullptr);
   StatusTuple detach_tracepoint(const std::string& tracepoint);
 
+  StatusTuple attach_perf_event(uint32_t ev_type, uint32_t ev_config,
+                                const std::string& probe_func,
+                                uint64_t sample_period, uint64_t sample_freq,
+                                pid_t pid = -1, int cpu = -1,
+                                int group_fd = -1);
+  StatusTuple detach_perf_event(uint32_t ev_type, uint32_t ev_config);
+
   template <class KeyType, class ValueType>
   BPFHashTable<KeyType, ValueType> get_hash_table(const std::string& name) {
     return BPFHashTable<KeyType, ValueType>(bpf_module_.get(), name);
@@ -105,6 +113,7 @@ private:
   StatusTuple detach_uprobe_event(const std::string& event, open_probe_t& attr);
   StatusTuple detach_tracepoint_event(const std::string& tracepoint,
                                       open_probe_t& attr);
+  StatusTuple detach_perf_event_all_cpu(open_probe_t& attr);
 
   std::string attach_type_debug(bpf_attach_type type) {
     switch (type) {
@@ -146,6 +155,7 @@ private:
   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_;
 };
 
 }  // namespace ebpf
-- 
2.7.4