Merge pull request #458 from zaafar/percpu_support

Percpu support

src/cc/frontends/clang/b_frontend_action.cc

@@ -570,6 +570,10 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       map_type = BPF_MAP_TYPE_HASH;
     } else if (A->getName() == "maps/array") {
       map_type = BPF_MAP_TYPE_ARRAY;
+    } else if (A->getName() == "maps/percpu_hash") {
+      map_type = BPF_MAP_TYPE_PERCPU_HASH;
+    } else if (A->getName() == "maps/percpu_array") {
+      map_type = BPF_MAP_TYPE_PERCPU_ARRAY;
     } else if (A->getName() == "maps/histogram") {
       if (table.key_desc == "\"int\"")
         map_type = BPF_MAP_TYPE_ARRAY;

src/python/bcc/__init__.py

@@ -260,7 +260,7 @@ class BPF(object):
         cls = type(str(desc[0]), (base,), dict(_fields_=fields))
         return cls
 
-    def get_table(self, name, keytype=None, leaftype=None):
+    def get_table(self, name, keytype=None, leaftype=None, reducer=None):
         map_id = lib.bpf_table_id(self.module, name.encode("ascii"))
         map_fd = lib.bpf_table_fd(self.module, name.encode("ascii"))
         if map_fd < 0:
@@ -275,7 +275,7 @@ class BPF(object):
             if not leaf_desc:
                 raise Exception("Failed to load BPF Table %s leaf desc" % name)
             leaftype = BPF._decode_table_type(json.loads(leaf_desc.decode()))
-        return Table(self, map_id, map_fd, keytype, leaftype)
+        return Table(self, map_id, map_fd, keytype, leaftype, reducer=reducer)
 
     def __getitem__(self, key):
         if key not in self.tables:

src/python/bcc/table.py

@@ -17,6 +17,7 @@ import ctypes as ct
 import multiprocessing
 
 from .libbcc import lib, _RAW_CB_TYPE
+from subprocess import check_output
 
 BPF_MAP_TYPE_HASH = 1
 BPF_MAP_TYPE_ARRAY = 2
@@ -73,8 +74,8 @@ def _print_log2_hist(vals, val_type):
                       _stars(val, val_max, stars)))
 
 
