mirror of https://github.com/aya-rs/aya
aya-ebpf: Add BTF map definitions for hash maps
Introduce BTF map definitions for hash maps, with custom structs indicating the metadata of the map, which end up in the `.maps` section.reviewable/pr1367/r2
Michal R
1 month ago
parent
70c77c5ea7
commit
131c0b3ca3
@ -0,0 +1,255 @@
|
|||||||
|
use core::{borrow::Borrow, cell::UnsafeCell};
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
bindings::bpf_map_type::{
|
||||||
|
BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_LRU_HASH, BPF_MAP_TYPE_LRU_PERCPU_HASH,
|
||||||
|
BPF_MAP_TYPE_PERCPU_HASH,
|
||||||
|
},
|
||||||
|
btf_map_def,
|
||||||
|
cty::{c_long, c_void},
|
||||||
|
insert, lookup, remove,
|
||||||
|
};
|
||||||
|
|
||||||
|
btf_map_def!(HashMapDef, BPF_MAP_TYPE_HASH);
|
||||||
|
|
||||||
|
#[repr(transparent)]
|
||||||
|
pub struct HashMap<K, V, const M: usize, const F: usize = 0>(UnsafeCell<HashMapDef<K, V, M, F>>);
|
||||||
|
|
||||||
|
unsafe impl<K: Sync, V: Sync, const M: usize, const F: usize> Sync for HashMap<K, V, M, F> {}
|
||||||
|
|
||||||
|
impl<K, V, const M: usize, const F: usize> HashMap<K, V, M, F> {
|
||||||
|
#[expect(
|
||||||
|
clippy::new_without_default,
|
||||||
|
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||||
|
)]
|
||||||
|
pub const fn new() -> Self {
|
||||||
|
Self(UnsafeCell::new(HashMapDef::new()))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the value associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub unsafe fn get(&self, key: impl Borrow<K>) -> Option<&V> {
|
||||||
|
unsafe { get(self.0.get().cast(), key.borrow()) }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr(&self, key: impl Borrow<K>) -> Option<*const V> {
|
||||||
|
get_ptr(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the mutable pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr_mut(&self, key: impl Borrow<K>) -> Option<*mut V> {
|
||||||
|
get_ptr_mut(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inserts a key-value pair into the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn insert(
|
||||||
|
&self,
|
||||||
|
key: impl Borrow<K>,
|
||||||
|
value: impl Borrow<V>,
|
||||||
|
flags: u64,
|
||||||
|
) -> Result<(), c_long> {
|
||||||
|
insert(self.0.get().cast(), key.borrow(), value.borrow(), flags)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Removes a key from the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn remove(&self, key: impl Borrow<K>) -> Result<(), c_long> {
|
||||||
|
remove(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
btf_map_def!(LruHashMapDef, BPF_MAP_TYPE_LRU_HASH);
|
||||||
|
|
||||||
|
#[repr(transparent)]
|
||||||
|
pub struct LruHashMap<K, V, const M: usize, const F: usize = 0>(
|
||||||
|
UnsafeCell<LruHashMapDef<K, V, M, F>>,
|
||||||
|
);
|
||||||
|
|
||||||
|
unsafe impl<K: Sync, V: Sync, const M: usize, const F: usize> Sync for LruHashMap<K, V, M, F> {}
|
||||||
|
|
||||||
|
impl<K, V, const M: usize, const F: usize> LruHashMap<K, V, M, F> {
|
||||||
|
#[expect(
|
||||||
|
clippy::new_without_default,
|
||||||
|
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||||
|
)]
|
||||||
|
pub const fn new() -> Self {
|
||||||
|
Self(UnsafeCell::new(LruHashMapDef::new()))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the value associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub unsafe fn get(&self, key: impl Borrow<K>) -> Option<&V> {
|
||||||
|
unsafe { get(self.0.get().cast(), key.borrow()) }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr(&self, key: impl Borrow<K>) -> Option<*const V> {
|
||||||
|
get_ptr(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the mutable pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr_mut(&self, key: impl Borrow<K>) -> Option<*mut V> {
|
||||||
|
get_ptr_mut(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inserts a key-value pair into the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn insert(
|
||||||
|
&self,
|
||||||
|
key: impl Borrow<K>,
|
||||||
|
value: impl Borrow<V>,
|
||||||
|
flags: u64,
|
||||||
|
) -> Result<(), c_long> {
|
||||||
|
insert(self.0.get().cast(), key.borrow(), value.borrow(), flags)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Removes a key from the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn remove(&self, key: impl Borrow<K>) -> Result<(), c_long> {
|
||||||
|
remove(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
btf_map_def!(PerCpuHashMapDef, BPF_MAP_TYPE_PERCPU_HASH);
|
||||||
|
|
||||||
|
#[repr(transparent)]
|
||||||
|
pub struct PerCpuHashMap<K, V, const M: usize, const F: usize>(
|
||||||
|
UnsafeCell<PerCpuHashMapDef<K, V, M, F>>,
|
||||||
|
);
|
||||||
|
|
||||||
|
unsafe impl<K: Sync, V: Sync, const M: usize, const F: usize> Sync for PerCpuHashMap<K, V, M, F> {}
|
||||||
|
|
||||||
|
impl<K, V, const M: usize, const F: usize> PerCpuHashMap<K, V, M, F> {
|
||||||
|
#[expect(
|
||||||
|
clippy::new_without_default,
|
||||||
|
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||||
|
)]
|
||||||
|
pub const fn new() -> Self {
|
||||||
|
Self(UnsafeCell::new(PerCpuHashMapDef::new()))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the value associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub unsafe fn get(&self, key: impl Borrow<K>) -> Option<&V> {
|
||||||
|
unsafe { get(self.0.get().cast(), key.borrow()) }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr(&self, key: impl Borrow<K>) -> Option<*const V> {
|
||||||
|
get_ptr(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the mutable pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr_mut(&self, key: impl Borrow<K>) -> Option<*mut V> {
|
||||||
|
get_ptr_mut(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inserts a key-value pair into the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn insert(
|
||||||
|
&self,
|
||||||
|
key: impl Borrow<K>,
|
||||||
|
value: impl Borrow<V>,
|
||||||
|
flags: u64,
|
||||||
|
) -> Result<(), c_long> {
|
||||||
|
insert(self.0.get().cast(), key.borrow(), value.borrow(), flags)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Removes a key from the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn remove(&self, key: impl Borrow<K>) -> Result<(), c_long> {
|
||||||
|
remove(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
btf_map_def!(LruPerCpuHashMapDef, BPF_MAP_TYPE_LRU_PERCPU_HASH);
|
||||||
|
|
||||||
|
#[repr(transparent)]
|
||||||
|
pub struct LruPerCpuHashMap<K, V, const M: usize, const F: usize = 0>(
|
||||||
|
UnsafeCell<LruPerCpuHashMapDef<K, V, M, F>>,
|
||||||
|
);
|
||||||
|
|
||||||
|
unsafe impl<K: Sync, V: Sync, const M: usize, const F: usize> Sync
|
||||||
|
for LruPerCpuHashMap<K, V, M, F>
|
||||||
|
{
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<K, V, const M: usize, const F: usize> LruPerCpuHashMap<K, V, M, F> {
|
||||||
|
#[expect(
|
||||||
|
clippy::new_without_default,
|
||||||
|
reason = "BPF maps are always used as static variables, therefore this method has to be `const`. `Default::default` is not `const`."
|
||||||
|
)]
|
||||||
|
pub const fn new() -> Self {
|
||||||
|
Self(UnsafeCell::new(LruPerCpuHashMapDef::new()))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the value associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub unsafe fn get(&self, key: impl Borrow<K>) -> Option<&V> {
|
||||||
|
unsafe { get(self.0.get().cast(), key.borrow()) }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr(&self, key: impl Borrow<K>) -> Option<*const V> {
|
||||||
|
get_ptr(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[doc = "Retrieves the mutable pointer associated with `key` from the map."]
|
||||||
|
#[doc = include_str!("../maps/map_safety.md")]
|
||||||
|
#[inline]
|
||||||
|
pub fn get_ptr_mut(&self, key: impl Borrow<K>) -> Option<*mut V> {
|
||||||
|
get_ptr_mut(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inserts a key-value pair into the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn insert(
|
||||||
|
&self,
|
||||||
|
key: impl Borrow<K>,
|
||||||
|
value: impl Borrow<V>,
|
||||||
|
flags: u64,
|
||||||
|
) -> Result<(), c_long> {
|
||||||
|
insert(self.0.get().cast(), key.borrow(), value.borrow(), flags)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Removes a key from the map.
|
||||||
|
#[inline]
|
||||||
|
pub fn remove(&self, key: impl Borrow<K>) -> Result<(), c_long> {
|
||||||
|
remove(self.0.get().cast(), key.borrow())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
unsafe fn get<'a, K, V>(def: *mut c_void, key: &K) -> Option<&'a V> {
|
||||||
|
get_ptr(def, key).map(|p| unsafe { &*p })
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn get_ptr_mut<K, V>(def: *mut c_void, key: &K) -> Option<*mut V> {
|
||||||
|
lookup(def, key).map(|p| p.as_ptr())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
fn get_ptr<K, V>(def: *mut c_void, key: &K) -> Option<*const V> {
|
||||||
|
lookup::<_, V>(def.cast(), key).map(|p| p.as_ptr().cast_const())
|
||||||
|
}
|
||||||
@ -0,0 +1,135 @@
|
|||||||
|
#![no_std]
|
||||||
|
#![no_main]
|
||||||
|
#![expect(unused_crate_dependencies, reason = "used in other bins")]
|
||||||
|
|
||||||
|
#[cfg(not(test))]
|
||||||
|
extern crate ebpf_panic;
|
||||||
|
|
||||||
|
use aya_ebpf::{
|
||||||
|
btf_maps::{Array, HashMap, LruHashMap, LruPerCpuHashMap, PerCpuHashMap},
|
||||||
|
cty::c_long,
|
||||||
|
macros::{btf_map, map, uprobe},
|
||||||
|
maps::{
|
||||||
|
Array as LegacyArray, HashMap as LegacyHashMap, LruHashMap as LegacyLruHashMap,
|
||||||
|
LruPerCpuHashMap as LegacyLruPerCpuHashMap, PerCpuHashMap as LegacyPerCpuHashMap,
|
||||||
|
},
|
||||||
|
programs::ProbeContext,
|
||||||
|
};
|
||||||
|
use integration_common::hash_map::GET_INDEX;
|
||||||
|
|
||||||
|
#[btf_map]
|
||||||
|
static RESULT: Array<u32, 3 /* max_elements */, 0> = Array::new();
|
||||||
|
#[btf_map]
|
||||||
|
static HASH_MAP: HashMap<u32, u32, 10 /* max_elements */, 0> = HashMap::new();
|
||||||
|
#[btf_map]
|
||||||
|
static LRU_HASH_MAP: LruHashMap<u32, u32, 10 /* max_elements */, 0> = LruHashMap::new();
|
||||||
|
#[btf_map]
|
||||||
|
static PER_CPU_HASH_MAP: PerCpuHashMap<u32, u32, 10 /* max_elements */, 0> = PerCpuHashMap::new();
|
||||||
|
#[btf_map]
|
||||||
|
static LRU_PER_CPU_HASH_MAP: LruPerCpuHashMap<u32, u32, 10 /* max_elements */, 0> =
|
||||||
|
LruPerCpuHashMap::new();
|
||||||
|
|
||||||
|
#[map]
|
||||||
|
static RESULT_LEGACY: LegacyArray<u32> = LegacyArray::with_max_entries(3, 0);
|
||||||
|
#[map]
|
||||||
|
static HASH_MAP_LEGACY: LegacyHashMap<u32, u32> = LegacyHashMap::with_max_entries(10, 0);
|
||||||
|
#[map]
|
||||||
|
static LRU_HASH_MAP_LEGACY: LegacyLruHashMap<u32, u32> = LegacyLruHashMap::with_max_entries(10, 0);
|
||||||
|
#[map]
|
||||||
|
static PER_CPU_HASH_MAP_LEGACY: LegacyPerCpuHashMap<u32, u32> =
|
||||||
|
LegacyPerCpuHashMap::with_max_entries(10, 0);
|
||||||
|
#[map]
|
||||||
|
static LRU_PER_CPU_HASH_MAP_LEGACY: LegacyLruPerCpuHashMap<u32, u32> =
|
||||||
|
LegacyLruPerCpuHashMap::with_max_entries(10, 0);
|
||||||
|
|
||||||
|
macro_rules! define_result_set {
|
||||||
|
(
|
||||||
|
$result_map:ident,
|
||||||
|
$result_set_fn:ident
|
||||||
|
) => {
|
||||||
|
#[inline(always)]
|
||||||
|
fn $result_set_fn(index: u32, value: u32) -> Result<(), c_long> {
|
||||||
|
let ptr = $result_map.get_ptr_mut(index).ok_or(-1)?;
|
||||||
|
let dst = unsafe { ptr.as_mut() };
|
||||||
|
let dst_res = dst.ok_or(-1)?;
|
||||||
|
*dst_res = value;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
define_result_set!(RESULT, result_set);
|
||||||
|
define_result_set!(RESULT_LEGACY, result_set_legacy);
|
||||||
|
|
||||||
|
macro_rules! define_hash_map_test {
|
||||||
|
(
|
||||||
|
$hash_map:ident,
|
||||||
|
$result_set_fn:ident,
|
||||||
|
$insert_prog:ident,
|
||||||
|
$get_prog:ident
|
||||||
|
$(,)?
|
||||||
|
) => {
|
||||||
|
#[uprobe]
|
||||||
|
fn $insert_prog(ctx: ProbeContext) -> Result<(), c_long> {
|
||||||
|
let key = ctx.arg(0).ok_or(-1)?;
|
||||||
|
let value = ctx.arg(1).ok_or(-1)?;
|
||||||
|
$hash_map.insert(&key, &value, 0)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
#[uprobe]
|
||||||
|
fn $get_prog(ctx: ProbeContext) -> Result<(), c_long> {
|
||||||
|
let key = ctx.arg(0).ok_or(-1)?;
|
||||||
|
let value = unsafe { $hash_map.get(&key).ok_or(-1)? };
|
||||||
|
$result_set_fn(GET_INDEX, *value)?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
define_hash_map_test!(HASH_MAP, result_set, hash_map_insert, hash_map_get);
|
||||||
|
define_hash_map_test!(
|
||||||
|
HASH_MAP_LEGACY,
|
||||||
|
result_set_legacy,
|
||||||
|
hash_map_insert_legacy,
|
||||||
|
hash_map_get_legacy,
|
||||||
|
);
|
||||||
|
|
||||||
|
define_hash_map_test!(
|
||||||
|
LRU_HASH_MAP,
|
||||||
|
result_set,
|
||||||
|
lru_hash_map_insert,
|
||||||
|
lru_hash_map_get
|
||||||
|
);
|
||||||
|
define_hash_map_test!(
|
||||||
|
LRU_HASH_MAP_LEGACY,
|
||||||
|
result_set_legacy,
|
||||||
|
lru_hash_map_insert_legacy,
|
||||||
|
lru_hash_map_get_legacy,
|
||||||
|
);
|
||||||
|
|
||||||
|
define_hash_map_test!(
|
||||||
|
PER_CPU_HASH_MAP,
|
||||||
|
result_set,
|
||||||
|
per_cpu_hash_map_insert,
|
||||||
|
per_cpu_hash_map_get,
|
||||||
|
);
|
||||||
|
define_hash_map_test!(
|
||||||
|
PER_CPU_HASH_MAP_LEGACY,
|
||||||
|
result_set_legacy,
|
||||||
|
per_cpu_hash_map_insert_legacy,
|
||||||
|
per_cpu_hash_map_get_legacy,
|
||||||
|
);
|
||||||
|
|
||||||
|
define_hash_map_test!(
|
||||||
|
LRU_PER_CPU_HASH_MAP,
|
||||||
|
result_set,
|
||||||
|
lru_per_cpu_hash_map_insert,
|
||||||
|
lru_per_cpu_hash_map_get,
|
||||||
|
);
|
||||||
|
define_hash_map_test!(
|
||||||
|
LRU_PER_CPU_HASH_MAP_LEGACY,
|
||||||
|
result_set_legacy,
|
||||||
|
lru_per_cpu_hash_map_insert_legacy,
|
||||||
|
lru_per_cpu_hash_map_get_legacy,
|
||||||
|
);
|
||||||
@ -0,0 +1,238 @@
|
|||||||
|
use std::thread;
|
||||||
|
|
||||||
|
use aya::{
|
||||||
|
Ebpf, EbpfLoader,
|
||||||
|
maps::{Array, HashMap, MapData, MapError, PerCpuHashMap},
|
||||||
|
programs::UProbe,
|
||||||
|
};
|
||||||
|
use integration_common::hash_map::GET_INDEX;
|
||||||
|
use nix::{
|
||||||
|
sched::{CpuSet, sched_setaffinity},
|
||||||
|
unistd::Pid,
|
||||||
|
};
|
||||||
|
|
||||||
|
/// Triggers the eBPF program that inserts the given `key` and `value` pair
|
||||||
|
/// into the hash map.
|
||||||
|
#[unsafe(no_mangle)]
|
||||||
|
#[inline(never)]
|
||||||
|
extern "C" fn hash_map_insert(key: u32, value: u32) {
|
||||||
|
std::hint::black_box((key, value));
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Triggers the eBPF program that retrieves the value associated with the
|
||||||
|
/// `key` and inserts it into the array.
|
||||||
|
#[unsafe(no_mangle)]
|
||||||
|
#[inline(never)]
|
||||||
|
extern "C" fn hash_map_get(key: u32) {
|
||||||
|
std::hint::black_box(key);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Loads the uprobe program and attaches it to the given `symbol`.
|
||||||
|
fn load_program(ebpf: &mut Ebpf, prog_name: &str, symbol: &str) {
|
||||||
|
let prog: &mut UProbe = ebpf.program_mut(prog_name).unwrap().try_into().unwrap();
|
||||||
|
prog.load().unwrap();
|
||||||
|
prog.attach(symbol, "/proc/self/exe", None, None).unwrap();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Loads the pair of programs:
|
||||||
|
///
|
||||||
|
/// * `insert_prog` that inserts key and value pairs into the `hash_map`.
|
||||||
|
/// * `get_prog` that retrieves values from the `hash_map` and inserts them
|
||||||
|
/// into `result_map`.
|
||||||
|
///
|
||||||
|
/// Returns the result array and the hash map.
|
||||||
|
fn load_programs_with_maps<'a>(
|
||||||
|
ebpf: &'a mut Ebpf,
|
||||||
|
result_array: &'a str,
|
||||||
|
hash_map: &'a str,
|
||||||
|
insert_prog: &'a str,
|
||||||
|
get_prog: &'a str,
|
||||||
|
) -> (Array<&'a MapData, u32>, HashMap<&'a MapData, u32, u32>) {
|
||||||
|
load_program(ebpf, insert_prog, "hash_map_insert");
|
||||||
|
load_program(ebpf, get_prog, "hash_map_get");
|
||||||
|
|
||||||
|
let result_array = ebpf.map(result_array).unwrap();
|
||||||
|
let result_array = Array::<_, u32>::try_from(result_array).unwrap();
|
||||||
|
|
||||||
|
let hash_map = ebpf.map(hash_map).unwrap();
|
||||||
|
let hash_map = HashMap::<_, u32, u32>::try_from(hash_map).unwrap();
|
||||||
|
|
||||||
|
(result_array, hash_map)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Loads the `insert_prog` program that inserts elements into the
|
||||||
|
/// `per_cpu_hash_map`. Returns the map.
|
||||||
|
fn load_program_with_per_cpu_map<'a>(
|
||||||
|
ebpf: &'a mut Ebpf,
|
||||||
|
per_cpu_hash_map: &'a str,
|
||||||
|
insert_prog: &'a str,
|
||||||
|
) -> PerCpuHashMap<&'a MapData, u32, u32> {
|
||||||
|
load_program(ebpf, insert_prog, "hash_map_insert");
|
||||||
|
|
||||||
|
let hash_map = ebpf.map(per_cpu_hash_map).unwrap();
|
||||||
|
PerCpuHashMap::<_, u32, u32>::try_from(hash_map).unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test_log::test]
|
||||||
|
fn test_hash_map() {
|
||||||
|
let mut ebpf = EbpfLoader::new().load(crate::HASH_MAP).unwrap();
|
||||||
|
for (result_map_name, hash_map_name, insert_prog_name, get_prog_name) in [
|
||||||
|
// BTF map definitions.
|
||||||
|
("RESULT", "HASH_MAP", "hash_map_insert", "hash_map_get"),
|
||||||
|
// Legacy map definitions.
|
||||||
|
(
|
||||||
|
"RESULT_LEGACY",
|
||||||
|
"HASH_MAP_LEGACY",
|
||||||
|
"hash_map_insert_legacy",
|
||||||
|
"hash_map_get_legacy",
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let (result_array, hash_map) = load_programs_with_maps(
|
||||||
|
&mut ebpf,
|
||||||
|
result_map_name,
|
||||||
|
hash_map_name,
|
||||||
|
insert_prog_name,
|
||||||
|
get_prog_name,
|
||||||
|
);
|
||||||
|
|
||||||
|
let seq = 0_u32..9;
|
||||||
|
for i in seq.clone() {
|
||||||
|
hash_map_insert(i.pow(2), i);
|
||||||
|
}
|
||||||
|
for i in seq.clone() {
|
||||||
|
// Assert the value returned by user-space API.
|
||||||
|
let key = i.pow(2);
|
||||||
|
let value = hash_map.get(&key, 0).unwrap();
|
||||||
|
assert_eq!(value, i);
|
||||||
|
// Assert the value returned by eBPF in-kernel API.
|
||||||
|
hash_map_get(key);
|
||||||
|
let result = result_array.get(&GET_INDEX, 0).unwrap();
|
||||||
|
assert_eq!(result, i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test_log::test]
|
||||||
|
fn test_lru_hash_map() {
|
||||||
|
let mut ebpf = EbpfLoader::new().load(crate::HASH_MAP).unwrap();
|
||||||
|
for (result_map_name, hash_map_name, insert_prog_name, get_prog_name) in [
|
||||||
|
// BTF map definitions.
|
||||||
|
(
|
||||||
|
"RESULT",
|
||||||
|
"LRU_HASH_MAP",
|
||||||
|
"lru_hash_map_insert",
|
||||||
|
"lru_hash_map_get",
|
||||||
|
),
|
||||||
|
// Legacy map definitions.
|
||||||
|
(
|
||||||
|
"RESULT_LEGACY",
|
||||||
|
"LRU_HASH_MAP_LEGACY",
|
||||||
|
"lru_hash_map_insert_legacy",
|
||||||
|
"lru_hash_map_get_legacy",
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let (result_array, hash_map) = load_programs_with_maps(
|
||||||
|
&mut ebpf,
|
||||||
|
result_map_name,
|
||||||
|
hash_map_name,
|
||||||
|
insert_prog_name,
|
||||||
|
get_prog_name,
|
||||||
|
);
|
||||||
|
|
||||||
|
// Insert elements over capacity.
|
||||||
|
let seq = 0_u32..15;
|
||||||
|
for i in seq.clone() {
|
||||||
|
hash_map_insert(i.pow(2), i);
|
||||||
|
}
|
||||||
|
// Check whether elements 0..5 got evicted.
|
||||||
|
for i in 0_u32..5 {
|
||||||
|
let key = i.pow(2);
|
||||||
|
assert!(matches!(hash_map.get(&key, 0), Err(MapError::KeyNotFound)));
|
||||||
|
}
|
||||||
|
// Check whether the newest 10 elements can be retrieved.
|
||||||
|
for i in 5_u32..15 {
|
||||||
|
// Assert the value returned by user-space API.
|
||||||
|
let key = i.pow(2);
|
||||||
|
let value = hash_map.get(&key, 0).unwrap();
|
||||||
|
assert_eq!(value, i);
|
||||||
|
// Assert the value returned by eBPF in-kernel API.
|
||||||
|
hash_map_get(key);
|
||||||
|
let result = result_array.get(&GET_INDEX, 0).unwrap();
|
||||||
|
assert_eq!(result, i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test_log::test]
|
||||||
|
fn test_per_cpu_hash_map() {
|
||||||
|
let mut ebpf = EbpfLoader::new().load(crate::HASH_MAP).unwrap();
|
||||||
|
for (hash_map_name, insert_prog_name) in [
|
||||||
|
// BTF map definitions.
|
||||||
|
("PER_CPU_HASH_MAP", "per_cpu_hash_map_insert"),
|
||||||
|
// Legacy map definitions.
|
||||||
|
("PER_CPU_HASH_MAP_LEGACY", "per_cpu_hash_map_insert_legacy"),
|
||||||
|
] {
|
||||||
|
let hash_map = load_program_with_per_cpu_map(&mut ebpf, hash_map_name, insert_prog_name);
|
||||||
|
|
||||||
|
let seq = 0_u32..9;
|
||||||
|
thread::scope(|s| {
|
||||||
|
let seq = seq.clone();
|
||||||
|
s.spawn(move || {
|
||||||
|
let mut cpu_set = CpuSet::new();
|
||||||
|
cpu_set.set(0).unwrap();
|
||||||
|
sched_setaffinity(Pid::from_raw(0), &cpu_set).unwrap();
|
||||||
|
|
||||||
|
for i in seq {
|
||||||
|
hash_map_insert(i.pow(2), i);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
});
|
||||||
|
for i in seq.clone() {
|
||||||
|
let key = i.pow(2);
|
||||||
|
let values = hash_map.get(&key, 0).unwrap();
|
||||||
|
assert_eq!(values.first().unwrap(), &i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test_log::test]
|
||||||
|
fn test_lru_per_cpu_hash_map() {
|
||||||
|
let mut ebpf = EbpfLoader::new().load(crate::HASH_MAP).unwrap();
|
||||||
|
for (hash_map_name, insert_prog_name) in [
|
||||||
|
// BTF map definitions.
|
||||||
|
("LRU_PER_CPU_HASH_MAP", "lru_per_cpu_hash_map_insert"),
|
||||||
|
// Legacy map definitions.
|
||||||
|
(
|
||||||
|
"LRU_PER_CPU_HASH_MAP_LEGACY",
|
||||||
|
"lru_per_cpu_hash_map_insert_legacy",
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let hash_map = load_program_with_per_cpu_map(&mut ebpf, hash_map_name, insert_prog_name);
|
||||||
|
|
||||||
|
// Insert elements over capacity.
|
||||||
|
let seq = 0_u32..15;
|
||||||
|
thread::scope(|s| {
|
||||||
|
let seq = seq.clone();
|
||||||
|
s.spawn(move || {
|
||||||
|
let mut cpu_set = CpuSet::new();
|
||||||
|
cpu_set.set(0).unwrap();
|
||||||
|
sched_setaffinity(Pid::from_raw(0), &cpu_set).unwrap();
|
||||||
|
|
||||||
|
for i in seq {
|
||||||
|
hash_map_insert(i.pow(2), i);
|
||||||
|
}
|
||||||
|
});
|
||||||
|
});
|
||||||
|
// Check whether elements 0..5 got evicted.
|
||||||
|
for i in 0_u32..5 {
|
||||||
|
let key = i.pow(2);
|
||||||
|
assert!(matches!(hash_map.get(&key, 0), Err(MapError::KeyNotFound)));
|
||||||
|
}
|
||||||
|
// Check whether the newest 10 elements can be retrieved.
|
||||||
|
for i in 5_u32..15 {
|
||||||
|
let key = i.pow(2);
|
||||||
|
let values = hash_map.get(&key, 0).unwrap();
|
||||||
|
assert_eq!(values.first().unwrap(), &i);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in New Issue