| 1 | use std::{ |
| 2 | collections::BTreeMap, |
| 3 | fs, io, |
| 4 | os::unix::net::{UnixListener, UnixStream}, |
| 5 | path::PathBuf, |
| 6 | process, |
| 7 | sync::{Arc, Mutex, MutexGuard}, |
| 8 | thread, |
| 9 | time::{Duration, Instant}, |
| 10 | }; |
| 11 | |
| 12 | use crate::{ |
| 13 | client, |
| 14 | config::Config, |
| 15 | paths, |
| 16 | proto::{self, Acquire, Event, LeaseId, LeaseInfo, Reader, Request, Status}, |
| 17 | sched::{self, Capacity, Disk, FsId, Need, Queued, Reason, Running, Verdict}, |
| 18 | sys, |
| 19 | units::Bytes, |
| 20 | }; |
| 21 | |
| 22 | struct Run { |
| 23 | started: Instant, |
| 24 | rss: Bytes, |
| 25 | reminded: Option<Instant>, |
| 26 | } |
| 27 | |
| 28 | struct Lease { |
| 29 | req: Acquire, |
| 30 | heavy: bool, |
| 31 | fs: Vec<FsId>, |
| 32 | queued_at: Instant, |
| 33 | run: Option<Run>, |
| 34 | out: UnixStream, |
| 35 | waiting: Option<Reason>, |
| 36 | last_pos: usize, |
| 37 | watchers: Vec<UnixStream>, |
| 38 | } |
| 39 | |
| 40 | impl Lease { |
| 41 | fn elapsed(&self) -> Duration { |
| 42 | self.run |
| 43 | .as_ref() |
| 44 | .map_or(self.queued_at, |r| r.started) |
| 45 | .elapsed() |
| 46 | } |
| 47 | |
| 48 | fn overdue(&self, grace: f64) -> bool { |
| 49 | self.run.is_some() && self.elapsed() > self.req.need.time.mul_f64(grace) |
| 50 | } |
| 51 | } |
| 52 | |
| 53 | struct State { |
| 54 | cfg: Config, |
| 55 | next: LeaseId, |
| 56 | leases: BTreeMap<LeaseId, Lease>, |
| 57 | cap: Capacity, |
| 58 | baselines: BTreeMap<FsId, Bytes>, |
| 59 | conns: usize, |
| 60 | idle_since: Instant, |
| 61 | } |
| 62 | |
| 63 | type Shared = Arc<Mutex<State>>; |
| 64 | |
| 65 | fn lock(shared: &Shared) -> MutexGuard<'_, State> { |
| 66 | shared |
| 67 | .lock() |
| 68 | .unwrap_or_else(std::sync::PoisonError::into_inner) |
| 69 | } |
| 70 | |
| 71 | impl State { |
| 72 | fn new(cfg: Config) -> Self { |
| 73 | let (total, live) = sys::memory(); |
| 74 | let cap = Capacity { |
| 75 | ram_budget: cfg.budget.ram.resolve(total), |
| 76 | ram_live: live, |
| 77 | ram_reserve: cfg.budget.ram_reserve, |
| 78 | cores: cfg.budget.cores.resolve(sys::cores()), |
| 79 | disks: BTreeMap::new(), |
| 80 | pressure: false, |
| 81 | }; |
| 82 | Self { |
| 83 | cfg, |
| 84 | next: 1, |
| 85 | leases: BTreeMap::new(), |
| 86 | cap, |
| 87 | baselines: BTreeMap::new(), |
| 88 | conns: 0, |
| 89 | idle_since: Instant::now(), |
| 90 | } |
| 91 | } |
| 92 | |
| 93 | fn heavy(&self, need: &Need) -> bool { |
| 94 | let h = &self.cfg.heavy; |
| 95 | need.ram >= h.ram || need.disk >= h.disk || need.cores >= h.cores.resolve(sys::cores()) |
| 96 | } |
| 97 | |
| 98 | /// capped at a tenth of the fs, or a 16G tmpfs /tmp with a 20G reserve |
| 99 | /// admits nothing ever |
| 100 | fn reserve(&self, size: Bytes) -> Bytes { |
| 101 | self.cfg.budget.disk_reserve.min(Bytes(size.0 / 10)) |
| 102 | } |
| 103 | |
| 104 | fn fs_for(&mut self, cwd: &PathBuf) -> Vec<FsId> { |
| 105 | let mut ids = Vec::new(); |
| 106 | for path in std::iter::once(cwd).chain(&self.cfg.budget.disk_paths) { |
| 107 | if let Some(fs) = sys::fs(path) { |
| 108 | if !ids.contains(&fs.id) { |
| 109 | ids.push(fs.id); |
| 110 | } |
| 111 | let reserve = self.reserve(fs.size); |
| 112 | self.cap.disks.entry(fs.id).or_insert(Disk { |
| 113 | free: fs.free, |
| 114 | reserve, |
| 115 | consumed: Bytes::ZERO, |
| 116 | }); |
| 117 | } |
| 118 | } |
| 119 | ids |
| 120 | } |
| 121 | |
| 122 | fn refresh(&mut self) { |
| 123 | let (_, live) = sys::memory(); |
| 124 | self.cap.ram_live = live; |
| 125 | self.cap.pressure = self |
| 126 | .cfg |
| 127 | .budget |
| 128 | .psi_max |
| 129 | .is_some_and(|max| sys::pressure().is_some_and(|p| p > max)); |
| 130 | |
| 131 | let mut paths: Vec<PathBuf> = self.leases.values().map(|l| l.req.cwd.clone()).collect(); |
| 132 | paths.extend(self.cfg.budget.disk_paths.iter().cloned()); |
| 133 | let mut disks = BTreeMap::new(); |
| 134 | for path in paths { |
| 135 | if let Some(fs) = sys::fs(&path) { |
| 136 | disks.insert(fs.id, (fs.free, self.reserve(fs.size))); |
| 137 | } |
| 138 | } |
| 139 | |
| 140 | let busy: Vec<FsId> = self |
| 141 | .leases |
| 142 | .values() |
| 143 | .filter(|l| l.run.is_some()) |
| 144 | .flat_map(|l| l.fs.iter().copied()) |
| 145 | .collect(); |
| 146 | self.baselines.retain(|id, _| busy.contains(id)); |
| 147 | self.cap.disks = disks |
| 148 | .into_iter() |
| 149 | .map(|(id, (free, reserve))| { |
| 150 | let consumed = if busy.contains(&id) { |
| 151 | self.baselines |
| 152 | .entry(id) |
| 153 | .or_insert(free) |
| 154 | .saturating_sub(free) |
| 155 | } else { |
| 156 | Bytes::ZERO |
| 157 | }; |
| 158 | ( |
| 159 | id, |
| 160 | Disk { |
| 161 | free, |
| 162 | reserve, |
| 163 | consumed, |
| 164 | }, |
| 165 | ) |
| 166 | }) |
| 167 | .collect(); |
| 168 | } |
| 169 | |
| 170 | fn running(&self) -> Vec<Running> { |
| 171 | self.leases |
| 172 | .values() |
| 173 | .filter_map(|l| { |
| 174 | let run = l.run.as_ref()?; |
| 175 | Some(Running { |
| 176 | need: l.req.need, |
| 177 | heavy: l.heavy, |
| 178 | fs: l.fs.clone(), |
| 179 | rss: run.rss, |
| 180 | remaining: l.req.need.time.saturating_sub(run.started.elapsed()), |
| 181 | }) |
| 182 | }) |
| 183 | .collect() |
| 184 | } |
| 185 | |
| 186 | fn reschedule(&mut self) { |
| 187 | let queue: Vec<Queued> = self |
| 188 | .leases |
| 189 | .iter() |
| 190 | .filter(|(_, l)| l.run.is_none()) |
| 191 | .map(|(id, l)| Queued { |
| 192 | id: *id, |
| 193 | need: l.req.need, |
| 194 | heavy: l.heavy, |
| 195 | fs: l.fs.clone(), |
| 196 | }) |
| 197 | .collect(); |
| 198 | if queue.is_empty() { |
| 199 | return; |
| 200 | } |
| 201 | let verdicts = sched::plan( |
| 202 | &queue, |
| 203 | &self.running(), |
| 204 | &self.cap, |
| 205 | self.cfg.queue.backfill_max_time.0, |
| 206 | ); |
| 207 | let cores = self.cap.cores; |
| 208 | let mut position = 0; |
| 209 | for (id, verdict) in verdicts { |
| 210 | let Some(lease) = self.leases.get_mut(&id) else { |
| 211 | continue; |
| 212 | }; |
| 213 | match verdict { |
| 214 | Verdict::Admit => { |
| 215 | lease.run = Some(Run { |
| 216 | started: Instant::now(), |
| 217 | rss: Bytes::ZERO, |
| 218 | reminded: None, |
| 219 | }); |
| 220 | let _ = proto::send( |
| 221 | &mut lease.out, |
| 222 | &Event::Granted { |
| 223 | lease: id, |
| 224 | cores: lease.req.need.cores.clamp(1, cores.max(1)), |
| 225 | heavy: lease.heavy, |
| 226 | }, |
| 227 | ); |
| 228 | } |
| 229 | Verdict::Wait(reason) => { |
| 230 | position += 1; |
| 231 | lease.waiting = Some(reason); |
| 232 | if lease.last_pos != position { |
| 233 | lease.last_pos = position; |
| 234 | let _ = proto::send( |
| 235 | &mut lease.out, |
| 236 | &Event::Queued { |
| 237 | lease: id, |
| 238 | position, |
| 239 | reason: Some(reason), |
| 240 | }, |
| 241 | ); |
| 242 | } |
| 243 | } |
| 244 | } |
| 245 | } |
| 246 | } |
| 247 | |
| 248 | fn remove(&mut self, id: LeaseId) { |
| 249 | if let Some(lease) = self.leases.remove(&id) { |
| 250 | for mut w in lease.watchers { |
| 251 | let _ = proto::send(&mut w, &Event::Fired { overdue: false }); |
| 252 | } |
| 253 | } |
| 254 | if self.leases.is_empty() { |
| 255 | self.idle_since = Instant::now(); |
| 256 | } |
| 257 | self.refresh(); |
| 258 | self.reschedule(); |
| 259 | } |
| 260 | |
| 261 | fn fire_overdue(&mut self) { |
| 262 | let grace = self.cfg.overrun.time_grace; |
| 263 | for lease in self.leases.values_mut() { |
| 264 | if !lease.watchers.is_empty() && lease.overdue(grace) { |
| 265 | for mut w in lease.watchers.drain(..) { |
| 266 | let _ = proto::send(&mut w, &Event::Fired { overdue: true }); |
| 267 | } |
| 268 | } |
| 269 | } |
| 270 | } |
| 271 | |
| 272 | fn info(&self) -> Vec<LeaseInfo> { |
| 273 | let grace = self.cfg.overrun.time_grace; |
| 274 | let mut position = 0; |
| 275 | self.leases |
| 276 | .iter() |
| 277 | .map(|(id, l)| { |
| 278 | if l.run.is_none() { |
| 279 | position += 1; |
| 280 | } |
| 281 | LeaseInfo { |
| 282 | id: *id, |
| 283 | cmd: l.req.cmd.clone(), |
| 284 | cwd: l.req.cwd.clone(), |
| 285 | session: l.req.session.clone(), |
| 286 | need: l.req.need, |
| 287 | heavy: l.heavy, |
| 288 | running: l.run.is_some(), |
| 289 | elapsed: l.elapsed().as_secs(), |
| 290 | rss: l.run.as_ref().map_or(Bytes::ZERO, |r| r.rss), |
| 291 | overdue: l.overdue(grace), |
| 292 | waiting: l.waiting.filter(|_| l.run.is_none()), |
| 293 | position: if l.run.is_none() { position } else { 0 }, |
| 294 | } |
| 295 | }) |
| 296 | .collect() |
| 297 | } |
| 298 | |
| 299 | fn status(&self) -> Status { |
| 300 | let (ram_free, cores_free) = sched::headroom(&self.running(), &self.cap); |
| 301 | Status { |
| 302 | leases: self.info(), |
| 303 | ram_free, |
| 304 | ram_budget: self.cap.ram_budget, |
| 305 | cores_free, |
| 306 | cores: self.cap.cores, |
| 307 | } |
| 308 | } |
| 309 | } |
| 310 | |
| 311 | fn serve(shared: &Shared, mut stream: UnixStream) -> io::Result<()> { |
| 312 | let mut reader = Reader::new(&stream)?; |
| 313 | let Some(req) = reader.recv::<Request>()? else { |
| 314 | return Ok(()); |
| 315 | }; |
| 316 | match req { |
| 317 | Request::Acquire(req) => { |
| 318 | let id = { |
| 319 | let mut st = lock(shared); |
| 320 | let id = st.next; |
| 321 | st.next += 1; |
| 322 | let fs = st.fs_for(&req.cwd); |
| 323 | let heavy = st.heavy(&req.need); |
| 324 | st.leases.insert( |
| 325 | id, |
| 326 | Lease { |
| 327 | req, |
| 328 | heavy, |
| 329 | fs, |
| 330 | queued_at: Instant::now(), |
| 331 | run: None, |
| 332 | out: stream.try_clone()?, |
| 333 | waiting: None, |
| 334 | last_pos: 0, |
| 335 | watchers: Vec::new(), |
| 336 | }, |
| 337 | ); |
| 338 | st.reschedule(); |
| 339 | id |
| 340 | }; |
| 341 | hold(shared, id, &mut reader); |
| 342 | lock(shared).remove(id); |
| 343 | Ok(()) |
| 344 | } |
| 345 | Request::Status => { |
| 346 | let status = lock(shared).status(); |
| 347 | proto::send(&mut stream, &Event::Status(status)) |
| 348 | } |
| 349 | Request::Kill { lease } => { |
| 350 | let mut st = lock(shared); |
| 351 | let reply = match st.leases.get_mut(&lease) { |
| 352 | Some(l) => { |
| 353 | let _ = proto::send(&mut l.out, &Event::Kill); |
| 354 | Event::Ok |
| 355 | } |
| 356 | None => Event::Err { |
| 357 | msg: format!("no lease {lease}"), |
| 358 | }, |
| 359 | }; |
| 360 | drop(st); |
| 361 | proto::send(&mut stream, &reply) |
| 362 | } |
| 363 | Request::Extend { lease, by } => { |
| 364 | let mut st = lock(shared); |
| 365 | let reply = match st.leases.get_mut(&lease) { |
| 366 | Some(l) => { |
| 367 | l.req.need.time += by.0; |
| 368 | if let Some(run) = &mut l.run { |
| 369 | run.reminded = None; |
| 370 | } |
| 371 | Event::Ok |
| 372 | } |
| 373 | None => Event::Err { |
| 374 | msg: format!("no lease {lease}"), |
| 375 | }, |
| 376 | }; |
| 377 | drop(st); |
| 378 | proto::send(&mut stream, &reply) |
| 379 | } |
| 380 | Request::Watch { lease } => { |
| 381 | let mut st = lock(shared); |
| 382 | let grace = st.cfg.overrun.time_grace; |
| 383 | match st.leases.get_mut(&lease) { |
| 384 | Some(l) if l.overdue(grace) => { |
| 385 | proto::send(&mut stream, &Event::Fired { overdue: true }) |
| 386 | } |
| 387 | Some(l) => { |
| 388 | l.watchers.push(stream); |
| 389 | Ok(()) |
| 390 | } |
| 391 | None => proto::send(&mut stream, &Event::Fired { overdue: false }), |
| 392 | } |
| 393 | } |
| 394 | Request::Overdue { session } => { |
| 395 | let mut st = lock(shared); |
| 396 | let grace = st.cfg.overrun.time_grace; |
| 397 | let every = st.cfg.overrun.remind_every.0; |
| 398 | let mut due = Vec::new(); |
| 399 | for (id, l) in &mut st.leases { |
| 400 | if l.req.session.as_deref() != Some(&session) || !l.overdue(grace) { |
| 401 | continue; |
| 402 | } |
| 403 | let Some(run) = &mut l.run else { continue }; |
| 404 | if run.reminded.is_some_and(|t| t.elapsed() < every) { |
| 405 | continue; |
| 406 | } |
| 407 | run.reminded = Some(Instant::now()); |
| 408 | due.push(*id); |
| 409 | } |
| 410 | let leases = st |
| 411 | .info() |
| 412 | .into_iter() |
| 413 | .filter(|i| due.contains(&i.id)) |
| 414 | .collect(); |
| 415 | drop(st); |
| 416 | proto::send(&mut stream, &Event::Overdue { leases }) |
| 417 | } |
| 418 | Request::Started { .. } | Request::Usage { .. } | Request::Done => proto::send( |
| 419 | &mut stream, |
| 420 | &Event::Err { |
| 421 | msg: "no lease on this connection".into(), |
| 422 | }, |
| 423 | ), |
| 424 | } |
| 425 | } |
| 426 | |
| 427 | /// a client that vanishes without saying done got sigkilled or similar, and |
| 428 | /// pdeathsig only reaches its direct child. take the rest of the group down |
| 429 | /// so a headless build doesn't keep eating ram off the books |
| 430 | fn hold(shared: &Shared, id: LeaseId, reader: &mut Reader) { |
| 431 | let mut pgid = None; |
| 432 | let clean = loop { |
| 433 | match reader.recv::<Request>() { |
| 434 | Ok(Some(Request::Usage { rss })) => { |
| 435 | if let Some(run) = lock(shared) |
| 436 | .leases |
| 437 | .get_mut(&id) |
| 438 | .and_then(|l| l.run.as_mut()) |
| 439 | { |
| 440 | run.rss = rss; |
| 441 | } |
| 442 | } |
| 443 | Ok(Some(Request::Started { pid })) => pgid = Some(pid), |
| 444 | Ok(Some(Request::Done)) => break true, |
| 445 | Ok(Some(_)) => {} |
| 446 | Ok(None) | Err(_) => break false, |
| 447 | } |
| 448 | }; |
| 449 | if !clean && let Some(pgid) = pgid.filter(|p| *p > 1) { |
| 450 | unsafe { libc::kill(-pgid, libc::SIGTERM) }; |
| 451 | } |
| 452 | } |
| 453 | |
| 454 | fn tick(shared: &Shared) { |
| 455 | loop { |
| 456 | thread::sleep(Duration::from_secs(1)); |
| 457 | let mut st = lock(shared); |
| 458 | st.refresh(); |
| 459 | st.reschedule(); |
| 460 | st.fire_overdue(); |
| 461 | if st.conns == 0 && st.leases.is_empty() && st.idle_since.elapsed() >= st.cfg.queue.linger.0 |
| 462 | { |
| 463 | // a spawner holds this while it clears a stale socket, so taking it |
| 464 | // here stops us unlinking our successor's socket on the way out |
| 465 | if let Some(_guard) = try_spawn_lock() { |
| 466 | let _ = fs::remove_file(paths::socket()); |
| 467 | process::exit(0); |
| 468 | } |
| 469 | } |
| 470 | } |
| 471 | } |
| 472 | |
| 473 | fn try_spawn_lock() -> Option<client::Flock> { |
| 474 | let file = fs::OpenOptions::new() |
| 475 | .write(true) |
| 476 | .open(paths::spawn_lock()) |
| 477 | .ok()?; |
| 478 | client::Flock::try_exclusive(file) |
| 479 | } |
| 480 | |
| 481 | pub fn main(cfg: Config) -> ! { |
| 482 | // fork once more so the client that spawned us can reap its child now |
| 483 | // instead of carrying a zombie for the length of its build |
| 484 | match unsafe { libc::fork() } { |
| 485 | -1 => { |
| 486 | eprintln!("roundabout: fork failed: {}", io::Error::last_os_error()); |
| 487 | process::exit(1); |
| 488 | } |
| 489 | 0 => {} |
| 490 | _ => process::exit(0), |
| 491 | } |
| 492 | |
| 493 | let listener = match UnixListener::bind(paths::socket()) { |
| 494 | Ok(l) => l, |
| 495 | Err(e) if e.kind() == io::ErrorKind::AddrInUse && client::connect().is_some() => { |
| 496 | process::exit(0) |
| 497 | } |
| 498 | Err(e) => { |
| 499 | eprintln!("roundabout: bind: {e}"); |
| 500 | process::exit(1); |
| 501 | } |
| 502 | }; |
| 503 | |
| 504 | let shared: Shared = Arc::new(Mutex::new(State::new(cfg))); |
| 505 | { |
| 506 | let shared = shared.clone(); |
| 507 | thread::spawn(move || tick(&shared)); |
| 508 | } |
| 509 | |
| 510 | for stream in listener.incoming() { |
| 511 | let Ok(stream) = stream else { continue }; |
| 512 | lock(&shared).conns += 1; |
| 513 | let shared = shared.clone(); |
| 514 | thread::spawn(move || { |
| 515 | if let Err(e) = serve(&shared, stream) { |
| 516 | eprintln!("roundabout: {e}"); |
| 517 | } |
| 518 | let mut st = lock(&shared); |
| 519 | st.conns -= 1; |
| 520 | if st.leases.is_empty() { |
| 521 | st.idle_since = Instant::now(); |
| 522 | } |
| 523 | }); |
| 524 | } |
| 525 | process::exit(0) |
| 526 | } |