1#!/usr/bin/env python
2# Copyright (c) PLUMgrid, Inc.
3# Licensed under the Apache License, Version 2.0 (the "License")
4
5import os
6import unittest
7from bcc import BPF
8import multiprocessing
9
10class TestPercpu(unittest.TestCase):
11
12    def setUp(self):
13        try:
14            b = BPF(text='BPF_TABLE("percpu_array", u32, u32, stub, 1);')
15        except:
16            raise unittest.SkipTest("PerCpu unsupported on this kernel")
17
18    def test_helper(self):
19        test_prog1 = """
20        BPF_PERCPU_ARRAY(stub_default);
21        BPF_PERCPU_ARRAY(stub_type, u64);
22        BPF_PERCPU_ARRAY(stub_full, u64, 1024);
23        """
24        BPF(text=test_prog1)
25
26    def test_u64(self):
27        test_prog1 = """
28        BPF_TABLE("percpu_hash", u32, u64, stats, 1);
29        int hello_world(void *ctx) {
30            u32 key=0;
31            u64 value = 0, *val;
32            val = stats.lookup_or_init(&key, &value);
33            *val += 1;
34            return 0;
35        }
36        """
37        bpf_code = BPF(text=test_prog1)
38        stats_map = bpf_code.get_table("stats")
39        event_name = bpf_code.get_syscall_fnname("clone")
40        bpf_code.attach_kprobe(event=event_name, fn_name="hello_world")
41        ini = stats_map.Leaf()
42        for i in range(0, multiprocessing.cpu_count()):
43            ini[i] = 0
44        stats_map[ stats_map.Key(0) ] = ini
45        f = os.popen("hostname")
46        f.close()
47        self.assertEqual(len(stats_map),1)
48        val = stats_map[ stats_map.Key(0) ]
49        sum = stats_map.sum(stats_map.Key(0))
50        avg = stats_map.average(stats_map.Key(0))
51        max = stats_map.max(stats_map.Key(0))
52        self.assertGreater(sum.value, int(0))
53        self.assertGreater(max.value, int(0))
54        bpf_code.detach_kprobe(event_name)
55
56    def test_u32(self):
57        test_prog1 = """
58        BPF_TABLE("percpu_array", u32, u32, stats, 1);
59        int hello_world(void *ctx) {
60            u32 key=0;
61            u32 value = 0, *val;
62            val = stats.lookup_or_init(&key, &value);
63            *val += 1;
64            return 0;
65        }
66        """
67        bpf_code = BPF(text=test_prog1)
68        stats_map = bpf_code.get_table("stats")
69        event_name = bpf_code.get_syscall_fnname("clone")
70        bpf_code.attach_kprobe(event=event_name, fn_name="hello_world")
71        ini = stats_map.Leaf()
72        for i in range(0, multiprocessing.cpu_count()):
73            ini[i] = 0
74        stats_map[ stats_map.Key(0) ] = ini
75        f = os.popen("hostname")
76        f.close()
77        self.assertEqual(len(stats_map),1)
78        val = stats_map[ stats_map.Key(0) ]
79        sum = stats_map.sum(stats_map.Key(0))
80        avg = stats_map.average(stats_map.Key(0))
81        max = stats_map.max(stats_map.Key(0))
82        self.assertGreater(sum.value, int(0))
83        self.assertGreater(max.value, int(0))
84        bpf_code.detach_kprobe(event_name)
85
86    def test_struct_custom_func(self):
87        test_prog2 = """
88        typedef struct counter {
89        u32 c1;
90        u32 c2;
91        } counter;
92        BPF_TABLE("percpu_hash", u32, counter, stats, 1);
93        int hello_world(void *ctx) {
94            u32 key=0;
95            counter value = {0,0}, *val;
96            val = stats.lookup_or_init(&key, &value);
97            val->c1 += 1;
98            val->c2 += 1;
99            return 0;
100        }
101        """
102        bpf_code = BPF(text=test_prog2)
103        stats_map = bpf_code.get_table("stats",
104                reducer=lambda x,y: stats_map.sLeaf(x.c1+y.c1))
105        event_name = bpf_code.get_syscall_fnname("clone")
106        bpf_code.attach_kprobe(event=event_name, fn_name="hello_world")
107        ini = stats_map.Leaf()
108        for i in ini:
109            i = stats_map.sLeaf(0,0)
110        stats_map[ stats_map.Key(0) ] = ini
111        f = os.popen("hostname")
112        f.close()
113        self.assertEqual(len(stats_map),1)
114        k = stats_map[ stats_map.Key(0) ]
115        self.assertGreater(k.c1, int(0))
116        bpf_code.detach_kprobe(event_name)
117
118
119if __name__ == "__main__":
120    unittest.main()
121