aya_friday/maps/hash_map/
per_cpu_hash_map.rs

1//! Per-CPU hash map.
2use std::{
3    borrow::{Borrow, BorrowMut},
4    marker::PhantomData,
5    os::fd::AsFd as _,
6};
7
8use crate::{
9    Pod,
10    maps::{
11        IterableMap, MapData, MapError, MapIter, MapKeys, PerCpuValues, check_kv_size, hash_map,
12    },
13    sys::{SyscallError, bpf_map_lookup_elem_per_cpu, bpf_map_update_elem_per_cpu},
14};
15
16/// Similar to [`HashMap`](crate::maps::HashMap) but each CPU holds a separate value for a given key. Typically used to
17/// minimize lock contention in eBPF programs.
18///
19/// This type can be used with eBPF maps of type `BPF_MAP_TYPE_PERCPU_HASH` and
20/// `BPF_MAP_TYPE_LRU_PERCPU_HASH`.
21///
22/// # Minimum kernel version
23///
24/// The minimum kernel version required to use this feature is 4.6.
25///
26/// # Examples
27///
28/// ```no_run
29/// # let mut bpf = aya::Ebpf::load(&[])?;
30/// use aya::maps::PerCpuHashMap;
31///
32/// const CPU_IDS: u8 = 1;
33/// const WAKEUPS: u8 = 2;
34///
35/// let mut hm = PerCpuHashMap::<_, u8, u32>::try_from(bpf.map_mut("PER_CPU_STORAGE").unwrap())?;
36/// let cpu_ids = unsafe { hm.get(&CPU_IDS, 0)? };
37/// let wakeups = unsafe { hm.get(&WAKEUPS, 0)? };
38/// for (cpu_id, wakeups) in cpu_ids.iter().zip(wakeups.iter()) {
39///     println!("cpu {} woke up {} times", cpu_id, wakeups);
40/// }
41/// # Ok::<(), aya::EbpfError>(())
42/// ```
43#[doc(alias = "BPF_MAP_TYPE_LRU_PERCPU_HASH")]
44#[doc(alias = "BPF_MAP_TYPE_PERCPU_HASH")]
45pub struct PerCpuHashMap<T, K: Pod, V: Pod> {
46    pub(crate) inner: T,
47    _kv: PhantomData<(K, V)>,
48}
49
50impl<T: Borrow<MapData>, K: Pod, V: Pod> PerCpuHashMap<T, K, V> {
51    pub(crate) fn new(map: T) -> Result<Self, MapError> {
52        let data = map.borrow();
53        check_kv_size::<K, V>(data)?;
54
55        Ok(Self {
56            inner: map,
57            _kv: PhantomData,
58        })
59    }
60
61    /// Returns a slice of values - one for each CPU - associated with the key.
62    pub fn get(&self, key: &K, flags: u64) -> Result<PerCpuValues<V>, MapError> {
63        let fd = self.inner.borrow().fd().as_fd();
64        let values =
65            bpf_map_lookup_elem_per_cpu(fd, key, flags).map_err(|io_error| SyscallError {
66                call: "bpf_map_lookup_elem",
67                io_error,
68            })?;
69        values.ok_or(MapError::KeyNotFound)
70    }
71
72    /// An iterator visiting all key-value pairs in arbitrary order. The
73    /// iterator item type is `Result<(K, PerCpuValues<V>), MapError>`.
74    pub fn iter(&self) -> MapIter<'_, K, PerCpuValues<V>, Self> {
75        MapIter::new(self)
76    }
77
78    /// An iterator visiting all keys in arbitrary order. The iterator element
79    /// type is `Result<K, MapError>`.
80    pub fn keys(&self) -> MapKeys<'_, K> {
81        MapKeys::new(self.inner.borrow())
82    }
83}
84
85impl<'a, T: Borrow<MapData>, K: Pod, V: Pod> IntoIterator for &'a PerCpuHashMap<T, K, V> {
86    type Item = Result<(K, PerCpuValues<V>), MapError>;
87    type IntoIter = MapIter<'a, K, PerCpuValues<V>, PerCpuHashMap<T, K, V>>;
88
89    fn into_iter(self) -> Self::IntoIter {
90        self.iter()
91    }
92}
93
94impl<T: BorrowMut<MapData>, K: Pod, V: Pod> PerCpuHashMap<T, K, V> {
95    /// Inserts a slice of values - one for each CPU - for the given key.
96    ///
97    /// # Examples
98    ///
99    /// ```no_run
100    /// # #[derive(thiserror::Error, Debug)]
101    /// # enum Error {
102    /// #     #[error(transparent)]
103    /// #     IO(#[from] std::io::Error),
104    /// #     #[error(transparent)]
105    /// #     Map(#[from] aya::maps::MapError),
106    /// #     #[error(transparent)]
107    /// #     Ebpf(#[from] aya::EbpfError)
108    /// # }
109    /// # let mut bpf = aya::Ebpf::load(&[])?;
110    /// use aya::maps::{PerCpuHashMap, PerCpuValues};
111    /// use aya::util::nr_cpus;
112    ///
113    /// const RETRIES: u8 = 1;
114    ///
115    /// let nr_cpus = nr_cpus().map_err(|(_, error)| error)?;
116    /// let mut hm = PerCpuHashMap::<_, u8, u32>::try_from(bpf.map_mut("PER_CPU_STORAGE").unwrap())?;
117    /// hm.insert(
118    ///     RETRIES,
119    ///     PerCpuValues::try_from(vec![3u32; nr_cpus])?,
120    ///     0,
121    /// )?;
122    /// # Ok::<(), Error>(())
123    /// ```
124    pub fn insert(
125        &mut self,
126        key: impl Borrow<K>,
127        values: PerCpuValues<V>,
128        flags: u64,
129    ) -> Result<(), MapError> {
130        let fd = self.inner.borrow_mut().fd().as_fd();
131        bpf_map_update_elem_per_cpu(fd, key.borrow(), &values, flags)
132            .map_err(|io_error| SyscallError {
133                call: "bpf_map_update_elem",
134                io_error,
135            })
136            .map_err(Into::into)
137    }
138
139    /// Removes a key from the map.
140    pub fn remove(&mut self, key: &K) -> Result<(), MapError> {
141        hash_map::remove(self.inner.borrow_mut(), key)
142    }
143}
144
145impl<T: Borrow<MapData>, K: Pod, V: Pod> IterableMap<K, PerCpuValues<V>>
146    for PerCpuHashMap<T, K, V>
147{
148    fn map(&self) -> &MapData {
149        self.inner.borrow()
150    }
151
152    fn get(&self, key: &K) -> Result<PerCpuValues<V>, MapError> {
153        Self::get(self, key, 0)
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use std::io;
160
161    use assert_matches::assert_matches;
162    use aya_obj::generated::bpf_map_type;
163    use libc::ENOENT;
164
165    use super::*;
166    use crate::{
167        maps::{Map, test_utils},
168        sys::{SysResult, override_syscall},
169    };
170
171    fn sys_error(value: i32) -> SysResult {
172        Err((-1, io::Error::from_raw_os_error(value)))
173    }
174
175    #[test]
176    fn test_try_from_ok() {
177        let map = Map::PerCpuHashMap(test_utils::new_map(test_utils::new_obj_map::<u32>(
178            bpf_map_type::BPF_MAP_TYPE_PERCPU_HASH,
179        )));
180        let _unused: PerCpuHashMap<_, u32, u32> = map.try_into().unwrap();
181    }
182    #[test]
183    fn test_try_from_ok_lru() {
184        let map_data = || {
185            test_utils::new_map(test_utils::new_obj_map::<u32>(
186                bpf_map_type::BPF_MAP_TYPE_LRU_PERCPU_HASH,
187            ))
188        };
189        let map = Map::PerCpuHashMap(map_data());
190        let _unused: PerCpuHashMap<_, u32, u32> = map.try_into().unwrap();
191        let map = Map::PerCpuLruHashMap(map_data());
192        let _unused: PerCpuHashMap<_, u32, u32> = map.try_into().unwrap();
193    }
194    #[test]
195    fn test_get_not_found() {
196        let map_data = || {
197            test_utils::new_map(test_utils::new_obj_map::<u32>(
198                bpf_map_type::BPF_MAP_TYPE_LRU_PERCPU_HASH,
199            ))
200        };
201        let map = Map::PerCpuHashMap(map_data());
202        let map = PerCpuHashMap::<_, u32, u32>::try_from(&map).unwrap();
203
204        override_syscall(|_| sys_error(ENOENT));
205
206        assert_matches!(map.get(&1, 0), Err(MapError::KeyNotFound));
207    }
208}
aya_friday/maps/hash_map/per_cpu_hash_map.rs

aya_friday/maps/hash_map/
per_cpu_hash_map.rs
