mirror of https://github.com/rust-bpf/rust-bcc
146 lines
3.4 KiB
Rust
146 lines
3.4 KiB
Rust
// Copyright 2021 Twitter, Inc.
|
|
// Licensed under the Apache License, Version 2.0
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
use bcc::perf_event::{Event, SoftwareEvent};
|
|
use bcc::BccError;
|
|
use bcc::{PerfEvent, PerfEventArray, BPF};
|
|
use clap::{App, Arg};
|
|
|
|
use core::sync::atomic::{AtomicBool, Ordering};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::{thread, time};
|
|
|
|
// Both consants are arbitrary
|
|
const DEFAULT_SAMPLE_FREQ: u64 = 99; // Hertz
|
|
const DEFAULT_DURATION: u64 = 10; // Seconds
|
|
|
|
#[cfg(any(
|
|
feature = "v0_4_0",
|
|
feature = "v0_5_0",
|
|
feature = "v0_6_0",
|
|
feature = "v0_6_1"
|
|
))]
|
|
fn do_main(_runnable: Arc<AtomicBool>) -> Result<(), BccError> {
|
|
panic!("requires bcc >= 0.7.0");
|
|
}
|
|
|
|
#[cfg(not(any(
|
|
feature = "v0_4_0",
|
|
feature = "v0_5_0",
|
|
feature = "v0_6_0",
|
|
feature = "v0_6_1"
|
|
)))]
|
|
fn do_main(runnable: Arc<AtomicBool>) -> Result<(), BccError> {
|
|
let matches = App::new("memory")
|
|
.arg(
|
|
Arg::with_name("duration")
|
|
.long("duration")
|
|
.short("d")
|
|
.help("How long to run this trace for (in seconds)")
|
|
.takes_value(true),
|
|
)
|
|
.get_matches();
|
|
|
|
let duration: u64 = matches
|
|
.value_of("duration")
|
|
.map(|v| v.parse().expect("Invalid duration"))
|
|
.unwrap_or(DEFAULT_DURATION);
|
|
|
|
let cpus = bcc::cpuonline::get()?.len() as u32;
|
|
|
|
let code = format!(
|
|
"{}\n{}",
|
|
format!("#define NUM_CPU {}", cpus),
|
|
include_str!("memory.c").to_string()
|
|
);
|
|
|
|
let mut bpf = BPF::new(&code)?;
|
|
PerfEventArray::new()
|
|
.table("loads_perf")
|
|
.event(Event::Raw {
|
|
event_code: 0x0E,
|
|
umask: 0x0F,
|
|
counter_mask: 0x00,
|
|
invert: false,
|
|
any_thread: false,
|
|
edge_detect: false,
|
|
})
|
|
.attach(&mut bpf)?;
|
|
|
|
println!("Running for {} seconds", duration);
|
|
|
|
let mut elapsed = 0;
|
|
while runnable.load(Ordering::SeqCst) {
|
|
thread::sleep(time::Duration::new(1, 0));
|
|
|
|
if elapsed == duration {
|
|
break;
|
|
}
|
|
elapsed += 1;
|
|
}
|
|
|
|
// Count misses
|
|
let mut loads = bpf.table("loads")?;
|
|
let loads_map = to_map(&mut loads);
|
|
|
|
let mut total_loads = 0;
|
|
|
|
for i in 0..cpus {
|
|
let loads = loads_map.get(&i).unwrap_or(&0);
|
|
|
|
total_loads += *loads;
|
|
}
|
|
|
|
println!("{:<-12}", "TOTAL_LOADS");
|
|
println!("{:<-12}", total_loads,);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn to_map(table: &mut bcc::table::Table) -> HashMap<u32, u64> {
|
|
let mut map = HashMap::new();
|
|
|
|
for entry in table.iter() {
|
|
let key = parse_u32(entry.key);
|
|
let value = parse_u64(entry.value);
|
|
|
|
map.insert(key, value);
|
|
}
|
|
|
|
map
|
|
}
|
|
|
|
fn parse_u32(x: Vec<u8>) -> u32 {
|
|
let mut v = [0_u8; 4];
|
|
for (i, byte) in v.iter_mut().enumerate() {
|
|
*byte = *x.get(i).unwrap_or(&0);
|
|
}
|
|
|
|
u32::from_ne_bytes(v)
|
|
}
|
|
|
|
fn parse_u64(x: Vec<u8>) -> u64 {
|
|
let mut v = [0_u8; 8];
|
|
for (i, byte) in v.iter_mut().enumerate() {
|
|
*byte = *x.get(i).unwrap_or(&0);
|
|
}
|
|
|
|
u64::from_ne_bytes(v)
|
|
}
|
|
|
|
fn main() {
|
|
let runnable = Arc::new(AtomicBool::new(true));
|
|
let r = runnable.clone();
|
|
ctrlc::set_handler(move || {
|
|
r.store(false, Ordering::SeqCst);
|
|
})
|
|
.expect("Failed to set handler for SIGINT / SIGTERM");
|
|
|
|
if let Err(x) = do_main(runnable) {
|
|
eprintln!("Error: {}", x);
|
|
std::process::exit(1);
|
|
}
|
|
}
|