per_cpu_hash_map.rs 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. //! Per-CPU hash map.
  2. use std::{
  3. borrow::{Borrow, BorrowMut},
  4. marker::PhantomData,
  5. };
  6. use crate::{
  7. maps::{
  8. check_kv_size, hash_map, IterableMap, MapData, MapError, MapIter, MapKeys, PerCpuValues,
  9. },
  10. sys::{bpf_map_lookup_elem_per_cpu, bpf_map_update_elem_per_cpu},
  11. Pod,
  12. };
  13. /// Similar to [`HashMap`](crate::maps::HashMap) but each CPU holds a separate value for a given key. Tipically used to
  14. /// minimize lock contention in eBPF programs.
  15. ///
  16. /// This type can be used with eBPF maps of type `BPF_MAP_TYPE_PERCPU_HASH` and
  17. /// `BPF_MAP_TYPE_LRU_PERCPU_HASH`.
  18. ///
  19. /// # Minimum kernel version
  20. ///
  21. /// The minimum kernel version required to use this feature is 4.6.
  22. ///
  23. /// # Examples
  24. ///
  25. /// ```no_run
  26. /// # let mut bpf = aya::Bpf::load(&[])?;
  27. /// use aya::maps::PerCpuHashMap;
  28. ///
  29. /// const CPU_IDS: u8 = 1;
  30. /// const WAKEUPS: u8 = 2;
  31. ///
  32. /// let mut hm = PerCpuHashMap::<_, u8, u32>::try_from(bpf.map_mut("PER_CPU_STORAGE").unwrap())?;
  33. /// let cpu_ids = unsafe { hm.get(&CPU_IDS, 0)? };
  34. /// let wakeups = unsafe { hm.get(&WAKEUPS, 0)? };
  35. /// for (cpu_id, wakeups) in cpu_ids.iter().zip(wakeups.iter()) {
  36. /// println!("cpu {} woke up {} times", cpu_id, wakeups);
  37. /// }
  38. /// # Ok::<(), aya::BpfError>(())
  39. /// ```
  40. #[doc(alias = "BPF_MAP_TYPE_LRU_PERCPU_HASH")]
  41. #[doc(alias = "BPF_MAP_TYPE_PERCPU_HASH")]
  42. pub struct PerCpuHashMap<T, K: Pod, V: Pod> {
  43. inner: T,
  44. _k: PhantomData<K>,
  45. _v: PhantomData<V>,
  46. }
  47. impl<T: Borrow<MapData>, K: Pod, V: Pod> PerCpuHashMap<T, K, V> {
  48. pub(crate) fn new(map: T) -> Result<PerCpuHashMap<T, K, V>, MapError> {
  49. let data = map.borrow();
  50. check_kv_size::<K, V>(data)?;
  51. let _ = data.fd_or_err()?;
  52. Ok(PerCpuHashMap {
  53. inner: map,
  54. _k: PhantomData,
  55. _v: PhantomData,
  56. })
  57. }
  58. /// Returns a slice of values - one for each CPU - associated with the key.
  59. pub fn get(&self, key: &K, flags: u64) -> Result<PerCpuValues<V>, MapError> {
  60. let fd = self.inner.borrow().fd_or_err()?;
  61. let values = bpf_map_lookup_elem_per_cpu(fd, key, flags).map_err(|(_, io_error)| {
  62. MapError::SyscallError {
  63. call: "bpf_map_lookup_elem".to_owned(),
  64. io_error,
  65. }
  66. })?;
  67. values.ok_or(MapError::KeyNotFound)
  68. }
  69. /// An iterator visiting all key-value pairs in arbitrary order. The
  70. /// iterator item type is `Result<(K, PerCpuValues<V>), MapError>`.
  71. pub fn iter(&self) -> MapIter<'_, K, PerCpuValues<V>, Self> {
  72. MapIter::new(self)
  73. }
  74. /// An iterator visiting all keys in arbitrary order. The iterator element
  75. /// type is `Result<K, MapError>`.
  76. pub fn keys(&self) -> MapKeys<'_, K> {
  77. MapKeys::new(self.inner.borrow())
  78. }
  79. }
  80. impl<T: BorrowMut<MapData>, K: Pod, V: Pod> PerCpuHashMap<T, K, V> {
  81. /// Inserts a slice of values - one for each CPU - for the given key.
  82. ///
  83. /// # Examples
  84. ///
  85. /// ```no_run
  86. /// # #[derive(thiserror::Error, Debug)]
  87. /// # enum Error {
  88. /// # #[error(transparent)]
  89. /// # IO(#[from] std::io::Error),
  90. /// # #[error(transparent)]
  91. /// # Map(#[from] aya::maps::MapError),
  92. /// # #[error(transparent)]
  93. /// # Bpf(#[from] aya::BpfError)
  94. /// # }
  95. /// # let mut bpf = aya::Bpf::load(&[])?;
  96. /// use aya::maps::{PerCpuHashMap, PerCpuValues};
  97. /// use aya::util::nr_cpus;
  98. ///
  99. /// const RETRIES: u8 = 1;
  100. ///
  101. /// let mut hm = PerCpuHashMap::<_, u8, u32>::try_from(bpf.map_mut("PER_CPU_STORAGE").unwrap())?;
  102. /// hm.insert(
  103. /// RETRIES,
  104. /// PerCpuValues::try_from(vec![3u32; nr_cpus()?])?,
  105. /// 0,
  106. /// )?;
  107. /// # Ok::<(), Error>(())
  108. /// ```
  109. pub fn insert(
  110. &mut self,
  111. key: impl Borrow<K>,
  112. values: PerCpuValues<V>,
  113. flags: u64,
  114. ) -> Result<(), MapError> {
  115. let fd = self.inner.borrow_mut().fd_or_err()?;
  116. bpf_map_update_elem_per_cpu(fd, key.borrow(), &values, flags).map_err(
  117. |(_, io_error)| MapError::SyscallError {
  118. call: "bpf_map_update_elem".to_owned(),
  119. io_error,
  120. },
  121. )?;
  122. Ok(())
  123. }
  124. /// Removes a key from the map.
  125. pub fn remove(&mut self, key: &K) -> Result<(), MapError> {
  126. hash_map::remove(self.inner.borrow_mut(), key)
  127. }
  128. }
  129. impl<T: Borrow<MapData>, K: Pod, V: Pod> IterableMap<K, PerCpuValues<V>>
  130. for PerCpuHashMap<T, K, V>
  131. {
  132. fn map(&self) -> &MapData {
  133. self.inner.borrow()
  134. }
  135. fn get(&self, key: &K) -> Result<PerCpuValues<V>, MapError> {
  136. PerCpuHashMap::get(self, key, 0)
  137. }
  138. }