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