Merge pull request #1183 from palmtenor/firstkey

Use new Kernel functionality to get first key of map

src/cc/BPFTable.h

@@ -70,6 +70,11 @@ class BPFTableBase {
                            static_cast<void*>(value)) >= 0;
   }
 
+  bool first(KeyType* key) {
+    return bpf_get_first_key(desc.fd, static_cast<void*>(key),
+                             desc.key_size) >= 0;
+  }
+
   bool next(KeyType* key, KeyType* next_key) {
     return bpf_get_next_key(desc.fd, static_cast<void*>(key),
                             static_cast<void*>(next_key)) >= 0;
@@ -174,17 +179,18 @@ class BPFHashTable : public BPFTableBase<KeyType, ValueType> {
 
   std::vector<std::pair<KeyType, ValueType>> get_table_offline() {
     std::vector<std::pair<KeyType, ValueType>> res;
-
-    KeyType cur, nxt;
+    KeyType cur;
     ValueType value;
 
+    if (!this->first(&cur))
+      return res;
+
     while (true) {
-      if (!this->next(&cur, &nxt))
+      if (!this->lookup(&cur, &value))
         break;
-      if (!this->lookup(&nxt, &value))
+      res.emplace_back(cur, value);
+      if (!this->next(&cur, &cur))
         break;
-      res.emplace_back(nxt, value);
-      std::swap(cur, nxt);
     }
 
     return res;

src/cc/libbpf.c

@@ -133,6 +133,44 @@ int bpf_delete_elem(int fd, void *key)
   return syscall(__NR_bpf, BPF_MAP_DELETE_ELEM, &attr, sizeof(attr));
 }
 
+int bpf_get_first_key(int fd, void *key, size_t key_size)
+{
+  union bpf_attr attr;
+  int i, res;
+
+  memset(&attr, 0, sizeof(attr));
+  attr.map_fd = fd;
+  attr.key = 0;
+  attr.next_key = ptr_to_u64(key);
+
+  // 4.12 and above kernel supports passing NULL to BPF_MAP_GET_NEXT_KEY
+  // to get first key of the map. For older kernels, the call will fail.
+  res = syscall(__NR_bpf, BPF_MAP_GET_NEXT_KEY, &attr, sizeof(attr));
+  if (res < 0 && errno == EFAULT) {
+    // Fall back to try to find a non-existing key.
+    static unsigned char try_values[3] = {0, 0xff, 0x55};
+    attr.key = ptr_to_u64(key);
+    for (i = 0; i < 3; i++) {
+      memset(key, try_values[i], key_size);
+      // We want to check the existence of the key but we don't know the size
+      // of map's value. So we pass an invalid pointer for value, expect
+      // the call to fail and check if the error is ENOENT indicating the
+      // key doesn't exist. If we use NULL for the invalid pointer, it might
+      // trigger a page fault in kernel and affect performence. Hence we use
+      // ~0 which will fail and return fast.
+      // This should fail since we pass an invalid pointer for value.
+      if (bpf_lookup_elem(fd, key, ~0) >= 0)
+        return -1;
+      // This means the key doesn't exist.
+      if (errno == ENOENT)
+        return syscall(__NR_bpf, BPF_MAP_GET_NEXT_KEY, &attr, sizeof(attr));
+    }
+    return -1;
+  } else {
+    return res;
+  }
+}
+
 int bpf_get_next_key(int fd, void *key, void *next_key)
 {
   union bpf_attr attr;

src/cc/libbpf.h

@@ -34,6 +34,7 @@ int bpf_create_map(enum bpf_map_type map_type, int key_size, int value_size,
 int bpf_update_elem(int fd, void *key, void *value, unsigned long long flags);
 int bpf_lookup_elem(int fd, void *key, void *value);
 int bpf_delete_elem(int fd, void *key);
+int bpf_get_first_key(int fd, void *key, size_t key_size);
 int bpf_get_next_key(int fd, void *key, void *next_key);
 
 int bpf_prog_load(enum bpf_prog_type prog_type,

src/python/bcc/libbcc.py

@@ -69,6 +69,8 @@ lib.bpf_table_leaf_sscanf.argtypes = [ct.c_void_p, ct.c_ulonglong,
 # keep in sync with libbpf.h
 lib.bpf_get_next_key.restype = ct.c_int
 lib.bpf_get_next_key.argtypes = [ct.c_int, ct.c_void_p, ct.c_void_p]
+lib.bpf_get_first_key.restype = ct.c_int
+lib.bpf_get_first_key.argtypes = [ct.c_int, ct.c_void_p, ct.c_uint]
 lib.bpf_lookup_elem.restype = ct.c_int
 lib.bpf_lookup_elem.argtypes = [ct.c_int, ct.c_void_p, ct.c_void_p]
 lib.bpf_update_elem.restype = ct.c_int

src/python/bcc/table.py

@@ -245,25 +245,15 @@ class TableBase(MutableMapping):
             self[k] = self.Leaf()
 
     def __iter__(self):
-        return TableBase.Iter(self, self.Key)
+        return TableBase.Iter(self)
 
     def iter(self): return self.__iter__()
     def keys(self): return self.__iter__()
 
     class Iter(object):
-        def __init__(self, table, keytype):
-            self.Key = keytype
+        def __init__(self, table):
             self.table = table
-            k = self.Key()
-            kp = ct.pointer(k)
-            # if 0 is a valid key, try a few alternatives
-            if k in table:
-                ct.memset(kp, 0xff, ct.sizeof(k))
-                if k in table:
-                    ct.memset(kp, 0x55, ct.sizeof(k))
-                    if k in table:
-                        raise Exception("Unable to allocate iterator")
-            self.key = k
+            self.key = None
         def __iter__(self):
             return self
         def __next__(self):
@@ -275,10 +265,17 @@ class TableBase(MutableMapping):
     def next(self, key):
         next_key = self.Key()
         next_key_p = ct.pointer(next_key)
-        key_p = ct.pointer(key)
-        res = lib.bpf_get_next_key(self.map_fd,
-                ct.cast(key_p, ct.c_void_p),
-                ct.cast(next_key_p, ct.c_void_p))
+
+        if key is None:
+            res = lib.bpf_get_first_key(self.map_fd,
+                    ct.cast(next_key_p, ct.c_void_p),
+                    ct.sizeof(self.Key))
+        else:
+            key_p = ct.pointer(key)
+            res = lib.bpf_get_next_key(self.map_fd,
+                    ct.cast(key_p, ct.c_void_p),
+                    ct.cast(next_key_p, ct.c_void_p))
+
         if res < 0:
             raise StopIteration()
         return next_key

tests/cc/test_hash_table.cc

@@ -71,4 +71,17 @@ TEST_CASE("test hash table", "[hash_table]") {
     res = t.get_value(k, v2);
     REQUIRE(res.code() != 0);
   }
+
+  SECTION("walk table") {
+    for (int i = 1; i <= 10; i++) {
+      res = t.update_value(i * 3, i);
+      REQUIRE(res.code() == 0);
+    }
+    auto offline = t.get_table_offline();
+    REQUIRE(offline.size() == 10);
+    for (const auto &pair : offline) {
+      REQUIRE(pair.first % 3 == 0);
+      REQUIRE(pair.first / 3 == pair.second);
+    }
+  }
 }