From 906af6fbf03d334c9d504fde6a01a3b769f38ffa Mon Sep 17 00:00:00 2001 From: Huapeng Zhou Date: Thu, 1 Dec 2016 20:01:46 -0800 Subject: [PATCH] Add lru_hash/lru_percpu_hash to python/lua --- src/lua/bcc/table.lua | 3 +++ src/lua/bpf/cdef.lua | 6 ++++- src/python/bcc/table.py | 14 +++++++++++ tests/python/test_lru.py | 59 +++++++++++++++++++++++++++++++++++++++++++++ tests/python/test_percpu.py | 9 +++---- 5 files changed, 84 insertions(+), 7 deletions(-) create mode 100644 tests/python/test_lru.py diff --git a/src/lua/bcc/table.lua b/src/lua/bcc/table.lua index 35d93ed..767402c 100644 --- a/src/lua/bcc/table.lua +++ b/src/lua/bcc/table.lua @@ -26,6 +26,9 @@ BaseTable.static.BPF_MAP_TYPE_PERF_EVENT_ARRAY = 4 BaseTable.static.BPF_MAP_TYPE_PERCPU_HASH = 5 BaseTable.static.BPF_MAP_TYPE_PERCPU_ARRAY = 6 BaseTable.static.BPF_MAP_TYPE_STACK_TRACE = 7 +BaseTable.static.BPF_MAP_TYPE_CGROUP_ARRAY = 8 +BaseTable.static.BPF_MAP_TYPE_LRU_HASH = 9 +BaseTable.static.BPF_MAP_TYPE_LRU_PERCPU_HASH = 10 function BaseTable:initialize(t_type, bpf, map_id, map_fd, key_type, leaf_type) assert(t_type == libbcc.bpf_table_type_id(bpf.module, map_id)) diff --git a/src/lua/bpf/cdef.lua b/src/lua/bpf/cdef.lua index 07749c1..d7be776 100644 --- a/src/lua/bpf/cdef.lua +++ b/src/lua/bpf/cdef.lua @@ -138,9 +138,13 @@ else S.c.BPF_MAP.PERCPU_ARRAY = 6 S.c.BPF_MAP.STACK_TRACE = 7 S.c.BPF_MAP.CGROUP_ARRAY = 8 + S.c.BPF_MAP.LRU_HASH = 9 + S.c.BPF_MAP.LRU_PERCPU_HASH = 10 end if not S.c.BPF_PROG.TRACEPOINT then S.c.BPF_PROG.TRACEPOINT = 5 + S.c.BPF_PROG.XDP = 6 + S.c.BPF_PROG.PERF_EVENT = 7 end end @@ -230,4 +234,4 @@ function M.tracepoint_type(tp) return string.format('struct { %s }', table.concat(fields)) end -return M \ No newline at end of file +return M diff --git a/src/python/bcc/table.py b/src/python/bcc/table.py index 002e9e3..49bf2a2 100644 --- a/src/python/bcc/table.py +++ b/src/python/bcc/table.py @@ -28,6 +28,9 @@ BPF_MAP_TYPE_PERF_EVENT_ARRAY = 4 BPF_MAP_TYPE_PERCPU_HASH = 5 BPF_MAP_TYPE_PERCPU_ARRAY = 6 BPF_MAP_TYPE_STACK_TRACE = 7 +BPF_MAP_TYPE_CGROUP_ARRAY = 8 +BPF_MAP_TYPE_LRU_HASH = 9 +BPF_MAP_TYPE_LRU_PERCPU_HASH = 10 stars_max = 40 @@ -97,6 +100,10 @@ def Table(bpf, map_id, map_fd, keytype, leaftype, **kwargs): 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) + elif ttype == BPF_MAP_TYPE_LRU_HASH: + t = LruHash(bpf, map_id, map_fd, keytype, leaftype) + elif ttype == BPF_MAP_TYPE_LRU_PERCPU_HASH: + t = LruPerCpuHash(bpf, map_id, map_fd, keytype, leaftype) if t == None: raise Exception("Unknown table type %d" % ttype) return t @@ -299,6 +306,9 @@ class HashTable(TableBase): if res < 0: raise KeyError +class LruHash(HashTable): + def __init__(self, *args, **kwargs): + super(LruHash, self).__init__(*args, **kwargs) class ArrayBase(TableBase): def __init__(self, *args, **kwargs): @@ -526,6 +536,10 @@ class PerCpuHash(HashTable): result.value/=self.total_cpu return result +class LruPerCpuHash(PerCpuHash): + def __init__(self, *args, **kwargs): + super(LruPerCpuHash, self).__init__(*args, **kwargs) + class PerCpuArray(ArrayBase): def __init__(self, *args, **kwargs): self.reducer = kwargs.pop("reducer", None) diff --git a/tests/python/test_lru.py b/tests/python/test_lru.py new file mode 100644 index 0000000..f900bcb --- /dev/null +++ b/tests/python/test_lru.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python +# Copyright (c) PLUMgrid, Inc. +# Licensed under the Apache License, Version 2.0 (the "License") + +import ctypes as ct +import os +import unittest +from bcc import BPF +import multiprocessing + +class TestLru(unittest.TestCase): + def test_lru_hash(self): + b = BPF(text="""BPF_TABLE("lru_hash", int, u64, lru, 1024);""") + t = b["lru"] + for i in range(1, 1032): + t[ct.c_int(i)] = ct.c_ulonglong(i) + for i, v in t.items(): + self.assertEqual(v.value, i.value) + # BPF_MAP_TYPE_LRU_HASH eviction happens in batch and we expect less + # items than specified size. + self.assertLess(len(t), 1024); + + def test_lru_percpu_hash(self): + test_prog1 = """ + BPF_TABLE("lru_percpu_hash", 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; + } + """ + b = BPF(text=test_prog1) + stats_map = b.get_table("stats") + b.attach_kprobe(event="sys_clone", fn_name="hello_world") + ini = stats_map.Leaf() + for i in range(0, multiprocessing.cpu_count()): + ini[i] = 0 + # First initialize with key 1 + stats_map[ stats_map.Key(1) ] = ini + # Then initialize with key 0 + stats_map[ stats_map.Key(0) ] = ini + # Key 1 should have been evicted + with self.assertRaises(KeyError): + val = stats_map[ stats_map.Key(1) ] + 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) + b.detach_kprobe("sys_clone") + +if __name__ == "__main__": + unittest.main() diff --git a/tests/python/test_percpu.py b/tests/python/test_percpu.py index 3f89a54..994ed76 100755 --- a/tests/python/test_percpu.py +++ b/tests/python/test_percpu.py @@ -20,7 +20,6 @@ class TestPercpu(unittest.TestCase): 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") @@ -37,6 +36,7 @@ class TestPercpu(unittest.TestCase): max = stats_map.max(stats_map.Key(0)) self.assertGreater(sum.value, 0L) self.assertGreater(max.value, 0L) + bpf_code.detach_kprobe("sys_clone") def test_u32(self): test_prog1 = """ @@ -49,7 +49,6 @@ class TestPercpu(unittest.TestCase): 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") @@ -66,6 +65,7 @@ class TestPercpu(unittest.TestCase): max = stats_map.max(stats_map.Key(0)) self.assertGreater(sum.value, 0L) self.assertGreater(max.value, 0L) + bpf_code.detach_kprobe("sys_clone") def test_struct_custom_func(self): test_prog2 = """ @@ -83,7 +83,6 @@ class TestPercpu(unittest.TestCase): 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)) @@ -97,9 +96,7 @@ class TestPercpu(unittest.TestCase): 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") + bpf_code.detach_kprobe("sys_clone") if __name__ == "__main__": -- 2.7.4