roundabout

use std::{

    collections::BTreeMap,

    fs, io,

    os::{

        fd::FromRawFd,

        unix::net::{UnixListener, UnixStream},

    },

    path::PathBuf,

    process,

    sync::{

        Arc, Mutex, MutexGuard,

        atomic::{AtomicBool, Ordering},

    },

    thread,

    time::{Duration, Instant},

};

use crate::{

    client,

    config::Config,

    paths,

    proto::{self, Acquire, Event, IntentInfo, LeaseId, LeaseInfo, Reader, Request, Status},

    sched::{self, Capacity, Disk, FsId, Need, Queued, Reason, Running, Verdict},

    sys,

    units::Bytes,

};

struct Run {

    started: Instant,

    rss: Bytes,

    reminded: Option<Instant>,

    /// asked-for plus bonus. only the asked-for part counts for admission,

    /// oversubscribed cores just make everyone slower

    cores: u32,

}

struct Intent {

    need: Need,

    until: Instant,

}

struct Lease {

    req: Acquire,

    heavy: bool,

    fs: Vec<FsId>,

    queued_at: Instant,

    run: Option<Run>,

    out: UnixStream,

    waiting: Option<Reason>,

    last_pos: usize,

    watchers: Vec<UnixStream>,

}

impl Lease {

    fn elapsed(&self) -> Duration {

        self.run

            .as_ref()

            .map_or(self.queued_at, |r| r.started)

            .elapsed()

    }

    fn overdue(&self, grace: f64) -> bool {

        self.run.is_some() && self.elapsed() > self.req.need.time.mul_f64(grace)

    }

}

struct State {

    cfg: Config,

    next: LeaseId,

    leases: BTreeMap<LeaseId, Lease>,

    intents: BTreeMap<String, Intent>,

    cap: Capacity,

    baselines: BTreeMap<FsId, Bytes>,

    conns: usize,

    idle_since: Instant,

}

type Shared = Arc<Mutex<State>>;

fn lock(shared: &Shared) -> MutexGuard<'_, State> {

    shared

        .lock()

        .unwrap_or_else(std::sync::PoisonError::into_inner)

}

impl State {

    fn new(cfg: Config) -> Self {

        let mut st = Self {

            cfg,

            next: 1,

            leases: BTreeMap::new(),

            intents: BTreeMap::new(),

            cap: Capacity {

                ram_budget: Bytes::ZERO,

                ram_live: Bytes::ZERO,

                ram_reserve: Bytes::ZERO,

                cores: 1,

                disks: BTreeMap::new(),

                pressure: false,

            },

            baselines: BTreeMap::new(),

            conns: 0,

            idle_since: Instant::now(),

        };

        st.apply_budget();

        st

    }

    fn apply_budget(&mut self) {

        let (total, live) = sys::memory();

        self.cap.ram_budget = self.cfg.budget.ram.resolve(total);

        self.cap.ram_live = live;

        self.cap.ram_reserve = self.cfg.budget.ram_reserve;

        self.cap.cores = self.cfg.budget.cores.resolve(sys::cores());

    }

    /// running leases keep the heaviness they started with, so a reload can't

    /// retroactively put two heavies side by side

    fn reload(&mut self, cfg: Config) {

        self.cfg = cfg;

        self.apply_budget();

        let queued: Vec<(LeaseId, Need)> = self

            .leases

            .iter()

            .filter(|(_, l)| l.run.is_none())

            .map(|(id, l)| (*id, l.req.need))

            .collect();

        for (id, need) in queued {

            let heavy = self.heavy(&need);

            if let Some(l) = self.leases.get_mut(&id) {

                l.heavy = heavy;

            }

        }

        self.refresh();

        self.reschedule();

    }

    fn heavy(&self, need: &Need) -> bool {

        let h = &self.cfg.heavy;

        need.ram >= h.ram || need.disk >= h.disk || need.cores >= h.cores.resolve(sys::cores())

    }

    /// capped at a tenth of the fs, or a 16G tmpfs /tmp with a 20G reserve

    /// admits nothing ever

    fn reserve(&self, size: Bytes) -> Bytes {

        self.cfg.budget.disk_reserve.min(Bytes(size.0 / 10))

    }

    fn fs_for(&mut self, cwd: &PathBuf) -> Vec<FsId> {

        let mut ids = Vec::new();

        for path in std::iter::once(cwd).chain(&self.cfg.budget.disk_paths) {

            if let Some(fs) = sys::fs(path) {

                if !ids.contains(&fs.id) {

                    ids.push(fs.id);

                }

                let reserve = self.reserve(fs.size);

                self.cap.disks.entry(fs.id).or_insert(Disk {

                    free: fs.free,

                    reserve,

                    consumed: Bytes::ZERO,

                });

            }

        }

        ids

    }

