Merge pull request #261 from iovisor/bblanco_dev

Fixes for table indexing and clear()

src/python/bcc/__init__.py

@@ -223,14 +223,9 @@ class BPF(object):
                     raise KeyError
 
         def clear(self):
-            if self.ttype in (BPF.ARRAY, BPF.PROG_ARRAY):
-                # Special case clear, since this class is currently behaving
-                # like a dict but popitem on an array causes an infinite loop.
-                # TODO: derive Table from array.array instead
-                for k in self.keys():
-                    self.__delitem__(k)
-            else:
-                super(BPF.Table, self).clear()
+            # default clear uses popitem, which can race with the bpf prog
+            for k in self.keys():
+                self.__delitem__(k)
 
         @staticmethod
         def _stars(val, val_max, width):
@@ -319,7 +314,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):

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)

tests/cc/test_array.py

+#!/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()

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()