honey-biscuit-workshop/workshop-baker/src/main.rs

use capctl;
use workshop_baker::bare::run_bare;
use workshop_baker::notify::queue::NotificationQueue;
use workshop_baker::types::debug::DebugFeature;
use workshop_baker::types::resource::ResourceRegistry;
use workshop_baker::{
    bake::bake,
    cli::{Cli, Commands, Parser},
    error::CliError,
    finalize::finalize,
    notify::notify,
    prebake::prebake,
};
use workshop_engine::{EventSender, ExecutionContext};

async fn worker_main(cli: Cli) -> Result<(), CliError> {
    // a magic closure hack
    let require = |name: &str, val: &Option<String>| -> String {
        val.clone().unwrap_or_else(|| {
            log::error!("--{} is required", name);
            std::process::exit(1);
        })
    };
    // worker depends on these information to work
    let pipeline = require("pipeline", &cli.pipeline);
    let build_id = require("build-id", &cli.build_id);
    let username = require("username", &cli.username);

    println!("── worker standalone mode ───────────────────────");
    println!("  {:<14} {:?}", "prebake:", cli.prebake);
    println!("  {:<14} {:?}", "bake:", cli.bake);
    println!("  {:<14} {:?}", "bake-base:", cli.bake_base);
    println!("  {:<14} {:?}", "finalize:", cli.finalize);
    println!("──────────────────────────────────────────────────");

    let debug_flags = DebugFeature::parse(cli.debug.as_deref());
    if !debug_flags.is_empty() {
        log::info!("Debug features enabled: {:?}", debug_flags);
    }

    let mut ctx = ExecutionContext {
        standalone: true,
        task_id: format!("{}-{}-worker", pipeline, build_id),
        pipeline_name: pipeline.clone(),
        build_id: build_id.clone(),
        username: username.clone(),
        dry_run: cli.dry_run,
        privileged: false,
        debug_flags: debug_flags.bits(),
        ..Default::default()
    };

    let (tx, event_rx) = tokio::sync::mpsc::channel(256);
    let event_tx: EventSender = Some(tx);
    let _event_handle = tokio::spawn(workshop_baker::prebake::event::event_receiver(
        event_rx,
        pipeline.clone(),
        build_id.clone(),
        ctx.debug_flags,
    ));

    let (to_tx, to_rx) = tokio::sync::mpsc::channel(256);
    let (retry_tx, retry_rx) = tokio::sync::mpsc::channel(256);

    // Resource registry — fetch plugins/templates before pipeline stages
    let mut registry = ResourceRegistry::new();
    if let Ok(finalize_cfg) = workshop_baker::finalize::parse(&cli.finalize) {
        let _ = workshop_baker::prefetch::prefetch_resources(&finalize_cfg, &mut registry).await;
    }
    let registry = registry; // freeze mut → immut

    let notify_finalize = cli.finalize.clone();
    let notify_ctx = ctx.clone();
    let notify_etx = event_tx.clone();
    let notify_registry = registry.clone();
    let notify_debug = ctx.debug_flags;
    let notify_handle = tokio::spawn(async move {
        notify(
            &notify_finalize,
            &mut notify_ctx.clone(),
            notify_etx,
            to_rx,
            retry_tx,
            notify_debug,
            &notify_registry,
        )
        .await
    });

    let nevent_tx = to_tx.clone();
    let mut queue = NotificationQueue::new(to_tx, retry_rx, notify_debug);

    let pipe_result: anyhow::Result<()> = async {
        prebake(
            &cli.prebake,
            &cli,
            None,
            &mut ctx,
            event_tx.clone(),
            Some(nevent_tx.clone()),
        )
        .await?;
        let bake_result = bake(
            &cli.bake,
            &cli.bake_base,
            &cli.prebake,
            &cli,
            &mut ctx,
            event_tx.clone(),
            Some(nevent_tx.clone()),
        )
        .await;
        if bake_result.is_err() {
            workshop_baker::notify::types::send_stage_event(
                &Some(nevent_tx.clone()),
                workshop_baker::types::buildstatus::StagePhase::Bake,
                "bake",
                workshop_baker::types::buildstatus::StageOutcome::Failure,
            )
            .await;
        }
        bake_result?;
        finalize(&cli.finalize, &cli, &mut ctx, event_tx.clone(), &registry).await?;
        Ok(())
    }
    .await;

    if let Err(e) = pipe_result {
        // Notify pipeline failure before exiting
        workshop_baker::notify::types::send_stage_event(
            &Some(nevent_tx.clone()),
            workshop_baker::types::buildstatus::StagePhase::Bake,
            "pipeline",
            workshop_baker::types::buildstatus::StageOutcome::PipelineFailure,
        )
        .await;
        log::error!("Pipeline stage failed: {}", e);
        return Err(CliError::General(e));
    }

    // Notify pipeline success
    workshop_baker::notify::types::send_stage_event(
        &Some(nevent_tx.clone()),
        workshop_baker::types::buildstatus::StagePhase::Bake,
        "pipeline",
        workshop_baker::types::buildstatus::StageOutcome::PipelineSuccess,
    )
    .await;
    log::info!("Pipeline complete. Draining notifications...");
    drop(nevent_tx);
    queue.tx = None;
    if let Err(e) = queue.drain().await {
        log::error!("Notification queue error: {}", e);
    }
    let _ = notify_handle.await;

    // Cleanup after notifications are fully drained
    let pipeline_ok = pipe_result.is_ok();
    let _ = workshop_baker::finalize::parse(&cli.finalize).map(|cfg| {
        let _ = workshop_baker::finalize::cleanup::cleanup(&cfg.cleanup, &ctx, pipeline_ok);
    });
    Ok(())
}

