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/r10
Michal R
1 month ago
parent
3a7d1bebb4
commit
92f914fd77
@ -0,0 +1,302 @@
|
||||
//! Hash map types that can be shared between eBPF programs and user-space.
|
||||
|
||||
#![deny(missing_docs)]
|
||||
|
||||
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,
|
||||
};
|
||||
|
||||
/// Generates a hash map definition with common methods.
|
||||
macro_rules! hash_map {
|
||||
(
|
||||
$map_doc:literal,
|
||||
$map_doc_examples:literal,
|
||||
$name:ident,
|
||||
$def:ident
|
||||
$(,)?
|
||||
) => {
|
||||
#[doc = include_str!($map_doc)]
|
||||
#[doc = $map_doc_examples]
|
||||
#[repr(transparent)]
|
||||
pub struct $name<K, V, const M: usize, const F: usize = 0>(UnsafeCell<$def<K, V, M, F>>);
|
||||
|
||||
unsafe impl<K: Sync, V: Sync, const M: usize, const F: usize> Sync for $name<K, V, M, F> {}
|
||||
|
||||
impl<K, V, const M: usize, const F: usize> $name<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`."
|
||||
)]
|
||||
#[doc = "Creates a new map."]
|
||||
pub const fn new() -> Self {
|
||||
Self(UnsafeCell::new($def::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())
|
||||
}
|
||||
|
||||
#[doc = "Inserts the given `key` and `value` pair to 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)
|
||||
}
|
||||
|
||||
#[doc = "Removes the given `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!(HashMapDef, BPF_MAP_TYPE_HASH);
|
||||
btf_map_def!(LruHashMapDef, BPF_MAP_TYPE_LRU_HASH);
|
||||
btf_map_def!(PerCpuHashMapDef, BPF_MAP_TYPE_PERCPU_HASH);
|
||||
btf_map_def!(LruPerCpuHashMapDef, BPF_MAP_TYPE_LRU_PERCPU_HASH);
|
||||
|
||||
hash_map!(
|
||||
"../maps/docs/hash_map.md",
|
||||
r#"# Examples
|
||||
|
||||
```rust,no_run
|
||||
use aya_ebpf::{
|
||||
maps::HashMap,
|
||||
macros::{map, tracepoint},
|
||||
programs::TracePointContext,
|
||||
EbpfContext as _,
|
||||
};
|
||||
|
||||
/// A hash map that counts syscalls issued by different processes.
|
||||
#[map]
|
||||
static COUNTER: HashMap<
|
||||
// PID.
|
||||
u32,
|
||||
// Count of syscalls issued by the given process.
|
||||
u32
|
||||
> = HashMap::with_max_entries(
|
||||
// Maximum number of elements. Reaching this capacity triggers an error.
|
||||
10,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
|
||||
/// A simple program attached to the `sys_enter` tracepoint that counts
|
||||
/// syscalls.
|
||||
#[tracepoint]
|
||||
fn sys_enter(ctx: TracePointContext) {
|
||||
let pid = ctx.pid();
|
||||
|
||||
if let Some(mut count) = COUNTER.get_ptr_mut(pid) {
|
||||
unsafe { *count += 1 };
|
||||
} else {
|
||||
COUNTER.insert(
|
||||
pid,
|
||||
// New value.
|
||||
1,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
}
|
||||
}
|
||||
```"#,
|
||||
HashMap,
|
||||
HashMapDef,
|
||||
);
|
||||
|
||||
hash_map!(
|
||||
"../maps/docs/lru_hash_map.md",
|
||||
r#"# Examples
|
||||
|
||||
```rust,no_run
|
||||
use aya_ebpf::{
|
||||
maps::LruHashMap,
|
||||
macros::{map, tracepoint},
|
||||
programs::TracePointContext,
|
||||
EbpfContext as _,
|
||||
};
|
||||
|
||||
/// A hash map that counts syscalls issued by different processes.
|
||||
#[map]
|
||||
static COUNTER: LruHashMap<
|
||||
// PID.
|
||||
u32,
|
||||
// Count of syscalls issued by the given process.
|
||||
u32,
|
||||
|
||||
> = LruHashMap::with_max_entries(
|
||||
// Maximum number of elements. Reaching this capacity triggers eviction of
|
||||
// the least used elements.
|
||||
10,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
|
||||
/// A simple program attached to the `sys_enter` tracepoint that counts
|
||||
/// syscalls.
|
||||
#[tracepoint]
|
||||
fn sys_enter(ctx: TracePointContext) {
|
||||
let pid = ctx.pid();
|
||||
|
||||
if let Some(mut count) = COUNTER.get_ptr_mut(pid) {
|
||||
unsafe { *count += 1 };
|
||||
} else {
|
||||
COUNTER.insert(
|
||||
pid,
|
||||
// New value.
|
||||
1,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
}
|
||||
}
|
||||
```"#,
|
||||
LruHashMap,
|
||||
LruHashMapDef,
|
||||
);
|
||||
|
||||
hash_map!(
|
||||
"../maps/docs/per_cpu_hash_map.md",
|
||||
r#"# Examples
|
||||
|
||||
```rust,no_run
|
||||
use aya_ebpf::{
|
||||
maps::PerCpuHashMap,
|
||||
macros::{map, tracepoint},
|
||||
programs::TracePointContext,
|
||||
EbpfContext as _,
|
||||
};
|
||||
|
||||
/// A hash map that counts syscalls issued by different processes.
|
||||
#[map]
|
||||
static COUNTER: PerCpuHashMap<
|
||||
// PID.
|
||||
u32,
|
||||
// Count of syscalls issued by the given process.
|
||||
u32
|
||||
> = PerCpuHashMap::with_max_entries(
|
||||
// Maximum number of elements. Reaching this capacity triggers an error.
|
||||
10,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
|
||||
/// A simple program attached to the `sys_enter` tracepoint that counts
|
||||
/// syscalls.
|
||||
#[tracepoint]
|
||||
fn sys_enter(ctx: TracePointContext) {
|
||||
let pid = ctx.pid();
|
||||
|
||||
if let Some(mut count) = COUNTER.get_ptr_mut(pid) {
|
||||
unsafe { *count += 1 };
|
||||
} else {
|
||||
COUNTER.insert(
|
||||
pid,
|
||||
// New value.
|
||||
1,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
}
|
||||
}
|
||||
```"#,
|
||||
PerCpuHashMap,
|
||||
PerCpuHashMapDef,
|
||||
);
|
||||
|
||||
hash_map!(
|
||||
"../maps/docs/lru_per_cpu_hash_map.md",
|
||||
r#"# Examples
|
||||
|
||||
```rust,no_run
|
||||
use aya_ebpf::{
|
||||
maps::LruPerCpuHashMap,
|
||||
macros::{map, tracepoint},
|
||||
programs::TracePointContext,
|
||||
EbpfContext as _,
|
||||
};
|
||||
|
||||
/// A hash map that counts syscalls issued by different processes.
|
||||
#[map]
|
||||
static COUNTER: LruPerCpuHashMap<
|
||||
// PID.
|
||||
u32,
|
||||
// Count of syscalls issued by the given process.
|
||||
u32,
|
||||
|
||||
> = LruPerCpuHashMap::with_max_entries(
|
||||
// Maximum number of elements. Reaching this capacity triggers eviction of
|
||||
// the least used elements.
|
||||
10,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
|
||||
/// A simple program attached to the `sys_enter` tracepoint that counts
|
||||
/// syscalls.
|
||||
#[tracepoint]
|
||||
fn sys_enter(ctx: TracePointContext) {
|
||||
let pid = ctx.pid();
|
||||
|
||||
if let Some(mut count) = COUNTER.get_ptr_mut(pid) {
|
||||
unsafe { *count += 1 };
|
||||
} else {
|
||||
COUNTER.insert(
|
||||
pid,
|
||||
// New value.
|
||||
1,
|
||||
// Optional flags.
|
||||
0
|
||||
);
|
||||
}
|
||||
}
|
||||
```"#,
|
||||
LruPerCpuHashMap,
|
||||
LruPerCpuHashMapDef,
|
||||
);
|
||||
|
||||
#[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,234 @@
|
||||
use std::thread;
|
||||
|
||||
use aya::{
|
||||
Ebpf, EbpfLoader,
|
||||
maps::{Array, HashMap, MapData, MapError, PerCpuHashMap},
|
||||
programs::UProbe,
|
||||
};
|
||||
use integration_common::hash_map::GET_INDEX;
|
||||
|
||||
/// 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 = nix::sched::CpuSet::new();
|
||||
cpu_set.set(0).unwrap();
|
||||
nix::sched::sched_setaffinity(nix::unistd::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 = nix::sched::CpuSet::new();
|
||||
cpu_set.set(0).unwrap();
|
||||
nix::sched::sched_setaffinity(nix::unistd::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