roundabout

//! fifo with easy backfill. the first job that can't start gets a reservation,

//! and anyone behind it only cuts in if they'll be gone before it starts or fit

//! in the leftovers once it does

use std::{collections::BTreeMap, time::Duration};

use serde::{Deserialize, Serialize};

use crate::units::Bytes;

pub type JobId = u64;

pub type FsId = u64;

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]

pub struct Need {

    pub ram: Bytes,

    pub disk: Bytes,

    pub cores: u32,

    pub time: Duration,

}

#[derive(Clone, Debug)]

pub struct Queued {

    pub id: JobId,

    pub need: Need,

    pub heavy: bool,

    pub fs: Vec<FsId>,

}

#[derive(Clone, Debug)]

pub struct Running {

    pub need: Need,

    pub heavy: bool,

    pub fs: Vec<FsId>,

    pub rss: Bytes,

    /// zero once overdue

    pub remaining: Duration,

}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]

pub struct Disk {

    pub free: Bytes,

    pub reserve: Bytes,

    /// written since the fs last went from idle to busy. free space already

    /// reflects this, so it's netted off promises to avoid counting it twice

    pub consumed: Bytes,

}

#[derive(Clone, Debug)]

pub struct Capacity {

    pub ram_budget: Bytes,

    pub ram_live: Bytes,

    pub ram_reserve: Bytes,

    pub cores: u32,

    pub disks: BTreeMap<FsId, Disk>,

    pub pressure: bool,

}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]

#[serde(rename_all = "snake_case", tag = "reason", content = "fs")]

pub enum Reason {

    HeavyRunning,

    Pressure,

    Ram,

    Cores,

    Disk(FsId),

    Reserved,

}

impl std::fmt::Display for Reason {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.write_str(match self {

            Self::HeavyRunning => "another heavy build is running",

            Self::Pressure => "memory pressure is high",

            Self::Ram => "not enough free RAM",

            Self::Cores => "not enough free cores",

            Self::Disk(_) => "not enough free disk",

            Self::Reserved => "waiting behind a reserved build",

        })

    }

}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]

pub enum Verdict {

    Admit,

    Wait(Reason),

}

#[derive(Clone, Debug)]

struct Free {

    ram: Bytes,

    cores: u32,

    disk: BTreeMap<FsId, Bytes>,

    heavy_slot: bool,

}

impl Free {

    fn shortfall(&self, need: &Need, fs: &[FsId], heavy: bool) -> Option<Reason> {

        if heavy && !self.heavy_slot {

            return Some(Reason::HeavyRunning);

        }

        if need.ram > self.ram {

            return Some(Reason::Ram);

        }

        if need.cores > self.cores {

            return Some(Reason::Cores);

        }

        fs.iter()

            .find(|id| self.disk.get(id).is_some_and(|free| need.disk > *free))

            .map(|id| Reason::Disk(*id))

    }

    fn take(&mut self, need: &Need, fs: &[FsId], heavy: bool) {

        self.ram = self.ram.saturating_sub(need.ram);

        self.cores = self.cores.saturating_sub(need.cores);

        for id in fs {

            if let Some(free) = self.disk.get_mut(id) {

                *free = free.saturating_sub(need.disk);

            }

        }

        if heavy {

            self.heavy_slot = false;

        }

    }

    /// no disk here. builds leave their output behind when they finish

    fn give_back(&mut self, ram: Bytes, cores: u32, heavy: bool) {

        self.ram += ram;

        self.cores += cores;

        if heavy {

            self.heavy_slot = true;

        }

    }

}

fn clamp(need: Need, cap: &Capacity) -> Need {

    Need {

        ram: need.ram.min(cap.ram_budget),

        cores: need.cores.clamp(1, cap.cores.max(1)),

        ..need

    }

}

fn free_now(running: &[Running], cap: &Capacity) -> Free {

    let charged: Bytes = running.iter().map(|r| r.need.ram.max(r.rss)).sum();

    let unrealised: Bytes = running

        .iter()

        .map(|r| r.need.ram.saturating_sub(r.rss))

        .sum();

    let by_budget = cap.ram_budget.saturating_sub(charged);

    let by_live = cap

        .ram_live

        .saturating_sub(cap.ram_reserve)

        .saturating_sub(unrealised);

    let disk = cap

        .disks

        .iter()

        .map(|(id, d)| {

            let promised: Bytes = running

                .iter()

                .filter(|r| r.fs.contains(id))

                .map(|r| r.need.disk)

                .sum();

            let outstanding = promised.saturating_sub(d.consumed);

            (

                *id,

                d.free.saturating_sub(d.reserve).saturating_sub(outstanding),

            )

        })

        .collect();

    Free {

        ram: by_budget.min(by_live),

        cores: cap

            .cores

            .saturating_sub(running.iter().map(|r| r.need.cores).sum()),

        disk,

        heavy_slot: !running.iter().any(|r| r.heavy),

    }

}

