ebpf_histogram_ebpf/
lib.rs

1#![no_std]
2use aya_ebpf::maps::PerCpuHashMap;
3
4#[inline(always)]
5fn bpf_log2(mut v: u32) -> u32 {
6    let mut r: u32;
7    let mut shift: u32;
8    r = ((v > 0xFFFF) as u32) << 4;
9    v >>= r;
10    shift = ((v > 0xFF) as u32) << 3;
11    v >>= shift;
12    r |= shift;
13    shift = ((v > 0xF) as u32) << 2;
14    v >>= shift;
15    r |= shift;
16    shift = ((v > 0x3) as u32) << 1;
17    v >>= shift;
18    r |= shift;
19    r |= v >> 1;
20    r
21}
22
23/// Return the log2(v) ceiled
24/// It should match the equivalent function in BCC
25fn bpf_log2l(v: u64) -> u32 {
26    let lo: u32 = (v & 0xFFFFFFFF) as u32;
27    let hi: u32 = (v >> 32) as u32;
28
29    if hi != 0 {
30        bpf_log2(hi) + 32 + 1
31    } else {
32        bpf_log2(lo) + 1
33    }
34}
35
36#[derive(Copy, Clone)]
37#[repr(C, packed)]
38pub struct Key<T: Sized> {
39    pub bucket: u32,
40    pub sub_key: T,
41}
42
43pub struct BpfHistogram<T> {
44    map: PerCpuHashMap<Key<T>, u64>,
45}
46
47impl<T> BpfHistogram<T> {
48    pub const fn with_max_entries(max_entries: u32, flags: u32) -> BpfHistogram<T> {
49        BpfHistogram {
50            map: PerCpuHashMap::with_max_entries(max_entries, flags),
51        }
52    }
53
54    #[inline(always)]
55    pub fn observe(&self, sub_key: T, value: u64) {
56        let bucket = bpf_log2l(value);
57        let key = Key { bucket, sub_key };
58
59        unsafe {
60            let counter = match self.map.get(&key) {
61                Some(i) => i + 1,
62                None => 1,
63            };
64            let _ = self.map.insert(&key, &counter, 0);
65        }
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use crate::bpf_log2l;
72
73    #[test]
74    fn bpf_log2_works() {
75        assert_eq!(bpf_log2l(10), 4);
76        assert_eq!(bpf_log2l(10000000), 24);
77        assert_eq!(bpf_log2l(1000000000), 30);
78        assert_eq!(bpf_log2l(100000000000), 37);
79        assert_eq!(bpf_log2l(10000000000000), 44);
80        assert_eq!(bpf_log2l(2043866223362000), 51);
81    }
82}
ebpf_histogram_ebpf/lib.rs

ebpf_histogram_ebpf/
lib.rs
