aya_metrics_ebpf/
lib.rs

1// #![no_std] except for tests and user space!
2#![cfg_attr(not(test), no_std)]
3#![deny(missing_docs)]
4#![deny(clippy::unwrap_used)]
5
6//! Provides counter functionality with testable no_std implementations for use in BPF.
7
8#[cfg(any(test, target_arch = "bpf"))]
9use aya_metrics_common::{Counter, Meter, BPF_COUNTERS_MAX_ENTRIES};
10
11// Module with implementations depending on the `aya-bpf` module.
12// The `bpf` module only compiles with the `bpf` feature enabled.
13#[cfg(target_arch = "bpf")]
14mod bpf {
15    use super::*;
16    use aya_ebpf::macros::map;
17    use aya_ebpf::maps::PerCpuArray;
18
19    // A BPF map to store counter metrics
20    #[map(name = "COUNTERS")]
21    pub static mut COUNTERS: PerCpuArray<u64> =
22        PerCpuArray::<u64>::with_max_entries(BPF_COUNTERS_MAX_ENTRIES as u32, 0);
23}
24
25// Include everything from the `bpf` module.
26// The `bpf` module itself only exists as a namespace which is hidden behind the `bpf` feature.
27#[cfg(target_arch = "bpf")]
28use bpf::*;
29
30/// Increments a counter.
31///
32/// Counters represent a single monotonic value, which means the value can only be incremented, not decremented, and
33/// always starts out with an initial value of zero.
34///
35/// # Arguments
36///
37/// * `counter` - An identifier for a counter metric. It is used as an index into the underlying BPF map.
38/// * `value`   - The amount by which the counter should be incremented.
39///
40#[cfg(any(test, target_arch = "bpf"))]
41#[inline(always)]
42pub fn counter<T: Counter>(counter: T, value: u64) {
43    // SAFETY: Instances of PerCpuArray are thread local in eBPF. We can therefore be sure that concurrent
44    // accesses will not happen on other threads and, within this function, counter is the sole reference to COUNTERS.
45    // It is not leaked from this function, so concurrent &mut references cannot be introduced by calling this function multiple times.
46    if let Some(counter) = unsafe { COUNTERS.get_ptr_mut(Meter::index(&counter)) } {
47        unsafe { *counter += value };
48    }
49}
50
51// Module containing mocks for the `bpf` module.
52// The `bpf` module only compiles with the `bpf` feature enabled. It contains dependencies from `aya-bpf` which either
53// do not compile or work correctly from user space. Defining mocks allows testing implementations that use `aya-bpf`.
54// Ideally these would live in the `elastic-flow-collector-ebpf` crate if it was possible to add tests there.
55#[cfg(test)]
56mod bpf_mocks {
57    use std::cell::Cell;
58
59    use super::BPF_COUNTERS_MAX_ENTRIES;
60
61    pub struct PerCpuArray<T, const N: usize> {
62        pub data: Cell<[T; N]>,
63        _t: core::marker::PhantomData<T>,
64    }
65
66    impl<const N: usize> PerCpuArray<u64, N> {
67        pub const fn new() -> PerCpuArray<u64, N> {
68            PerCpuArray {
69                data: Cell::new([0u64; N]),
70                _t: core::marker::PhantomData,
71            }
72        }
73
74        pub fn get_ptr_mut(&mut self, index: u32) -> Option<*mut u64> {
75            let data = self.data.get_mut();
76            let ptr = data as *mut u64;
77            let ptr_at = unsafe { ptr.offset(index.try_into().unwrap()) };
78            Some(ptr_at)
79        }
80    }
81
82    pub static mut COUNTERS: PerCpuArray<u64, BPF_COUNTERS_MAX_ENTRIES> =
83        PerCpuArray::<u64, BPF_COUNTERS_MAX_ENTRIES>::new();
84}
85
86// Include everything from the `bpf_mocks` module for tests.
87#[cfg(test)]
88use bpf_mocks::*;
89
90#[cfg(test)]
91mod test {
92    use super::*;
93
94    #[derive(Copy, Clone, Debug)]
95    enum MockCounter {
96        Test1,
97        Test2,
98    }
99
100    impl Counter for MockCounter {
101        fn name(self) -> String {
102            match self {
103                MockCounter::Test1 => "test1".to_string(),
104                MockCounter::Test2 => "test2".to_string(),
105            }
106        }
107
108        fn index(&self) -> u32 {
109            match self {
110                MockCounter::Test1 => 0,
111                MockCounter::Test2 => BPF_COUNTERS_MAX_ENTRIES as u32 - 1,
112            }
113        }
114    }
115
116    #[test]
117    fn test_counter() {
118        let mut expected = [0u64; BPF_COUNTERS_MAX_ENTRIES];
119
120        let actual = unsafe { COUNTERS.data.get() };
121        assert_eq!(actual, expected);
122
123        // test adding some numbers
124        counter(MockCounter::Test1, 1);
125        counter(MockCounter::Test2, 42);
126        let actual = unsafe { COUNTERS.data.get() };
127        *expected.first_mut().unwrap() = 1;
128        *expected.last_mut().unwrap() = 42;
129        assert_eq!(actual, expected);
130
131        // test adding zero
132        counter(MockCounter::Test1, 0);
133        counter(MockCounter::Test2, 0);
134        let actual = unsafe { COUNTERS.data.get() };
135        assert_eq!(actual, expected);
136
137        // test adding again increments existing values
138        counter(MockCounter::Test1, 1);
139        counter(MockCounter::Test2, 1);
140        let actual = unsafe { COUNTERS.data.get() };
141        *expected.first_mut().unwrap() = 2;
142        *expected.last_mut().unwrap() = 43;
143        assert_eq!(actual, expected);
144    }
145}
aya_metrics_ebpf/lib.rs

aya_metrics_ebpf/
lib.rs
