1#!/usr/bin/env python3 2# Copyright (c) PLUMgrid, Inc. 3# Licensed under the Apache License, Version 2.0 (the "License") 4 5from ctypes import c_uint, c_ulong, Structure 6from bcc import BPF 7from time import sleep 8import sys 9from unittest import main, TestCase 10 11text = b""" 12#include <linux/ptrace.h> 13struct Ptr { u64 ptr; }; 14struct Counters { char unused; __int128 stat1; }; 15BPF_HASH(stats, struct Ptr, struct Counters, 1024); 16 17int count_sched(struct pt_regs *ctx) { 18 struct Ptr key = {.ptr=PT_REGS_PARM1(ctx)}; 19 struct Counters zleaf; 20 21 memset(&zleaf, 0, sizeof(zleaf)); 22 struct Counters *val = stats.lookup_or_try_init(&key, &zleaf); 23 if (val) { 24 val->stat1++; 25 } 26 return 0; 27} 28""" 29 30class TestTracingEvent(TestCase): 31 def setUp(self): 32 b = BPF(text=text, debug=0) 33 self.stats = b.get_table(b"stats") 34 b.attach_kprobe(event_re=b"^finish_task_switch$|^finish_task_switch\.isra\.\d$", 35 fn_name=b"count_sched") 36 37 def test_sched1(self): 38 for i in range(0, 100): 39 sleep(0.01) 40 for key, leaf in self.stats.items(): 41 print("ptr %x:" % key.ptr, "stat1 (%d %d)" % (leaf.stat1[1], leaf.stat1[0])) 42 43if __name__ == "__main__": 44 main() 45