refactor(bake/baker): tidy up bake.rs

This commit is contained in:
Catty Steve
2026-06-09 23:36:58 +08:00
parent 9ff9073fce
commit 6d9ddbca56
5 changed files with 324 additions and 202 deletions
+20 -193
View File
@@ -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<String>, Vec<String>) = 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<Box<dyn Future<Output = Result<(), BakeError>> + 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
+38 -8
View File
@@ -38,7 +38,9 @@ pub enum Decorator {
/// Chain functions via pipe with comma-separated names.
Pipe(Vec<String>),
/// Run concurrently with other parallel functions.
Parallel,
/// `None` = wildcard ("*" group, compatible with any group).
/// `Some(name)` = restricted to the named group.
Parallel(Option<String>),
/// 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<Option<Decorator>,
}
}),
"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());
}
+196
View File
@@ -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<Box<dyn Future<Output = Result<(), BakeError>> + 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<String, Function>,
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<BakeError> = 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(())
}
+54 -1
View File
@@ -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<String, Function>,
) -> (Vec<String>, HashMap<String, Vec<String>>) {
let mut star_nodes = Vec::new();
let mut named_groups: HashMap<String, Vec<String>> = 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<String, Vec<String>>,
) -> Option<Vec<String>> {
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::*;
+16
View File
@@ -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))
}