Improve C++ API perf buffer polling

examples/cpp/RandomRead.cc

@@ -30,6 +30,7 @@ struct urandom_read_args {
 struct event_t {
   int pid;
   char comm[16];
+  int cpu;
   int got_bits;
 };
 
@@ -39,6 +40,7 @@ int on_urandom_read(struct urandom_read_args* attr) {
   struct event_t event = {};
   event.pid = bpf_get_current_pid_tgid();
   bpf_get_current_comm(&event.comm, sizeof(event.comm));
+  event.cpu = bpf_get_smp_processor_id();
   event.got_bits = attr->got_bits;
 
   events.perf_submit(attr, &event, sizeof(event));
@@ -50,13 +52,15 @@ int on_urandom_read(struct urandom_read_args* attr) {
 struct event_t {
   int pid;
   char comm[16];
+  int cpu;
   int got_bits;
 };
 
 void handle_output(void* cb_cookie, void* data, int data_size) {
   auto event = static_cast<event_t*>(data);
-  std::cout << "PID: " << event->pid << " (" << event->comm << ") "
-            << "Read " << event->got_bits << " bits" << std::endl;
+  std::cout << "PID: " << event->pid << " (" << event->comm << ") on CPU "
+            << event->cpu << " read " << event->got_bits << " bits"
+            << std::endl;
 }
 
 ebpf::BPF* bpf;

src/cc/BPF.cc

@@ -98,7 +98,7 @@ StatusTuple BPF::detach_all() {
   }
 
   for (auto it : perf_buffers_) {
-    auto res = it.second->close();
+    auto res = it.second->close_all_cpu();
     if (res.code() != 0) {
       error_msg += "Failed to close perf buffer " + it.first + ": ";
       error_msg += res.msg() + "\n";
@@ -279,7 +279,7 @@ StatusTuple BPF::open_perf_buffer(const std::string& name,
   if (perf_buffers_.find(name) == perf_buffers_.end())
     perf_buffers_[name] = new BPFPerfBuffer(bpf_module_.get(), name);
   auto table = perf_buffers_[name];
-  TRY2(table->open(cb, cb_cookie));
+  TRY2(table->open_all_cpu(cb, cb_cookie));
   return StatusTuple(0);
 }
 
@@ -287,7 +287,7 @@ StatusTuple BPF::close_perf_buffer(const std::string& name) {
   auto it = perf_buffers_.find(name);
   if (it == perf_buffers_.end())
     return StatusTuple(-1, "Perf buffer for %s not open", name.c_str());
-  TRY2(it->second->close());
+  TRY2(it->second->close_all_cpu());
   return StatusTuple(0);
 }
 

src/cc/BPFTable.cc

@@ -79,13 +79,23 @@ StatusTuple BPFPerfBuffer::open_on_cpu(perf_reader_raw_cb cb, int cpu,
     return StatusTuple(-1, "Unable to open perf buffer on CPU %d: %s", cpu,
                        strerror(errno));
   }
-  cpu_readers_[cpu] = static_cast<perf_reader*>(reader);
+  cpu_readers_[cpu] = reader;
+  readers_.push_back(reader);
   return StatusTuple(0);
 }
 
-StatusTuple BPFPerfBuffer::open(perf_reader_raw_cb cb, void* cb_cookie) {
-  for (int i = 0; i < sysconf(_SC_NPROCESSORS_ONLN); i++)
-    TRY2(open_on_cpu(cb, i, cb_cookie));
+StatusTuple BPFPerfBuffer::open_all_cpu(perf_reader_raw_cb cb,
+                                        void* cb_cookie) {
+  if (cpu_readers_.size() != 0 || readers_.size() != 0)
+    return StatusTuple(-1, "Previously opened perf buffer not cleaned");
+
+  for (int i = 0; i < sysconf(_SC_NPROCESSORS_ONLN); i++) {
+    auto res = open_on_cpu(cb, i, cb_cookie);
+    if (res.code() != 0) {
+      TRY2(close_all_cpu());
+      return res;
+    }
+  }
   return StatusTuple(0);
 }
 
@@ -100,7 +110,7 @@ StatusTuple BPFPerfBuffer::close_on_cpu(int cpu) {
   return StatusTuple(0);
 }
 
-StatusTuple BPFPerfBuffer::close() {
+StatusTuple BPFPerfBuffer::close_all_cpu() {
   std::string errors;
   bool has_error = false;
   for (int i = 0; i < sysconf(_SC_NPROCESSORS_ONLN); i++) {
@@ -111,21 +121,20 @@ StatusTuple BPFPerfBuffer::close() {
       has_error = true;
     }
   }
+  readers_.clear();
   if (has_error)
     return StatusTuple(-1, errors);
   return StatusTuple(0);
 }
 
 void BPFPerfBuffer::poll(int timeout) {
-  perf_reader* readers[cpu_readers_.size()];
-  int i = 0;
-  for (auto it : cpu_readers_)
-    readers[i++] = it.second;
-  perf_reader_poll(cpu_readers_.size(), readers, timeout);
+  if (readers_.empty())
+    return;
+  perf_reader_poll(readers_.size(), readers_.data(), timeout);
 }
 
 BPFPerfBuffer::~BPFPerfBuffer() {
-  auto res = close();
+  auto res = close_all_cpu();
   if (res.code() != 0)
     std::cerr << "Failed to close all perf buffer on destruction: "
               << res.msg() << std::endl;

src/cc/BPFTable.h

@@ -126,8 +126,8 @@ public:
       : BPFTableBase<int, int>(bpf_module, name) {}
   ~BPFPerfBuffer();
 
-  StatusTuple open(perf_reader_raw_cb cb, void* cb_cookie);
-  StatusTuple close();
+  StatusTuple open_all_cpu(perf_reader_raw_cb cb, void* cb_cookie);
+  StatusTuple close_all_cpu();
   void poll(int timeout);
 
 private:
@@ -135,6 +135,7 @@ private:
   StatusTuple close_on_cpu(int cpu);
 
   std::map<int, perf_reader*> cpu_readers_;
+  std::vector<perf_reader*> readers_;
 };
 
 }  // namespace ebpf