From 6c33a5230ef5e34303de992f7b87b99656127a48 Mon Sep 17 00:00:00 2001
From: Teng Qin <qinteng@fb.com>
Date: Fri, 9 Mar 2018 17:03:31 -0800
Subject: [PATCH] C++ interface for attaching to perf event with raw
 perf_event_attr argument

---
 src/cc/api/BPF.cc | 51 +++++++++++++++++++++++++++++++++++++++++++++--
 src/cc/api/BPF.h  |  7 ++++++-
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/src/cc/api/BPF.cc b/src/cc/api/BPF.cc
index b3c8e242c46d..c1ef66a1fa61 100644
--- a/src/cc/api/BPF.cc
+++ b/src/cc/api/BPF.cc
@@ -15,6 +15,7 @@
  */
 
 #include <linux/bpf.h>
+#include <linux/perf_event.h>
 #include <unistd.h>
 #include <cstdio>
 #include <cstring>
@@ -292,12 +293,13 @@ StatusTuple BPF::attach_perf_event(uint32_t ev_type, uint32_t ev_config,
   int probe_fd;
   TRY2(load_func(probe_func, BPF_PROG_TYPE_PERF_EVENT, probe_fd));
 
-  auto fds = new std::map<int, int>();
   std::vector<int> cpus;
   if (cpu >= 0)
     cpus.push_back(cpu);
   else
     cpus = get_online_cpus();
+  auto fds = new std::vector<std::pair<int, int>>();
+  fds->reserve(cpus.size());
   for (int i : cpus) {
     int fd = bpf_attach_perf_event(probe_fd, ev_type, ev_config, sample_period,
                                    sample_freq, pid, i, group_fd);
@@ -309,7 +311,7 @@ StatusTuple BPF::attach_perf_event(uint32_t ev_type, uint32_t ev_config,
       return StatusTuple(-1, "Failed to attach perf event type %d config %d",
                          ev_type, ev_config);
     }
-    fds->emplace(i, fd);
+    fds->emplace_back(i, fd);
   }
 
   open_probe_t p = {};
@@ -319,6 +321,46 @@ StatusTuple BPF::attach_perf_event(uint32_t ev_type, uint32_t ev_config,
   return StatusTuple(0);
 }
 
+StatusTuple BPF::attach_perf_event_raw(void* perf_event_attr,
+                                       const std::string& probe_func,
+                                       pid_t pid, int cpu, int group_fd) {
+  auto attr = static_cast<struct perf_event_attr*>(perf_event_attr);
+  auto ev_pair = std::make_pair(attr->type, attr->config);
+  if (perf_events_.find(ev_pair) != perf_events_.end())
+    return StatusTuple(-1, "Perf event type %d config %d already attached",
+                       attr->type, attr->config);
+
+  int probe_fd;
+  TRY2(load_func(probe_func, BPF_PROG_TYPE_PERF_EVENT, probe_fd));
+
+  std::vector<int> cpus;
+  if (cpu >= 0)
+    cpus.push_back(cpu);
+  else
+    cpus = get_online_cpus();
+  auto fds = new std::vector<std::pair<int, int>>();
+  fds->reserve(cpus.size());
+  for (int i : cpus) {
+    int fd = bpf_attach_perf_event_raw(probe_fd, attr, pid, i, group_fd);
+    if (fd < 0) {
+      for (const 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",
+                         attr->type, attr->config);
+    }
+    fds->emplace_back(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_probe_attach_type attach_type) {
   std::string event = get_kprobe_event(kernel_func, attach_type);
@@ -394,6 +436,11 @@ StatusTuple BPF::detach_perf_event(uint32_t ev_type, uint32_t ev_config) {
   return StatusTuple(0);
 }
 
+StatusTuple BPF::detach_perf_event_raw(void* perf_event_attr) {
+  auto attr = static_cast<struct perf_event_attr*>(perf_event_attr);
+  return detach_perf_event(attr->type, attr->config);
+}
+
 StatusTuple BPF::open_perf_event(const std::string& name, uint32_t type,
                                  uint64_t config) {
   if (perf_event_arrays_.find(name) == perf_event_arrays_.end()) {
diff --git a/src/cc/api/BPF.h b/src/cc/api/BPF.h
index 3f2556f13262..c529721e42ee 100644
--- a/src/cc/api/BPF.h
+++ b/src/cc/api/BPF.h
@@ -36,7 +36,7 @@ namespace ebpf {
 struct open_probe_t {
   int perf_event_fd;
   std::string func;
-  std::map<int, int>* per_cpu_fd;
+  std::vector<std::pair<int, int>>* per_cpu_fd;
 };
 
 class USDT;
@@ -84,7 +84,12 @@ class BPF {
                                 uint64_t sample_period, uint64_t sample_freq,
                                 pid_t pid = -1, int cpu = -1,
                                 int group_fd = -1);
+  StatusTuple attach_perf_event_raw(void* perf_event_attr,
+                                    const std::string& probe_func,
+                                    pid_t pid = -1, int cpu = -1,
+                                    int group_fd = -1);
   StatusTuple detach_perf_event(uint32_t ev_type, uint32_t ev_config);
+  StatusTuple detach_perf_event_raw(void* perf_event_attr);
 
   BPFTable get_table(const std::string& name) {
     TableStorage::iterator it;