Support iterating from a 0-filled table key

In the case that 0-filled keys are valid in the table, the previous
algorithm did not properly iterate.

The API of the bpf_get_next_key routine is such that the iteration
over a map should start with an invalid key. When a 0 key is valid, this
causes iteration to start anywhere inside the hash table, skipping some
entries. So, add logic to the Iter object to test if the init key is
invalid. If otherwise, try a few alternatives until an invalid key is
found. If none found, raise an exception.

Also adds a test for indexing arrays from 0, which nows works with this
too.

Fixes: #260
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

diff --git a/src/python/bcc/__init__.py b/src/python/bcc/__init__.py
index 62eb104..1197ce8 100644 (file)
--- a/src/python/bcc/__init__.py
+++ b/src/python/bcc/__init__.py
@@ -319,7 +319,16 @@ class BPF(object):
             def __init__(self, table, keytype):
                 self.Key = keytype
                 self.table = table
-                self.key = self.Key()
+                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
             def __iter__(self):
                 return self
             def __next__(self):

diff --git a/tests/cc/CMakeLists.txt b/tests/cc/CMakeLists.txt
index a5533a8..8fd5715 100644 (file)
--- a/tests/cc/CMakeLists.txt
+++ b/tests/cc/CMakeLists.txt
@@ -44,3 +44,5 @@ add_test(NAME py_test_histogram WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_histogram sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_histogram.py)
 add_test(NAME py_test_callchain WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_callchain sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_callchain.py)
+add_test(NAME py_array WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_array sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_array.py)

diff --git a/tests/cc/test_array.py b/tests/cc/test_array.py
new file mode 100755 (executable)
index 0000000..f83203f
--- /dev/null
+++ b/tests/cc/test_array.py
@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+# Copyright (c) PLUMgrid, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+from bcc import BPF
+from ctypes import c_int, c_ulonglong
+import random
+import time
+from unittest import main, TestCase
+
+class TestArray(TestCase):
+    def test_simple(self):
+        b = BPF(text="""BPF_TABLE("array", int, u64, table1, 128);""")
+        t1 = b["table1"]
+        t1[c_int(0)] = c_ulonglong(100)
+        t1[c_int(127)] = c_ulonglong(1000)
+        for i, v in t1.items():
+            if i.value == 0:
+                self.assertEqual(v.value, 100)
+            if i.value == 127:
+                self.assertEqual(v.value, 1000)
+        self.assertEqual(len(t1), 128)
+
+if __name__ == "__main__":
+    main()

diff --git a/tests/cc/test_stat1.py b/tests/cc/test_stat1.py
index 98ff756..23b3a29 100755 (executable)
--- a/tests/cc/test_stat1.py
+++ b/tests/cc/test_stat1.py
@@ -56,5 +56,24 @@ class TestBPFSocket(TestCase):
         self.stats.clear()
         self.assertEqual(len(self.stats), 0)
 
+    def test_empty_key(self):
+        # test with a 0 key
+        self.stats.clear()
+        self.stats[self.stats.Key()] = self.stats.Leaf(100, 200)
+        x = self.stats.popitem()
+        self.stats[self.stats.Key(10, 20)] = self.stats.Leaf(300, 400)
+        with self.assertRaises(KeyError):
+            x = self.stats[self.stats.Key()]
+        (_, x) = self.stats.popitem()
+        self.assertEqual(x.rx_pkts, 300)
+        self.assertEqual(x.tx_pkts, 400)
+        self.stats.clear()
+        self.assertEqual(len(self.stats), 0)
+        self.stats[self.stats.Key()] = x
+        self.stats[self.stats.Key(0, 1)] = x
+        self.stats[self.stats.Key(0, 2)] = x
+        self.stats[self.stats.Key(0, 3)] = x
+        self.assertEqual(len(self.stats), 4)
+
 if __name__ == "__main__":
     main()