-def Table(bpf, map_id, map_fd, keytype, leaftype):
-    """Table(bpf, map_id, map_fd, keytype, leaftype)
+def Table(bpf, map_id, map_fd, keytype, leaftype, **kwargs):
+    """Table(bpf, map_id, map_fd, keytype, leaftype, **kwargs)
 
     Create a python object out of a reference to a bpf table handle"""
 
@@ -89,9 +90,9 @@ def Table(bpf, map_id, map_fd, keytype, leaftype):
     elif ttype == BPF_MAP_TYPE_PERF_EVENT_ARRAY:
         t = PerfEventArray(bpf, map_id, map_fd, keytype, leaftype)
     elif ttype == BPF_MAP_TYPE_PERCPU_HASH:
-        t = PerCpuHashTable(bpf, map_id, map_fd, keytype, leaftype)
+        t = PerCpuHash(bpf, map_id, map_fd, keytype, leaftype, **kwargs)
     elif ttype == BPF_MAP_TYPE_PERCPU_ARRAY:
-        t = PerCpuArray(bpf, map_id, map_fd, keytype, leaftype)
+        t = PerCpuArray(bpf, map_id, map_fd, keytype, leaftype, **kwargs)
     elif ttype == BPF_MAP_TYPE_STACK_TRACE:
         t = StackTrace(bpf, map_id, map_fd, keytype, leaftype)
     if t == None:
@@ -403,13 +404,111 @@ class PerfEventArray(ArrayBase):
             del(self.bpf.open_kprobes()[(id(self), key)])
         del self._cbs[key]
 
-class PerCpuHashTable(TableBase):
+class PerCpuHash(HashTable):
     def __init__(self, *args, **kwargs):
-        raise Exception("Unsupported")
+        self.reducer = kwargs.pop("reducer", None)
+        super(PerCpuHash, self).__init__(*args, **kwargs)
+        self.sLeaf = self.Leaf
+        self.total_cpu = multiprocessing.cpu_count()
+        # This needs to be 8 as hard coded into the linux kernel.
+        self.alignment = ct.sizeof(self.sLeaf) % 8
+        if self.alignment is 0:
+            self.Leaf = self.sLeaf * self.total_cpu
+        else:
+            # Currently Float, Char, un-aligned structs are not supported
+            if self.sLeaf == ct.c_uint:
+                self.Leaf = ct.c_uint64 * self.total_cpu
+            elif self.sLeaf == ct.c_int:
+                self.Leaf = ct.c_int64 * self.total_cpu
+            else:
+                raise IndexError("Leaf must be aligned to 8 bytes")
+
+    def getvalue(self, key):
+        result = super(PerCpuHash, self).__getitem__(key)
+        if self.alignment is 0:
+            ret = result
+        else:
+            ret = (self.sLeaf * self.total_cpu)()
+            for i in range(0, self.total_cpu):
+                ret[i] = result[i]
+        return ret
+
+    def __getitem__(self, key):
+        if self.reducer:
+            return reduce(self.reducer, self.getvalue(key))
+        else:
+            return self.getvalue(key)
+
+    def __setitem__(self, key, leaf):
+        super(PerCpuHash, self).__setitem__(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)))
+
+    def max(self, key):
+        if isinstance(self.Leaf(), ct.Structure):
+            raise IndexError("Leaf must be an integer type for default max functions")
+        return self.sLeaf(max(self.getvalue(key)))
+
+    def average(self, key):
+        result = self.sum(key)
+        result.value/=self.total_cpu
+        return result
 
 class PerCpuArray(ArrayBase):
     def __init__(self, *args, **kwargs):
-        raise Exception("Unsupported")
+        self.reducer = kwargs.pop("reducer", None)
+        super(PerCpuArray, self).__init__(*args, **kwargs)
+        self.sLeaf = self.Leaf
+        self.total_cpu = multiprocessing.cpu_count()
+        # This needs to be 8 as hard coded into the linux kernel.
+        self.alignment = ct.sizeof(self.sLeaf) % 8
+        if self.alignment is 0:
+            self.Leaf = self.sLeaf * self.total_cpu
+        else:
+            # Currently Float, Char, un-aligned structs are not supported
+            if self.sLeaf == ct.c_uint:
+                self.Leaf = ct.c_uint64 * self.total_cpu
+            elif self.sLeaf == ct.c_int:
+                self.Leaf = ct.c_int64 * self.total_cpu
+            else:
+                raise IndexError("Leaf must be aligned to 8 bytes")
+
+    def getvalue(self, key):
+        result = super(PerCpuArray, self).__getitem__(key)
+        if self.alignment is 0:
+            ret = result
+        else:
+            ret = (self.sLeaf * self.total_cpu)()
+            for i in range(0, self.total_cpu):
+                ret[i] = result[i]
+        return ret
+
+    def __getitem__(self, key):
+        if (self.reducer):
+            return reduce(self.reducer, self.getvalue(key))
+        else:
+            return self.getvalue(key)
+
+    def __setitem__(self, key, leaf):
+        super(PerCpuArray, self).__setitem__(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)))
+
+    def max(self, key):
+        if isinstance(self.Leaf(), ct.Structure):
+            raise IndexError("Leaf must be an integer type for default max functions")
+        return self.sLeaf(max(self.getvalue(key)))
+
+    def average(self, key):
+        result = self.sum(key)
+        result.value/=self.total_cpu
+        return result
 
 class StackTrace(TableBase):
     MAX_DEPTH = 127

tests/python/test_percpu.py

+#!/usr/bin/env python
+# Copyright (c) PLUMgrid, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+import os
+import unittest
+from bcc import BPF
+import multiprocessing
+
+class TestPercpu(unittest.TestCase):
+
+    def test_u64(self):
+        test_prog1 = """
+        BPF_TABLE("percpu_hash", u32, u64, stats, 1);
+        int hello_world(void *ctx) {
+            u32 key=0;
+            u64 value = 0, *val;
+            val = stats.lookup_or_init(&key, &value);
+            *val += 1;
+            return 0;
+        }
+        """
+        self.addCleanup(self.cleanup)
+        bpf_code = BPF(text=test_prog1)
+        stats_map = bpf_code.get_table("stats")
+        bpf_code.attach_kprobe(event="sys_clone", fn_name="hello_world")
+        ini = stats_map.Leaf()
+        for i in range(0, multiprocessing.cpu_count()):
+            ini[i] = 0
+        stats_map[ stats_map.Key(0) ] = ini
+        f = os.popen("hostname")
+        f.close()
+        self.assertEqual(len(stats_map),1)
+        val = stats_map[ stats_map.Key(0) ]
+        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)
+
+    def test_u32(self):
+        test_prog1 = """
+        BPF_TABLE("percpu_array", u32, u32, stats, 1);
+        int hello_world(void *ctx) {
+            u32 key=0;
+            u32 value = 0, *val;
+            val = stats.lookup_or_init(&key, &value);
+            *val += 1;
+            return 0;
+        }
+        """
+        self.addCleanup(self.cleanup)
+        bpf_code = BPF(text=test_prog1)
+        stats_map = bpf_code.get_table("stats")
+        bpf_code.attach_kprobe(event="sys_clone", fn_name="hello_world")
+        ini = stats_map.Leaf()
+        for i in range(0, multiprocessing.cpu_count()):
+            ini[i] = 0
+        stats_map[ stats_map.Key(0) ] = ini
+        f = os.popen("hostname")
+        f.close()
+        self.assertEqual(len(stats_map),1)
+        val = stats_map[ stats_map.Key(0) ]
+        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)
+
+    def test_struct_custom_func(self):
+        test_prog2 = """
+        typedef struct counter {
+        u32 c1;
+        u32 c2;
+        } counter;
+        BPF_TABLE("percpu_hash", u32, counter, stats, 1);
+        int hello_world(void *ctx) {
+            u32 key=0;
+            counter value = {0,0}, *val;
+            val = stats.lookup_or_init(&key, &value);
+            val->c1 += 1;
+            val->c2 += 1;
+            return 0;
+        }
+        """
+        self.addCleanup(self.cleanup)
+        bpf_code = BPF(text=test_prog2)
+        stats_map = bpf_code.get_table("stats",
+                reducer=lambda x,y: stats_map.sLeaf(x.c1+y.c1))
+        bpf_code.attach_kprobe(event="sys_clone", fn_name="hello_world")
+        ini = stats_map.Leaf()
+        for i in ini:
+            i = stats_map.sLeaf(0,0)
+        stats_map[ stats_map.Key(0) ] = ini
+        f = os.popen("hostname")
+        f.close()
+        self.assertEqual(len(stats_map),1)
+        k = stats_map[ stats_map.Key(0) ]
+        self.assertGreater(k.c1, 0L)
+
+    def cleanup(self):
+        BPF.detach_kprobe("sys_clone")
+
+
+if __name__ == "__main__":
+    unittest.main()