1#!/usr/bin/env python
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="""
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    return 0;
21}
22""")
23        for i in range(0, 32):
24            for j in range(0, random.randint(1, 10)):
25                try: del b["stub"][c_ulonglong(1 << i)]
26                except: pass
27        b["hist1"].print_log2_hist()
28
29        for i in range(32, 64):
30            for j in range(0, random.randint(1, 10)):
31                try: del b["stub"][c_ulonglong(1 << i)]
32                except: pass
33        b["hist1"].print_log2_hist()
34        b.cleanup()
35
36    def test_struct(self):
37        b = BPF(text="""
38#include <uapi/linux/ptrace.h>
39struct bpf_map;
40typedef struct { void *map; u64 slot; } Key;
41BPF_HISTOGRAM(hist1, Key, 1024);
42BPF_HASH(stub1);
43BPF_HASH(stub2);
44int kprobe__htab_map_delete_elem(struct pt_regs *ctx, struct bpf_map *map, u64 *k) {
45    hist1.increment((Key){map, bpf_log2l(*k)});
46    return 0;
47}
48""")
49        for i in range(0, 64):
50            for j in range(0, random.randint(1, 10)):
51                try: del b["stub1"][c_ulonglong(1 << i)]
52                except: pass
53                try: del b["stub2"][c_ulonglong(1 << i)]
54                except: pass
55        b["hist1"].print_log2_hist()
56        b.cleanup()
57
58    def test_chars(self):
59        b = BPF(text="""
60#include <uapi/linux/ptrace.h>
61#include <linux/sched.h>
62typedef struct { char name[TASK_COMM_LEN]; u64 slot; } Key;
63BPF_HISTOGRAM(hist1, Key, 1024);
64int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) {
65    Key k = {.slot = bpf_log2l(prev->real_start_time)};
66    if (!bpf_get_current_comm(&k.name, sizeof(k.name)))
67        hist1.increment(k);
68    return 0;
69}
70""")
71        for i in range(0, 100): time.sleep(0.01)
72        b["hist1"].print_log2_hist()
73        b.cleanup()
74
75    def test_multiple_key(self):
76        b = BPF(text="""
77#include <uapi/linux/ptrace.h>
78#include <uapi/linux/fs.h>
79struct hist_s_key {
80    u64 key_1;
81    u64 key_2;
82};
83struct hist_key {
84    struct hist_s_key s_key;
85    u64 slot;
86};
87BPF_HISTOGRAM(mk_hist, struct hist_key, 1024);
88int kprobe__vfs_read(struct pt_regs *ctx, struct file *file,
89        char __user *buf, size_t count) {
90    struct hist_key key = {.slot = bpf_log2l(count)};
91    key.s_key.key_1 = (unsigned long)buf & 0x70;
92    key.s_key.key_2 = (unsigned long)buf & 0x7;
93    mk_hist.increment(key);
94    return 0;
95}
96""")
97        def bucket_sort(buckets):
98            buckets.sort()
99            return buckets
100
101        for i in range(0, 100): time.sleep(0.01)
102        b["mk_hist"].print_log2_hist("size", "k_1 & k_2",
103                section_print_fn=lambda bucket: "%3d %d" % (bucket[0], bucket[1]),
104                bucket_fn=lambda bucket: (bucket.key_1, bucket.key_2),
105                strip_leading_zero=True,
106                bucket_sort_fn=bucket_sort)
107        b.cleanup()
108
109if __name__ == "__main__":
110    main()
111