diff --git a/examples/example-probes/Cargo.toml b/examples/example-probes/Cargo.toml index bea65eb4..9bfb7df0 100644 --- a/examples/example-probes/Cargo.toml +++ b/examples/example-probes/Cargo.toml @@ -73,3 +73,8 @@ required-features = ["probes"] name = "tasks" path = "src/tasks/main.rs" required-features = ["probes", "kernel5_8"] + +[[bin]] +name = "hashmaps" +path = "src/hashmaps/main.rs" +required-features = ["probes"] diff --git a/examples/example-probes/src/hashmaps/main.rs b/examples/example-probes/src/hashmaps/main.rs new file mode 100644 index 00000000..97c47615 --- /dev/null +++ b/examples/example-probes/src/hashmaps/main.rs @@ -0,0 +1,68 @@ +//! This is an example of showing difference between `PerCpuHashMap` and +//! `HashMap`. The former is per-cpu data structure and users don't need to +//! worry about race condition. The latter is global data structure so it has +//! race condition problems. +//! +//! `PerCpuArray` can be used instead of bpf stack to hold temporary values +//! that exceeds the maximum size of bpf stack (512 bytes). +#![no_std] +#![no_main] +use example_probes::hashmaps::*; +use redbpf_probes::kprobe::prelude::*; + +program!(0xFFFFFFFE, "GPL"); + +#[map] +static mut ALT_STACK: PerCpuArray = PerCpuArray::with_max_entries(1); + +#[map] +static mut BIG_STRUCT: LruHashMap = LruHashMap::with_max_entries(16); + +#[map] +static mut PCPU_MEM_ALLOC: PerCpuHashMap = PerCpuHashMap::with_max_entries(16); + +#[map] +static mut MEM_ALLOC: HashMap = HashMap::with_max_entries(16); + +#[kprobe] +unsafe fn sched_fork(_regs: Registers) { + let rnd_key = (bpf_get_prandom_u32() & 0xff) as i8; + if let Some(bigstruct) = BIG_STRUCT.get_mut(&rnd_key) { + bigstruct.f2[99] = 99; + BIG_STRUCT.set(&rnd_key, bigstruct); + } else { + // maximum size of bpf stack is 512 bytes. BigStructure struct is 808 + // bytes. So it can not be located in stack. Use percpu array to hold + // temporary BigStructure value. Note that if percpu array is used for + // this purpose, the size of percpu array must be 1. This is checked by + // BPF verifier. + let bigstruct = ALT_STACK.get_mut(0).unwrap(); + for x in 0..=99 { + bigstruct.f2[x] = x; + } + + BIG_STRUCT.set(&rnd_key, bigstruct); + } +} + +#[kprobe] +unsafe fn __kmalloc(regs: Registers) { + let mut size = regs.parm1() as usize; + let mut max: usize = 9999; + for x in 1..=12 { + size >>= 1; + if size == 0 { + max = usize::pow(2, x) - 1; + break; + } + } + if let Some(count) = PCPU_MEM_ALLOC.get_mut(&max) { + *count += 1; + let count = MEM_ALLOC.get_mut(&max).unwrap(); + *count += 1; + } else { + let count = 1; + PCPU_MEM_ALLOC.set(&max, &count); + MEM_ALLOC.set(&max, &count); + } +} diff --git a/examples/example-probes/src/hashmaps/mod.rs b/examples/example-probes/src/hashmaps/mod.rs new file mode 100644 index 00000000..bae5abb4 --- /dev/null +++ b/examples/example-probes/src/hashmaps/mod.rs @@ -0,0 +1,15 @@ +#[repr(C)] +#[derive(Clone, Debug)] +pub struct BigStructure { + pub f1: usize, + pub f2: [usize; 100], +} + +impl Default for BigStructure { + fn default() -> Self { + BigStructure { + f1: 0, + f2: [0; 100], + } + } +} diff --git a/examples/example-probes/src/lib.rs b/examples/example-probes/src/lib.rs index ea9df2e3..83f72982 100644 --- a/examples/example-probes/src/lib.rs +++ b/examples/example-probes/src/lib.rs @@ -12,7 +12,8 @@ pub mod bindings; pub mod echo; +pub mod hashmaps; pub mod mallocstacks; +pub mod tasks; pub mod tcp_lifetime; pub mod vfsreadlat; -pub mod tasks; diff --git a/examples/example-userspace/examples/hashmaps.rs b/examples/example-userspace/examples/hashmaps.rs new file mode 100644 index 00000000..8bc0f796 --- /dev/null +++ b/examples/example-userspace/examples/hashmaps.rs @@ -0,0 +1,77 @@ +//! This example shows usage of HashMap, PerCpuHashMap and LruHashMap. And +//! also it confirms you that hashmap has race condition problems. You should +//! consider PerCpuHashMap if your program needs to store accurate map data. + +use libc; +use std::process; +use std::time::Duration; +use tokio::{signal::ctrl_c, time::sleep}; +use tracing::{error, subscriber, Level}; +use tracing_subscriber::FmtSubscriber; + +use probes::hashmaps::BigStructure; +use redbpf::{load::Loader, HashMap, LruHashMap, PerCpuHashMap}; + +#[tokio::main(flavor = "current_thread")] +async fn main() { + let subscriber = FmtSubscriber::builder() + .with_max_level(Level::TRACE) + .finish(); + subscriber::set_global_default(subscriber).unwrap(); + if unsafe { libc::geteuid() != 0 } { + error!("You must be root to use eBPF!"); + process::exit(1); + } + + let mut loaded = Loader::load(probe_code()).expect("error loading probe"); + for kp in loaded.kprobes_mut() { + kp.attach_kprobe(kp.name().as_str(), 0) + .expect(format!("error on attach_kprobe to {}", kp.name()).as_str()); + } + + let big_struct = + LruHashMap::::new(loaded.map("BIG_STRUCT").expect("map not found")) + .expect("error on LruHashMap::new"); + let pcpu_mem_alloc = + PerCpuHashMap::::new(loaded.map("PCPU_MEM_ALLOC").expect("map not found")) + .expect("error on PerCpuHashMap::new"); + let mem_alloc = HashMap::::new(loaded.map("MEM_ALLOC").expect("map not found")) + .expect("error on HashMap::new"); + println!("Hit Ctrl-C to quit"); + loop { + tokio::select! { + _ = sleep(Duration::from_secs(1)) => {} + _ = ctrl_c() => break + } + + let mut alloc_stats = mem_alloc.iter().collect::>(); + alloc_stats.sort(); + println!("[allocation size upto XXX bytes] => [number of __kmalloc call]"); + + for (size, total_cnt) in alloc_stats { + let pcpu_vals = pcpu_mem_alloc.get(size).unwrap(); + let exact_cnt: usize = pcpu_vals.iter().sum(); + if total_cnt != exact_cnt { + println!( + "{} => {} != {} (hashmap != pcpu hashmap)", + size, total_cnt, exact_cnt + ); + } else { + println!("{} => {}", size, total_cnt); + } + } + } + + println!(""); + println!("iterate over big structures!"); + for (_, bigstruct) in big_struct.iter() { + println!("{:?}", bigstruct); + } +} + +fn probe_code() -> &'static [u8] { + include_bytes!(concat!( + env!("OUT_DIR"), + "/target/bpf/programs/hashmaps/hashmaps.elf" + )) +}