    fn refresh(&mut self) {

        let (_, live) = sys::memory();

        self.cap.ram_live = live;

        self.cap.pressure = self

            .cfg

            .budget

            .psi_max

            .is_some_and(|max| sys::pressure().is_some_and(|p| p > max));

        let mut paths: Vec<PathBuf> = self.leases.values().map(|l| l.req.cwd.clone()).collect();

        paths.extend(self.cfg.budget.disk_paths.iter().cloned());

        let mut disks = BTreeMap::new();

        for path in paths {

            if let Some(fs) = sys::fs(&path) {

                disks.insert(fs.id, (fs.free, self.reserve(fs.size)));

            }

        }

        let busy: Vec<FsId> = self

            .leases

            .values()

            .filter(|l| l.run.is_some())

            .flat_map(|l| l.fs.iter().copied())

            .collect();

        self.baselines.retain(|id, _| busy.contains(id));

        self.cap.disks = disks

            .into_iter()

            .map(|(id, (free, reserve))| {

                let consumed = if busy.contains(&id) {

                    self.baselines

                        .entry(id)

                        .or_insert(free)

                        .saturating_sub(free)

                } else {

                    Bytes::ZERO

                };

                (

                    id,

                    Disk {

                        free,

                        reserve,

                        consumed,

                    },

                )

            })

            .collect();

    }

    fn running(&self) -> Vec<Running> {

        self.leases

            .values()

            .filter_map(|l| {

                let run = l.run.as_ref()?;

                Some(Running {

                    need: l.req.need,

                    heavy: l.heavy,

                    fs: l.fs.clone(),

                    rss: run.rss,

                    remaining: l.req.need.time.saturating_sub(run.started.elapsed()),

                })

            })

            .collect()

    }

    fn reschedule(&mut self) {

        let queue: Vec<Queued> = self

            .leases

            .iter()

            .filter(|(_, l)| l.run.is_none())

            .map(|(id, l)| Queued {

                id: *id,

                need: l.req.need,

                heavy: l.heavy,

                fs: l.fs.clone(),

            })

            .collect();

        if queue.is_empty() {

            return;

        }

        let verdicts = sched::plan(

            &queue,

            &self.running(),

            &self.cap,

            self.cfg.queue.backfill_max_time.0,

        );

        let cores = self.cap.cores;

        let mut position = 0;

        for (id, verdict) in verdicts {

            match verdict {

                Verdict::Admit => {

                    let (extra, extra_ram) = self.bonus(id);

                    let Some(lease) = self.leases.get_mut(&id) else {

                        continue;

                    };

                    lease.req.need.ram += extra_ram;

                    let granted = lease.req.need.cores.clamp(1, cores.max(1)) + extra;

                    lease.run = Some(Run {

                        started: Instant::now(),

                        rss: Bytes::ZERO,

                        reminded: None,

                        cores: granted,

                    });

                    let _ = proto::send(

                        &mut lease.out,

                        &Event::Granted {

                            lease: id,

                            cores: granted,

                            ram: lease.req.need.ram,

                            heavy: lease.heavy,

                        },

                    );

                }

                Verdict::Wait(reason) => {

                    let Some(lease) = self.leases.get_mut(&id) else {

                        continue;

                    };

                    position += 1;

                    lease.waiting = Some(reason);

                    if lease.last_pos != position {

                        lease.last_pos = position;

                        let _ = proto::send(

                            &mut lease.out,

                            &Event::Queued {

                                lease: id,

                                position,

                                reason: Some(reason),

                            },

                        );

                    }

                }

            }

        }

    }

    fn held_by(&self, session: &str) -> (u32, Bytes) {

        self.leases

            .values()

            .filter(|l| l.req.session.as_deref() == Some(session))

            .fold((0, Bytes::ZERO), |(c, r), l| {

                (

                    c + l.run.as_ref().map_or(l.req.need.cores, |run| run.cores),

                    r + l.req.need.ram,

                )

            })

    }

    /// idle cores the lease can have on top of its ask, minus whatever other

    /// sessions said they'll want and don't hold yet. capped by ram at the

    /// lease's own ram per core, since every extra job eats its share