pub fn headroom(running: &[Running], cap: &Capacity) -> (Bytes, u32) {

    let free = free_now(running, cap);

    (free.ram, free.cores)

}

struct Reservation {

    /// `None` when nothing finishing would free enough, e.g. disk or someone

    /// else's ram. backfill falls back to the time cap then

    shadow: Option<Duration>,

    extra: Option<Free>,

}

fn reserve(

    job: &Queued,

    need: &Need,

    free: &Free,

    ends: &[(Duration, Bytes, u32, bool)],

) -> Reservation {

    let mut sim = free.clone();

    let mut ends = ends.to_vec();

    ends.sort_by_key(|e| e.0);

    for (at, ram, cores, heavy) in ends {

        sim.give_back(ram, cores, heavy);

        if sim.shortfall(need, &job.fs, job.heavy).is_none() {

            sim.take(need, &job.fs, job.heavy);

            return Reservation {

                shadow: Some(at),

                extra: Some(sim),

            };

        }

    }

    Reservation {

        shadow: None,

        extra: None,

    }

}

pub fn plan(

    queue: &[Queued],

    running: &[Running],

    cap: &Capacity,

    backfill_max: Duration,

) -> Vec<(JobId, Verdict)> {

    let mut free = free_now(running, cap);

    let mut ends: Vec<_> = running

        .iter()

        .map(|r| (r.remaining, r.need.ram.max(r.rss), r.need.cores, r.heavy))

        .collect();

    let mut reservation: Option<Reservation> = None;

    let mut out = Vec::with_capacity(queue.len());

    for job in queue {

        let need = clamp(job.need, cap);

        if job.heavy && cap.pressure {

            out.push((job.id, Verdict::Wait(Reason::Pressure)));

            continue;

        }

        if let Some(why) = free.shortfall(&need, &job.fs, job.heavy) {

            if reservation.is_none() {

                reservation = Some(reserve(job, &need, &free, &ends));

            }

            out.push((job.id, Verdict::Wait(why)));

            continue;

        }

        if let Some(res) = &mut reservation {

            let finishes_first = match res.shadow {

                Some(shadow) => need.time <= shadow,

                None => need.time <= backfill_max,

            };

            let fits_extra = res

                .extra

                .as_ref()

                .is_some_and(|x| x.shortfall(&need, &job.fs, job.heavy).is_none());

            if !finishes_first && !fits_extra {

                out.push((job.id, Verdict::Wait(Reason::Reserved)));

                continue;

            }

            if !finishes_first && let Some(extra) = &mut res.extra {

                extra.take(&need, &job.fs, job.heavy);

            }

        }

        free.take(&need, &job.fs, job.heavy);

        ends.push((need.time, need.ram, need.cores, job.heavy));

        out.push((job.id, Verdict::Admit));

    }

    out

}

#[cfg(test)]

mod tests {

    use super::*;

    const FS: FsId = 1;

    fn need(ram_gib: u64, cores: u32, disk_gib: u64, mins: u64) -> Need {

        Need {

            ram: Bytes::gib(ram_gib),

            disk: Bytes::gib(disk_gib),

            cores,

            time: Duration::from_secs(mins * 60),

        }

    }

    fn q(id: JobId, need: Need, heavy: bool) -> Queued {

        Queued {

            id,

            need,

            heavy,

            fs: vec![FS],

        }

    }

    fn cap() -> Capacity {

        Capacity {

            ram_budget: Bytes::gib(24),

            ram_live: Bytes::gib(28),

            ram_reserve: Bytes::gib(4),

            cores: 16,

            disks: BTreeMap::from([(

                FS,

                Disk {

                    free: Bytes::gib(100),

                    reserve: Bytes::gib(20),

                    consumed: Bytes::ZERO,

                },

            )]),

            pressure: false,

        }

    }

    fn verdicts(plan: &[(JobId, Verdict)]) -> Vec<Verdict> {

        plan.iter().map(|(_, v)| *v).collect()

    }

    const BF: Duration = Duration::from_secs(120);

    #[test]

    fn two_heavies_never_collide() {

        let out = plan(

            &[q(1, need(8, 8, 5, 10), true), q(2, need(8, 4, 5, 10), true)],

            &[],

            &cap(),

            BF,

        );

        assert_eq!(

            verdicts(&out),

            [Verdict::Admit, Verdict::Wait(Reason::HeavyRunning)]

        );

    }

    #[test]

    fn disk_never_double_booked() {

        let out = plan(

            &[

                q(1, need(1, 1, 50, 10), false),

                q(2, need(1, 1, 50, 10), false),

            ],

            &[],

            &cap(),

            BF,

        );

        assert_eq!(

            verdicts(&out),

            [Verdict::Admit, Verdict::Wait(Reason::Disk(FS))]

        );

    }

}