diff --git a/workshop-baker/src/bake.rs b/workshop-baker/src/bake.rs index e49c2f0..73bf36d 100644 --- a/workshop-baker/src/bake.rs +++ b/workshop-baker/src/bake.rs @@ -1,32 +1,22 @@ use crate::bake::builder::RenderMode; use crate::bake::constant::DEFAULT_WORKSPACE; use crate::bake::error::BakeError; -use crate::bake::parser::Function; +use crate::bake::util::is_parallel; use crate::cli::Cli; use crate::prebake; use crate::prebake::PrebakeConfig; -use crate::prebake::security::get_drop_after; -use crate::prebake::stage::PrebakeStage; -use std::future::Future; use std::path::{Path, PathBuf}; -use std::pin::Pin; -use std::time::Duration; -use tokio::sync::mpsc; -use workshop_engine::{Engine, EventSender, ExecutionContext, ExecutionError}; +use workshop_engine::{Engine, EventSender, ExecutionContext}; mod builder; pub mod constant; mod decorator; pub mod error; +mod execute; pub mod parser; mod schedule; pub mod trivial; - -/// Completion event sent from a spawned task to the batch consumer loop. -struct TaskCompletion { - name: String, - result: Result<(), BakeError>, -} +pub mod util; pub async fn bake( script_path: &Path, @@ -83,20 +73,11 @@ pub async fn bake( return Err(BakeError::CircularDependency(remaining_nodes)); } - // Separate @parallel from sequential (consume ready for owned names) - let (sequential, parallel): (Vec, Vec) = ready.into_iter() - .partition(|name| { - let f = &func_map[name.as_str()]; - !f.decorators - .iter() - .any(|d| matches!(d, decorator::Decorator::Parallel)) - }); - - // Execute sequential functions first - for name in &sequential { - execute_one( + // Rule 1: sequential always goes first + if let Some(seq_name) = ready.iter().find(|n| !is_parallel(&func_map[n.as_str()])).cloned() { + execute::execute_one( &engine, - &func_map[name.as_str()], + &func_map[seq_name.as_str()], ctx, &event_tx, &workspace, @@ -104,54 +85,23 @@ pub async fn bake( bake_base_path, ) .await?; + dag.pop(); + continue; } - // Execute @parallel functions concurrently - if !parallel.is_empty() { - let handles: Vec<_> = parallel - .iter() - .map(|name| { - let func = func_map[name.as_str()].clone(); - let mut ctx_local = ctx.clone(); - let et = event_tx.clone(); - let ws = workspace.clone(); - let rem = remaining.clone(); - let bb = bake_base_path.to_path_buf(); - let n = name.clone(); - tokio::spawn(async move { - let eng = Engine::new(); - ( - n, - execute_one(&eng, &func, &mut ctx_local, &et, &ws, &rem, &bb) - .await, - ) - }) - }) - .collect(); - - for handle in handles { - let (name, result) = handle.await.map_err(|_| { - BakeError::ScriptExecutionError(ExecutionError::ExecutionFailed( - "parallel task panicked".to_string(), - )) - })?; - let func = &func_map[name.as_str()]; - let is_fallible = func - .decorators - .iter() - .any(|d| matches!(d, decorator::Decorator::Fallible)); - match result { - Err(e) if !is_fallible => return Err(e), - Err(e) => { - log::warn!("'{}' failed (fallible, non-fatal): {}", name, e); - } - Ok(_) => {} - } + // All remaining are parallel — group and pick a batch + let (star_nodes, groups) = schedule::partition_by_group(&ready, &func_map); + if let Some(batch) = schedule::pick_batch(&star_nodes, &groups) { + execute::fire_parallel_batch( + &batch, &engine, &func_map, ctx, &event_tx, + &workspace, &remaining, bake_base_path, + ).await?; + for _name in &batch { + dag.pop(); } } - // Commit the entire batch — all nodes completed (or fallible-failed) - dag.pop_all(); + // No complete group and no sequential — nothing fireable } } @@ -173,128 +123,5 @@ fn get_workspace(prebake_config: &PrebakeConfig) -> PathBuf { .unwrap_or_else(|| PathBuf::from(DEFAULT_WORKSPACE)) } -fn privileged(prebake_config: &PrebakeConfig) -> bool { - if let Some(bootstrap) = &prebake_config.bootstrap - && bootstrap.user == "root" - { - return true; - } - let drop_after = get_drop_after(&prebake_config.security).unwrap_or_default(); - PrebakeStage::Never == drop_after -} - -// ── Execution helpers ── - -async fn execute_one( - engine: &Engine, - func: &Function, - ctx: &mut ExecutionContext, - event_tx: &EventSender, - workspace: &Path, - remaining: &str, - bake_base: &Path, -) -> Result<(), BakeError> { - // Extract combinators preserving original order - let combos: Vec<&decorator::Decorator> = func - .decorators - .iter() - .filter(|d| decorator::is_combinator(d)) - .collect(); - - // Build the script - let mut modified = func.clone(); - modified.body = remaining.to_string() + &func.body; - let script = builder::build_script( - &modified, bake_base, workspace, None, RenderMode::Full, - )?; - - // Execute with combinator chain - execute_with_combos(engine, &script, ctx, event_tx, &combos).await -} - -async fn execute_with_combos( - engine: &Engine, - script: &PathBuf, - ctx: &ExecutionContext, - event_tx: &EventSender, - combos: &[&decorator::Decorator], -) -> Result<(), BakeError> { - build_combinator_chain(engine, script, ctx, event_tx, combos, 0).await -} - -/// Builds a combinator execution chain using `Box::pin` to avoid infinite-sized async fn recursion. -/// Each combinator wraps the inner future: processed outermost (idx=0) → innermost (baseline). -fn build_combinator_chain<'a>( - engine: &'a Engine, - script: &'a PathBuf, - ctx: &'a ExecutionContext, - event_tx: &'a EventSender, - combos: &'a [&'a decorator::Decorator], - idx: usize, -) -> Pin> + Send + 'a>> { - use decorator::Decorator; - - // Baseline: no more combinators — execute the script once - if idx >= combos.len() { - return Box::pin(async move { - engine.execute_script(script, ctx, event_tx).await?; - Ok(()) - }); - } - - match combos[idx] { - Decorator::Loop(count) => Box::pin(async move { - for i in 0..*count { - log::debug!("Loop {}/{}", i + 1, count); - build_combinator_chain(engine, script, ctx, event_tx, combos, idx + 1).await?; - } - Ok(()) - }), - Decorator::Timeout(secs) => Box::pin(async move { - tokio::time::timeout( - Duration::from_secs(*secs as u64), - build_combinator_chain(engine, script, ctx, event_tx, combos, idx + 1), - ) - .await - .map_err(|_| BakeError::ScriptExecutionError(ExecutionError::Timeout))? - }), - Decorator::Retry(retries, delay_secs) => Box::pin(async move { - let max_attempts = *retries as usize + 1; - let delay = Duration::from_secs(*delay_secs as u64); - let mut last_err = None; - for attempt in 0..max_attempts { - match build_combinator_chain( - engine, script, ctx, event_tx, combos, idx + 1, - ) - .await - { - Ok(_) => return Ok(()), - Err(e) => { - // Timeout should not be retried - if matches!(&e, BakeError::ScriptExecutionError(ExecutionError::Timeout)) { - return Err(e); - } - if attempt < max_attempts - 1 { - last_err = Some(e); - log::warn!( - "Retry {}/{} after {}s", - attempt + 1, - retries, - delay_secs - ); - tokio::time::sleep(delay).await; - } else { - return Err(e); - } - } - } - } - Err(last_err.unwrap()) - }), - // Non-combinator decorators: skip (shouldn't reach here if filtered correctly) - _ => build_combinator_chain(engine, script, ctx, event_tx, combos, idx + 1), - } -} - // ── AI Coding Agent marker ── // Code in this module PARTIALLY or FULLY utilized AI Coding Agent diff --git a/workshop-baker/src/bake/decorator.rs b/workshop-baker/src/bake/decorator.rs index 4d75cec..656f93e 100644 --- a/workshop-baker/src/bake/decorator.rs +++ b/workshop-baker/src/bake/decorator.rs @@ -38,7 +38,9 @@ pub enum Decorator { /// Chain functions via pipe with comma-separated names. Pipe(Vec), /// Run concurrently with other parallel functions. - Parallel, + /// `None` = wildcard ("*" group, compatible with any group). + /// `Some(name)` = restricted to the named group. + Parallel(Option), /// Run as a background daemon process. Daemon, /// Health check endpoint for daemon verification. @@ -54,7 +56,7 @@ pub fn is_flag(d: &Decorator) -> bool { | Decorator::Export | Decorator::If(_) | Decorator::After(_) - | Decorator::Parallel + | Decorator::Parallel(_) | Decorator::Health(_) | Decorator::Unknown ) @@ -70,6 +72,14 @@ pub fn is_mode(d: &Decorator) -> bool { matches!(d, Decorator::Pipe(_) | Decorator::Daemon) } +/// Extract the parallel group name. Returns "*" for wildcard (no group specified). +pub fn parallel_group(d: &Decorator) -> &str { + match d { + Decorator::Parallel(Some(g)) => g.as_str(), + _ => "*", + } +} + impl Display for Decorator { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { @@ -83,7 +93,10 @@ impl Display for Decorator { Decorator::Retry(count, delay) => write!(f, "@retry({}, {})", count, delay), Decorator::Export => write!(f, "@export"), Decorator::Pipe(funcs) => write!(f, "@pipe({})", funcs.join(", ")), - Decorator::Parallel => write!(f, "@parallel"), + Decorator::Parallel(group) => match group { + None => write!(f, "@parallel"), + Some(g) => write!(f, "@parallel({})", g), + }, Decorator::Daemon => write!(f, "@daemon"), Decorator::Health(check) => write!(f, "@health({})", check), } @@ -185,7 +198,14 @@ pub fn parse_decorator(line: &str, line_num: usize) -> Result, } }), "export" => decorator_noarg(Decorator::Export), - "parallel" => decorator_noarg(Decorator::Parallel), + "parallel" => { + if has_arg { + let group = required_arg()?; + Ok(Some(Decorator::Parallel(Some(group.to_string())))) + } else { + Ok(Some(Decorator::Parallel(None))) + } + } "daemon" => decorator_noarg(Decorator::Daemon), "health" => required_arg().map(|a| Some(Decorator::Health(a.to_string()))), "pipe" => required_arg().and_then(|a| { @@ -314,7 +334,11 @@ mod tests { #[test] fn test_decorator_display_parallel() { - assert_eq!(format!("{}", Decorator::Parallel), "@parallel"); + assert_eq!(format!("{}", Decorator::Parallel(None)), "@parallel"); + assert_eq!( + format!("{}", Decorator::Parallel(Some("build".to_string()))), + "@parallel(build)" + ); } #[test] @@ -364,7 +388,13 @@ mod tests { #[test] fn test_parse_decorator_parallel() { let result = parse_decorator("# @parallel", 1).unwrap(); - assert_eq!(result, Some(Decorator::Parallel)); + assert_eq!(result, Some(Decorator::Parallel(None))); + } + + #[test] + fn test_parse_decorator_parallel_with_group() { + let result = parse_decorator("# @parallel(build)", 1).unwrap(); + assert_eq!(result, Some(Decorator::Parallel(Some("build".to_string())))); } #[test] @@ -600,8 +630,8 @@ mod tests { } #[test] - fn test_parse_decorator_parallel_with_arg() { - let result = parse_decorator("# @parallel(true)", 1); + fn test_parse_decorator_parallel_empty_args() { + let result = parse_decorator("# @parallel()", 1); assert!(result.is_err()); } diff --git a/workshop-baker/src/bake/execute.rs b/workshop-baker/src/bake/execute.rs new file mode 100644 index 0000000..309e68b --- /dev/null +++ b/workshop-baker/src/bake/execute.rs @@ -0,0 +1,196 @@ +use crate::bake::builder::{self, RenderMode}; +use crate::bake::decorator::{self, Decorator}; +use crate::bake::error::BakeError; +use crate::bake::parser::Function; +use crate::bake::util::is_fallible; +use std::collections::HashMap; +use std::future::Future; +use std::path::{Path, PathBuf}; +use std::pin::Pin; +use std::time::Duration; +use workshop_engine::{Engine, EventSender, ExecutionContext, ExecutionError}; + +/// Execute a single function with combinator wrapping. +pub async fn execute_one( + engine: &Engine, + func: &Function, + ctx: &mut ExecutionContext, + event_tx: &EventSender, + workspace: &Path, + remaining: &str, + bake_base: &Path, +) -> Result<(), BakeError> { + let combos: Vec<&Decorator> = func + .decorators + .iter() + .filter(|d| decorator::is_combinator(d)) + .collect(); + + let mut modified = func.clone(); + modified.body = remaining.to_string() + &func.body; + let script = builder::build_script( + &modified, bake_base, workspace, None, RenderMode::Full, + )?; + + execute_with_combos(engine, &script, ctx, event_tx, &combos).await +} + +async fn execute_with_combos( + engine: &Engine, + script: &PathBuf, + ctx: &ExecutionContext, + event_tx: &EventSender, + combos: &[&Decorator], +) -> Result<(), BakeError> { + build_combinator_chain(engine, script, ctx, event_tx, combos, 0).await +} + +/// Builds a combinator execution chain using `Box::pin` to avoid +/// infinite-sized async fn recursion. +fn build_combinator_chain<'a>( + engine: &'a Engine, + script: &'a PathBuf, + ctx: &'a ExecutionContext, + event_tx: &'a EventSender, + combos: &'a [&'a Decorator], + idx: usize, +) -> Pin> + Send + 'a>> { + if idx >= combos.len() { + return Box::pin(async move { + engine.execute_script(script, ctx, event_tx).await?; + Ok(()) + }); + } + + match combos[idx] { + Decorator::Loop(count) => Box::pin(async move { + for i in 0..*count { + log::debug!("Loop {}/{}", i + 1, count); + build_combinator_chain( + engine, script, ctx, event_tx, combos, idx + 1, + ) + .await?; + } + Ok(()) + }), + Decorator::Timeout(secs) => Box::pin(async move { + tokio::time::timeout( + Duration::from_secs(*secs as u64), + build_combinator_chain( + engine, script, ctx, event_tx, combos, idx + 1, + ), + ) + .await + .map_err(|_| { + BakeError::ScriptExecutionError(ExecutionError::Timeout) + })? + }), + Decorator::Retry(retries, delay_secs) => Box::pin(async move { + let max_attempts = *retries as usize + 1; + let delay = Duration::from_secs(*delay_secs as u64); + let mut last_err = None; + for attempt in 0..max_attempts { + match build_combinator_chain( + engine, script, ctx, event_tx, combos, idx + 1, + ) + .await + { + Ok(()) => return Ok(()), + Err(e) => { + if matches!( + &e, + BakeError::ScriptExecutionError(ExecutionError::Timeout) + ) { + return Err(e); + } + if attempt < max_attempts - 1 { + last_err = Some(e); + log::warn!( + "Retry {}/{} after {}s", + attempt + 1, + retries, + delay_secs + ); + tokio::time::sleep(delay).await; + } else { + return Err(e); + } + } + } + } + Err(last_err.unwrap()) + }), + _ => build_combinator_chain( + engine, script, ctx, event_tx, combos, idx + 1, + ), + } +} + +/// Fire a batch of parallel nodes concurrently. +/// +/// All nodes in the batch are waited on before returning — matching the +/// GitHub/GitLab/Concourse convention of collecting full failure information. +/// A non-`@fallible` failure is surfaced only after all nodes complete. +pub async fn fire_parallel_batch( + batch: &[String], + _engine: &Engine, + func_map: &HashMap, + ctx: &mut ExecutionContext, + event_tx: &EventSender, + workspace: &Path, + remaining: &str, + bake_base: &Path, +) -> Result<(), BakeError> { + let handles: Vec<_> = batch + .iter() + .map(|name| { + let function = func_map[name.as_str()].clone(); + let mut task_ctx = ctx.clone(); + let task_events = event_tx.clone(); + let task_workspace = workspace.to_path_buf(); + let task_remaining = remaining.to_string(); + let task_bake_base = bake_base.to_path_buf(); + let task_name = name.clone(); + tokio::spawn(async move { + let task_engine = Engine::new(); + ( + task_name, + execute_one( + &task_engine, + &function, + &mut task_ctx, + &task_events, + &task_workspace, + &task_remaining, + &task_bake_base, + ) + .await, + ) + }) + }) + .collect(); + + let mut first_fatal: Option = None; + for handle in handles { + let (name, result) = handle.await.map_err(|_| { + BakeError::ScriptExecutionError(ExecutionError::ExecutionFailed( + "parallel task panicked".to_string(), + )) + })?; + match result { + Ok(()) => {} + Err(ref err) if is_fallible(&func_map[&name]) => { + log::warn!("'{}' failed (fallible, non-fatal): {}", name, err); + } + Err(err) => { + log::error!("'{}' failed (fatal): {}", name, err); + first_fatal.get_or_insert(err); + } + } + } + + if let Some(err) = first_fatal { + return Err(err); + } + Ok(()) +} diff --git a/workshop-baker/src/bake/schedule.rs b/workshop-baker/src/bake/schedule.rs index 9de16b5..2b4b2d8 100644 --- a/workshop-baker/src/bake/schedule.rs +++ b/workshop-baker/src/bake/schedule.rs @@ -1,4 +1,4 @@ -use super::{decorator::Decorator, parser::Function}; +use super::{decorator::Decorator, decorator::parallel_group, parser::Function}; use crate::bake::error::BakeError; use std::collections::HashMap; use topological_sort::TopologicalSort; @@ -68,6 +68,59 @@ pub fn build_dag( Ok((ts, owned_map)) } +/// Group parallel nodes by their declared group. +/// Returns `(star_nodes, named_groups)`. +/// Star nodes (`group == "*"`) are compatible with any group. +pub fn partition_by_group( + names: &[String], + func_map: &HashMap, +) -> (Vec, HashMap>) { + let mut star_nodes = Vec::new(); + let mut named_groups: HashMap> = HashMap::new(); + for name in names { + let func = &func_map[name.as_str()]; + let group = func + .decorators + .iter() + .find_map(|d| { + if let Decorator::Parallel(_) = d { + Some(parallel_group(d)) + } else { + None + } + }) + .unwrap_or("*"); + if group == "*" { + star_nodes.push(name.clone()); + } else { + named_groups + .entry(group.to_string()) + .or_default() + .push(name.clone()); + } + } + (star_nodes, named_groups) +} + +/// Pick a complete group to execute. +/// - If any named group has ready nodes, pick the first (star nodes join). +/// - If only star nodes exist, fire them together. +/// - Otherwise, return None. +pub fn pick_batch( + star_nodes: &[String], + groups: &HashMap>, +) -> Option> { + for nodes in groups.values() { + let mut batch = nodes.clone(); + batch.extend_from_slice(star_nodes); + return Some(batch); + } + if !star_nodes.is_empty() && groups.is_empty() { + return Some(star_nodes.to_vec()); + } + None +} + #[cfg(test)] mod tests { use super::*; diff --git a/workshop-baker/src/bake/util.rs b/workshop-baker/src/bake/util.rs new file mode 100644 index 0000000..e5b0323 --- /dev/null +++ b/workshop-baker/src/bake/util.rs @@ -0,0 +1,16 @@ +use crate::bake::decorator::Decorator; +use crate::bake::parser::Function; + +/// True if the function is decorated with `@parallel`. +pub fn is_parallel(func: &Function) -> bool { + func.decorators + .iter() + .any(|d| matches!(d, Decorator::Parallel(_))) +} + +/// True if the function is decorated with `@fallible`. +pub fn is_fallible(func: &Function) -> bool { + func.decorators + .iter() + .any(|d| matches!(d, Decorator::Fallible)) +}