Merge pull request #913 from iovisor/python23_percpu

Fix python2/3 incompatible percpu helpers

src/python/bcc/table.py

@@ -14,12 +14,14 @@
 
 from collections import MutableMapping
 import ctypes as ct
+from functools import reduce
 import multiprocessing
 import os
 
 from .libbcc import lib, _RAW_CB_TYPE
 from .perf import Perf
 from .utils import get_online_cpus
+from .utils import get_possible_cpus
 from subprocess import check_output
 
 BPF_MAP_TYPE_HASH = 1
@@ -560,7 +562,7 @@ class PerCpuHash(HashTable):
         self.reducer = kwargs.pop("reducer", None)
         super(PerCpuHash, self).__init__(*args, **kwargs)
         self.sLeaf = self.Leaf
-        self.total_cpu = multiprocessing.cpu_count()
+        self.total_cpu = len(get_possible_cpus())
         # This needs to be 8 as hard coded into the linux kernel.
         self.alignment = ct.sizeof(self.sLeaf) % 8
         if self.alignment is 0:
@@ -596,7 +598,7 @@ class PerCpuHash(HashTable):
     def sum(self, key):
         if isinstance(self.Leaf(), ct.Structure):
             raise IndexError("Leaf must be an integer type for default sum functions")
-        return self.sLeaf(reduce(lambda x,y: x+y, self.getvalue(key)))
+        return self.sLeaf(sum(self.getvalue(key)))
 
     def max(self, key):
         if isinstance(self.Leaf(), ct.Structure):
@@ -605,8 +607,7 @@ class PerCpuHash(HashTable):
 
     def average(self, key):
         result = self.sum(key)
-        result.value/=self.total_cpu
-        return result
+        return result.value / self.total_cpu
 
 class LruPerCpuHash(PerCpuHash):
     def __init__(self, *args, **kwargs):
@@ -617,7 +618,7 @@ class PerCpuArray(ArrayBase):
         self.reducer = kwargs.pop("reducer", None)
         super(PerCpuArray, self).__init__(*args, **kwargs)
         self.sLeaf = self.Leaf
-        self.total_cpu = multiprocessing.cpu_count()
+        self.total_cpu = len(get_possible_cpus())
         # This needs to be 8 as hard coded into the linux kernel.
         self.alignment = ct.sizeof(self.sLeaf) % 8
         if self.alignment is 0:
@@ -653,7 +654,7 @@ class PerCpuArray(ArrayBase):
     def sum(self, key):
         if isinstance(self.Leaf(), ct.Structure):
             raise IndexError("Leaf must be an integer type for default sum functions")
-        return self.sLeaf(reduce(lambda x,y: x+y, self.getvalue(key)))
+        return self.sLeaf(sum(self.getvalue(key)))
 
     def max(self, key):
         if isinstance(self.Leaf(), ct.Structure):
@@ -662,8 +663,7 @@ class PerCpuArray(ArrayBase):
 
     def average(self, key):
         result = self.sum(key)
-        result.value/=self.total_cpu
-        return result
+        return result.value / self.total_cpu
 
 class StackTrace(TableBase):
     MAX_DEPTH = 127

tests/python/CMakeLists.txt

@@ -58,6 +58,8 @@ add_test(NAME py_test_perf_event WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test_perf_event sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_perf_event.py)
 add_test(NAME py_test_utils WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test_utils sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_utils.py)
+add_test(NAME py_test_percpu WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test_percpu sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_percpu.py)
 
 add_test(NAME py_test_dump_func WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_dump_func simple ${CMAKE_CURRENT_SOURCE_DIR}/test_dump_func.py)

tests/python/test_percpu.py

@@ -9,6 +9,12 @@ import multiprocessing
 
 class TestPercpu(unittest.TestCase):
 
+    def setUp(self):
+        try:
+            b = BPF(text='BPF_TABLE("percpu_array", u32, u32, stub, 1);')
+        except:
+            raise unittest.SkipTest("PerCpu unsupported on this kernel")
+
     def test_u64(self):
         test_prog1 = """
         BPF_TABLE("percpu_hash", u32, u64, stats, 1);
@@ -34,8 +40,8 @@ class TestPercpu(unittest.TestCase):
         sum = stats_map.sum(stats_map.Key(0))
         avg = stats_map.average(stats_map.Key(0))
         max = stats_map.max(stats_map.Key(0))
-        self.assertGreater(sum.value, 0L)
-        self.assertGreater(max.value, 0L)
+        self.assertGreater(sum.value, int(0))
+        self.assertGreater(max.value, int(0))
         bpf_code.detach_kprobe("sys_clone")
 
     def test_u32(self):
@@ -63,8 +69,8 @@ class TestPercpu(unittest.TestCase):
         sum = stats_map.sum(stats_map.Key(0))
         avg = stats_map.average(stats_map.Key(0))
         max = stats_map.max(stats_map.Key(0))
-        self.assertGreater(sum.value, 0L)
-        self.assertGreater(max.value, 0L)
+        self.assertGreater(sum.value, int(0))
+        self.assertGreater(max.value, int(0))
         bpf_code.detach_kprobe("sys_clone")
 
     def test_struct_custom_func(self):
@@ -95,7 +101,7 @@ class TestPercpu(unittest.TestCase):
         f.close()
         self.assertEqual(len(stats_map),1)
         k = stats_map[ stats_map.Key(0) ]
-        self.assertGreater(k.c1, 0L)
+        self.assertGreater(k.c1, int(0))
         bpf_code.detach_kprobe("sys_clone")