Misc improvements

Add helpers to generate random values for workload.
Use BTF to attach fentry/fexit to BPF programs (WIP).
Add SIGTERM hook to finish workers processes.

Author: erthalion

Cargo.lock

@@ -27,3 +27,4 @@ pest = "2.8.1"
 pest_derive = "2.8.1"
 llvm-sys = "201.0.1"
 docopt = "1.1.1"
+signal-hook = "0.3.18"

Cargo.toml

@@ -3,10 +3,16 @@ use crate::script::ast::MachineInstruction;
 use log::{debug, trace};
 use std::{
     io::{prelude::*, BufReader},
+    mem,
     net::TcpListener,
     thread,
 };
 
+use libc::SYS_bpf;
+
+use aya::{include_bytes_aligned, programs::FEntry, Btf, Ebpf, Endianness};
+use aya_obj::generated::{bpf_attr, bpf_btf_info, bpf_cmd};
+
 #[derive(Debug)]
 pub enum MachineError {
     Internal,
@@ -62,10 +68,80 @@ fn start_server(addr: String, target_port: u16) -> Result<(), MachineError> {
     Ok(())
 }
 
+fn start_bpf_profiling() -> Result<(), MachineError> {
+    debug!("Starting eBPF profiling");
+    let btf_fd;
+    let mut buf = vec![0u8; 4096];
+
+    // Load the BPF program
+    unsafe {
+        let mut fd_attr = mem::zeroed::<bpf_attr>();
+        fd_attr.__bindgen_anon_6.__bindgen_anon_1.btf_id = 293;
+
+        btf_fd = libc::syscall(
+            SYS_bpf,
+            bpf_cmd::BPF_BTF_GET_FD_BY_ID,
+            &fd_attr,
+            mem::size_of::<bpf_attr>(),
+        );
+
+        let mut info_attr = mem::zeroed::<bpf_attr>();
+        let mut info = mem::zeroed::<bpf_btf_info>();
+
+        info.btf = buf.as_mut_ptr() as _;
+        info.btf_size = buf.len() as _;
+
+        info_attr.info.bpf_fd = btf_fd as u32;
+        info_attr.info.info = &info as *const _ as u64;
+        info_attr.info.info_len = mem::size_of_val(&info) as u32;
+
+        libc::syscall(
+            SYS_bpf,
+            bpf_cmd::BPF_OBJ_GET_INFO_BY_FD,
+            &info_attr,
+            mem::size_of::<bpf_attr>(),
+        );
+    }
+
+    let btf = Btf::parse(&buf, Endianness::default()).unwrap();
+
+    debug!("Got btf {:?}", btf);
+
+    debug!("Loading eBPF program");
+    let mut bpf = match Ebpf::load(include_bytes_aligned!("../fentry.bpf.o")) {
+        Ok(prog) => {
+            debug!("Loaded prog");
+            prog
+        }
+        Err(e) => {
+            debug!("Cannot load eBPF program, {e}");
+            panic!("AAA")
+        }
+    };
+
+    debug!("Loaded eBPF program");
+    //let btf = Btf::from_sys_fs().unwrap();
+    let btf = match Btf::parse_file("/tmp/btf", Endianness::default()) {
+        Ok(data) => data,
+        Err(e) => panic!("Cannot parse BTF {e}"),
+    };
+    let fentry: &mut FEntry =
+        bpf.program_mut("fentry_XXX").unwrap().try_into().unwrap();
+
+    fentry.load("handle_tp", &btf).unwrap();
+    debug!("Loaded fentry program");
+
+    fentry.attach().unwrap();
+    debug!("Attached fentry");
+
+    Ok(())
+}
+
 pub fn apply(instr: MachineInstruction) -> Result<(), MachineError> {
     match instr {
         MachineInstruction::Server { port } => {
             start_server("127.0.0.1".to_string(), port)
         }
+        MachineInstruction::Profile { target: _ } => start_bpf_profiling(),
     }
 }

src/machine.rs

@@ -31,6 +31,9 @@ use serde::Deserialize;
 use std::time::SystemTime;
 use std::{thread, time};
 
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+
 use berserker::machine::apply;
 use berserker::script::{ast::Node, parser::parse_instructions};
 use berserker::{
@@ -53,8 +56,9 @@ struct Args {
     flag_f: Option<String>,
 }
 
-fn run_script(script_path: String) -> Vec<Option<i32>> {
+fn run_script(script_path: String) -> Vec<(i32, u64)> {
     let mut handles = vec![];
+    info!("Loading script: {:?}", script_path);
 
     let ast: Vec<Node> =
         parse_instructions(&std::fs::read_to_string(script_path).unwrap())
@@ -72,8 +76,8 @@ fn run_script(script_path: String) -> Vec<Option<i32>> {
         };
 
         for instr in m_instructions {
-            debug!("INSTR {:?}", instr);
-            thread::spawn(move || apply(instr.clone()));
+            //thread::spawn(move || apply(instr.clone()));
+            let _ = apply(instr.clone());
         }
     };
 
@@ -96,14 +100,22 @@ fn run_script(script_path: String) -> Vec<Option<i32>> {
             .unwrap_or(String::from("0"))
             .parse()
             .unwrap();
+
+        let duration: u64 = args
+            .get("duration")
+            .cloned()
+            .unwrap_or(String::from("0"))
+            .parse()
+            .unwrap();
+
         let h: Vec<_> = (0..workers)
             .map(|_| {
                 let worker = new_script_worker(node.clone());
 
                 match fork() {
                     Ok(Fork::Parent(child)) => {
                         info!("Child {}", child);
-                        Some(child)
+                        Some((child, duration))
                     }
                     Ok(Fork::Child) => {
                         worker.run_payload().unwrap();
@@ -120,10 +132,10 @@ fn run_script(script_path: String) -> Vec<Option<i32>> {
         handles.extend(h);
     });
 
-    handles
+    handles.iter().filter_map(|i| *i).collect()
 }
 
-fn run_workload(config: WorkloadConfig) -> Vec<Option<i32>> {
+fn run_workload(config: WorkloadConfig) -> Vec<(i32, u64)> {
     let mut lower = 1024;
     let mut upper = 1024;
 
@@ -142,7 +154,7 @@ fn run_workload(config: WorkloadConfig) -> Vec<Option<i32>> {
             match fork() {
                 Ok(Fork::Parent(child)) => {
                     info!("Child {}", child);
-                    Some(child)
+                    Some((child, config.duration))
                 }
                 Ok(Fork::Child) => {
                     if config.per_core {
@@ -162,83 +174,101 @@ fn run_workload(config: WorkloadConfig) -> Vec<Option<i32>> {
         .collect();
 
     info!("In total: {}", upper);
-    handles
+    handles.iter().filter_map(|i| *i).collect()
 }
 
 fn main() {
     env_logger::init();
 
+    let terminating = Arc::new(AtomicBool::new(false));
+    signal_hook::flag::register(
+        signal_hook::consts::SIGTERM,
+        Arc::clone(&terminating),
+    )
+    .unwrap();
+
     let args: Args = Docopt::new(USAGE)
         .and_then(|d| d.deserialize())
         .unwrap_or_else(|e| e.exit());
 
     debug!("ARGS {:?}", args);
 
-    let default_config = String::from("workload.toml");
     let duration_timer = SystemTime::now();
     let script_path = args.flag_f;
-    let config_path = args.flag_c.unwrap_or(default_config);
-
-    let config = Config::builder()
-        // Add in `./Settings.toml`
-        .add_source(
-            config::File::with_name("/etc/berserker/workload.toml")
-                .required(false),
-        )
-        .add_source(
-            config::File::with_name(config_path.as_str()).required(false),
-        )
-        // Add in settings from the environment (with a prefix of APP)
-        // Eg.. `BERSERKER__WORKLOAD__ARRIVAL_RATE=1` would set the `arrival_rate` key
-        .add_source(
-            config::Environment::with_prefix("BERSERKER")
-                .try_parsing(true)
-                .separator("__"),
-        )
-        .build()
-        .unwrap()
-        .try_deserialize::<WorkloadConfig>()
-        .unwrap();
-
-    info!("Config: {:?}", config);
 
     let handles = match script_path {
         Some(path) => run_script(path),
-        None => run_workload(config),
+        None => {
+            let default_config = String::from("workload.toml");
+            let config_path = args.flag_c.unwrap_or(default_config);
+
+            let config = Config::builder()
+                // Add in `./Settings.toml`
+                .add_source(
+                    config::File::with_name("/etc/berserker/workload.toml")
+                        .required(false),
+                )
+                .add_source(
+                    config::File::with_name(config_path.as_str())
+                        .required(false),
+                )
+                // Add in settings from the environment (with a prefix of APP)
+                // Eg.. `BERSERKER__WORKLOAD__ARRIVAL_RATE=1` would set the `arrival_rate` key
+                .add_source(
+                    config::Environment::with_prefix("BERSERKER")
+                        .try_parsing(true)
+                        .separator("__"),
+                )
+                .build()
+                .unwrap()
+                .try_deserialize::<WorkloadConfig>()
+                .unwrap();
+
+            info!("Config: {:?}", config);
+            run_workload(config)
+        }
     };
 
     let processes = &handles.clone();
 
     thread::scope(|s| {
-        if config.duration != 0 {
-            // Spin a watcher thread
-            s.spawn(move || loop {
-                thread::sleep(time::Duration::from_secs(1));
-                let elapsed = duration_timer.elapsed().unwrap().as_secs();
-
-                if elapsed > config.duration {
-                    for handle in processes.iter().flatten() {
-                        info!("Terminating: {}", *handle);
-                        let _ = kill(Pid::from_raw(*handle), Signal::SIGTERM);
-                    }
+        // Spin a watcher thread
+        s.spawn(move || loop {
+            thread::sleep(time::Duration::from_secs(1));
+            let elapsed = duration_timer.elapsed().unwrap().as_secs();
+
+            // Find all processes with expired duration. If we've received
+            // SIGTERM, get all processes.
+            let expired = processes
+                .iter()
+                .filter(|(_, duration)| {
+                    (*duration > 0 && *duration < elapsed)
+                        || terminating.load(Ordering::Relaxed)
+                })
+                .collect::<Vec<_>>();
+
+            for (handle, _) in &expired {
+                info!("Terminating: {}", *handle);
+                let _ = kill(Pid::from_raw(*handle), Signal::SIGKILL);
+            }
 
-                    break;
-                }
-            });
-        }
+            if expired.len() == processes.len() {
+                break;
+            }
+        });
 
         s.spawn(move || {
-            for handle in processes.iter().flatten() {
-                info!("waitpid: {}", *handle);
-                match waitpid(Pid::from_raw(*handle), None) {
+            for (handle, _) in handles {
+                info!("waitpid: {}", handle);
+                match waitpid(Pid::from_raw(handle), None) {
                     Ok(_) => {
-                        info!("{:?} stopped", *handle)
+                        info!("{handle:?} stopped")
                     }
                     Err(Errno::ECHILD) => {
-                        info! {"no process {:?} found", *handle}
+                        info!("no process {handle:?} found")
                     }
                     Err(e) => {
-                        panic! {"cannot wait for {:?}: {:?} ", *handle, e}
+                        panic!("cannot wait for {handle:?}: {e:?}")
                     }
                 };
             }

src/main.rs

@@ -3,20 +3,21 @@ use std::collections::HashMap;
 #[derive(Debug, Clone)]
 pub enum Arg {
     Const { text: String },
-
     Var { name: String },
+    Dynamic { name: String, args: Vec<Arg> },
 }
 
 #[derive(Debug, Clone)]
 pub enum Instruction {
-    Task { name: String, args: Vec<Arg> },
-    Open { path: String },
-    Debug { text: String },
+    Task { name: Arg, args: Vec<Arg> },
+    Open { path: Arg },
+    Debug { text: Arg },
 }
 
 #[derive(Debug, Clone)]
 pub enum MachineInstruction {
     Server { port: u16 },
+    Profile { target: String },
 }
 
 #[derive(Debug, Clone)]