Add hashmaps example

Signed-off-by: Junyeong Jeong <rhdxmr@gmail.com>

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"]

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<BigStructure> = PerCpuArray::with_max_entries(1);
+
+#[map]
+static mut BIG_STRUCT: LruHashMap<i8, BigStructure> = LruHashMap::with_max_entries(16);
+
+#[map]
+static mut PCPU_MEM_ALLOC: PerCpuHashMap<usize, usize> = PerCpuHashMap::with_max_entries(16);
+
+#[map]
+static mut MEM_ALLOC: HashMap<usize, usize> = 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);
+    }
+}

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],
+        }
+    }
+}

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;

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::<i8, BigStructure>::new(loaded.map("BIG_STRUCT").expect("map not found"))
+            .expect("error on LruHashMap::new");
+    let pcpu_mem_alloc =
+        PerCpuHashMap::<usize, usize>::new(loaded.map("PCPU_MEM_ALLOC").expect("map not found"))
+            .expect("error on PerCpuHashMap::new");
+    let mem_alloc = HashMap::<usize, usize>::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::<Vec<(usize, usize)>>();
+        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"
+    ))
+}