Skip to content
This repository has been archived by the owner on Jul 11, 2023. It is now read-only.

Commit

Permalink
Add Module::global and Loaded::global methods
Browse files Browse the repository at this point in the history
  • Loading branch information
rhdxmr committed Feb 10, 2022
1 parent 823f2c4 commit f8d3019
Show file tree
Hide file tree
Showing 6 changed files with 81 additions and 27 deletions.
13 changes: 4 additions & 9 deletions examples/example-probes/src/global_var/main.rs
Original file line number Diff line number Diff line change
@@ -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 {
Expand Down
22 changes: 12 additions & 10 deletions examples/example-probes/src/global_var/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
16 changes: 10 additions & 6 deletions examples/example-userspace/examples/global-var.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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! {
Expand Down
17 changes: 17 additions & 0 deletions redbpf-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
32 changes: 32 additions & 0 deletions redbpf/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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>,
Expand Down
8 changes: 6 additions & 2 deletions redbpf/src/load/loader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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)
}
}

0 comments on commit f8d3019

Please sign in to comment.