diff --git a/src/main.rs b/src/main.rs index 7914d1e..ec10aa4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -294,17 +294,37 @@ mod tests { } #[test] - fn test_task_script() { + fn test_task_script_random_arg() { + let input = r#" + main () { + task(stub, random_string()); + } + "#; + + let nodes: Vec = parse_instructions(input).unwrap(); + assert_eq!(nodes.len(), 1); + let prepared_nodes = apply_rules(nodes); + + new_script_worker(prepared_nodes[0].clone()) + .run_payload() + .unwrap(); + } + + #[test] + fn test_task_script_no_arg() { let input = r#" main () { task(stub); } "#; - let ast: Vec = parse_instructions(input).unwrap(); - assert_eq!(ast.len(), 1); + let nodes: Vec = parse_instructions(input).unwrap(); + assert_eq!(nodes.len(), 1); + let prepared_nodes = apply_rules(nodes); - new_script_worker(ast[0].clone()).run_payload().unwrap(); + new_script_worker(prepared_nodes[0].clone()) + .run_payload() + .unwrap(); } #[test] diff --git a/src/script/ast.rs b/src/script/ast.rs index e72cd2d..69c7a83 100644 --- a/src/script/ast.rs +++ b/src/script/ast.rs @@ -2,6 +2,9 @@ use std::collections::HashMap; #[derive(Debug, Clone, PartialEq)] pub enum Arg { + /// Null constant + Null {}, + /// Simple constant Const { text: String }, @@ -14,9 +17,16 @@ pub enum Arg { #[derive(Debug, Clone, PartialEq)] pub enum Instruction { + /// Execute a binary with specified name and arguments Task { name: Arg, args: Vec }, + + /// Open a file at specified path Open { path: Arg }, + + /// Print a debugging message (subject to configured log level) Debug { text: Arg }, + + /// Send a message to a server at specified address Ping { server: Arg }, } diff --git a/src/script/grammar.peg b/src/script/grammar.peg index 41ff05b..51e02fa 100644 --- a/src/script/grammar.peg +++ b/src/script/grammar.peg @@ -27,7 +27,7 @@ arg = { } args = {"(" ~ (arg ~ ("," ~ arg)* ~ ","?)? ~ ")"} -value = {(ASCII_ALPHANUMERIC| "." | " " | "/" | ":")*} +value = {(ASCII_ALPHANUMERIC | "." | " " | "/" | ":" | "_")*} param = {ident ~ "=" ~ value} params = {"(" ~ (param ~ ("," ~ param)* ~ ","?)? ~ ")"} diff --git a/src/script/rules.rs b/src/script/rules.rs index f8d1608..baccec5 100644 --- a/src/script/rules.rs +++ b/src/script/rules.rs @@ -1,13 +1,39 @@ use log::debug; -use crate::script::ast::{Instruction, Node}; +use crate::script::ast::{Arg, Instruction, Node}; use std::collections::HashMap; -fn apply_instruction_rules( +/// Apply following transformations to task instruction: +/// * If no arguments provided for the task, add an empty argument +fn apply_task_rules(task: &Instruction, _node: &Node) -> Instruction { + let Instruction::Task { name, args } = task else { + unreachable!() + }; + + let new_args = if args.is_empty() { + vec![Arg::Null {}] + } else { + args.to_vec() + }; + + Instruction::Task { + name: name.clone(), + args: new_args, + } +} + +fn apply_instructions_rules( instructions: &[Instruction], - _node: &Node, + node: &Node, ) -> Vec { - instructions.to_vec() + instructions + .iter() + .map(|i| match i { + Instruction::Task { .. } => apply_task_rules(i, node), + _ => i.clone(), + }) + .collect::>() + .to_vec() } fn apply_arg_rules( @@ -43,7 +69,7 @@ fn apply_work_rules(work: Node) -> Node { Node::Work { name: name.clone(), args: apply_arg_rules(args, &work), - instructions: apply_instruction_rules(instructions, &work), + instructions: apply_instructions_rules(instructions, &work), dist: dist.clone(), } } diff --git a/src/worker/script.rs b/src/worker/script.rs index 4020f52..ce20ae8 100644 --- a/src/worker/script.rs +++ b/src/worker/script.rs @@ -50,6 +50,7 @@ struct BuildContext<'a> { ee: LLVMExecutionEngineRef, builder: LLVMBuilderRef, module: LLVMModuleRef, + context: LLVMContextRef, module_state: &'a HashMap, module_runtime: &'a HashMap, } @@ -116,23 +117,45 @@ pub unsafe extern "C" fn ping(addr: *const i8) -> u64 { /// The caller must ensure the pointer is valid and points to a null /// terminated C-string. #[unsafe(no_mangle)] -pub unsafe extern "C" fn task(name: *const i8, random: bool) -> u64 { +pub unsafe extern "C" fn task(name: *const i8, args: *const i8) -> u64 { let name = unsafe { CStr::from_ptr(name) }; - debug!("Task {:?} {:?}", name, random); - let uniq_arg: String = rand::thread_rng() + let mut task = Command::new(name.to_str().unwrap()); + + if !args.is_null() { + let args = unsafe { CStr::from_ptr(args) }; + debug!("Task {:?} {:?}", name, args); + + task.args(args.to_str().unwrap().split(' ')); + } else { + debug!("Task {:?}, null", name); + } + + let status = task.status().expect("Failed to execute task"); + + status.code().unwrap_or(0).try_into().unwrap() +} + +thread_local! { + static POINTERS: RefCell> = const { RefCell::new(vec![]) }; +} + +/// Return a randomly generated string. +/// +/// # Safety +/// The caller must ensure the pointer is valid and points to a null +/// terminated C-string. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn random_string() -> *const i8 { + let rand: String = rand::thread_rng() .sample_iter(&Alphanumeric) .take(7) .map(char::from) .collect(); - let _res = Command::new(name.to_str().unwrap()) - .arg(uniq_arg) - .output() - .unwrap(); - 0 -} -thread_local! { - static POINTERS: RefCell> = const { RefCell::new(vec![]) }; + let result = CString::new(rand).unwrap().into_raw(); + + POINTERS.with(|ps| ps.borrow_mut().push(result)); + result } /// Return a randomly generated path. @@ -187,8 +210,8 @@ pub static RUNTIME: LazyLock> = "task".to_string(), RuntimeFunc { func: task as *const () as usize, - param_count: 1, - param_types: &[RuntimeType::Pointer], + param_count: 2, + param_types: &[RuntimeType::Pointer, RuntimeType::Pointer], return_type: RuntimeType::Int, }, ), @@ -229,6 +252,15 @@ pub static RUNTIME: LazyLock> = return_type: RuntimeType::Pointer, }, ), + ( + "random_string".to_string(), + RuntimeFunc { + func: random_string as *const () as usize, + param_count: 0, + param_types: &[], + return_type: RuntimeType::Pointer, + }, + ), // utils ( "cleanup".to_string(), @@ -243,8 +275,12 @@ pub static RUNTIME: LazyLock> = }); impl ScriptWorker { - fn jit_instruction(name: &CStr, arg: Arg, ctx: &BuildContext) { - let mut arg_ptr = Self::get_arg_value(arg, ctx); + fn jit_instruction(name: &CStr, args: Vec, ctx: &BuildContext) { + let (args_ref, args_len, args_cap) = args + .iter() + .map(|a| Self::get_arg_value(a.clone(), ctx)) + .collect::>() + .into_raw_parts(); let (func, func_type) = ctx .module_runtime @@ -256,16 +292,30 @@ impl ScriptWorker { ctx.builder, *func_type, *func, - &mut arg_ptr, - 1, + args_ref, + args_len.try_into().unwrap(), name.as_ptr() as *const _, ); + + let _ = Vec::from_raw_parts(args_ref, args_len, args_cap); } } fn get_arg_value(arg: Arg, ctx: &BuildContext) -> LLVMValueRef { match arg { + Arg::Null {} => unsafe { + let td = LLVMGetExecutionEngineTargetData(ctx.ee); + let iptr = LLVMIntPtrTypeInContext(ctx.context, td); + LLVMConstNull(iptr) + }, Arg::Const { text } => unsafe { + // The name of all constants created this way will be "const", + // which is ugly, but not a problem as LLVM modifies this to + // make sure uniqueness, i.e. they will be: + // + // @const, @const.1, @const.2, ... + // + // in the jited code. LLVMBuildGlobalString( ctx.builder, format!("{text}\0").as_ptr() as *const _, @@ -285,12 +335,11 @@ impl ScriptWorker { .get(&name) .expect("No dynamic variable in the static runtime"); - let text = match &args[0] { - Arg::Const { text } => text, - unknown => { - panic!("Unknown dynamic variable argument: {unknown:?}") - } - }; + let (args_ref, args_len, args_cap) = args + .iter() + .map(|a| Self::get_arg_value(a.clone(), ctx)) + .collect::>() + .into_raw_parts(); unsafe { trace!("Add mapping to {:?}", name); @@ -305,20 +354,17 @@ impl ScriptWorker { runtime_func.func as *mut c_void, ); - let mut helper_ptr = LLVMBuildGlobalString( - ctx.builder, - format!("{text}\0").as_ptr() as *const _, - c"const".as_ptr() as *const _, - ); - - LLVMBuildCall2( + let call = LLVMBuildCall2( ctx.builder, *func_type, *func, - &mut helper_ptr, - 1, + args_ref, + args_len.try_into().unwrap(), c"{name}".as_ptr() as *const _, - ) + ); + + let _ = Vec::from_raw_parts(args_ref, args_len, args_cap); + call } } } @@ -371,6 +417,7 @@ impl ScriptWorker { // get a type for main function let i64t = LLVMInt64TypeInContext(context); + let boolt = LLVMInt1TypeInContext(context); let iptr = LLVMIntPtrTypeInContext(context, td); // Insert runtime functions into the module @@ -440,6 +487,12 @@ impl ScriptWorker { ); module_state.insert(String::from("stub"), stub_ptr); + let true_ptr = LLVMConstInt(boolt, 1, 0); + module_state.insert(String::from("true"), true_ptr); + + let false_ptr = LLVMConstInt(boolt, 0, 0); + module_state.insert(String::from("false"), false_ptr); + let Node::Work { ref instructions, .. } = node @@ -453,6 +506,7 @@ impl ScriptWorker { ee, builder, module, + context, module_state: &module_state, module_runtime: &module_runtime, }; @@ -461,23 +515,26 @@ impl ScriptWorker { for instr in instructions { // JIT the instruction and collect it's name let name = match instr.clone() { - Instruction::Task { name, args: _ } => { - Self::jit_instruction(c"task", name, &ctx); + Instruction::Task { name, args } => { + let mut task_args = vec![name]; + task_args.extend_from_slice(&args); + + Self::jit_instruction(c"task", task_args, &ctx); "task" } Instruction::Open { path } => { - Self::jit_instruction(c"open", path, &ctx); + Self::jit_instruction(c"open", vec![path], &ctx); "open" } Instruction::Ping { server } => { - Self::jit_instruction(c"ping", server, &ctx); + Self::jit_instruction(c"ping", vec![server], &ctx); "ping" } Instruction::Debug { text } => { - Self::jit_instruction(c"debug", text, &ctx); + Self::jit_instruction(c"debug", vec![text], &ctx); "debug" } }; @@ -502,13 +559,7 @@ impl ScriptWorker { } // Final instruction to clear dangling pointers - Self::jit_instruction( - c"cleanup", - Arg::Const { - text: "".to_string(), - }, - &ctx, - ); + Self::jit_instruction(c"cleanup", vec![], &ctx); let module_func = LLVMGetNamedFunction( module, diff --git a/workloads/example.ber b/workloads/example.ber index ff6fd79..7b5a690 100644 --- a/workloads/example.ber +++ b/workloads/example.ber @@ -10,7 +10,7 @@ main (workers = 2, duration = 10) { // debug(text) -- log with DEBUG level // open(path) -- open file by path, create if needed and write something to it debug("run task stub"); - task(stub); + task(stub, random_string()); debug("open file /tmp/test"); open("/tmp/test"); } : exp { diff --git a/workloads/example.short.ber b/workloads/example.short.ber index 4095746..b4fcafa 100644 --- a/workloads/example.short.ber +++ b/workloads/example.short.ber @@ -5,7 +5,7 @@ machine { main (workers = 1) { debug("run task stub"); - task(stub); + task(stub, random_string()); debug("open file /tmp/test"); open("/tmp/test"); debug("ping server");