RedBear-OS/init/src/scheduler.rs

use std::collections::{BTreeMap, VecDeque};

use crate::InitConfig;
use crate::color::{init_debug, init_error, init_info, status_ok, status_skip};
use crate::unit::{Unit, UnitId, UnitKind, UnitStore};

pub struct Scheduler {
    pending: VecDeque<Job>,
}

struct Job {
    unit: UnitId,
    kind: JobKind,
}

enum JobKind {
    Start,
}

impl Scheduler {
    pub fn new() -> Scheduler {
        Scheduler {
            pending: VecDeque::new(),
        }
    }

    pub fn schedule_start_and_report_errors(
        &mut self,
        unit_store: &mut UnitStore,
        unit_id: UnitId,
    ) {
        let mut errors = vec![];
        self.schedule_start(unit_store, unit_id, &mut errors);
        for error in errors {
            init_error(&format!("{}", error));
        }
    }

    pub fn schedule_start(
        &mut self,
        unit_store: &mut UnitStore,
        unit_id: UnitId,
        errors: &mut Vec<String>,
    ) {
        let loaded_units = unit_store.load_units(unit_id.clone(), errors);
        for unit_id in loaded_units {
            if !unit_store.unit(&unit_id).conditions_met() {
                continue;
            }

            self.pending.push_back(Job {
                unit: unit_id,
                kind: JobKind::Start,
            });
        }
    }

    pub fn step(&mut self, unit_store: &mut UnitStore, init_config: &mut InitConfig) {
        let mut defer_count: BTreeMap<String, usize> = BTreeMap::new();
        'a: loop {
            let Some(job) = self.pending.pop_front() else {
                return;
            };

            match job.kind {
                JobKind::Start => {
                    let unit = unit_store.unit_mut(&job.unit);

                    let mut blocked = false;
                    for dep in &unit.info.requires_weak {
                        for pending_job in &self.pending {
                            if &pending_job.unit == dep {
                                blocked = true;
                                break;
                            }
                        }
                        if blocked { break; }
                    }

                    if blocked {
                        let cnt = defer_count.entry(job.unit.0.clone()).or_insert(0);
                        *cnt += 1;
                        if *cnt <= 3 {
                            self.pending.push_back(job);
                            continue 'a;
                        }
                    }

                    run(unit, init_config);
                }
            }
        }
    }
}

fn run(unit: &mut Unit, config: &mut InitConfig) {
    match &unit.kind {
        UnitKind::LegacyScript { script } => {
            for cmd in script.clone() {
                if config.log_debug {
                    init_debug(&format!("running: {cmd:?}"));
                }
                cmd.run(config);
            }
        }
        UnitKind::Service { service } => {
            let desc = unit.info.description.as_ref().unwrap_or(&unit.id.0);
            if config.skip_cmd.contains(&service.cmd) {
                status_skip(&format!("Skipping {} ({})", desc, service.cmd));
                return;
            }
            if !service.condition_met() {
                status_skip(&format!(
                    "Skipping {} (ConditionPathExists not met)",
                    desc
                ));
                return;
            }
            if config.log_debug {
                init_info(&format!("Starting {} ({})", desc, service.cmd));
            } else {
                status_ok(&format!("Started {}", desc));
            }
            service.spawn(&config.envs);
        }
        UnitKind::Target {} => {
            if config.log_debug {
                init_info(&format!("Reached target {}",
                    unit.info.description.as_ref().unwrap_or(&unit.id.0)));
            }
        }
    }
}