Use bpf_get_first_key in Python API

diff --git a/src/python/bcc/libbcc.py b/src/python/bcc/libbcc.py
index f2168905bc86cd5df48588ea99bcd2fd489e911c..0bba2c6587944bde491e4d9af1635ba835cab7a0 100644 (file)
--- a/src/python/bcc/libbcc.py
+++ b/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

diff --git a/src/python/bcc/table.py b/src/python/bcc/table.py
index c03a674fbff68746370537dadfd0b138ccbe5fdc..99c4fbf9b88fd821e980b961d8a4c40e1cefe431 100644 (file)
--- a/src/python/bcc/table.py
+++ b/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