Merge pull request #1711 from palmtenor/cgrouparray

Support Cgroup Array in BCC

examples/cpp/RandomRead.cc

@@ -19,6 +19,10 @@ const std::string BPF_PROGRAM = R"(
 #include <linux/sched.h>
 #include <uapi/linux/ptrace.h>
 
+#ifndef CGROUP_FILTER
+#define CGROUP_FILTER 0
+#endif
+
 struct urandom_read_args {
   // See /sys/kernel/debug/tracing/events/random/urandom_read/format
   uint64_t common__unused;
@@ -35,8 +39,12 @@ struct event_t {
 };
 
 BPF_PERF_OUTPUT(events);
+BPF_CGROUP_ARRAY(cgroup, 1);
 
 int on_urandom_read(struct urandom_read_args* attr) {
+  if (CGROUP_FILTER && (cgroup.check_current_task(0) != 1))
+    return 0;
+
   struct event_t event = {};
   event.pid = bpf_get_current_pid_tgid();
   bpf_get_current_comm(&event.comm, sizeof(event.comm));
@@ -72,12 +80,29 @@ void signal_handler(int s) {
 }
 
 int main(int argc, char** argv) {
+  if (argc != 1 && argc != 2) {
+    std::cerr << "USAGE: RandomRead [cgroup2_path]" << std::endl;
+    return 1;
+  }
+
+  std::vector<std::string> cflags = {};
+  if (argc == 2)
+    cflags.emplace_back("-DCGROUP_FILTER=1");
+
   bpf = new ebpf::BPF();
-  auto init_res = bpf->init(BPF_PROGRAM);
+  auto init_res = bpf->init(BPF_PROGRAM, cflags, {});
   if (init_res.code() != 0) {
     std::cerr << init_res.msg() << std::endl;
     return 1;
   }
+  if (argc == 2) {
+    auto cgroup_array = bpf->get_cgroup_array("cgroup");
+    auto update_res = cgroup_array.update_value(0, argv[1]);
+    if (update_res.code() != 0) {
+      std::cerr << update_res.msg() << std::endl;
+      return 1;
+    }
+  }
 
   auto attach_res =
       bpf->attach_tracepoint("random:urandom_read", "on_urandom_read");

src/cc/api/BPF.cc

@@ -585,6 +585,13 @@ BPFProgTable BPF::get_prog_table(const std::string& name) {
   return BPFProgTable({});
 }
 
+BPFCgroupArray BPF::get_cgroup_array(const std::string& name) {
+  TableStorage::iterator it;
+  if (bpf_module_->table_storage().Find(Path({bpf_module_->id(), name}), it))
+    return BPFCgroupArray(it->second);
+  return BPFCgroupArray({});
+}
+
 BPFStackTable BPF::get_stack_table(const std::string& name, bool use_debug_file,
                                    bool check_debug_file_crc) {
   TableStorage::iterator it;

src/cc/api/BPF.h

@@ -132,6 +132,8 @@ class BPF {
 
   BPFProgTable get_prog_table(const std::string& name);
 
+  BPFCgroupArray get_cgroup_array(const std::string& name);
+
   BPFStackTable get_stack_table(const std::string& name,
                                 bool use_debug_file = true,
                                 bool check_debug_file_crc = true);

src/cc/api/BPFTable.cc

@@ -17,6 +17,7 @@
 #include <linux/elf.h>
 #include <linux/perf_event.h>
 #include <sys/epoll.h>
+#include <fcntl.h>
 #include <unistd.h>
 #include <cerrno>
 #include <cinttypes>
@@ -30,6 +31,7 @@
 #include "bcc_exception.h"
 #include "bcc_syms.h"
 #include "common.h"
+#include "file_desc.h"
 #include "libbpf.h"
 #include "perf_reader.h"
 
@@ -433,4 +435,39 @@ BPFPerfEventArray::~BPFPerfEventArray() {
               << std::endl;
   }
 }
+
+StatusTuple BPFProgTable::update_value(const int& index, const int& prog_fd) {
+  if (!this->update(const_cast<int*>(&index), const_cast<int*>(&prog_fd)))
+    return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFProgTable::remove_value(const int& index) {
+  if (!this->remove(const_cast<int*>(&index)))
+    return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFCgroupArray::update_value(const int& index,
+                                         const int& cgroup2_fd) {
+  if (!this->update(const_cast<int*>(&index), const_cast<int*>(&cgroup2_fd)))
+    return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFCgroupArray::update_value(const int& index,
+                                         const std::string& cgroup2_path) {
+  FileDesc f(::open(cgroup2_path.c_str(), O_RDONLY | O_CLOEXEC));
+  if ((int)f < 0)
+    return StatusTuple(-1, "Unable to open %s", cgroup2_path.c_str());
+  TRY2(update_value(index, (int)f));
+  return StatusTuple(0);
+}
+
+StatusTuple BPFCgroupArray::remove_value(const int& index) {
+  if (!this->remove(const_cast<int*>(&index)))
+    return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
+  return StatusTuple(0);
+}
+
 }  // namespace ebpf

src/cc/api/BPFTable.h

@@ -343,18 +343,21 @@ public:
       throw std::invalid_argument("Table '" + desc.name + "' is not a prog table");
   }
 
-  // updates an element
-  StatusTuple update_value(const int& index, const int& value) {
-    if (!this->update(const_cast<int*>(&index), const_cast<int*>(&value)))
-      return StatusTuple(-1, "Error updating value: %s", std::strerror(errno));
-    return StatusTuple(0);
-  }
+  StatusTuple update_value(const int& index, const int& prog_fd);
+  StatusTuple remove_value(const int& index);
+};
 
-  StatusTuple remove_value(const int& index) {
-    if (!this->remove(const_cast<int*>(&index)))
-      return StatusTuple(-1, "Error removing value: %s", std::strerror(errno));
-    return StatusTuple(0);
+class BPFCgroupArray : public BPFTableBase<int, int> {
+public:
+  BPFCgroupArray(const TableDesc& desc)
+      : BPFTableBase<int, int>(desc) {
+    if (desc.type != BPF_MAP_TYPE_CGROUP_ARRAY)
+      throw std::invalid_argument("Table '" + desc.name + "' is not a cgroup array");
   }
+
+  StatusTuple update_value(const int& index, const int& cgroup2_fd);
+  StatusTuple update_value(const int& index, const std::string& cgroup2_path);
+  StatusTuple remove_value(const int& index);
 };
 
 }  // namespace ebpf

src/cc/export/helpers.h

@@ -103,6 +103,17 @@ struct _name##_table_t { \
 __attribute__((section("maps/perf_array"))) \
 struct _name##_table_t _name = { .max_entries = (_max_entries) }
 
+// Table for cgroup file descriptors
+#define BPF_CGROUP_ARRAY(_name, _max_entries) \
+struct _name##_table_t { \
+  int key; \
+  u32 leaf; \
+  int (*check_current_task) (int); \
+  u32 max_entries; \
+}; \
+__attribute__((section("maps/cgroup_array"))) \
+struct _name##_table_t _name = { .max_entries = (_max_entries) }
+
 #define BPF_HASH1(_name) \
   BPF_TABLE("hash", u64, u64, _name, 10240)
 #define BPF_HASH2(_name, _key_type) \

src/cc/frontends/clang/b_frontend_action.cc

@@ -459,6 +459,9 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
           } else if (memb_name == "perf_counter_value") {
             prefix = "bpf_perf_event_read_value";
             suffix = ")";
+          } else if (memb_name == "check_current_task") {
+            prefix = "bpf_current_task_under_cgroup";
+            suffix = ")";
           } else {
             error(Call->getLocStart(), "invalid bpf_table operation %0") << memb_name;
             return false;
@@ -762,6 +765,8 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       table.max_entries = numcpu;
     } else if (A->getName() == "maps/perf_array") {
       map_type = BPF_MAP_TYPE_PERF_EVENT_ARRAY;
+    } else if (A->getName() == "maps/cgroup_array") {
+      map_type = BPF_MAP_TYPE_CGROUP_ARRAY;
     } else if (A->getName() == "maps/stacktrace") {
       map_type = BPF_MAP_TYPE_STACK_TRACE;
     } else if (A->getName() == "maps/extern") {

src/python/bcc/table.py

@@ -142,6 +142,8 @@ def Table(bpf, map_id, map_fd, keytype, leaftype, **kwargs):
         t = LruHash(bpf, map_id, map_fd, keytype, leaftype)
     elif ttype == BPF_MAP_TYPE_LRU_PERCPU_HASH:
         t = LruPerCpuHash(bpf, map_id, map_fd, keytype, leaftype)
+    elif ttype == BPF_MAP_TYPE_CGROUP_ARRAY:
+        t = CgroupArray(bpf, map_id, map_fd, keytype, leaftype)
     if t == None:
         raise Exception("Unknown table type %d" % ttype)
     return t
@@ -199,8 +201,7 @@ class TableBase(MutableMapping):
         return leaf
 
     def __setitem__(self, key, leaf):
-        res = lib.bpf_update_elem(self.map_fd, ct.byref(key), ct.byref(leaf),
-                                  0)
+        res = lib.bpf_update_elem(self.map_fd, ct.byref(key), ct.byref(leaf), 0)
         if res < 0:
             errstr = os.strerror(ct.get_errno())
             raise Exception("Could not update table: %s" % errstr)
@@ -477,6 +478,40 @@ class ProgArray(ArrayBase):
             leaf = self.Leaf(leaf.fd)
         super(ProgArray, self).__setitem__(key, leaf)
 
+class FileDesc:
+    def __init__(self, fd):
+        if (fd is None) or (fd < 0):
+            raise Exception("Invalid file descriptor")
+        self.fd = fd
+
+    def clean_up(self):
+        if (self.fd is not None) and (self.fd >= 0):
+            os.close(self.fd)
+            self.fd = None
+
+    def __del__(self):
+        self.clean_up()
+
+    def __enter__(self, *args, **kwargs):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.clean_up()
+
+class CgroupArray(ArrayBase):
+    def __init__(self, *args, **kwargs):
+        super(CgroupArray, self).__init__(*args, **kwargs)
+
+    def __setitem__(self, key, leaf):
+        if isinstance(leaf, int):
+            super(CgroupArray, self).__setitem__(key, self.Leaf(leaf))
+        elif isinstance(leaf, str):
+            # TODO: Add os.O_CLOEXEC once we move to Python version >3.3
+            with FileDesc(os.open(leaf, os.O_RDONLY)) as f:
+                super(CgroupArray, self).__setitem__(key, self.Leaf(f.fd))
+        else:
+            raise Exception("Cgroup array key must be either FD or cgroup path")
+
 class PerfEventArray(ArrayBase):
 
     def __init__(self, *args, **kwargs):