xref: /aosp_15_r20/external/bcc/tests/python/test_histogram.py (revision 387f9dfdfa2baef462e92476d413c7bc2470293e) 
1#!/usr/bin/env python3
2# Copyright (c) PLUMgrid, Inc.
3# Licensed under the Apache License, Version 2.0 (the "License")
4
5from bcc import BPF
6from ctypes import c_int, c_ulonglong
7import random
8import time
9from unittest import main, TestCase
10
11class TestHistogram(TestCase):
12    def test_simple(self):
13        b = BPF(text=b"""
14#include <uapi/linux/ptrace.h>
15struct bpf_map;
16BPF_HISTOGRAM(hist1);
17BPF_HASH(stub);
18int kprobe__htab_map_delete_elem(struct pt_regs *ctx, struct bpf_map *map, u64 *k) {
19    hist1.increment(bpf_log2l(*k));
20    hist1.atomic_increment(bpf_log2l(*k));
21    return 0;
22}
23""")
24        for i in range(0, 32):
25            for j in range(0, random.randint(1, 10)):
26                try: del b[b"stub"][c_ulonglong(1 << i)]
27                except: pass
28        b[b"hist1"].print_log2_hist()
29
30        for i in range(32, 64):
31            for j in range(0, random.randint(1, 10)):
32                try: del b[b"stub"][c_ulonglong(1 << i)]
33                except: pass
34        b[b"hist1"].print_log2_hist()
35        b.cleanup()
36
37    def test_struct(self):
38        b = BPF(text=b"""
39#include <uapi/linux/ptrace.h>
40struct bpf_map;
41typedef struct { void *map; u64 slot; } Key;
42BPF_HISTOGRAM(hist1, Key, 1024);
43BPF_HASH(stub1);
44BPF_HASH(stub2);
45int kprobe__htab_map_delete_elem(struct pt_regs *ctx, struct bpf_map *map, u64 *k) {
46    hist1.increment((Key){map, bpf_log2l(*k)});
47    hist1.atomic_increment((Key){map, bpf_log2l(*k)});
48    return 0;
49}
50""")
51        for i in range(0, 64):
52            for j in range(0, random.randint(1, 10)):
53                try: del b[b"stub1"][c_ulonglong(1 << i)]
54                except: pass
55                try: del b[b"stub2"][c_ulonglong(1 << i)]
56                except: pass
57        b[b"hist1"].print_log2_hist()
58        b.cleanup()
59
60    def test_chars(self):
61        b = BPF(text=b"""
62#include <uapi/linux/ptrace.h>
63#include <linux/sched.h>
64#include <linux/version.h>
65typedef struct { char name[TASK_COMM_LEN]; u64 slot; } Key;
66BPF_HISTOGRAM(hist1, Key, 1024);
67int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) {
68#if LINUX_VERSION_CODE < KERNEL_VERSION(5,5,0)
69    Key k = {.slot = bpf_log2l(prev->real_start_time)};
70#else
71    Key k = {.slot = bpf_log2l(prev->start_boottime)};
72#endif
73    if (!bpf_get_current_comm(&k.name, sizeof(k.name))) {
74        hist1.increment(k);
75        hist1.atomic_increment(k);
76    }
77    return 0;
78}
79""")
80        for i in range(0, 100): time.sleep(0.01)
81        b[b"hist1"].print_log2_hist()
82        b.cleanup()
83
84    def test_multiple_key(self):
85        b = BPF(text=b"""
86#include <uapi/linux/ptrace.h>
87#include <uapi/linux/fs.h>
88struct hist_s_key {
89    u64 key_1;
90    u64 key_2;
91};
92struct hist_key {
93    struct hist_s_key s_key;
94    u64 slot;
95};
96BPF_HISTOGRAM(mk_hist, struct hist_key, 1024);
97int kprobe__vfs_read(struct pt_regs *ctx, struct file *file,
98        char __user *buf, size_t count) {
99    struct hist_key key = {.slot = bpf_log2l(count)};
100    key.s_key.key_1 = (unsigned long)buf & 0x70;
101    key.s_key.key_2 = (unsigned long)buf & 0x7;
102    mk_hist.increment(key);
103    return 0;
104}
105""")
106        def bucket_sort(buckets):
107            buckets.sort()
108            return buckets
109
110        for i in range(0, 100): time.sleep(0.01)
111        b[b"mk_hist"].print_log2_hist("size", "k_1 & k_2",
112                section_print_fn=lambda bucket: "%3d %d" % (bucket[0], bucket[1]),
113                bucket_fn=lambda bucket: (bucket.key_1, bucket.key_2),
114                strip_leading_zero=True,
115                bucket_sort_fn=bucket_sort)
116        b.cleanup()
117
118if __name__ == "__main__":
119    main()
120