Add Module::global and Loaded::global methods

Author: rhdxmr

examples/example-probes/src/global_var/main.rs

@@ -1,21 +1,16 @@
 #![no_std]
 #![no_main]
-use core::sync::atomic::{AtomicU64, Ordering};
+use core::sync::atomic::Ordering;
+
 use redbpf_probes::kprobe::prelude::*;
 
+use example_probes::global_var::{GLOBAL_VAR, GLOBAL_VAR_INCORRECT};
+
 program!(0xFFFFFFFE, "GPL");
 
 #[map]
 static mut PERCPU_MAP: PerCpuArray<u64> = PerCpuArray::with_max_entries(1);
 
-// global variable is shared between multiple cores so proper synchronization
-// should be involved carefully.
-static GLOBAL_VAR: AtomicU64 = AtomicU64::new(0);
-
-// global variable without any synchronization mechanism. This results in wrong
-// statistics.
-static mut GLOBAL_VAR_INCORRECT: u64 = 0;
-
 #[kprobe]
 fn incr_write_count(_regs: Registers) {
     unsafe {

examples/example-probes/src/global_var/mod.rs

@@ -1,11 +1,13 @@
-use cty::*;
+use core::sync::atomic::AtomicU64;
 
-// This is where you should define the types shared by the kernel and user
-// space, eg:
-//
-// #[repr(C)]
-// #[derive(Debug)]
-// pub struct SomeEvent {
-//     pub pid: u64,
-//     ...
-// }
+use redbpf_macros::global;
+
+/// global variable is shared between multiple cores so proper synchronization
+/// should be involved carefully.
+#[global]
+pub static GLOBAL_VAR: AtomicU64 = AtomicU64::new(0);
+
+/// global variable without any synchronization mechanism. This results in wrong
+/// statistics.
+#[global]
+pub static mut GLOBAL_VAR_INCORRECT: u64 = 0;

examples/example-userspace/examples/global-var.rs

@@ -15,6 +15,7 @@ use tracing_subscriber::FmtSubscriber;
 
 use redbpf::{load::Loader, Array, PerCpuArray};
 
+use probes::global_var::{GLOBAL_VAR, GLOBAL_VAR_INCORRECT};
 #[repr(C)]
 #[derive(Debug, Clone)]
 struct Data {
@@ -39,19 +40,22 @@ async fn main() {
         .expect("kprobe_mut error")
         .attach_kprobe("ksys_write", 0)
         .expect("error attach_kprobe");
-    let global = Array::<Data>::new(loaded.map(".bss").expect("map not found"))
-        .expect("can not initialize Array");
+
+    let gvar = loaded
+        .global::<u64>("GLOBAL_VAR")
+        .expect("error on accessing gvar");
+    let gvar_wo_sync = loaded
+        .global::<u64>("GLOBAL_VAR_INCORRECT")
+        .expect("error on accessing gvar-wo-sync ");
     let percpu_map =
         PerCpuArray::<u64>::new(loaded.map("PERCPU_MAP").expect("PERCPU_MAP not found"))
             .expect("can not initialize PerCpuArray");
-
     loop {
-        let gval = global.get(0).expect("global var value");
         let pcpu_val = percpu_map.get(0).expect("percpu value");
         println!(
             "w/ sync, w/o sync, pcpu = {}, {}, {}",
-            gval.var,
-            gval.var_wo_sync,
+            gvar.load().unwrap(),
+            gvar_wo_sync.load().unwrap(),
             pcpu_val.iter().sum::<u64>()
         );
         select! {

redbpf-macros/src/lib.rs

@@ -770,3 +770,20 @@ pub fn task_iter(attrs: TokenStream, item: TokenStream) -> TokenStream {
 
     probe_impl("task_iter", attrs, wrapper, name)
 }
+
+/// Attribute macro for defining global variables in probes
+///
+/// **NOTE** It is not required to use it when global variables are accessed
+/// only inside probes. But it has to used if global variables are also
+/// accessed by userspace program.
+#[proc_macro_attribute]
+pub fn global(_attrs: TokenStream, item: TokenStream) -> TokenStream {
+    let item = parse_macro_input!(item as ItemStatic);
+    let section_name = format!("globals/{}", item.ident.to_string());
+    quote! {
+        #[no_mangle]
+        #[link_section = #section_name]
+        #item
+    }
+    .into()
+}

redbpf/src/lib.rs

@@ -1383,6 +1383,15 @@ impl Module {
     pub fn task_iter_mut(&mut self, name: &str) -> Option<&mut TaskIter> {
         self.task_iters_mut().find(|p| p.common.name == name)
     }
+
+    pub fn global<T: Clone>(&self, name: &str) -> Result<GlobalVariable<'_, T>> {
+        let map = self.map(name).ok_or_else(|| {
+            error!("map not found: {}", name);
+            Error::Map
+        })?;
+
+        GlobalVariable::new(map)
+    }
 }
 
 impl<'a> ModuleBuilder<'a> {
@@ -1486,6 +1495,10 @@ impl<'a> ModuleBuilder<'a> {
                         symval_to_map_builders.insert(sym.st_value, map_builder);
                     }
                 }
+                (hdr::SHT_PROGBITS, Some("globals"), Some(name)) => {
+                    let map_builder = MapBuilder::with_section_data(name, &content)?;
+                    map_builders.insert(shndx, map_builder);
+                }
                 (hdr::SHT_PROGBITS, Some(kind @ "kprobe"), Some(name))
                 | (hdr::SHT_PROGBITS, Some(kind @ "kretprobe"), Some(name))
                 | (hdr::SHT_PROGBITS, Some(kind @ "uprobe"), Some(name))
@@ -2744,6 +2757,25 @@ impl Drop for TaskIter {
     }
 }
 
+pub struct GlobalVariable<'a, T: Clone> {
+    array: Array<'a, T>,
+}
+
+impl<'a, T: Clone> GlobalVariable<'a, T> {
+    fn new(map: &Map) -> Result<GlobalVariable<T>> {
+        let array = Array::<T>::new(map)?;
+        Ok(GlobalVariable { array })
+    }
+
+    pub fn load(&self) -> Option<T> {
+        self.array.get(0)
+    }
+
+    pub fn store(&self, val: T) -> Result<()> {
+        self.array.set(0, val)
+    }
+}
+
 #[inline]
 fn add_relocation(
     rels: &mut Vec<RelocationInfo>,

redbpf/src/load/loader.rs

@@ -15,8 +15,8 @@ use std::path::Path;
 use crate::load::map_io::PerfMessageStream;
 use crate::{cpus, Program};
 use crate::{
-    Error, KProbe, Map, Module, PerfMap, SkLookup, SocketFilter, StreamParser, StreamVerdict,
-    TaskIter, UProbe, XDP,
+    Error, GlobalVariable, KProbe, Map, Module, PerfMap, SkLookup, SocketFilter, StreamParser,
+    StreamVerdict, TaskIter, UProbe, XDP,
 };
 
 #[derive(Debug)]
@@ -188,4 +188,8 @@ impl Loaded {
     pub fn task_iter_mut(&mut self, name: &str) -> Option<&mut TaskIter> {
         self.module.task_iter_mut(name)
     }
+
+    pub fn global<T: Clone>(&self, name: &str) -> Result<GlobalVariable<'_, T>, Error> {
+        self.module.global(name)
+    }
 }