1use std::collections::{BTreeMap, BTreeSet};
6use std::time::{Duration, SystemTime};
7
8use crate::api::Waiting;
9use crate::github::QueuedJob;
10
11pub mod engine;
12pub mod warm;
13
14pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
17
18pub const CPU_PRIORITY_AFTER: Duration = Duration::from_secs(120);
21
22#[derive(Debug, Clone, Copy)]
24pub struct Hold {
25 pub until: SystemTime,
27 pub cause: crate::api::HoldCause,
29}
30
31#[derive(Debug, Clone)]
33pub struct Repository {
34 pub name: String,
36 pub profile: String,
38 pub labels: Vec<String>,
40 pub paused: bool,
42 pub held: Option<Hold>,
44 pub queued: Vec<QueuedJob>,
46 pub runner_memory: Option<u64>,
49 pub runner_cpu: Option<u64>,
52 pub cpu_waiting_since: Option<SystemTime>,
54 pub warm: usize,
56}
57
58impl Repository {
59 fn starved_of_cpu(&self, now: SystemTime) -> bool {
61 self.cpu_waiting_since
62 .and_then(|since| now.duration_since(since).ok())
63 .is_some_and(|waited| waited >= CPU_PRIORITY_AFTER)
64 }
65}
66
67#[derive(Debug, Clone)]
69pub struct Profile {
70 pub memory: u64,
72 pub cpus: u64,
75 pub preparing: Option<Preparing>,
77 pub upkeep_due: bool,
80 pub count: usize,
82 pub slots: Vec<Slot>,
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum Preparing {
89 Reseeding,
91 PullingImage,
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
97pub enum Slot {
98 Free,
100 Runner(RunnerSlot),
102 Closed,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq)]
109pub struct RunnerSlot {
110 pub repository: String,
112 pub busy: bool,
114 pub started_at: SystemTime,
116 pub memory: u64,
118 pub cpu: u64,
120}
121
122#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
124pub struct SlotRef {
125 pub profile: String,
127 pub index: usize,
129}
130
131#[derive(Debug, Clone, PartialEq, Eq)]
133pub enum Decision {
134 Start {
136 slot: SlotRef,
138 repository: String,
140 memory: u64,
142 cpu: u64,
144 warm: bool,
146 replacing: Vec<SlotRef>,
149 },
150 HeldForCpu {
152 repository: String,
154 },
155 Stop {
158 slot: SlotRef,
160 },
161}
162
163#[derive(Debug, Clone, Default, PartialEq, Eq)]
165pub struct Plan {
166 pub decisions: Vec<Decision>,
168 pub waiting: BTreeMap<u64, Waiting>,
172 pub warm: BTreeSet<SlotRef>,
174}
175
176#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
178pub struct Budget {
179 pub memory: Option<u64>,
181 pub cpu: Option<u64>,
183 pub host: Option<u64>,
185 pub held: bool,
188}
189
190pub fn matches(job: &QueuedJob, labels: &[String]) -> bool {
193 job.labels.iter().all(|wanted| {
194 labels
195 .iter()
196 .any(|label| label.eq_ignore_ascii_case(wanted))
197 })
198}
199
200pub fn plan(
213 repositories: &[Repository],
214 profiles: &BTreeMap<String, Profile>,
215 budget: Budget,
216 now: SystemTime,
217) -> Plan {
218 let mut decided = Plan::default();
219 let mut wanted = BTreeMap::new();
220 let mut warm_wanted = Vec::new();
221 let mut jobs = Vec::new();
222 for repository in repositories {
223 if let Some(sorted) = sort_out(repository, profiles, now, &mut decided) {
224 wanted.insert(repository.name.as_str(), sorted.serve.len());
225 jobs.extend(sorted.serve.into_iter().map(|job| (job, repository)));
226 warm_wanted.push((repository, sorted.warm));
227 }
228 }
229 jobs.sort_by_key(|(id, _)| *id);
230
231 let mut warm = Warm::new(&decided.warm, profiles, budget.held);
232 let mut taken = taken_slots(profiles);
233 let mut left = Left {
234 memory: budget.memory,
235 cpu: budget.cpu,
236 host: budget.host,
237 };
238 let mut started = Started::new(&taken, left);
239 let mut first_for_cpu = None;
241 for (job, repository) in jobs {
242 let Some(missing) = wanted.get_mut(repository.name.as_str()) else {
243 continue;
244 };
245 if *missing == 0 {
246 continue;
247 }
248 let (Some(profile), Some(taken)) = (
249 profiles.get(&repository.profile),
250 taken.get_mut(repository.profile.as_str()),
251 ) else {
252 continue;
253 };
254 if let Some(older) = first_for_cpu {
255 decided
256 .waiting
257 .insert(job, Waiting::BehindCpu { job: older });
258 continue;
259 }
260 let memory = repository.runner_memory.unwrap_or(profile.memory);
261 let cpu = repository.runner_cpu.unwrap_or(profile.cpus);
262 let free_index = taken.iter().position(|taken| !taken);
263 let fits = free_index.is_some() && left.fits(memory, cpu);
264 let room = if fits || profile.preparing.is_some() {
265 None
266 } else {
267 warm.make_room(repository, free_index.is_none(), left.short(memory, cpu))
268 };
269 let index = match (room.as_ref(), free_index) {
270 (Some(room), free) => free.or(room.slot).unwrap_or_default(),
271 (None, Some(index)) if fits => index,
272 (None, free) => {
273 let (reason, short) = why_waits(
274 &repository.profile,
275 profile,
276 free.is_some(),
277 memory,
278 cpu,
279 left,
280 &started,
281 );
282 decided.waiting.insert(job, reason);
283 if short == Short::Cpu {
284 let held = Decision::HeldForCpu {
285 repository: repository.name.clone(),
286 };
287 if !decided.decisions.contains(&held) {
288 decided.decisions.push(held);
289 }
290 if repository.starved_of_cpu(now) {
292 first_for_cpu = Some(job);
293 }
294 }
295 continue;
296 }
297 };
298 let replacing = room.map_or_else(Vec::new, |room| {
300 left.give(room.memory, room.cpu);
301 for slot in &room.stopped {
302 decided.warm.remove(slot);
303 decided
304 .decisions
305 .push(Decision::Stop { slot: slot.clone() });
306 }
307 room.stopped
308 });
309 taken[index] = true;
310 *missing -= 1;
311 started.note(&repository.profile, job);
312 left.take(memory, cpu);
313 decided
314 .decisions
315 .push(start(repository, index, memory, cpu, replacing));
316 }
317 if !budget.held {
318 start_warm(&warm_wanted, profiles, &mut taken, left, &mut decided);
319 }
320 decided
321}
322
323#[derive(Debug, Clone, Copy, PartialEq, Eq)]
325enum Short {
326 Slot,
327 Memory,
328 Cpu,
329}
330
331fn why_waits(
334 name: &str,
335 profile: &Profile,
336 free_slot: bool,
337 memory: u64,
338 cpu: u64,
339 left: Left,
340 started: &Started<'_>,
341) -> (Waiting, Short) {
342 let profile_name = || name.to_owned();
343 if !free_slot {
344 let reason = match profile.preparing {
345 Some(Preparing::Reseeding) => Waiting::Reseeding {
346 profile: profile_name(),
347 },
348 Some(Preparing::PullingImage) => Waiting::PullingImage {
349 profile: profile_name(),
350 },
351 None => started.for_slot(
352 name,
353 Waiting::NoSlot {
354 profile: profile_name(),
355 },
356 ),
357 };
358 return (reason, Short::Slot);
359 }
360 match (left.memory, left.host, left.cpu) {
361 (Some(free), _, _) if free < memory => (
362 started.for_memory(
363 memory,
364 Waiting::Memory {
365 needs: memory,
366 free,
367 },
368 ),
369 Short::Memory,
370 ),
371 (_, Some(free), _) if free < memory => (
372 started.for_host(
373 memory,
374 Waiting::HostMemory {
375 needs: memory,
376 free,
377 },
378 ),
379 Short::Memory,
380 ),
381 (_, _, Some(free)) => (
382 started.for_cpu(cpu, Waiting::Cpu { needs: cpu, free }),
383 Short::Cpu,
384 ),
385 (_, _, None) => (
387 Waiting::NoSlot {
388 profile: profile_name(),
389 },
390 Short::Slot,
391 ),
392 }
393}
394
395fn start(
398 repository: &Repository,
399 index: usize,
400 memory: u64,
401 cpu: u64,
402 replacing: Vec<SlotRef>,
403) -> Decision {
404 Decision::Start {
405 slot: SlotRef {
406 profile: repository.profile.clone(),
407 index,
408 },
409 repository: repository.name.clone(),
410 memory,
411 cpu,
412 warm: false,
413 replacing,
414 }
415}
416
417fn start_warm(
421 wanted: &[(&Repository, usize)],
422 profiles: &BTreeMap<String, Profile>,
423 taken: &mut BTreeMap<&str, Vec<bool>>,
424 mut left: Left,
425 decided: &mut Plan,
426) {
427 let room_wanted = decided.waiting.values().any(|reason| {
428 matches!(
429 reason,
430 Waiting::NoSlot { .. }
431 | Waiting::Memory { .. }
432 | Waiting::HostMemory { .. }
433 | Waiting::Cpu { .. }
434 | Waiting::BehindCpu { .. }
435 | Waiting::BehindStart { .. }
436 )
437 });
438 if room_wanted {
439 return;
440 }
441 for round in 0.. {
442 let mut any = false;
443 for (repository, wants) in wanted {
444 if *wants <= round {
445 continue;
446 }
447 let (Some(profile), Some(taken)) = (
448 profiles.get(&repository.profile),
449 taken.get_mut(repository.profile.as_str()),
450 ) else {
451 continue;
452 };
453 let (memory, cpu) = (profile.memory, profile.cpus);
454 let Some(index) = taken.iter().position(|taken| !taken) else {
455 continue;
456 };
457 if !left.fits(memory, cpu) {
458 continue;
459 }
460 taken[index] = true;
461 left.take(memory, cpu);
462 any = true;
463 decided.decisions.push(Decision::Start {
464 slot: SlotRef {
465 profile: repository.profile.clone(),
466 index,
467 },
468 repository: repository.name.clone(),
469 memory,
470 cpu,
471 warm: true,
472 replacing: Vec::new(),
473 });
474 }
475 if !any {
476 break;
477 }
478 }
479}
480
481#[derive(Debug, Clone, Copy)]
483struct Left {
484 memory: Option<u64>,
485 cpu: Option<u64>,
486 host: Option<u64>,
488}
489
490impl Left {
491 fn fits(self, memory: u64, cpu: u64) -> bool {
492 self.memory.is_none_or(|free| free >= memory)
493 && self.host.is_none_or(|free| free >= memory)
494 && self.cpu.is_none_or(|free| free >= cpu)
495 }
496
497 fn short(self, memory: u64, cpu: u64) -> Left {
500 let lacks = |free: Option<u64>| free.map(|free| memory.saturating_sub(free));
501 Left {
502 memory: lacks(self.memory).max(lacks(self.host)),
503 cpu: self.cpu.map(|free| cpu.saturating_sub(free)),
504 host: None,
505 }
506 }
507
508 fn take(&mut self, memory: u64, cpu: u64) {
509 self.memory = self.memory.map(|free| free.saturating_sub(memory));
510 self.host = self.host.map(|free| free.saturating_sub(memory));
511 self.cpu = self.cpu.map(|free| free.saturating_sub(cpu));
512 }
513
514 fn give(&mut self, memory: u64, cpu: u64) {
515 self.memory = self.memory.map(|free| free + memory);
516 self.host = self.host.map(|free| free + memory);
517 self.cpu = self.cpu.map(|free| free + cpu);
518 }
519}
520
521struct Sorted {
523 serve: Vec<u64>,
524 warm: usize,
525}
526
527fn sort_out(
531 repository: &Repository,
532 profiles: &BTreeMap<String, Profile>,
533 now: SystemTime,
534 decided: &mut Plan,
535) -> Option<Sorted> {
536 let mut matching: Vec<&QueuedJob> = repository
537 .queued
538 .iter()
539 .filter(|job| matches(job, &repository.labels))
540 .collect();
541 let profile = profiles.get(&repository.profile);
542 let mut warm = repository.warm;
543 if repository.paused {
544 decided
545 .waiting
546 .extend(matching.iter().map(|job| (job.id, Waiting::Paused)));
547 matching.clear();
549 warm = 0;
550 }
551 let upkeep_due = profile.is_some_and(|profile| profile.upkeep_due);
554 let stop_now = upkeep_due || repository.paused;
555 if upkeep_due {
556 warm = 0;
557 }
558 let (idle, closing): (Vec<(SlotRef, &RunnerSlot)>, Vec<_>) = profile
559 .into_iter()
560 .flat_map(|profile| profile.slots.iter().enumerate())
561 .filter_map(|(index, slot)| match slot {
562 Slot::Runner(runner) if runner.repository == repository.name && !runner.busy => Some((
563 SlotRef {
564 profile: repository.profile.clone(),
565 index,
566 },
567 runner,
568 )),
569 _ => None,
570 })
571 .partition(|(slot, _)| profile.is_none_or(|profile| slot.index < profile.count));
572 decided
574 .decisions
575 .extend(closing.into_iter().map(|(slot, _)| Decision::Stop { slot }));
576 let surplus = idle.len().saturating_sub(matching.len());
577 let stops: Vec<SlotRef> = idle
578 .iter()
579 .filter(|(_, runner)| {
580 stop_now
581 || now
582 .duration_since(runner.started_at)
583 .is_ok_and(|waited| waited >= IDLE_TIMEOUT)
584 })
585 .take(surplus.saturating_sub(warm))
586 .map(|(slot, _)| slot.clone())
587 .collect();
588 let kept: Vec<&SlotRef> = idle
590 .iter()
591 .map(|(slot, _)| slot)
592 .filter(|slot| !stops.contains(slot))
593 .collect();
594 decided
595 .warm
596 .extend(kept.iter().skip(matching.len()).map(|slot| (*slot).clone()));
597 let warm = warm.saturating_sub(kept.len().saturating_sub(matching.len()));
598 decided
599 .decisions
600 .extend(stops.into_iter().map(|slot| Decision::Stop { slot }));
601 let missing = matching.len().saturating_sub(idle.len());
602 let skip = matching.len() - missing;
604 decided.waiting.extend(
605 matching
606 .iter()
607 .take(skip)
608 .map(|job| (job.id, Waiting::RunnerStarting)),
609 );
610 let serve = matching.iter().skip(skip).map(|job| job.id);
611 match repository.held.filter(|held| held.until > now) {
612 Some(held) => {
613 let until = crate::time::ms(held.until) / 1000;
614 decided.waiting.extend(serve.map(|job| {
615 (
616 job,
617 Waiting::Backoff {
618 until,
619 cause: held.cause,
620 },
621 )
622 }));
623 None
624 }
625 None => Some(Sorted {
626 serve: serve.collect(),
627 warm,
628 }),
629 }
630}
631
632struct Warm<'plan> {
634 runners: Vec<Option<(SlotRef, &'plan RunnerSlot)>>,
636}
637
638struct Room {
640 stopped: Vec<SlotRef>,
642 slot: Option<usize>,
644 memory: u64,
645 cpu: u64,
646}
647
648impl<'plan> Warm<'plan> {
649 fn new(
651 warm: &BTreeSet<SlotRef>,
652 profiles: &'plan BTreeMap<String, Profile>,
653 held: bool,
654 ) -> Warm<'plan> {
655 let runners = warm
656 .iter()
657 .filter(|_| !held)
658 .filter_map(
659 |slot| match profiles.get(&slot.profile)?.slots.get(slot.index)? {
660 Slot::Runner(runner) => Some(Some((slot.clone(), runner))),
661 Slot::Free | Slot::Closed => None,
662 },
663 )
664 .collect();
665 Warm { runners }
666 }
667
668 fn make_room(&mut self, repository: &Repository, slot: bool, short: Left) -> Option<Room> {
672 let (memory, cpu) = (short.memory.unwrap_or(0), short.cpu.unwrap_or(0));
673 let mut candidates: Vec<usize> = self
674 .runners
675 .iter()
676 .enumerate()
677 .filter_map(|(at, runner)| {
678 runner
679 .as_ref()
680 .filter(|(_, runner)| runner.repository != repository.name)
681 .map(|_| at)
682 })
683 .collect();
684 candidates.sort_by_key(|at| {
686 self.runners[*at]
687 .as_ref()
688 .is_some_and(|(slot, _)| slot.profile != repository.profile)
689 });
690 let mut room = Room {
691 stopped: Vec::new(),
692 slot: None,
693 memory: 0,
694 cpu: 0,
695 };
696 let mut picked = Vec::new();
697 for at in candidates {
698 let short = (slot && room.slot.is_none()) || room.memory < memory || room.cpu < cpu;
699 if !short {
700 break;
701 }
702 let Some((ref_, runner)) = &self.runners[at] else {
703 continue;
704 };
705 let gives_slot = ref_.profile == repository.profile;
706 let helps = (gives_slot && slot && room.slot.is_none())
708 || (room.memory < memory && runner.memory > 0)
709 || (room.cpu < cpu && runner.cpu > 0);
710 if !helps {
711 continue;
712 }
713 if gives_slot && room.slot.is_none() {
714 room.slot = Some(ref_.index);
715 }
716 room.memory += runner.memory;
717 room.cpu += runner.cpu;
718 room.stopped.push(ref_.clone());
719 picked.push(at);
720 }
721 if (slot && room.slot.is_none()) || room.memory < memory || room.cpu < cpu {
722 return None;
723 }
724 for at in picked {
725 self.runners[at] = None;
726 }
727 Some(room)
728 }
729}
730
731struct Started<'plan> {
734 free_slots: BTreeSet<&'plan str>,
736 before: Left,
738 by_profile: BTreeMap<&'plan str, u64>,
740 first: Option<u64>,
741}
742
743impl<'plan> Started<'plan> {
744 fn new(taken: &BTreeMap<&'plan str, Vec<bool>>, before: Left) -> Started<'plan> {
745 Started {
746 free_slots: taken
747 .iter()
748 .filter(|(_, slots)| slots.contains(&false))
749 .map(|(name, _)| *name)
750 .collect(),
751 before,
752 by_profile: BTreeMap::new(),
753 first: None,
754 }
755 }
756
757 fn note(&mut self, profile: &'plan str, job: u64) {
759 self.by_profile.entry(profile).or_insert(job);
760 self.first.get_or_insert(job);
761 }
762
763 fn for_slot(&self, profile: &str, reason: Waiting) -> Waiting {
765 match self.by_profile.get(profile) {
766 Some(older) if self.free_slots.contains(profile) => {
767 Waiting::BehindStart { job: *older }
768 }
769 _ => reason,
770 }
771 }
772
773 fn for_memory(&self, needs: u64, reason: Waiting) -> Waiting {
775 Self::behind(self.first, self.before.memory, needs).unwrap_or(reason)
776 }
777
778 fn for_host(&self, needs: u64, reason: Waiting) -> Waiting {
780 Self::behind(self.first, self.before.host, needs).unwrap_or(reason)
781 }
782
783 fn for_cpu(&self, needs: u64, reason: Waiting) -> Waiting {
785 Self::behind(self.first, self.before.cpu, needs).unwrap_or(reason)
786 }
787
788 fn behind(first: Option<u64>, free: Option<u64>, needs: u64) -> Option<Waiting> {
789 first
790 .filter(|_| free.is_some_and(|free| free >= needs))
791 .map(|job| Waiting::BehindStart { job })
792 }
793}
794
795fn taken_slots(profiles: &BTreeMap<String, Profile>) -> BTreeMap<&str, Vec<bool>> {
797 profiles
798 .iter()
799 .map(|(name, profile)| {
800 let taken = profile
801 .slots
802 .iter()
803 .map(|slot| profile.preparing.is_some() || *slot != Slot::Free)
804 .collect();
805 (name.as_str(), taken)
806 })
807 .collect()
808}
809
810#[cfg(test)]
811mod tests;