1 #!/usr/bin/env python 2 # Copyright (c) PLUMgrid, Inc. 3 # Licensed under the Apache License, Version 2.0 (the "License") 4 5 import os 6 import unittest 7 from bcc import BPF 8 import multiprocessing 9 10 class 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 119 if __name__ == "__main__": 120 unittest.main() 121
