roundabout

public

Star0

Watch0

6 commits 1 branch 0 tags 0 releases158.8 KiB

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, 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,
            Ok(Some(_)) => {}
            Ok(None) | Err(_) => break false,
        }
    };
    if !clean && let Some(pgid) = pgid.filter(|p| *p > 1) {
        unsafe { libc::kill(-pgid, libc::SIGTERM) };
    }
}

fn tick(shared: &Shared) {
    loop {
        thread::sleep(Duration::from_secs(1));
        let mut st = lock(shared);
        st.refresh();
        st.reschedule();
        st.fire_overdue();
        if st.conns == 0 && st.leases.is_empty() && st.idle_since.elapsed() >= st.cfg.queue.linger.0
        {
            // systemd owns the socket and queues whoever knocks while we're
            // leaving, so there's nothing to tidy
            if ACTIVATED.load(Ordering::Relaxed) {
                process::exit(0);
            }
            // a spawner holds this while it clears a stale socket, so taking it
            // here stops us unlinking our successor's socket on the way out
            if let Some(_guard) = try_spawn_lock() {
                let _ = fs::remove_file(paths::socket());
                process::exit(0);
            }
        }
    }
}

/// watches the directory, not the file. editors save by renaming over it,
/// which would orphan a watch on the old inode after the first edit
#[cfg(target_os = "linux")]
fn watch_config(shared: &Shared) {
    use std::os::unix::ffi::OsStrExt;

    let path = paths::config_file();
    let (Some(dir), Some(name)) = (path.parent(), path.file_name()) else {
        return;
    };
    let _ = fs::create_dir_all(dir);
    let Ok(cdir) = std::ffi::CString::new(dir.as_os_str().as_bytes()) else {
        return;
    };
    let fd = unsafe { libc::inotify_init1(libc::IN_CLOEXEC) };
    if fd < 0 {
        return;
    }
    let mask = libc::IN_CLOSE_WRITE
        | libc::IN_MOVED_TO
        | libc::IN_MOVED_FROM
        | libc::IN_CREATE
        | libc::IN_DELETE;
    if unsafe { libc::inotify_add_watch(fd, cdir.as_ptr(), mask) } < 0 {
        return;
    }

    let header = std::mem::size_of::<libc::inotify_event>();
    let mut buf = [0u8; 4096];
    loop {
        let n = unsafe { libc::read(fd, buf.as_mut_ptr().cast(), buf.len()) };
        let Ok(n) = usize::try_from(n) else {
            return;
        };
        let mut hit = false;
        let mut off = 0;
        while off + header <= n {
            let ev: libc::inotify_event =
                unsafe { std::ptr::read_unaligned(buf.as_ptr().add(off).cast()) };
            let len = ev.len as usize;
            let raw = &buf[off + header..(off + header + len).min(n)];
            let ev_name = raw.split(|b| *b == 0).next().unwrap_or_default();
            hit |= ev_name == name.as_bytes();
            off += header + len;
        }
        if !hit {
            continue;
        }
        thread::sleep(Duration::from_millis(100));
        match Config::load() {
            Ok(cfg) => {
                lock(shared).reload(cfg);
                eprintln!("roundabout: config reloaded");
            }
            Err(e) => eprintln!("roundabout: kept the old config, {e}"),
        }
    }
}

/// creates the file if it's gone, else a wiped runtime dir leaves us idling
/// forever unable to take the lock we need to exit
fn try_spawn_lock() -> Option<client::Flock> {
    let file = fs::OpenOptions::new()
        .create(true)
        .truncate(false)
        .write(true)
        .open(paths::spawn_lock())
        .ok()?;
    client::Flock::try_exclusive(file)
}

static ACTIVATED: AtomicBool = AtomicBool::new(false);

fn inherited() -> Option<UnixListener> {
    let ours = std::env::var("LISTEN_PID").ok()?.parse::<u32>().ok()? == process::id();
    if !ours || std::env::var("LISTEN_FDS").ok()? != "1" {
        return None;
    }
    // SAFETY: sd_listen_fds protocol, fd 3 is the listening socket
    Some(unsafe { UnixListener::from_raw_fd(3) })
}

fn bind() -> UnixListener {
    // fork once more so the client that spawned us can reap its child now
    // instead of carrying a zombie for the length of its build
    match unsafe { libc::fork() } {
        -1 => {
            eprintln!("roundabout: fork failed: {}", io::Error::last_os_error());
            process::exit(1);
        }
        0 => {}
        _ => process::exit(0),
    }
    match UnixListener::bind(paths::socket()) {
        Ok(l) => l,
        Err(e) if e.kind() == io::ErrorKind::AddrInUse && client::connect().is_some() => {
            process::exit(0)
        }
        Err(e) => {
            eprintln!("roundabout: bind: {e}");
            process::exit(1);
        }
    }
}

