Merge pull request iovisor#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 @@ def __init__(self, *args, **kwargs):
  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 @@ def __setitem__(self, key, leaf):
  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 @@ def max(self, key):
 
  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 @@ def __init__(self, *args, **kwargs):
  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 @@ def __setitem__(self, key, leaf):
  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 @@ def max(self, key):
 
  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 @@
 
 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 @@ def test_u64(self):
  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 @@ def test_u32(self):
  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 @@ def test_struct_custom_func(self):
  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")