    fn bonus(&self, id: LeaseId) -> (u32, Bytes) {

        let Some(lease) = self.leases.get(&id) else {

            return (0, Bytes::ZERO);

        };

        if lease.req.exact || !self.cfg.queue.elastic {

            return (0, Bytes::ZERO);

        }

        let need = lease.req.need;

        let (taken, queued_ram) = self.leases.iter().filter(|(k, _)| **k != id).fold(

            (need.cores, need.ram),

            |(c, r), (_, l)| match &l.run {

                Some(run) => (c + run.cores, r),

                None => (c + l.req.need.cores, r + l.req.need.ram),

            },

        );

        let mut spare = self.cap.cores.saturating_sub(taken);

        let mut ram_spare = sched::headroom(&self.running(), &self.cap)

            .0

            .saturating_sub(queued_ram);

        for (session, intent) in &self.intents {

            if lease.req.session.as_deref() == Some(session) {

                continue;

            }

            let (cores, ram) = self.held_by(session);

            spare = spare.saturating_sub(intent.need.cores.saturating_sub(cores));

            ram_spare = ram_spare.saturating_sub(intent.need.ram.saturating_sub(ram));

        }

        let per_core = need.ram.0 / u64::from(need.cores.max(1));

        let by_ram = ram_spare

            .0

            .checked_div(per_core)

            .map_or(spare, |n| n as u32);

        let extra = spare.min(by_ram);

        (extra, Bytes(per_core * u64::from(extra)))

    }

    /// explicit intents set the size, builds only keep one alive or start a

    /// default one sized like themselves

    fn intend(&mut self, session: &str, need: Need, explicit: bool) {

        let until = Instant::now() + self.cfg.queue.intent_ttl.0;

        match self.intents.get_mut(session) {

            Some(i) => {

                i.until = until;

                if explicit {

                    i.need = need;

                }

            }

            None => {

                self.intents

                    .insert(session.to_owned(), Intent { need, until });

            }

        }

    }

    fn expire_intents(&mut self) {

        let had = !self.intents.is_empty();

        let now = Instant::now();

        self.intents.retain(|_, i| i.until > now);

        if had && self.intents.is_empty() && self.leases.is_empty() {

            self.idle_since = now;

        }

    }

    fn remove(&mut self, id: LeaseId) {

        if let Some(lease) = self.leases.remove(&id) {

            if let Some(s) = &lease.req.session

                && let Some(i) = self.intents.get_mut(s)

            {

                i.until = Instant::now() + self.cfg.queue.intent_ttl.0;

            }

            for mut w in lease.watchers {

                let _ = proto::send(&mut w, &Event::Fired { overdue: false });

            }

        }

        if self.leases.is_empty() {

            self.idle_since = Instant::now();

        }

        self.refresh();

        self.reschedule();

    }

    fn fire_overdue(&mut self) {

        let grace = self.cfg.overrun.time_grace;

        for lease in self.leases.values_mut() {

            if !lease.watchers.is_empty() && lease.overdue(grace) {

                for mut w in lease.watchers.drain(..) {

                    let _ = proto::send(&mut w, &Event::Fired { overdue: true });

                }

            }

        }

    }

    fn info(&self) -> Vec<LeaseInfo> {

        let grace = self.cfg.overrun.time_grace;

        let mut position = 0;

        self.leases

            .iter()

            .map(|(id, l)| {

                if l.run.is_none() {

                    position += 1;

                }

                LeaseInfo {

                    id: *id,

                    cmd: l.req.cmd.clone(),

                    cwd: l.req.cwd.clone(),

                    session: l.req.session.clone(),

                    need: l.req.need,

                    heavy: l.heavy,

                    running: l.run.is_some(),

                    granted: l.run.as_ref().map_or(0, |r| r.cores),

                    elapsed: l.elapsed().as_secs(),

                    rss: l.run.as_ref().map_or(Bytes::ZERO, |r| r.rss),

                    overdue: l.overdue(grace),

                    waiting: l.waiting.filter(|_| l.run.is_none()),

                    position: if l.run.is_none() { position } else { 0 },

                }

            })

            .collect()

    }

    fn status(&self) -> Status {

        let (ram_free, cores_free) = sched::headroom(&self.running(), &self.cap);

        let now = Instant::now();

        Status {

            leases: self.info(),

            intents: self

                .intents

                .iter()

                .map(|(session, i)| IntentInfo {

                    session: session.clone(),

                    need: i.need,

                    expires_in: i.until.saturating_duration_since(now).as_secs(),

                })

                .collect(),

            ram_free,

            ram_budget: self.cap.ram_budget,

            cores_free,

            cores: self.cap.cores,

        }

    }

}

fn serve(shared: &Shared, mut stream: UnixStream) -> io::Result<()> {

    let mut reader = Reader::new(&stream)?;

    let Some(req) = reader.recv::<Request>()? else {

        return Ok(());

    };

    match req {

        Request::Acquire(req) => {

            let id = {

                let mut st = lock(shared);

                let id = st.next;

                st.next += 1;

                let fs = st.fs_for(&req.cwd);