add bpf_map_lookup_batch and bpf_map_delete_batch in bcc

docs/reference_guide.md

@@ -105,12 +105,14 @@ This guide is incomplete. If something feels missing, check the bcc and kernel s
         - [4. values()](#4-values)
         - [5. clear()](#5-clear)
         - [6. items_lookup_and_delete_batch()](#6-items_lookup_and_delete_batch)
-        - [7. print_log2_hist()](#7-print_log2_hist)
-        - [8. print_linear_hist()](#8-print_linear_hist)
-        - [9. open_ring_buffer()](#9-open_ring_buffer)
-        - [10. push()](#10-push)
-        - [11. pop()](#11-pop)
-        - [12. peek()](#12-peek)
+        - [7. items_lookup_batch()](#7-items_lookup_batch)
+        - [8. items_delete_batch()](#8-items_delete_batch)
+        - [9. print_log2_hist()](#9-print_log2_hist)
+        - [10. print_linear_hist()](#10-print_linear_hist)
+        - [11. open_ring_buffer()](#11-open_ring_buffer)
+        - [12. push()](#12-push)
+        - [13. pop()](#13-pop)
+        - [14. peek()](#14-peek)
     - [Helpers](#helpers)
         - [1. ksym()](#1-ksym)
         - [2. ksymname()](#2-ksymname)
@@ -1961,7 +1963,7 @@ Examples in situ:
 Syntax: ```table.items_lookup_and_delete_batch()```
 
 Returns an array of the keys in a table with a single call to BPF syscall. This can be used with BPF_HASH maps to fetch, and iterate, over the keys. It also clears the table: deletes all entries.
-You should rather use table.items_lookup_and_delete_batch() than table.items() followed by table.clear().
+You should rather use table.items_lookup_and_delete_batch() than table.items() followed by table.clear(). It requires kernel v5.6.
 
 Example:
 
@@ -1974,7 +1976,36 @@ while True:
     sleep(1)
 ```
 
-### 7. print_log2_hist()
+### 7. items_lookup_batch()
+
+Syntax: ```table.items_lookup_batch()```
+
+Returns an array of the keys in a table with a single call to BPF syscall. This can be used with BPF_HASH maps to fetch, and iterate, over the keys.
+You should rather use table.items_lookup_batch() than table.items(). It requires kernel v5.6.
+
+Example:
+
+```Python
+# print current value of map:
+print("%9s-%9s-%8s-%9s" % ("PID", "COMM", "fname", "counter"))
+while True:
+    for k, v in sorted(b['map'].items_lookup_batch(), key=lambda kv: (kv[0]).pid):
+        print("%9s-%9s-%8s-%9d" % (k.pid, k.comm, k.fname, v.counter))
+```
+
+### 8. items_delete_batch()
+
+Syntax: ```table.items_delete_batch(keys)```
+
+Arguments:
+
+- keys is optional and by default is None.
+In that case, it clears all entries of a BPF_HASH map when keys is None. It is more efficient than table.clear() since it generates only one system call.
+If a list of keys is given then only those keys and their associated values will be deleted.
+It requires kernel v5.6.
+
+
+### 9. print_log2_hist()
 
 Syntax: ```table.print_log2_hist(val_type="value", section_header="Bucket ptr", section_print_fn=None)```
 
@@ -2025,7 +2056,7 @@ Examples in situ:
 [search /examples](https://github.com/iovisor/bcc/search?q=print_log2_hist+path%3Aexamples+language%3Apython&type=Code),
 [search /tools](https://github.com/iovisor/bcc/search?q=print_log2_hist+path%3Atools+language%3Apython&type=Code)
 
-### 8. print_linear_hist()
+### 10. print_linear_hist()
 
 Syntax: ```table.print_linear_hist(val_type="value", section_header="Bucket ptr", section_print_fn=None)```
 
@@ -2084,7 +2115,7 @@ Examples in situ:
 [search /examples](https://github.com/iovisor/bcc/search?q=print_linear_hist+path%3Aexamples+language%3Apython&type=Code),
 [search /tools](https://github.com/iovisor/bcc/search?q=print_linear_hist+path%3Atools+language%3Apython&type=Code)
 
-### 9. open_ring_buffer()
+### 11. open_ring_buffer()
 
 Syntax: ```table.open_ring_buffer(callback, ctx=None)```
 
@@ -2146,7 +2177,7 @@ def print_event(ctx, data, size):
 Examples in situ:
 [search /examples](https://github.com/iovisor/bcc/search?q=open_ring_buffer+path%3Aexamples+language%3Apython&type=Code),
 
-### 10. push()
+### 12. push()
 
 Syntax: ```table.push(leaf, flags=0)```
 
@@ -2156,7 +2187,7 @@ Passing QueueStack.BPF_EXIST as a flag causes the Queue or Stack to discard the
 Examples in situ:
 [search /tests](https://github.com/iovisor/bcc/search?q=push+path%3Atests+language%3Apython&type=Code),
 
-### 11. pop()
+### 13. pop()
 
 Syntax: ```leaf = table.pop()```
 
@@ -2167,7 +2198,7 @@ Raises a KeyError exception if the operation does not succeed.
 Examples in situ:
 [search /tests](https://github.com/iovisor/bcc/search?q=pop+path%3Atests+language%3Apython&type=Code),
 
-### 12. peek()
+### 14. peek()
 
 Syntax: ```leaf = table.peek()```
 

src/cc/libbpf.c

@@ -358,8 +358,20 @@ int bpf_lookup_and_delete(int fd, void *key, void *value)
   return bpf_map_lookup_and_delete_elem(fd, key, value);
 }
 
-int bpf_lookup_and_delete_batch(int fd, __u32 *in_batch, __u32 *out_batch, void *keys,
-                                void *values, __u32 *count)
+int bpf_lookup_batch(int fd, __u32 *in_batch, __u32 *out_batch, void *keys,
+                     void *values, __u32 *count)
+{
+  return bpf_map_lookup_batch(fd, in_batch, out_batch, keys, values, count,
+                              NULL);
+}
+
+int bpf_delete_batch(int fd,  void *keys, __u32 *count)
+{
+  return bpf_map_delete_batch(fd, keys, count, NULL);
+}
+
+int bpf_lookup_and_delete_batch(int fd, __u32 *in_batch, __u32 *out_batch,
+                                void *keys, void *values, __u32 *count)
 {
   return bpf_map_lookup_and_delete_batch(fd, in_batch, out_batch, keys, values,
                                          count, NULL);

src/python/bcc/libbcc.py

@@ -83,9 +83,14 @@
         ct.c_ulonglong]
 lib.bpf_delete_elem.restype = ct.c_int
 lib.bpf_delete_elem.argtypes = [ct.c_int, ct.c_void_p]
+lib.bpf_delete_batch.restype = ct.c_int
+lib.bpf_delete_batch.argtypes = [ct.c_int, ct.c_void_p, ct.c_void_p]
+lib.bpf_lookup_batch.restype = ct.c_int
+lib.bpf_lookup_batch.argtypes = [ct.c_int, ct.POINTER(ct.c_uint32),
+        ct.POINTER(ct.c_uint32), ct.c_void_p, ct.c_void_p, ct.c_void_p]
 lib.bpf_lookup_and_delete_batch.restype = ct.c_int
 lib.bpf_lookup_and_delete_batch.argtypes = [ct.c_int, ct.POINTER(ct.c_uint32),
-        ct.POINTER(ct.c_uint32),ct.c_void_p, ct.c_void_p, ct.c_void_p]
+        ct.POINTER(ct.c_uint32), ct.c_void_p, ct.c_void_p, ct.c_void_p]
 lib.bpf_open_raw_sock.restype = ct.c_int
 lib.bpf_open_raw_sock.argtypes = [ct.c_char_p]
 lib.bpf_attach_socket.restype = ct.c_int

src/python/bcc/table.py

@@ -397,6 +397,67 @@ def clear(self):
         for k in self.keys():
             self.__delitem__(k)
 
+    def items_lookup_batch(self):
+        # batch size is set to the maximum
+        batch_size = self.max_entries
+        out_batch = ct.c_uint32(0)
+        keys = (type(self.Key()) * batch_size)()
+        values = (type(self.Leaf()) * batch_size)()
+        count = ct.c_uint32(batch_size)
+        res = lib.bpf_lookup_batch(self.map_fd,
+                                   None,
+                                   ct.byref(out_batch),
+                                   ct.byref(keys),
+                                   ct.byref(values),
+                                   ct.byref(count)
+                                   )
+
+        errcode = ct.get_errno()
+        if (errcode == errno.EINVAL):
+            raise Exception("BPF_MAP_LOOKUP_BATCH is invalid.")
+
+        if (res != 0 and errcode != errno.ENOENT):
+            raise Exception("BPF_MAP_LOOKUP_BATCH has failed")
+
+        for i in range(0, count.value):
+            yield (keys[i], values[i])
+
+    def items_delete_batch(self, keys=None):
+        """Delete all the key-value pairs in the map if no key are given.
+        In that case, it is faster to call lib.bpf_lookup_and_delete_batch than
+        create keys list and then call lib.bpf_delete_batch on these keys.
+        If a list of keys is given then it deletes the related key-value.
+        """
+        if keys is not None:
+            # a list is expected
+            if type(keys) != list:
+                raise TypeError
+
+            batch_size = len(keys)
+            if batch_size < 1 and batch_size > self.max_entries:
+                raise KeyError
+
+            count = ct.c_uint32(batch_size)
+
+            # build the array aka list of key_t that will be deleted
+            keylist = (type(self.Key()) * batch_size)()
+            for i, k in enumerate(keys):
+                keylist[i] = k
+
+            res = lib.bpf_delete_batch(self.map_fd,
+                                       ct.byref(keylist),
+                                       ct.byref(count)
+                                       )
+            errcode = ct.get_errno()
+            if (errcode == errno.EINVAL):
+                raise Exception("BPF_MAP_DELETE_BATCH is invalid.")
+
+            if (res != 0 and errcode != errno.ENOENT):
+                raise Exception("BPF_MAP_DELETE_BATCH has failed")
+        else:
+            for _ in self.items_lookup_and_delete_batch():
+                return
+
     def items_lookup_and_delete_batch(self):
         # batch size is set to the maximum
         batch_size = self.max_entries

tests/python/test_map_batch_ops.py

@@ -27,22 +27,73 @@ def kernel_version_ge(major, minor):
 
 @skipUnless(kernel_version_ge(5, 6), "requires kernel >= 5.6")
 class TestMapBatch(TestCase):
-    def test_lookup_and_delete_batch(self):
-        b = BPF(text="""BPF_HASH(map, int, int, 1024);""")
-        hmap = b["map"]
-        for i in range(0, 1024):
+    MAPSIZE = 1024
+
+    def fill_hashmap(self):
+        b = BPF(text=b"""BPF_HASH(map, int, int, %d);""" % self.MAPSIZE)
+        hmap = b[b"map"]
+        for i in range(0, self.MAPSIZE):
             hmap[ct.c_int(i)] = ct.c_int(i)
+        return hmap
 
-        # check the lookup
+    def check_hashmap_values(self, it):
         i = 0
-        for k, v in sorted(hmap.items_lookup_and_delete_batch()):
+        for k, v in sorted(it):
             self.assertEqual(k, i)
             self.assertEqual(v, i)
             i += 1
-        # and check the delete has workd, i.e map is empty
-        count = sum(1 for _ in hmap.items_lookup_and_delete_batch())
+        return i
+
+    def test_lookup_and_delete_batch(self):
+        # fill the hashmap
+        hmap = self.fill_hashmap()
+
+        # check values and count them
+        count = self.check_hashmap_values(hmap.items_lookup_and_delete_batch())
+        self.assertEqual(count, self.MAPSIZE)
+
+        # and check the delete has worked, i.e map is now empty
+        count = sum(1 for _ in hmap.items_lookup_batch())
+        self.assertEqual(count, 0)
+
+    def test_lookup_batch(self):
+        # fill the hashmap
+        hmap = self.fill_hashmap()
+
+        # check values and count them
+        count = self.check_hashmap_values(hmap.items_lookup_batch())
+        self.assertEqual(count, self.MAPSIZE)
+
+    def test_delete_batch_all_keysp(self):
+        # Delete all key/value in the map
+        # fill the hashmap
+        hmap = self.fill_hashmap()
+        hmap.items_delete_batch()
+
+        # check the delete has worked, i.e map is now empty
+        count = sum(1 for _ in hmap.items())
         self.assertEqual(count, 0)
 
+    def test_delete_batch_subset(self):
+        # Delete only a subset of key/value in the map
+        # fill the hashmap
+        hmap = self.fill_hashmap()
+        # Get 4 keys in this map.
+        subset_size = 32
+        keys = [None] * subset_size
+        i = 0
+        for k, v in hmap.items_lookup_batch():
+            if i < subset_size:
+                keys[i] = k
+                i += 1
+            else:
+                break
+
+        hmap.items_delete_batch(keys)
+        # check the delete has worked, i.e map is now empty
+        count = sum(1 for _ in hmap.items())
+        self.assertEqual(count, self.MAPSIZE - subset_size)
+
 
 if __name__ == "__main__":
     main()