roundabout

use std::{

    collections::BTreeMap,

    fs, io,

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

    path::PathBuf,

    process,

    sync::{Arc, Mutex, MutexGuard},

    thread,

    time::{Duration, Instant},

};

use crate::{

    client,

    config::Config,

    paths,

    proto::{self, Acquire, Event, 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>,

}

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>,

    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(),

            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 {

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

                continue;

            };

            match verdict {

                Verdict::Admit => {

                    lease.run = Some(Run {

                        started: Instant::now(),

                        rss: Bytes::ZERO,

                        reminded: None,

                    });

                    let _ = proto::send(

                        &mut lease.out,

                        &Event::Granted {

                            lease: id,

                            cores: lease.req.need.cores.clamp(1, cores.max(1)),

                            heavy: lease.heavy,

                        },

                    );

                }

                Verdict::Wait(reason) => {

                    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 remove(&mut self, id: LeaseId) {

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

            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(),

                    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);

        Status {

            leases: self.info(),

            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);

                let heavy = st.heavy(&req.need);

                st.leases.insert(

                    id,

                    Lease {

                        req,

                        heavy,

                        fs,

                        queued_at: Instant::now(),

                        run: None,

                        out: stream.try_clone()?,

                        waiting: None,

                        last_pos: 0,

                        watchers: Vec::new(),

                    },

                );

                st.reschedule();

                id

            };

            hold(shared, id, &mut reader);

            lock(shared).remove(id);

            Ok(())

        }

        Request::Status => {

            let status = lock(shared).status();

            proto::send(&mut stream, &Event::Status(status))

        }

        Request::Kill { lease } => {

            let mut st = lock(shared);

            let reply = match st.leases.get_mut(&lease) {

                Some(l) => {

                    let _ = proto::send(&mut l.out, &Event::Kill);

                    Event::Ok

                }

                None => Event::Err {

                    msg: format!("no lease {lease}"),

                },

            };

            drop(st);

            proto::send(&mut stream, &reply)

        }

        Request::Extend { lease, by } => {

            let mut st = lock(shared);

            let reply = match st.leases.get_mut(&lease) {

                Some(l) => {

                    l.req.need.time += by.0;

                    if let Some(run) = &mut l.run {

                        run.reminded = None;

                    }

                    Event::Ok

                }

                None => Event::Err {

                    msg: format!("no lease {lease}"),

                },

            };

            drop(st);

            proto::send(&mut stream, &reply)

        }

        Request::Watch { lease } => {

            let mut st = lock(shared);

            let grace = st.cfg.overrun.time_grace;

            match st.leases.get_mut(&lease) {

                Some(l) if l.overdue(grace) => {

                    proto::send(&mut stream, &Event::Fired { overdue: true })

                }

                Some(l) => {

                    l.watchers.push(stream);

                    Ok(())

                }

                None => proto::send(&mut stream, &Event::Fired { overdue: false }),

            }

        }

        Request::Overdue { session } => {

            let mut st = lock(shared);

            let grace = st.cfg.overrun.time_grace;

            let every = st.cfg.overrun.remind_every.0;

            let mut due = Vec::new();

            for (id, l) in &mut st.leases {

                if l.req.session.as_deref() != Some(&session) || !l.overdue(grace) {

                    continue;

                }

                let Some(run) = &mut l.run else { continue };

                if run.reminded.is_some_and(|t| t.elapsed() < every) {

                    continue;

                }

                run.reminded = Some(Instant::now());

                due.push(*id);

            }

            let leases = st

                .info()

                .into_iter()

                .filter(|i| due.contains(&i.id))

                .collect();

            drop(st);

            proto::send(&mut stream, &Event::Overdue { leases })

        }

        Request::Started { .. } | Request::Usage { .. } | Request::Done => proto::send(

            &mut stream,

            &Event::Err {

                msg: "no lease on this connection".into(),

            },

        ),

    }

}

/// a client that vanishes without saying done got sigkilled or similar, and

/// pdeathsig only reaches its direct child. take the rest of the group down

/// so a headless build doesn't keep eating ram off the books

fn hold(shared: &Shared, id: LeaseId, reader: &mut Reader) {

    let mut pgid = None;

    let clean = loop {

        match reader.recv::<Request>() {

            Ok(Some(Request::Usage { rss })) => {

                if let Some(run) = lock(shared)

                    .leases

                    .get_mut(&id)

                    .and_then(|l| l.run.as_mut())

                {

                    run.rss = rss;

                }

            }

            Ok(Some(Request::Started { pid })) => pgid = Some(pid),

            Ok(Some(Request::Done)) => break true,