use std::collections::{BTreeMap, BTreeSet};
use std::time::{Duration, SystemTime};
use crate::api::Waiting;
use crate::github::QueuedJob;
pub mod engine;
pub mod warm;
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>,
pub warm: usize,
}
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: Option<Preparing>,
pub upkeep_due: bool,
pub slots: Vec<Slot>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Preparing {
Reseeding,
PullingImage,
}
#[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,
pub memory: u64,
pub cpu: u64,
}
#[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,
warm: bool,
replacing: Vec<SlotRef>,
},
HeldForCpu {
repository: String,
},
Stop {
slot: SlotRef,
},
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Plan {
pub decisions: Vec<Decision>,
pub waiting: BTreeMap<u64, Waiting>,
pub warm: BTreeSet<SlotRef>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct Budget {
pub memory: Option<u64>,
pub cpu: Option<u64>,
pub held: bool,
}
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>,
budget: Budget,
now: SystemTime,
) -> Plan {
let mut decided = Plan::default();
let mut wanted = BTreeMap::new();
let mut warm_wanted = Vec::new();
let mut jobs = Vec::new();
for repository in repositories {
if let Some(sorted) = sort_out(repository, profiles, now, &mut decided) {
wanted.insert(repository.name.as_str(), sorted.serve.len());
jobs.extend(sorted.serve.into_iter().map(|job| (job, repository)));
warm_wanted.push((repository, sorted.warm));
}
}
jobs.sort_by_key(|(id, _)| *id);
let mut warm = Warm::new(&decided.warm, profiles, budget.held);
let mut taken = taken_slots(profiles);
let mut started = Started::new(&taken, budget.memory, budget.cpu);
let mut left = Left {
memory: budget.memory,
cpu: budget.cpu,
};
let mut first_for_cpu = None;
for (job, 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;
};
if let Some(older) = first_for_cpu {
decided
.waiting
.insert(job, Waiting::BehindCpu { job: older });
continue;
}
let memory = repository.runner_memory.unwrap_or(profile.memory);
let cpu = repository.runner_cpu.unwrap_or(profile.cpus);
let free_index = taken.iter().position(|taken| !taken);
let fits = free_index.is_some() && left.fits(memory, cpu);
let room = if fits || profile.preparing.is_some() {
None
} else {
warm.make_room(repository, free_index.is_none(), left.short(memory, cpu))
};
let index = match (room.as_ref(), free_index) {
(Some(room), free) => free.or(room.slot).unwrap_or_default(),
(None, Some(index)) if fits => index,
(None, free) => {
let (reason, short) = why_waits(
&repository.profile,
profile,
free.is_some(),
memory,
cpu,
left,
&started,
);
decided.waiting.insert(job, reason);
if short == Short::Cpu {
let held = Decision::HeldForCpu {
repository: repository.name.clone(),
};
if !decided.decisions.contains(&held) {
decided.decisions.push(held);
}
if repository.starved_of_cpu(now) {
first_for_cpu = Some(job);
}
}
continue;
}
};
let replacing = room.map_or_else(Vec::new, |room| {
left.give(room.memory, room.cpu);
for slot in &room.stopped {
decided.warm.remove(slot);
decided
.decisions
.push(Decision::Stop { slot: slot.clone() });
}
room.stopped
});
taken[index] = true;
*missing -= 1;
started.note(&repository.profile, job);
left.take(memory, cpu);
decided
.decisions
.push(start(repository, index, memory, cpu, replacing));
}
if !budget.held {
start_warm(&warm_wanted, profiles, &mut taken, left, &mut decided);
}
decided
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Short {
Slot,
Memory,
Cpu,
}
fn why_waits(
name: &str,
profile: &Profile,
free_slot: bool,
memory: u64,
cpu: u64,
left: Left,
started: &Started<'_>,
) -> (Waiting, Short) {
let profile_name = || name.to_owned();
if !free_slot {
let reason = match profile.preparing {
Some(Preparing::Reseeding) => Waiting::Reseeding {
profile: profile_name(),
},
Some(Preparing::PullingImage) => Waiting::PullingImage {
profile: profile_name(),
},
None => started.for_slot(
name,
Waiting::NoSlot {
profile: profile_name(),
},
),
};
return (reason, Short::Slot);
}
match (left.memory, left.cpu) {
(Some(free), _) if free < memory => (
started.for_memory(
memory,
Waiting::Memory {
needs: memory,
free,
},
),
Short::Memory,
),
(_, Some(free)) => (
started.for_cpu(cpu, Waiting::Cpu { needs: cpu, free }),
Short::Cpu,
),
(_, None) => (
Waiting::NoSlot {
profile: profile_name(),
},
Short::Slot,
),
}
}
fn start(
repository: &Repository,
index: usize,
memory: u64,
cpu: u64,
replacing: Vec<SlotRef>,
) -> Decision {
Decision::Start {
slot: SlotRef {
profile: repository.profile.clone(),
index,
},
repository: repository.name.clone(),
memory,
cpu,
warm: false,
replacing,
}
}
fn start_warm(
wanted: &[(&Repository, usize)],
profiles: &BTreeMap<String, Profile>,
taken: &mut BTreeMap<&str, Vec<bool>>,
mut left: Left,
decided: &mut Plan,
) {
let room_wanted = decided.waiting.values().any(|reason| {
matches!(
reason,
Waiting::NoSlot { .. }
| Waiting::Memory { .. }
| Waiting::Cpu { .. }
| Waiting::BehindCpu { .. }
| Waiting::BehindStart { .. }
)
});
if room_wanted {
return;
}
for round in 0.. {
let mut any = false;
for (repository, wants) in wanted {
if *wants <= round {
continue;
}
let (Some(profile), Some(taken)) = (
profiles.get(&repository.profile),
taken.get_mut(repository.profile.as_str()),
) else {
continue;
};
let (memory, cpu) = (profile.memory, profile.cpus);
let Some(index) = taken.iter().position(|taken| !taken) else {
continue;
};
if !left.fits(memory, cpu) {
continue;
}
taken[index] = true;
left.take(memory, cpu);
any = true;
decided.decisions.push(Decision::Start {
slot: SlotRef {
profile: repository.profile.clone(),
index,
},
repository: repository.name.clone(),
memory,
cpu,
warm: true,
replacing: Vec::new(),
});
}
if !any {
break;
}
}
}
#[derive(Debug, Clone, Copy)]
struct Left {
memory: Option<u64>,
cpu: Option<u64>,
}
impl Left {
fn fits(self, memory: u64, cpu: u64) -> bool {
self.memory.is_none_or(|free| free >= memory) && self.cpu.is_none_or(|free| free >= cpu)
}
fn short(self, memory: u64, cpu: u64) -> Left {
Left {
memory: self.memory.map(|free| memory.saturating_sub(free)),
cpu: self.cpu.map(|free| cpu.saturating_sub(free)),
}
}
fn take(&mut self, memory: u64, cpu: u64) {
self.memory = self.memory.map(|free| free.saturating_sub(memory));
self.cpu = self.cpu.map(|free| free.saturating_sub(cpu));
}
fn give(&mut self, memory: u64, cpu: u64) {
self.memory = self.memory.map(|free| free + memory);
self.cpu = self.cpu.map(|free| free + cpu);
}
}
struct Sorted {
serve: Vec<u64>,
warm: usize,
}
fn sort_out(
repository: &Repository,
profiles: &BTreeMap<String, Profile>,
now: SystemTime,
decided: &mut Plan,
) -> Option<Sorted> {
let mut matching: Vec<&QueuedJob> = repository
.queued
.iter()
.filter(|job| matches(job, &repository.labels))
.collect();
let profile = profiles.get(&repository.profile);
let mut warm = repository.warm;
if repository.paused {
decided
.waiting
.extend(matching.iter().map(|job| (job.id, Waiting::Paused)));
matching.clear();
warm = 0;
}
let upkeep_due = profile.is_some_and(|profile| profile.upkeep_due);
if upkeep_due {
warm = 0;
}
let idle: Vec<(SlotRef, &RunnerSlot)> = 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());
let stops: Vec<SlotRef> = idle
.iter()
.filter(|(_, runner)| {
upkeep_due
|| now
.duration_since(runner.started_at)
.is_ok_and(|waited| waited >= IDLE_TIMEOUT)
})
.take(surplus.saturating_sub(warm))
.map(|(slot, _)| slot.clone())
.collect();
let kept: Vec<&SlotRef> = idle
.iter()
.map(|(slot, _)| slot)
.filter(|slot| !stops.contains(slot))
.collect();
decided
.warm
.extend(kept.iter().skip(matching.len()).map(|slot| (*slot).clone()));
let warm = warm.saturating_sub(kept.len().saturating_sub(matching.len()));
decided
.decisions
.extend(stops.into_iter().map(|slot| Decision::Stop { slot }));
let missing = matching.len().saturating_sub(idle.len());
let skip = matching.len() - missing;
decided.waiting.extend(
matching
.iter()
.take(skip)
.map(|job| (job.id, Waiting::RunnerStarting)),
);
let serve = matching.iter().skip(skip).map(|job| job.id);
match repository.held_until.filter(|until| *until > now) {
Some(until) => {
let until = crate::time::ms(until) / 1000;
decided
.waiting
.extend(serve.map(|job| (job, Waiting::Backoff { until })));
None
}
None => Some(Sorted {
serve: serve.collect(),
warm,
}),
}
}
struct Warm<'plan> {
runners: Vec<Option<(SlotRef, &'plan RunnerSlot)>>,
}
struct Room {
stopped: Vec<SlotRef>,
slot: Option<usize>,
memory: u64,
cpu: u64,
}
impl<'plan> Warm<'plan> {
fn new(
warm: &BTreeSet<SlotRef>,
profiles: &'plan BTreeMap<String, Profile>,
held: bool,
) -> Warm<'plan> {
let runners = warm
.iter()
.filter(|_| !held)
.filter_map(
|slot| match profiles.get(&slot.profile)?.slots.get(slot.index)? {
Slot::Runner(runner) => Some(Some((slot.clone(), runner))),
Slot::Free | Slot::Closed => None,
},
)
.collect();
Warm { runners }
}
fn make_room(&mut self, repository: &Repository, slot: bool, short: Left) -> Option<Room> {
let (memory, cpu) = (short.memory.unwrap_or(0), short.cpu.unwrap_or(0));
let mut candidates: Vec<usize> = self
.runners
.iter()
.enumerate()
.filter_map(|(at, runner)| {
runner
.as_ref()
.filter(|(_, runner)| runner.repository != repository.name)
.map(|_| at)
})
.collect();
candidates.sort_by_key(|at| {
self.runners[*at]
.as_ref()
.is_some_and(|(slot, _)| slot.profile != repository.profile)
});
let mut room = Room {
stopped: Vec::new(),
slot: None,
memory: 0,
cpu: 0,
};
let mut picked = Vec::new();
for at in candidates {
let short = (slot && room.slot.is_none()) || room.memory < memory || room.cpu < cpu;
if !short {
break;
}
let Some((ref_, runner)) = &self.runners[at] else {
continue;
};
let gives_slot = ref_.profile == repository.profile;
let helps = (gives_slot && slot && room.slot.is_none())
|| (room.memory < memory && runner.memory > 0)
|| (room.cpu < cpu && runner.cpu > 0);
if !helps {
continue;
}
if gives_slot && room.slot.is_none() {
room.slot = Some(ref_.index);
}
room.memory += runner.memory;
room.cpu += runner.cpu;
room.stopped.push(ref_.clone());
picked.push(at);
}
if (slot && room.slot.is_none()) || room.memory < memory || room.cpu < cpu {
return None;
}
for at in picked {
self.runners[at] = None;
}
Some(room)
}
}
struct Started<'plan> {
free_slots: BTreeSet<&'plan str>,
memory: Option<u64>,
cpu: Option<u64>,
by_profile: BTreeMap<&'plan str, u64>,
first: Option<u64>,
}
impl<'plan> Started<'plan> {
fn new(
taken: &BTreeMap<&'plan str, Vec<bool>>,
memory: Option<u64>,
cpu: Option<u64>,
) -> Started<'plan> {
Started {
free_slots: taken
.iter()
.filter(|(_, slots)| slots.contains(&false))
.map(|(name, _)| *name)
.collect(),
memory,
cpu,
by_profile: BTreeMap::new(),
first: None,
}
}
fn note(&mut self, profile: &'plan str, job: u64) {
self.by_profile.entry(profile).or_insert(job);
self.first.get_or_insert(job);
}
fn for_slot(&self, profile: &str, reason: Waiting) -> Waiting {
match self.by_profile.get(profile) {
Some(older) if self.free_slots.contains(profile) => {
Waiting::BehindStart { job: *older }
}
_ => reason,
}
}
fn for_memory(&self, needs: u64, reason: Waiting) -> Waiting {
Self::behind(self.first, self.memory, needs).unwrap_or(reason)
}
fn for_cpu(&self, needs: u64, reason: Waiting) -> Waiting {
Self::behind(self.first, self.cpu, needs).unwrap_or(reason)
}
fn behind(first: Option<u64>, free: Option<u64>, needs: u64) -> Option<Waiting> {
first
.filter(|_| free.is_some_and(|free| free >= needs))
.map(|job| Waiting::BehindStart { job })
}
}
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.is_some() || *slot != Slot::Free)
.collect();
(name.as_str(), taken)
})
.collect()
}
#[cfg(test)]
mod tests;