Merge pull request iovisor#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 @@
 # 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 @@ def zero(self):
  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 @@ def next(self):
  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);
+ }
+ }
 }