Support native integer indexing in table.Array type

Improve the indexing in the Array class to be more like native python
list/array types. No need to use `t[c_int(0)]`, instead `t[0]` is
sufficient, for instance.

Add tests for the above. Relies on a new bpf_module function for
exposing the max_entries property of a table in order to range-check the
indices.

In one case, array was using a struct key type.
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/cc/bpf_common.cc

@@ -134,6 +134,18 @@ int bpf_table_type_id(void *program, size_t id) {
   return mod->table_type(id);
 }
 
+size_t bpf_table_max_entries(void *program, const char *table_name) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->table_max_entries(table_name);
+}
+
+size_t bpf_table_max_entries_id(void *program, size_t id) {
+  auto mod = static_cast<ebpf::BPFModule *>(program);
+  if (!mod) return 0;
+  return mod->table_max_entries(id);
+}
+
 const char * bpf_table_name(void *program, size_t id) {
   auto mod = static_cast<ebpf::BPFModule *>(program);
   if (!mod) return nullptr;

src/cc/bpf_common.h

@@ -42,6 +42,8 @@ int bpf_table_fd(void *program, const char *table_name);
 int bpf_table_fd_id(void *program, size_t id);
 int bpf_table_type(void *program, const char *table_name);
 int bpf_table_type_id(void *program, size_t id);
+size_t bpf_table_max_entries(void *program, const char *table_name);
+size_t bpf_table_max_entries_id(void *program, size_t id);
 const char * bpf_table_name(void *program, size_t id);
 const char * bpf_table_key_desc(void *program, const char *table_name);
 const char * bpf_table_key_desc_id(void *program, size_t id);

src/cc/bpf_module.cc

@@ -522,6 +522,15 @@ int BPFModule::table_type(size_t id) const {
   return (*tables_)[id].type;
 }
 
+size_t BPFModule::table_max_entries(const string &name) const {
+  return table_max_entries(table_id(name));
+}
+
+size_t BPFModule::table_max_entries(size_t id) const {
+  if (id >= tables_->size()) return 0;
+  return (*tables_)[id].max_entries;
+}
+
 const char * BPFModule::table_name(size_t id) const {
   if (id >= tables_->size()) return nullptr;
   return (*tables_)[id].name.c_str();

src/cc/bpf_module.h

@@ -70,6 +70,8 @@ class BPFModule {
   const char * table_name(size_t id) const;
   int table_type(const std::string &name) const;
   int table_type(size_t id) const;
+  size_t table_max_entries(const std::string &name) const;
+  size_t table_max_entries(size_t id) const;
   const char * table_key_desc(size_t id) const;
   const char * table_key_desc(const std::string &name) const;
   size_t table_key_size(size_t id) const;

src/python/bcc/libbcc.py

@@ -45,6 +45,8 @@ lib.bpf_table_fd.restype = ct.c_int
 lib.bpf_table_fd.argtypes = [ct.c_void_p, ct.c_char_p]
 lib.bpf_table_type_id.restype = ct.c_int
 lib.bpf_table_type_id.argtypes = [ct.c_void_p, ct.c_ulonglong]
+lib.bpf_table_max_entries_id.restype = ct.c_ulonglong
+lib.bpf_table_max_entries_id.argtypes = [ct.c_void_p, ct.c_ulonglong]
 lib.bpf_table_key_desc.restype = ct.c_char_p
 lib.bpf_table_key_desc.argtypes = [ct.c_void_p, ct.c_char_p]
 lib.bpf_table_leaf_desc.restype = ct.c_char_p

src/python/bcc/table.py

@@ -37,6 +37,7 @@ def _stars(val, val_max, width):
         text = text[:-1] + "+"
     return text
 
+
 def _print_log2_hist(vals, val_type):
     global stars_max
     log2_dist_max = 64
@@ -87,6 +88,7 @@ def Table(bpf, map_id, map_fd, keytype, leaftype):
         raise Exception("Unknown table type %d" % ttype)
     return t
 
+
 class TableBase(MutableMapping):
 
     def __init__(self, bpf, map_id, map_fd, keytype, leaftype):
@@ -154,14 +156,6 @@ class TableBase(MutableMapping):
         if res < 0:
             raise Exception("Could not update table")
 
-    def __len__(self):
-        i = 0
-        for k in self: i += 1
-        return i
-
-    def __delitem__(self, key):
-        raise Exception("__delitem__ not implemented, abstract base class")
-
     # override the MutableMapping's implementation of these since they
     # don't handle KeyError nicely
     def itervalues(self):
@@ -191,7 +185,6 @@ class TableBase(MutableMapping):
         for k in self.keys():
             self.__delitem__(k)
 
-
     def __iter__(self):
         return TableBase.Iter(self, self.Key)
 
@@ -270,17 +263,48 @@ class HashTable(TableBase):
     def __init__(self, *args, **kwargs):
         super(HashTable, self).__init__(*args, **kwargs)
 
+    def __len__(self):
+        i = 0
+        for k in self: i += 1
+        return i
+
     def __delitem__(self, key):
         key_p = ct.pointer(key)
         res = lib.bpf_delete_elem(self.map_fd, ct.cast(key_p, ct.c_void_p))
         if res < 0:
             raise KeyError
 
+
 class ArrayBase(TableBase):
     def __init__(self, *args, **kwargs):
         super(ArrayBase, self).__init__(*args, **kwargs)
+        self.max_entries = int(lib.bpf_table_max_entries_id(self.bpf.module,
+                self.map_id))
+
+    def _normalize_key(self, key):
+        if isinstance(key, int):
+            if key < 0:
+                key = len(self) + key
+            key = self.Key(key)
+        if not isinstance(key, ct._SimpleCData):
+            raise IndexError("Array index must be an integer type")
+        if key.value >= len(self):
+            raise IndexError("Array index out of range")
+        return key
+
+    def __len__(self):
+        return self.max_entries
+
+    def __getitem__(self, key):
+        key = self._normalize_key(key)
+        return super(ArrayBase, self).__getitem__(key)
+
+    def __setitem__(self, key, leaf):
+        key = self._normalize_key(key)
+        super(ArrayBase, self).__setitem__(key, leaf)
 
     def __delitem__(self, key):
+        key = self._normalize_key(key)
         key_p = ct.pointer(key)
 
         # Deleting from array type maps does not have an effect, so
@@ -292,16 +316,42 @@ class ArrayBase(TableBase):
         if res < 0:
             raise Exception("Could not clear item")
 
+    def __iter__(self):
+        return ArrayBase.Iter(self, self.Key)
+
+    class Iter(object):
+        def __init__(self, table, keytype):
+            self.Key = keytype
+            self.table = table
+            self.i = -1
+
+        def __iter__(self):
+            return self
+        def __next__(self):
+            return self.next()
+        def next(self):
+            self.i += 1
+            if self.i == len(self.table):
+                raise StopIteration()
+            return self.Key(self.i)
+
 class Array(ArrayBase):
     def __init__(self, *args, **kwargs):
         super(Array, self).__init__(*args, **kwargs)
 
 
-
 class ProgArray(ArrayBase):
     def __init__(self, *args, **kwargs):
         super(ProgArray, self).__init__(*args, **kwargs)
 
+    def __setitem__(self, key, leaf):
+        if isinstance(leaf, int):
+            leaf = self.Leaf(leaf)
+        if isinstance(leaf, self.bpf.Function):
+            leaf = self.Leaf(leaf.fd)
+        super(ProgArray, self).__setitem__(key, leaf)
+
+
 class PerfEventArray(ArrayBase):
     def __init__(self, *args, **kwargs):
         super(PerfEventArray, self).__init__(*args, **kwargs)

tests/cc/test_array.py

@@ -21,5 +21,20 @@ class TestArray(TestCase):
                 self.assertEqual(v.value, 1000)
         self.assertEqual(len(t1), 128)
 
+    def test_native_type(self):
+        b = BPF(text="""BPF_TABLE("array", int, u64, table1, 128);""")
+        t1 = b["table1"]
+        t1[0] = c_ulonglong(100)
+        t1[-2] = c_ulonglong(37)
+        t1[127] = c_ulonglong(1000)
+        for i, v in t1.items():
+            if i.value == 0:
+                self.assertEqual(v.value, 100)
+            if i.value == 127:
+                self.assertEqual(v.value, 1000)
+        self.assertEqual(len(t1), 128)
+        self.assertEqual(t1[-2].value, 37)
+        self.assertEqual(t1[-1].value, t1[127].value)
+
 if __name__ == "__main__":
     main()

tests/cc/test_trace3.py

@@ -35,6 +35,7 @@ class TestBlkRequest(TestCase):
         for key, leaf in self.latency.items():
             print("latency %u:" % key.value, "count %u" % leaf.value)
         sys.stdout.flush()
+        self.assertEqual(len(list(self.latency.keys())), len(self.latency))
 
 if __name__ == "__main__":
     main()

tests/cc/test_trace4.py

@@ -12,7 +12,7 @@ class TestKprobeRgx(TestCase):
         self.b = BPF(text="""
         typedef struct { int idx; } Key;
         typedef struct { u64 val; } Val;
-        BPF_TABLE("array", Key, Val, stats, 3);
+        BPF_TABLE("hash", Key, Val, stats, 3);
         int hello(void *ctx) {
           stats.lookup_or_init(&(Key){1}, &(Val){0})->val++;
           return 0;