use std::collections::BTreeMap;
use std::time::{Duration, SystemTime};
use crate::github::QueuedJob;
pub mod engine;
pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
pub const CPU_PRIORITY_AFTER: Duration = Duration::from_secs(120);
#[derive(Debug, Clone)]
pub struct Repository {
pub name: String,
pub profile: String,
pub labels: Vec<String>,
pub paused: bool,
pub held_until: Option<SystemTime>,
pub queued: Vec<QueuedJob>,
pub runner_memory: Option<u64>,
pub runner_cpu: Option<u64>,
pub cpu_waiting_since: Option<SystemTime>,
}
impl Repository {
fn starved_of_cpu(&self, now: SystemTime) -> bool {
self.cpu_waiting_since
.and_then(|since| now.duration_since(since).ok())
.is_some_and(|waited| waited >= CPU_PRIORITY_AFTER)
}
}
#[derive(Debug, Clone)]
pub struct Profile {
pub memory: u64,
pub cpus: u64,
pub preparing: bool,
pub slots: Vec<Slot>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Slot {
Free,
Runner(RunnerSlot),
Closed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RunnerSlot {
pub repository: String,
pub busy: bool,
pub started_at: SystemTime,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SlotRef {
pub profile: String,
pub index: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Decision {
Start {
slot: SlotRef,
repository: String,
memory: u64,
cpu: u64,
},
HeldForCpu {
repository: String,
},
Stop {
slot: SlotRef,
},
}
pub fn matches(job: &QueuedJob, labels: &[String]) -> bool {
job.labels.iter().all(|wanted| {
labels
.iter()
.any(|label| label.eq_ignore_ascii_case(wanted))
})
}
pub fn plan(
repositories: &[Repository],
profiles: &BTreeMap<String, Profile>,
free_memory: Option<u64>,
free_cpu: Option<u64>,
now: SystemTime,
) -> Vec<Decision> {
let mut decisions = Vec::new();
let mut wanted = BTreeMap::new();
let mut jobs = Vec::new();
for repository in repositories {
let held = repository.held_until.is_some_and(|until| until > now);
let matching: Vec<&QueuedJob> = if repository.paused {
Vec::new()
} else {
repository
.queued
.iter()
.filter(|job| matches(job, &repository.labels))
.collect()
};
let idle: Vec<(SlotRef, &RunnerSlot)> = profiles
.get(&repository.profile)
.into_iter()
.flat_map(|profile| profile.slots.iter().enumerate())
.filter_map(|(index, slot)| match slot {
Slot::Runner(runner) if runner.repository == repository.name && !runner.busy => {
Some((
SlotRef {
profile: repository.profile.clone(),
index,
},
runner,
))
}
_ => None,
})
.collect();
let surplus = idle.len().saturating_sub(matching.len());
decisions.extend(
idle.iter()
.filter(|(_, runner)| {
now.duration_since(runner.started_at)
.is_ok_and(|waited| waited >= IDLE_TIMEOUT)
})
.take(surplus)
.map(|(slot, _)| Decision::Stop { slot: slot.clone() }),
);
if !held {
let missing = matching.len().saturating_sub(idle.len());
wanted.insert(repository.name.as_str(), missing);
let skip = matching.len() - missing;
jobs.extend(matching.iter().skip(skip).map(|job| (job.id, repository)));
}
}
jobs.sort_by_key(|(id, _)| *id);
let mut taken = taken_slots(profiles);
let mut free_memory = free_memory;
let mut free_cpu = free_cpu;
for (_, repository) in jobs {
let Some(missing) = wanted.get_mut(repository.name.as_str()) else {
continue;
};
if *missing == 0 {
continue;
}
let (Some(profile), Some(taken)) = (
profiles.get(&repository.profile),
taken.get_mut(repository.profile.as_str()),
) else {
continue;
};
let memory = repository.runner_memory.unwrap_or(profile.memory);
if free_memory.is_some_and(|free| free < memory) {
continue;
}
let Some(index) = taken.iter().position(|taken| !taken) else {
continue;
};
let cpu = repository.runner_cpu.unwrap_or(profile.cpus);
if free_cpu.is_some_and(|free| free < cpu) {
let held = Decision::HeldForCpu {
repository: repository.name.clone(),
};
if !decisions.contains(&held) {
decisions.push(held);
}
if repository.starved_of_cpu(now) {
break;
}
continue;
}
taken[index] = true;
*missing -= 1;
free_memory = free_memory.map(|free| free - memory);
free_cpu = free_cpu.map(|free| free - cpu);
decisions.push(Decision::Start {
slot: SlotRef {
profile: repository.profile.clone(),
index,
},
repository: repository.name.clone(),
memory,
cpu,
});
}
decisions
}
fn taken_slots(profiles: &BTreeMap<String, Profile>) -> BTreeMap<&str, Vec<bool>> {
profiles
.iter()
.map(|(name, profile)| {
let taken = profile
.slots
.iter()
.map(|slot| profile.preparing || *slot != Slot::Free)
.collect();
(name.as_str(), taken)
})
.collect()
}
#[cfg(test)]
mod tests;