refactor: centralize ExecutionContext and event handling in main

- Lift ctx creation and event_receiver spawning to main
- Inject ctx and event_tx into prebake/bake/finalize
- Add standalone flag to ExecutionContext
- Wire up finalize command in main
This commit is contained in:
Catty Steve
2026-04-21 11:38:02 +08:00
parent e9128adfa0
commit d8ba166a3b
7 changed files with 87 additions and 59 deletions
+5 -20
View File
@@ -4,7 +4,7 @@ use crate::constant::DEFAULT_WORKSPACE;
use crate::constant::TRIVIAL_SCRIPT_NAME;
use crate::engine::EventSender;
use crate::error::BakeError;
use crate::prebake::event;
use crate::ExecutionContext;
use crate::prebake::PrebakeConfig;
use crate::prebake::security::get_drop_after;
use crate::prebake::stage::PrebakeStage;
@@ -21,6 +21,8 @@ pub async fn bake(
bake_base_path: &Path,
prebake_path: &Path,
cli: &Cli,
ctx: &mut ExecutionContext,
event_tx: EventSender,
) -> Result<(), BakeError> {
let script_content = std::fs::read_to_string(script_path)?;
let prebake_content = std::fs::read_to_string(prebake_path)?;
@@ -41,26 +43,12 @@ pub async fn bake(
let use_template = true;
let engine = Engine::new();
let mut ctx = crate::ExecutionContext {
task_id: format!("{}-{}-bake", cli.pipeline, cli.build_id),
pipeline_name: cli.pipeline.clone(),
username: cli.username.clone(),
dry_run: cli.dry_run,
privileged: privileged(&prebake),
..Default::default()
};
if let Some(env) = &prebake.envvars {
if let Some(prebake) = &env.bake {
ctx.env_vars.extend(prebake.clone());
if let Some(prebake_env) = &env.bake {
ctx.env_vars.extend(prebake_env.clone());
}
}
let (tx, event_rx) = tokio::sync::mpsc::channel(256);
let event_tx: EventSender = Some(tx);
let pipeline_name = cli.pipeline.clone();
let build_id = cli.build_id.clone();
let receiver_handle = tokio::spawn(event::event_receiver(event_rx, pipeline_name, build_id));
if !has_pipeline || trivial {
let function = Function{
name: TRIVIAL_SCRIPT_NAME.to_string(),
@@ -92,9 +80,6 @@ pub async fn bake(
}
}
drop(event_tx);
let _ = receiver_handle.await;
Ok(())
}
+1
View File
@@ -456,6 +456,7 @@ mod tests {
shell: "bash".to_string(),
dry_run: false,
privileged: false,
standalone: false,
}
}
+3 -1
View File
@@ -5,7 +5,7 @@ use std::time::Duration;
use tokio::sync::mpsc;
pub use super::cgroups::CgroupManager;
pub use crate::error::{EXITCODE_EXECUTION_ERROR, EXITCODE_IO_ERROR, ExecutionError};
pub use crate::error::{ExecutionError, EXITCODE_EXECUTION_ERROR, EXITCODE_IO_ERROR};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceUsage {
@@ -27,6 +27,7 @@ pub struct ExecutionContext {
pub shell: String,
pub dry_run: bool,
pub privileged: bool,
pub standalone: bool,
}
#[derive(Debug, Clone)]
@@ -49,6 +50,7 @@ impl Default for ExecutionContext {
shell: String::from("bash"),
dry_run: false,
privileged: false,
standalone: false,
}
}
}
+3 -15
View File
@@ -6,6 +6,7 @@ pub mod template;
pub mod event;
use crate::cli::Cli;
use crate::engine::EventSender;
use std::path::Path;
pub use config::*;
pub use error::FinalizeError;
@@ -19,8 +20,9 @@ fn parse(path: &Path) -> Result<FinalizeConfig, FinalizeError> {
pub async fn finalize(
finalize_path: &Path,
cli: &Cli,
ctx: &mut crate::ExecutionContext,
event_tx: EventSender,
) -> Result<(), FinalizeError> {
// Initialize
let finalize = parse(finalize_path)?;
let validate_result = finalize.validate();
if let Err(errors) = validate_result {
@@ -30,20 +32,6 @@ pub async fn finalize(
return Err(FinalizeError::ValidateError(errors));
};
let mut ctx = crate::ExecutionContext {
pipeline_name: cli.pipeline.clone(),
task_id: format!("{}-{}-finalize", cli.pipeline, cli.build_id),
username: cli.username.clone(),
dry_run: cli.dry_run,
privileged: false, // really?
..Default::default()
};
let (tx, event_rx) = tokio::sync::mpsc::channel(256);
let event_tx: crate::engine::EventSender = Some(tx);
let pipeline_name = cli.pipeline.clone();
let build_id = cli.build_id.clone();
let receiver_handle = tokio::spawn(event::event_receiver(event_rx, pipeline_name, build_id));
// Register Plugins
todo!();
}
+71 -5
View File
@@ -1,10 +1,14 @@
use std::path::Path;
use env_logger;
use workshop_baker::{
cli::{Cli, Commands, Parser},
daemon,
engine::{EventSender, ExecutionContext},
error::PrebakeError,
prebake::prebake,
bake::bake,
finalize::finalize,
};
#[tokio::main]
@@ -18,10 +22,30 @@ async fn main() {
Some(Commands::Monitor { group: _ }) => {
unimplemented!("executor-monitor should be moved to daemon instead.")
}
Some(Commands::Client { config:_ }) => {}
Some(Commands::Client { config:_ }) => {
unimplemented!("Client mode not yet implemented.")
}
Some(Commands::Prebake { config }) => {
log::info!("Running in standalone mode.");
let result = prebake(config, &cli, None).await;
let mut ctx = ExecutionContext {
standalone: true,
task_id: format!("{}-{}-init", cli.pipeline, cli.build_id),
pipeline_name: cli.pipeline.clone(),
username: cli.username.clone(),
dry_run: cli.dry_run,
privileged: true,
..Default::default()
};
let (tx, event_rx) = tokio::sync::mpsc::channel(256);
let event_tx: EventSender = Some(tx);
let pipeline_name = cli.pipeline.clone();
let build_id = cli.build_id.clone();
let receiver_handle = tokio::spawn(
workshop_baker::prebake::event::event_receiver(event_rx, pipeline_name, build_id)
);
let result = prebake(config, &cli, None, &mut ctx, event_tx).await;
drop(ctx);
let _ = receiver_handle.await;
dbg!(&result);
if let Err(ref e) = result {
log::error!("Prebake stage failed: {}", e);
@@ -36,15 +60,57 @@ async fn main() {
script, bake_base, prebake
}) => {
log::info!("Running in standalone mode.");
let result = bake(script, bake_base, prebake, &cli).await;
let mut ctx = ExecutionContext {
standalone: true,
task_id: format!("{}-{}-bake", cli.pipeline, cli.build_id),
pipeline_name: cli.pipeline.clone(),
username: cli.username.clone(),
dry_run: cli.dry_run,
privileged: false,
..Default::default()
};
let (tx, event_rx) = tokio::sync::mpsc::channel(256);
let event_tx: EventSender = Some(tx);
let pipeline_name = cli.pipeline.clone();
let build_id = cli.build_id.clone();
let receiver_handle = tokio::spawn(
workshop_baker::prebake::event::event_receiver(event_rx, pipeline_name, build_id)
);
let result = bake(script, bake_base, prebake, &cli, &mut ctx, event_tx).await;
drop(ctx);
let _ = receiver_handle.await;
dbg!(&result);
if result.is_err() {
log::error!("Bake stage failed! Exiting...");
std::process::exit(1);
}
}
Some(Commands::Finalize { config: _ }) => {
unimplemented!("finalize not implemented.")
Some(Commands::Finalize { config }) => {
log::info!("Running in standalone mode.");
let mut ctx = ExecutionContext {
standalone: true,
task_id: format!("{}-{}-finalize", cli.pipeline, cli.build_id),
pipeline_name: cli.pipeline.clone(),
username: cli.username.clone(),
dry_run: cli.dry_run,
privileged: false,
..Default::default()
};
let (tx, event_rx) = tokio::sync::mpsc::channel(256);
let event_tx: EventSender = Some(tx);
let pipeline_name = cli.pipeline.clone();
let build_id = cli.build_id.clone();
let receiver_handle = tokio::spawn(
workshop_baker::prebake::event::event_receiver(event_rx, pipeline_name, build_id)
);
let result = finalize(Path::new(config), &cli, &mut ctx, event_tx).await;
drop(ctx);
let _ = receiver_handle.await;
dbg!(&result);
if let Err(ref e) = result {
log::error!("Finalize stage failed: {}", e);
std::process::exit(1);
}
}
None => {
// Running in daemon mode
+3 -17
View File
@@ -30,6 +30,7 @@ use crate::{
cli::Cli,
engine::EventSender,
error::PrebakeError,
ExecutionContext,
prebake::{security::get_drop_after, stage::PrebakeStage},
};
use std::path::Path;
@@ -88,6 +89,8 @@ pub async fn prebake(
prebake_path: &Path,
cli: &Cli,
stage: Option<PrebakeStage>,
mut ctx: &mut ExecutionContext,
event_tx: EventSender,
) -> anyhow::Result<()> {
// Initialize
let prebake_content = std::fs::read_to_string(prebake_path).inspect_err(|e| {
@@ -109,14 +112,6 @@ pub async fn prebake(
PrebakeError::ConfigError(e)
})?;
let target_user = "vulcan";
let mut ctx = crate::ExecutionContext {
task_id: format!("{}-{}-init", cli.pipeline, cli.build_id),
pipeline_name: cli.pipeline.clone(),
username: cli.username.clone(),
dry_run: cli.dry_run,
privileged: true,
..Default::default()
};
// Wait for bakerd to finish Fetch-1
let stage = stage.unwrap_or(PrebakeStage::Bootstrap);
if let Some(env) = &prebake.envvars
@@ -125,12 +120,6 @@ pub async fn prebake(
ctx.env_vars.extend(prebake.clone());
}
let (tx, event_rx) = tokio::sync::mpsc::channel(256);
let event_tx: EventSender = Some(tx);
let pipeline_name = cli.pipeline.clone();
let build_id = cli.build_id.clone();
let receiver_handle = tokio::spawn(event::event_receiver(event_rx, pipeline_name, build_id));
// Bootstrap stage
if stage <= PrebakeStage::Bootstrap {
@@ -244,9 +233,6 @@ pub async fn prebake(
// Ready stage
ctx.task_id = format!("{}-{}-ready", cli.pipeline, cli.build_id);
drop(event_tx);
let _ = receiver_handle.await;
log::info!("Prebake done!");
Ok(())
}
+1 -1
View File
@@ -1,4 +1,4 @@
use serde::{de::Visitor, Deserialize, Deserializer};
use serde::{de::Visitor, Deserializer};
pub const INT_MAX_MS: u64 = i32::MAX as u64;