pub fn main(cfg: Config) -> ! {
    let listener = if let Some(l) = inherited() {
        ACTIVATED.store(true, Ordering::Relaxed);
        l
    } else {
        bind()
    };

    let shared: Shared = Arc::new(Mutex::new(State::new(cfg)));
    {
        let shared = shared.clone();
        thread::spawn(move || tick(&shared));
    }
    #[cfg(target_os = "linux")]
    {
        let shared = shared.clone();
        thread::spawn(move || watch_config(&shared));
    }

    for stream in listener.incoming() {
        let Ok(stream) = stream else { continue };
        lock(&shared).conns += 1;
        let shared = shared.clone();
        thread::spawn(move || {
            if let Err(e) = serve(&shared, stream) {
                eprintln!("roundabout: {e}");
            }
            let mut st = lock(&shared);
            st.conns -= 1;
            if st.leases.is_empty() {
                st.idle_since = Instant::now();
            }
        });
    }
    process::exit(0)
}

1	use std::{
2	collections::BTreeMap,
3	fs, io,
4	os::{
5	fd::FromRawFd,
6	unix::net::{UnixListener, UnixStream},
7	},
8	path::PathBuf,
9	process,
10	sync::{
11	Arc, Mutex, MutexGuard,
12	atomic::{AtomicBool, Ordering},
13	},
14	thread,
15	time::{Duration, Instant},
16	};
17
18	use crate::{
19	client,
20	config::Config,
21	paths,
22	proto::{self, Acquire, Event, LeaseId, LeaseInfo, Reader, Request, Status},
23	sched::{self, Capacity, Disk, FsId, Need, Queued, Reason, Running, Verdict},
24	sys,
25	units::Bytes,
26	};
27
28	struct Run {
29	started: Instant,
30	rss: Bytes,
31	reminded: Option<Instant>,
32	}
33
34	struct Lease {
35	req: Acquire,
36	heavy: bool,
37	fs: Vec<FsId>,
38	queued_at: Instant,
39	run: Option<Run>,
40	out: UnixStream,
41	waiting: Option<Reason>,
42	last_pos: usize,
43	watchers: Vec<UnixStream>,
44	}
45
46	impl Lease {
47	fn elapsed(&self) -> Duration {
48	self.run
49	.as_ref()
50	.map_or(self.queued_at, \|r\| r.started)
51	.elapsed()
52	}
53
54	fn overdue(&self, grace: f64) -> bool {
55	self.run.is_some() && self.elapsed() > self.req.need.time.mul_f64(grace)
56	}
57	}
58
59	struct State {
60	cfg: Config,
61	next: LeaseId,
62	leases: BTreeMap<LeaseId, Lease>,
63	cap: Capacity,
64	baselines: BTreeMap<FsId, Bytes>,
65	conns: usize,
66	idle_since: Instant,
67	}
68
69	type Shared = Arc<Mutex<State>>;
70
71	fn lock(shared: &Shared) -> MutexGuard<'_, State> {
72	shared
73	.lock()
74	.unwrap_or_else(std::sync::PoisonError::into_inner)
75	}
76
77	impl State {
78	fn new(cfg: Config) -> Self {
79	let mut st = Self {
80	cfg,
81	next: 1,
82	leases: BTreeMap::new(),
83	cap: Capacity {
84	ram_budget: Bytes::ZERO,
85	ram_live: Bytes::ZERO,
86	ram_reserve: Bytes::ZERO,
87	cores: 1,
88	disks: BTreeMap::new(),
89	pressure: false,
90	},
91	baselines: BTreeMap::new(),
92	conns: 0,
93	idle_since: Instant::now(),
94	};
95	st.apply_budget();
96	st
97	}
98
99	fn apply_budget(&mut self) {
100	let (total, live) = sys::memory();
101	self.cap.ram_budget = self.cfg.budget.ram.resolve(total);
102	self.cap.ram_live = live;
103	self.cap.ram_reserve = self.cfg.budget.ram_reserve;
104	self.cap.cores = self.cfg.budget.cores.resolve(sys::cores());
105	}
106
107	/// running leases keep the heaviness they started with, so a reload can't
108	/// retroactively put two heavies side by side
109	fn reload(&mut self, cfg: Config) {
110	self.cfg = cfg;
111	self.apply_budget();
112	let queued: Vec<(LeaseId, Need)> = self
113	.leases
114	.iter()
115	.filter(\|(_, l)\| l.run.is_none())
116	.map(\|(id, l)\| (*id, l.req.need))
117	.collect();
118	for (id, need) in queued {
119	let heavy = self.heavy(&need);
120	if let Some(l) = self.leases.get_mut(&id) {
121	l.heavy = heavy;
122	}
123	}
124	self.refresh();
125	self.reschedule();
126	}
127
128	fn heavy(&self, need: &Need) -> bool {
129	let h = &self.cfg.heavy;
130	need.ram >= h.ram \|\| need.disk >= h.disk \|\| need.cores >= h.cores.resolve(sys::cores())
131	}
132
133	/// capped at a tenth of the fs, or a 16G tmpfs /tmp with a 20G reserve
134	/// admits nothing ever
135	fn reserve(&self, size: Bytes) -> Bytes {
136	self.cfg.budget.disk_reserve.min(Bytes(size.0 / 10))
137	}
138
139	fn fs_for(&mut self, cwd: &PathBuf) -> Vec<FsId> {
140	let mut ids = Vec::new();
141	for path in std::iter::once(cwd).chain(&self.cfg.budget.disk_paths) {
142	if let Some(fs) = sys::fs(path) {
143	if !ids.contains(&fs.id) {
144	ids.push(fs.id);
145	}
146	let reserve = self.reserve(fs.size);
147	self.cap.disks.entry(fs.id).or_insert(Disk {
148	free: fs.free,
149	reserve,
150	consumed: Bytes::ZERO,
151	});
152	}
153	}
154	ids
155	}
156
157	fn refresh(&mut self) {
158	let (_, live) = sys::memory();
159	self.cap.ram_live = live;
160	self.cap.pressure = self
161	.cfg
162	.budget
163	.psi_max
164	.is_some_and(\|max\| sys::pressure().is_some_and(\|p\| p > max));
165
166	let mut paths: Vec<PathBuf> = self.leases.values().map(\|l\| l.req.cwd.clone()).collect();
167	paths.extend(self.cfg.budget.disk_paths.iter().cloned());
168	let mut disks = BTreeMap::new();
169	for path in paths {
170	if let Some(fs) = sys::fs(&path) {
171	disks.insert(fs.id, (fs.free, self.reserve(fs.size)));
172	}
173	}
174
175	let busy: Vec<FsId> = self
176	.leases
177	.values()
178	.filter(\|l\| l.run.is_some())
179	.flat_map(\|l\| l.fs.iter().copied())
180	.collect();
181	self.baselines.retain(\|id, _\| busy.contains(id));
182	self.cap.disks = disks
183	.into_iter()
184	.map(\|(id, (free, reserve))\| {
185	let consumed = if busy.contains(&id) {
186	self.baselines
187	.entry(id)
188	.or_insert(free)
189	.saturating_sub(free)
190	} else {
191	Bytes::ZERO
192	};
193	(
194	id,
195	Disk {
196	free,
197	reserve,
198	consumed,
199	},
200	)
201	})
202	.collect();
203	}
204
205	fn running(&self) -> Vec<Running> {
206	self.leases
207	.values()
208	.filter_map(\|l\| {
209	let run = l.run.as_ref()?;
210	Some(Running {
211	need: l.req.need,
212	heavy: l.heavy,
213	fs: l.fs.clone(),
214	rss: run.rss,
215	remaining: l.req.need.time.saturating_sub(run.started.elapsed()),
216	})
217	})
218	.collect()
219	}
220
221	fn reschedule(&mut self) {
222	let queue: Vec<Queued> = self
223	.leases
224	.iter()
225	.filter(\|(_, l)\| l.run.is_none())
226	.map(\|(id, l)\| Queued {
227	id: *id,
228	need: l.req.need,
229	heavy: l.heavy,
230	fs: l.fs.clone(),
231	})
232	.collect();
233	if queue.is_empty() {
234	return;
235	}
236	let verdicts = sched::plan(
237	&queue,
238	&self.running(),
239	&self.cap,
240	self.cfg.queue.backfill_max_time.0,
241	);
242	let cores = self.cap.cores;
243	let mut position = 0;
244	for (id, verdict) in verdicts {
245	let Some(lease) = self.leases.get_mut(&id) else {
246	continue;
247	};
248	match verdict {
249	Verdict::Admit => {
250	lease.run = Some(Run {
251	started: Instant::now(),
252	rss: Bytes::ZERO,
253	reminded: None,
254	});
255	let _ = proto::send(
256	&mut lease.out,
257	&Event::Granted {
258	lease: id,
259	cores: lease.req.need.cores.clamp(1, cores.max(1)),
260	heavy: lease.heavy,
261	},
262	);
263	}
264	Verdict::Wait(reason) => {
265	position += 1;
266	lease.waiting = Some(reason);
267	if lease.last_pos != position {
268	lease.last_pos = position;
269	let _ = proto::send(
270	&mut lease.out,
271	&Event::Queued {
272	lease: id,
273	position,
274	reason: Some(reason),
275	},
276	);
277	}
278	}
279	}
280	}
281	}
282
283	fn remove(&mut self, id: LeaseId) {
284	if let Some(lease) = self.leases.remove(&id) {
285	for mut w in lease.watchers {
286	let _ = proto::send(&mut w, &Event::Fired { overdue: false });
287	}
288	}
289	if self.leases.is_empty() {
290	self.idle_since = Instant::now();
291	}
292	self.refresh();
293	self.reschedule();
294	}
295
296	fn fire_overdue(&mut self) {
297	let grace = self.cfg.overrun.time_grace;
298	for lease in self.leases.values_mut() {
299	if !lease.watchers.is_empty() && lease.overdue(grace) {
300	for mut w in lease.watchers.drain(..) {
301	let _ = proto::send(&mut w, &Event::Fired { overdue: true });
302	}
303	}
304	}
305	}
306
307	fn info(&self) -> Vec<LeaseInfo> {
308	let grace = self.cfg.overrun.time_grace;
309	let mut position = 0;
310	self.leases
311	.iter()
312	.map(\|(id, l)\| {
313	if l.run.is_none() {
314	position += 1;
315	}
316	LeaseInfo {
317	id: *id,
318	cmd: l.req.cmd.clone(),
319	cwd: l.req.cwd.clone(),
320	session: l.req.session.clone(),
321	need: l.req.need,
322	heavy: l.heavy,
323	running: l.run.is_some(),
324	elapsed: l.elapsed().as_secs(),
325	rss: l.run.as_ref().map_or(Bytes::ZERO, \|r\| r.rss),
326	overdue: l.overdue(grace),
327	waiting: l.waiting.filter(\|_\| l.run.is_none()),
328	position: if l.run.is_none() { position } else { 0 },
329	}
330	})
331	.collect()
332	}
333
334	fn status(&self) -> Status {
335	let (ram_free, cores_free) = sched::headroom(&self.running(), &self.cap);
336	Status {
337	leases: self.info(),
338	ram_free,
339	ram_budget: self.cap.ram_budget,
340	cores_free,
341	cores: self.cap.cores,
342	}
343	}
344	}
345
346	fn serve(shared: &Shared, mut stream: UnixStream) -> io::Result<()> {
347	let mut reader = Reader::new(&stream)?;
348	let Some(req) = reader.recv::<Request>()? else {
349	return Ok(());
350	};
351	match req {
352	Request::Acquire(req) => {
353	let id = {
354	let mut st = lock(shared);
355	let id = st.next;
356	st.next += 1;
357	let fs = st.fs_for(&req.cwd);
358	let heavy = st.heavy(&req.need);
359	st.leases.insert(
360	id,
361	Lease {
362	req,
363	heavy,
364	fs,
365	queued_at: Instant::now(),
366	run: None,
367	out: stream.try_clone()?,
368	waiting: None,
369	last_pos: 0,
370	watchers: Vec::new(),
371	},
372	);
373	st.reschedule();
374	id
375	};
376	hold(shared, id, &mut reader);
377	lock(shared).remove(id);
378	Ok(())
379	}
380	Request::Status => {
381	let status = lock(shared).status();
382	proto::send(&mut stream, &Event::Status(status))
383	}
384	Request::Kill { lease } => {
385	let mut st = lock(shared);
386	let reply = match st.leases.get_mut(&lease) {
387	Some(l) => {
388	let _ = proto::send(&mut l.out, &Event::Kill);
389	Event::Ok
390	}
391	None => Event::Err {
392	msg: format!("no lease {lease}"),
393	},
394	};
395	drop(st);
396	proto::send(&mut stream, &reply)
397	}
398	Request::Extend { lease, by } => {
399	let mut st = lock(shared);
400	let reply = match st.leases.get_mut(&lease) {
401	Some(l) => {
402	l.req.need.time += by.0;
403	if let Some(run) = &mut l.run {
404	run.reminded = None;
405	}
406	Event::Ok
407	}
408	None => Event::Err {
409	msg: format!("no lease {lease}"),
410	},
411	};
412	drop(st);
413	proto::send(&mut stream, &reply)
414	}
415	Request::Watch { lease } => {
416	let mut st = lock(shared);
417	let grace = st.cfg.overrun.time_grace;
418	match st.leases.get_mut(&lease) {
419	Some(l) if l.overdue(grace) => {
420	proto::send(&mut stream, &Event::Fired { overdue: true })
421	}
422	Some(l) => {
423	l.watchers.push(stream);
424	Ok(())
425	}
426	None => proto::send(&mut stream, &Event::Fired { overdue: false }),
427	}
428	}
429	Request::Overdue { session } => {
430	let mut st = lock(shared);
431	let grace = st.cfg.overrun.time_grace;
432	let every = st.cfg.overrun.remind_every.0;
433	let mut due = Vec::new();
434	for (id, l) in &mut st.leases {
435	if l.req.session.as_deref() != Some(&session) \|\| !l.overdue(grace) {
436	continue;
437	}
438	let Some(run) = &mut l.run else { continue };
439	if run.reminded.is_some_and(\|t\| t.elapsed() < every) {
440	continue;
441	}
442	run.reminded = Some(Instant::now());
443	due.push(*id);
444	}
445	let leases = st
446	.info()
447	.into_iter()
448	.filter(\|i\| due.contains(&i.id))
449	.collect();
450	drop(st);
451	proto::send(&mut stream, &Event::Overdue { leases })
452	}
453	Request::Started { .. } \| Request::Usage { .. } \| Request::Done => proto::send(
454	&mut stream,
455	&Event::Err {
456	msg: "no lease on this connection".into(),
457	},
458	),
459	}
460	}
461
462	/// a client that vanishes without saying done got sigkilled or similar, and
463	/// pdeathsig only reaches its direct child. take the rest of the group down
464	/// so a headless build doesn't keep eating ram off the books
465	fn hold(shared: &Shared, id: LeaseId, reader: &mut Reader) {
466	let mut pgid = None;
467	let clean = loop {
468	match reader.recv::<Request>() {
469	Ok(Some(Request::Usage { rss })) => {
470	if let Some(run) = lock(shared)
471	.leases
472	.get_mut(&id)
473	.and_then(\|l\| l.run.as_mut())
474	{
475	run.rss = rss;
476	}
477	}
478	Ok(Some(Request::Started { pid })) => pgid = Some(pid),
479	Ok(Some(Request::Done)) => break true,
480	Ok(Some(_)) => {}
481	Ok(None) \| Err(_) => break false,
482	}
483	};
484	if !clean && let Some(pgid) = pgid.filter(\|p\| *p > 1) {
485	unsafe { libc::kill(-pgid, libc::SIGTERM) };
486	}
487	}
488
489	fn tick(shared: &Shared) {
490	loop {
491	thread::sleep(Duration::from_secs(1));
492	let mut st = lock(shared);
493	st.refresh();
494	st.reschedule();
495	st.fire_overdue();
496	if st.conns == 0 && st.leases.is_empty() && st.idle_since.elapsed() >= st.cfg.queue.linger.0
497	{
498	// systemd owns the socket and queues whoever knocks while we're
499	// leaving, so there's nothing to tidy
500	if ACTIVATED.load(Ordering::Relaxed) {
501	process::exit(0);
502	}
503	// a spawner holds this while it clears a stale socket, so taking it
504	// here stops us unlinking our successor's socket on the way out
505	if let Some(_guard) = try_spawn_lock() {
506	let _ = fs::remove_file(paths::socket());
507	process::exit(0);
508	}
509	}
510	}
511	}
512
513	/// watches the directory, not the file. editors save by renaming over it,
514	/// which would orphan a watch on the old inode after the first edit
515	#[cfg(target_os = "linux")]
516	fn watch_config(shared: &Shared) {
517	use std::os::unix::ffi::OsStrExt;
518
519	let path = paths::config_file();
520	let (Some(dir), Some(name)) = (path.parent(), path.file_name()) else {
521	return;
522	};
523	let _ = fs::create_dir_all(dir);
524	let Ok(cdir) = std::ffi::CString::new(dir.as_os_str().as_bytes()) else {
525	return;
526	};
527	let fd = unsafe { libc::inotify_init1(libc::IN_CLOEXEC) };
528	if fd < 0 {
529	return;
530	}
531	let mask = libc::IN_CLOSE_WRITE
532	\| libc::IN_MOVED_TO
533	\| libc::IN_MOVED_FROM
534	\| libc::IN_CREATE
535	\| libc::IN_DELETE;
536	if unsafe { libc::inotify_add_watch(fd, cdir.as_ptr(), mask) } < 0 {
537	return;
538	}
539
540	let header = std::mem::size_of::<libc::inotify_event>();
541	let mut buf = [0u8; 4096];
542	loop {
543	let n = unsafe { libc::read(fd, buf.as_mut_ptr().cast(), buf.len()) };
544	let Ok(n) = usize::try_from(n) else {
545	return;
546	};
547	let mut hit = false;
548	let mut off = 0;
549	while off + header <= n {
550	let ev: libc::inotify_event =
551	unsafe { std::ptr::read_unaligned(buf.as_ptr().add(off).cast()) };
552	let len = ev.len as usize;
553	let raw = &buf[off + header..(off + header + len).min(n)];
554	let ev_name = raw.split(\|b\| *b == 0).next().unwrap_or_default();
555	hit \|= ev_name == name.as_bytes();
556	off += header + len;
557	}
558	if !hit {
559	continue;
560	}
561	thread::sleep(Duration::from_millis(100));
562	match Config::load() {
563	Ok(cfg) => {
564	lock(shared).reload(cfg);
565	eprintln!("roundabout: config reloaded");
566	}
567	Err(e) => eprintln!("roundabout: kept the old config, {e}"),
568	}
569	}
570	}
571
572	/// creates the file if it's gone, else a wiped runtime dir leaves us idling
573	/// forever unable to take the lock we need to exit
574	fn try_spawn_lock() -> Option<client::Flock> {
575	let file = fs::OpenOptions::new()
576	.create(true)
577	.truncate(false)
578	.write(true)
579	.open(paths::spawn_lock())
580	.ok()?;
581	client::Flock::try_exclusive(file)
582	}
583
584	static ACTIVATED: AtomicBool = AtomicBool::new(false);
585
586	fn inherited() -> Option<UnixListener> {
587	let ours = std::env::var("LISTEN_PID").ok()?.parse::<u32>().ok()? == process::id();
588	if !ours \|\| std::env::var("LISTEN_FDS").ok()? != "1" {
589	return None;
590	}
591	// SAFETY: sd_listen_fds protocol, fd 3 is the listening socket
592	Some(unsafe { UnixListener::from_raw_fd(3) })
593	}
594
595	fn bind() -> UnixListener {
596	// fork once more so the client that spawned us can reap its child now
597	// instead of carrying a zombie for the length of its build
598	match unsafe { libc::fork() } {
599	-1 => {
600	eprintln!("roundabout: fork failed: {}", io::Error::last_os_error());
601	process::exit(1);
602	}
603	0 => {}
604	_ => process::exit(0),
605	}
606	match UnixListener::bind(paths::socket()) {
607	Ok(l) => l,
608	Err(e) if e.kind() == io::ErrorKind::AddrInUse && client::connect().is_some() => {
609	process::exit(0)
610	}
611	Err(e) => {
612	eprintln!("roundabout: bind: {e}");
613	process::exit(1);
614	}
615	}
616	}
617
618	pub fn main(cfg: Config) -> ! {
619	let listener = if let Some(l) = inherited() {
620	ACTIVATED.store(true, Ordering::Relaxed);
621	l
622	} else {
623	bind()
624	};
625
626	let shared: Shared = Arc::new(Mutex::new(State::new(cfg)));
627	{
628	let shared = shared.clone();
629	thread::spawn(move \|\| tick(&shared));
630	}
631	#[cfg(target_os = "linux")]
632	{
633	let shared = shared.clone();
634	thread::spawn(move \|\| watch_config(&shared));
635	}
636
637	for stream in listener.incoming() {
638	let Ok(stream) = stream else { continue };
639	lock(&shared).conns += 1;
640	let shared = shared.clone();
641	thread::spawn(move \|\| {
642	if let Err(e) = serve(&shared, stream) {
643	eprintln!("roundabout: {e}");
644	}
645	let mut st = lock(&shared);
646	st.conns -= 1;
647	if st.leases.is_empty() {
648	st.idle_since = Instant::now();
649	}
650	});
651	}
652	process::exit(0)
653	}