#[tokio::main]
async fn main() {
    env_logger::init();
    let cli = Cli::parse();
    if std::process::id() == 1 {
        log::error!("Baker is not meant to be run as PID=1");
        std::process::exit(1);
    }
    if let Err(e) = capctl::prctl::set_subreaper(true){
        log::error!("Failed to set self as subreaper: {}", e);
        log::error!("finalize.cleanup might not work properly");
    }

    // --debug help: print help and exit immediately
    if let Some(ref raw) = cli.debug {
        if DebugFeature::print_help(raw) {
            return;
        }
    }
    let debug_flags = DebugFeature::parse(cli.debug.as_deref());
    if cli.dry_run {
        log::warn!("Dry run enabled, no changes will be made.");
    }

    // --resource-registry is only allowed in bare mode
    if cli.resource_registry.is_some() && !matches!(cli.command, Some(Commands::Bare { .. })) {
        log::error!("--resource-registry is only supported in bare mode");
        std::process::exit(1);
    }
    match &cli.command {
        Some(Commands::Bare { command }) => {
            // --pipeline, --build-id, --username are optional in bare mode
            // and they will be defaulted to "BarePipeline", "0", "bare" respectively
            let pipeline = cli
                .pipeline
                .as_deref()
                .unwrap_or("BarePipeline")
                .to_string();
            let build_id = cli.build_id.as_deref().unwrap_or("0").to_string();
            let username = cli.username.as_deref().unwrap_or("bare").to_string();

            let mut ctx = ExecutionContext {
                standalone: true,
                task_id: format!("{}-{}-bare", pipeline, build_id),
                pipeline_name: pipeline.clone(),
                build_id: build_id.clone(),
                username: username.clone(),
                dry_run: cli.dry_run,
                privileged: false,
                debug_flags: debug_flags.bits(),
                ..Default::default()
            };
            let (tx, event_rx) = tokio::sync::mpsc::channel(256);
            let event_tx: EventSender = Some(tx);
            let event_handle = tokio::spawn(workshop_baker::prebake::event::event_receiver(
                event_rx,
                pipeline.clone(),
                build_id.clone(),
                ctx.debug_flags,
            ));

            let registry = if let Some(ref json_str) = cli.resource_registry {
                match workshop_baker::types::resource::ResourceRegistry::from_json_str(json_str) {
                    Ok(r) => {
                        log::info!("Loaded {} resources from --resource-registry", r.len());
                        r
                    }
                    Err(e) => {
                        log::error!("Failed to parse --resource-registry: {}", e);
                        std::process::exit(1);
                    }
                }
            } else {
                ResourceRegistry::new()
            };
            let result = run_bare(
                &cli, command, &mut ctx, event_tx, &pipeline, &build_id, &registry,
            )
            .await;

            drop(ctx);
            let _ = event_handle.await;

            if let Err(e) = result {
                log::error!("{}", e);
                std::process::exit(e.exit_code());
            }
        }
        Some(Commands::Daemon {}) => {
            unimplemented!("Daemon mode not implemented");
        }
        None => {
            if let Err(e) = worker_main(cli).await {
                log::error!("{}", e);
                std::process::exit(e.exit_code());
            }
        }
    }
}