added bpf_update_batch() API support for Python Maps

This commit aims at introducing items_update_batch, batch operation to update multiple key-value pairs at the same time. Doc has been updated accordingly, and a test is provided. Signed-off-by: Simone Magnani <[email protected]>
nanshuaibo · Apr 29, 2021 · 19df7ee · 19df7ee
1 parent 8034be6
commit 19df7ee
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 14 deletions.
diff --git a/docs/reference_guide.md b/docs/reference_guide.md
@@ -107,12 +107,13 @@ This guide is incomplete. If something feels missing, check the bcc and kernel s
  - [6. items_lookup_and_delete_batch()](#6-items_lookup_and_delete_batch)
  - [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)
+ - [9. items_update_batch()](#9-items_update_batch)
+ - [10. print_log2_hist()](#10-print_log2_hist)
+ - [11. print_linear_hist()](#11-print_linear_hist)
+ - [12. open_ring_buffer()](#12-open_ring_buffer)
+ - [13. push()](#13-push)
+ - [14. pop()](#14-pop)
+ - [15. peek()](#15-peek)
  - [Helpers](#helpers)
  - [1. ksym()](#1-ksym)
  - [2. ksymname()](#2-ksymname)
@@ -2005,7 +2006,19 @@ If a list of keys is given then only those keys and their associated values will
 It requires kernel v5.6.
 
 
-### 9. print_log2_hist()
+### 9. items_update_batch()
+
+Syntax: ```table.items_update_batch(keys, values)```
+
+Update all the provided keys with new values. The two arguments must be the same length and within the map limits (between 1 and the maximum entries).
+
+Arguments:
+
+- keys is the list of keys to be updated
+- values is the list containing the new values.
+
+
+### 10. print_log2_hist()
 
 Syntax: ```table.print_log2_hist(val_type="value", section_header="Bucket ptr", section_print_fn=None)```
 
@@ -2056,7 +2069,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)
 
-### 10. print_linear_hist()
+### 11. print_linear_hist()
 
 Syntax: ```table.print_linear_hist(val_type="value", section_header="Bucket ptr", section_print_fn=None)```
 
@@ -2115,7 +2128,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)
 
-### 11. open_ring_buffer()
+### 12. open_ring_buffer()
 
 Syntax: ```table.open_ring_buffer(callback, ctx=None)```
 
@@ -2177,7 +2190,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),
 
-### 12. push()
+### 13. push()
 
 Syntax: ```table.push(leaf, flags=0)```
 
@@ -2187,7 +2200,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),
 
-### 13. pop()
+### 14. pop()
 
 Syntax: ```leaf = table.pop()```
 
@@ -2198,7 +2211,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),
 
-### 14. peek()
+### 15. peek()
 
 Syntax: ```leaf = table.peek()```
 

diff --git a/src/cc/libbpf.c b/src/cc/libbpf.c
@@ -370,6 +370,11 @@ int bpf_delete_batch(int fd, void *keys, __u32 *count)
  return bpf_map_delete_batch(fd, keys, count, NULL);
 }
 
+int bpf_update_batch(int fd, void *keys, void *values, __u32 *count)
+{
+ return bpf_map_update_batch(fd, keys, values, count, NULL);
+}
+
 int bpf_lookup_and_delete_batch(int fd, __u32 *in_batch, __u32 *out_batch,
  void *keys, void *values, __u32 *count)
 {

diff --git a/src/python/bcc/libbcc.py b/src/python/bcc/libbcc.py
@@ -85,6 +85,9 @@
 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_update_batch.restype = ct.c_int
+lib.bpf_update_batch.argtypes = [ct.c_int, ct.c_void_p, ct.c_void_p,
+ ct.POINTER(ct.c_uint32)]
 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]

diff --git a/src/python/bcc/table.py b/src/python/bcc/table.py
@@ -460,6 +460,34 @@ def items_delete_batch(self, keys=None):
  for _ in self.items_lookup_and_delete_batch():
  return
 
+ def items_update_batch(self, keys, values):
+ """Update all the key-value pairs in the map provided.
+ The lists must be the same length, between 1 and the maximum number of entries.
+ """
+ # two ct.Array are expected
+ if not isinstance(keys, ct.Array) or not isinstance(values, ct.Array):
+ raise TypeError
+
+ batch_size = len(keys)
+
+ # check that batch between limits and coherent with the provided values
+ if batch_size < 1 or batch_size > self.max_entries or batch_size != len(values):
+ raise KeyError
+
+ count = ct.c_uint32(batch_size)
+ res = lib.bpf_update_batch(self.map_fd,
+ ct.byref(keys),
+ ct.byref(values),
+ ct.byref(count)
+ )
+
+ errcode = ct.get_errno()
+ if (errcode == errno.EINVAL):
+ raise Exception("BPF_MAP_UPDATE_BATCH is invalid.")
+
+ if (res != 0 and errcode != errno.ENOENT):
+ raise Exception("BPF_MAP_UPDATE_BATCH has failed")
+
  def items_lookup_and_delete_batch(self):
  # batch size is set to the maximum
  batch_size = self.max_entries

diff --git a/tests/python/test_map_batch_ops.py b/tests/python/test_map_batch_ops.py
@@ -6,11 +6,12 @@
 # Licensed under the Apache License, Version 2.0 (the "License")
 
 from __future__ import print_function
+from unittest import main, skipUnless, TestCase
 from bcc import BPF
+
+import os
 import distutils.version
-from unittest import main, skipUnless, TestCase
 import ctypes as ct
-import os
 
 
 def kernel_version_ge(major, minor):
@@ -94,6 +95,21 @@ def test_delete_batch_subset(self):
  count = sum(1 for _ in hmap.items())
  self.assertEqual(count, self.MAPSIZE - subset_size)
 
+ def test_update_batch(self):
+ hmap = self.fill_hashmap()
+
+ # preparing keys and new values arrays
+ keys = (hmap.Key * self.MAPSIZE)()
+ new_values = (hmap.Leaf * self.MAPSIZE)()
+ for i in range(self.MAPSIZE):
+ keys[i] = ct.c_int(i)
+ new_values[i] = ct.c_int(-1)
+ hmap.items_update_batch(keys, new_values)
+
+ # check the update has worked, i.e sum of values is -NUM_KEYS
+ count = sum(v.value for v in hmap.values())
+ self.assertEqual(count, -1*self.MAPSIZE)
+
 
 if __name__ == "__main__":
  main()