1use std::fmt;
36use std::num::NonZeroU16;
37
38use crate::attempt::RunnerAttempt;
39use crate::model::{Host, HostId, PolicyId};
40use crate::policy::ScalePolicy;
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum LimitingFactor {
49 Demand,
51 MinCapacity,
55 MaxCapacity,
57 HostCapacity,
59 MonitorOnly,
61 NotReconciling,
65 ForeignHost,
67}
68
69impl fmt::Display for LimitingFactor {
70 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71 f.write_str(match self {
72 LimitingFactor::Demand => "demand",
73 LimitingFactor::MinCapacity => "min_capacity",
74 LimitingFactor::MaxCapacity => "max_capacity",
75 LimitingFactor::HostCapacity => "host_capacity",
76 LimitingFactor::MonitorOnly => "monitor_only",
77 LimitingFactor::NotReconciling => "not_reconciling",
78 LimitingFactor::ForeignHost => "foreign_host",
79 })
80 }
81}
82
83#[derive(Debug, Clone, PartialEq, Eq)]
85pub struct Allocation {
86 pub policy_id: PolicyId,
87 pub demand: u32,
89 pub desired: u16,
91 pub active_owned: u16,
93 pub headroom_before: u16,
95 pub to_start: u16,
97 pub limiting_factor: LimitingFactor,
98}
99
100impl Allocation {
101 #[must_use]
102 pub const fn starts_nothing(&self) -> bool {
103 self.to_start == 0
104 }
105}
106
107#[derive(Debug, Clone, PartialEq, Eq)]
125pub struct HostAllocator<'a> {
126 host_id: HostId,
127 host_capacity: NonZeroU16,
128 attempts: Vec<&'a RunnerAttempt>,
130 active_total: u16,
131}
132
133impl<'a> HostAllocator<'a> {
134 #[must_use]
147 pub fn from_attempts(
148 host: &Host,
149 attempts: impl IntoIterator<Item = &'a RunnerAttempt>,
150 ) -> Self {
151 let attempts: Vec<&'a RunnerAttempt> = attempts.into_iter().collect();
152 let active_total = crate::attempt::active_count(attempts.iter().copied());
153 Self {
154 host_id: host.id,
155 host_capacity: host.host_capacity,
156 attempts,
157 active_total,
158 }
159 }
160
161 #[must_use]
162 pub fn host_capacity(&self) -> u16 {
163 self.host_capacity.get()
164 }
165
166 #[must_use]
168 pub const fn active_total(&self) -> u16 {
169 self.active_total
170 }
171
172 #[must_use]
174 pub fn headroom(&self) -> u16 {
175 self.host_capacity.get().saturating_sub(self.active_total)
176 }
177
178 pub fn allocate(&mut self, policy: &ScalePolicy, demand: u32) -> Allocation {
197 let active_owned =
198 crate::attempt::active_count_for(policy.id, self.attempts.iter().copied());
199 let headroom_before = self.headroom();
200
201 let refuse = |limiting_factor| Allocation {
202 policy_id: policy.id,
203 demand,
204 desired: 0,
205 active_owned,
206 headroom_before,
207 to_start: 0,
208 limiting_factor,
209 };
210
211 if !policy.is_owned_by(self.host_id) {
214 return refuse(LimitingFactor::ForeignHost);
215 }
216 if !policy.owns_runners() {
220 return refuse(LimitingFactor::MonitorOnly);
221 }
222 if !policy.may_start_runners() {
224 return refuse(LimitingFactor::NotReconciling);
225 }
226
227 let min = policy.min_capacity();
228 let max = policy
229 .max_capacity()
230 .expect("an Autoscale policy always has a max_capacity (D19)")
231 .get();
232
233 debug_assert!(min <= max, "PolicyMode invariant");
237 let desired = demand.clamp(u32::from(min), u32::from(max)) as u16;
238
239 let limiting_factor = if demand > u32::from(max) {
240 LimitingFactor::MaxCapacity
241 } else if demand < u32::from(min) {
242 LimitingFactor::MinCapacity
243 } else {
244 LimitingFactor::Demand
245 };
246
247 let wanted = desired.saturating_sub(active_owned);
248 let to_start = wanted.min(headroom_before);
249 let limiting_factor = if to_start < wanted {
250 LimitingFactor::HostCapacity
252 } else {
253 limiting_factor
254 };
255
256 self.active_total = self.active_total.saturating_add(to_start);
257
258 Allocation {
259 policy_id: policy.id,
260 demand,
261 desired,
262 active_owned,
263 headroom_before,
264 to_start,
265 limiting_factor,
266 }
267 }
268}
269
270#[cfg(test)]
271mod tests {
272 use super::*;
273 use crate::attempt::{
274 AttemptOutcome, AttemptState, FailureReason, PersistedAttempt, RunnerAttempt,
275 };
276 use crate::model::{
277 Arch, AttemptId, CachePolicy, HostLabel, Os, PolicyId, ScaleTarget, Timestamp,
278 };
279 use crate::policy::{PolicyMode, RoutingLabels, RunsOn, ScalePolicy};
280 use crate::workspace::WorkspaceKind;
281
282 fn ts(secs: i64) -> Timestamp {
283 chrono::DateTime::from_timestamp(secs, 0).expect("valid timestamp")
284 }
285
286 fn nz(v: u16) -> NonZeroU16 {
287 NonZeroU16::new(v).expect("non-zero")
288 }
289
290 const HOST: HostId = HostId::from_u128(7);
291
292 const NO_ATTEMPTS: &[RunnerAttempt] = &[];
297
298 fn host(capacity: u16) -> Host {
299 Host::new(HOST, "home-pc", Os::Windows, Arch::X64, nz(capacity), ts(0)).expect("valid host")
300 }
301
302 fn labels(name: &str) -> RoutingLabels {
303 RoutingLabels::derive(&HostLabel::new(name).unwrap(), Os::Windows, Arch::X64)
304 }
305
306 fn active_policy(id: u128, host_label: &str, max: u16) -> ScalePolicy {
309 let mut policy = ScalePolicy::new(
310 PolicyId::from_u128(id),
311 ScaleTarget::repository("o/r").unwrap(),
312 1,
313 HOST,
314 PolicyMode::autoscale(labels(host_label), 0, nz(max)).unwrap(),
315 CachePolicy::default(),
316 );
317 policy.activate().expect("pending -> active");
318 policy
319 }
320
321 fn attempt_in(state: AttemptState, id: u128, policy: u128) -> RunnerAttempt {
325 let outcome = state.is_terminal().then(|| match state {
326 AttemptState::Failed => {
327 AttemptOutcome::failed(FailureReason::ProcessExitedUnexpectedly)
328 }
329 AttemptState::Orphaned => AttemptOutcome::Orphaned,
330 _ => AttemptOutcome::CompletedJob,
331 });
332 RunnerAttempt::from_persisted(PersistedAttempt {
333 id: AttemptId::from_u128(id),
334 policy_id: PolicyId::from_u128(policy),
335 github_runner_id: None,
336 state,
337 outcome,
338 process_id: None,
339 runtime_path: "runtime/p/a".into(),
340 workspace_kind: WorkspaceKind::Ephemeral,
341 workspace_slot: None,
342 created_at: ts(0),
343 terminal_at: state.is_terminal().then(|| ts(0)),
344 last_state_change_at: ts(0),
345 })
346 .expect("a state/outcome pair the domain accepts")
347 }
348
349 #[test]
354 fn desired_clamps_above_and_below() {
355 let host = host(100);
356 let policy = active_policy(1, "home", 3);
357
358 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
360 let above = alloc.allocate(&policy, 10);
361 assert_eq!(above.desired, 3, "max_capacity beats reported demand");
362 assert_eq!(above.to_start, 3);
363 assert_eq!(above.limiting_factor, LimitingFactor::MaxCapacity);
364
365 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
367 let inside = alloc.allocate(&policy, 2);
368 assert_eq!(inside.desired, 2);
369 assert_eq!(inside.to_start, 2);
370 assert_eq!(inside.limiting_factor, LimitingFactor::Demand);
371
372 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
375 let none = alloc.allocate(&policy, 0);
376 assert_eq!(none.desired, 0);
377 assert_eq!(none.to_start, 0);
378 assert!(none.starts_nothing());
379 }
380
381 #[test]
382 fn a_non_zero_min_capacity_raises_desired_above_demand() {
383 let host = host(10);
387 let mut policy = ScalePolicy::new(
388 PolicyId::from_u128(1),
389 ScaleTarget::organization("acme").unwrap(),
390 1,
391 HOST,
392 PolicyMode::autoscale(labels("home"), 2, nz(5)).unwrap(),
393 CachePolicy::default(),
394 );
395 policy.activate().unwrap();
396
397 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
398 let got = alloc.allocate(&policy, 0);
399 assert_eq!(got.desired, 2);
400 assert_eq!(got.to_start, 2);
401 assert_eq!(got.limiting_factor, LimitingFactor::MinCapacity);
402 }
403
404 #[test]
409 fn the_same_queued_job_on_two_polls_yields_one_attempt_not_two() {
410 let host = host(4);
419 let policy = active_policy(1, "home", 4);
420 let queued = vec![RunsOn::Single("rm-home-win-x64".into())];
421
422 let demand = policy.tally(&queued).demand();
424 assert_eq!(demand, 1);
425 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
426 let first = alloc.allocate(&policy, demand);
427 assert_eq!(first.to_start, 1);
428
429 let attempts = vec![attempt_in(AttemptState::Starting, 1, 1)];
431
432 for poll in 2..=3 {
434 let demand = policy.tally(&queued).demand();
435 assert_eq!(demand, 1, "poll {poll}: the job has not left the queue");
436
437 let mut alloc = HostAllocator::from_attempts(&host, &attempts);
438 let again = alloc.allocate(&policy, demand);
439 assert_eq!(
440 again.to_start, 0,
441 "poll {poll} started another runner for a job already being \
442 served; the `- active_owned_runners` term was dropped from the \
443 formula"
444 );
445 assert_eq!(again.desired, 1);
446 assert_eq!(again.active_owned, 1);
447 }
448 }
449
450 #[test]
451 fn mutant_ignoring_in_flight_attempts_is_detected() {
452 let host = host(4);
453 let policy = active_policy(1, "home", 4);
454 let attempts = vec![attempt_in(AttemptState::Starting, 1, 1)];
455 let mut allocator = HostAllocator::from_attempts(&host, &attempts);
456 let protected = allocator.allocate(&policy, 1);
457 assert_eq!(protected.active_owned, 1);
458 assert_eq!(protected.to_start, 0);
459
460 let mutant_active_owned = 0_u16;
463 let mutant_to_start = protected
464 .desired
465 .saturating_sub(mutant_active_owned)
466 .min(protected.headroom_before);
467 assert_eq!(
468 mutant_to_start, 1,
469 "removing the in-flight term must make the duplicate-poll gate red"
470 );
471 }
472
473 #[test]
474 fn an_attempt_stops_counting_once_it_is_terminal() {
475 let host = host(4);
476 let policy = active_policy(1, "home", 4);
477
478 let in_flight = vec![
479 attempt_in(AttemptState::Allocated, 1, 1),
480 attempt_in(AttemptState::Starting, 2, 1),
481 attempt_in(AttemptState::Busy, 3, 1),
482 ];
483 let mut alloc = HostAllocator::from_attempts(&host, &in_flight);
484 assert_eq!(alloc.active_total(), 3);
485 assert_eq!(alloc.headroom(), 1);
486 assert_eq!(alloc.allocate(&policy, 4).to_start, 1);
487
488 let done = vec![
490 attempt_in(AttemptState::Finished, 1, 1),
491 attempt_in(AttemptState::Failed, 2, 1),
492 attempt_in(AttemptState::Cleaned, 3, 1),
493 ];
494 let mut alloc = HostAllocator::from_attempts(&host, &done);
495 assert_eq!(alloc.active_total(), 0);
496 assert_eq!(alloc.headroom(), 4);
497 assert_eq!(alloc.allocate(&policy, 4).to_start, 4);
498 }
499
500 #[test]
501 fn one_attempt_set_answers_both_ceilings() {
502 let host = host(10);
513 let mine = active_policy(1, "home", 9);
514 let theirs = active_policy(2, "office", 9);
515
516 let on_the_machine = vec![
517 attempt_in(AttemptState::Busy, 1, 1),
518 attempt_in(AttemptState::Starting, 2, 1),
519 attempt_in(AttemptState::Idle, 3, 2),
520 attempt_in(AttemptState::Finished, 4, 1),
522 ];
523
524 let mut alloc = HostAllocator::from_attempts(&host, &on_the_machine);
525 assert_eq!(alloc.active_total(), 3, "host-wide (D9), from the one set");
526
527 let got = alloc.allocate(&mine, 9);
528 assert_eq!(
529 got.active_owned, 2,
530 "per-policy (D7), from the same set and with no second argument that \
531 could have said otherwise"
532 );
533 assert_eq!(got.headroom_before, 7);
534 assert_eq!(got.to_start, 7, "9 wanted, 2 already in flight, 7 free");
535
536 let got = alloc.allocate(&theirs, 9);
537 assert_eq!(got.active_owned, 1);
538 assert_eq!(
539 got.to_start, 0,
540 "the first grant spent the headroom the second would have used"
541 );
542 }
543
544 #[test]
549 fn the_host_ceiling_binds_across_two_policies_whose_max_capacities_sum_higher() {
550 let host = host(3);
553 let a = active_policy(1, "home", 3);
554 let b = active_policy(2, "home", 3);
555
556 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
557 let first = alloc.allocate(&a, 10);
558 let second = alloc.allocate(&b, 10);
559
560 assert_eq!(first.to_start, 3, "the first policy takes the whole host");
561 assert_eq!(
562 second.to_start, 0,
563 "the second gets nothing; each policy is individually within its own \
564 max_capacity of 3, and 3 + 3 > host_capacity of 3"
565 );
566 assert_eq!(second.limiting_factor, LimitingFactor::HostCapacity);
567 assert_eq!(
568 first.to_start + second.to_start,
569 3,
570 "the sum across policies must never exceed host_capacity"
571 );
572 assert_eq!(alloc.headroom(), 0);
573 }
574
575 #[test]
576 fn the_host_ceiling_splits_headroom_between_policies_in_call_order() {
577 let host = host(5);
578 let a = active_policy(1, "home", 4);
579 let b = active_policy(2, "home", 4);
580 let c = active_policy(3, "home", 4);
581
582 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
583 let first = alloc.allocate(&a, 4);
584 let second = alloc.allocate(&b, 4);
585 let third = alloc.allocate(&c, 4);
586
587 assert_eq!(first.to_start, 4);
588 assert_eq!(second.to_start, 1, "one slot of headroom left");
589 assert_eq!(second.limiting_factor, LimitingFactor::HostCapacity);
590 assert_eq!(third.to_start, 0);
591 assert_eq!(
592 first.to_start + second.to_start + third.to_start,
593 5,
594 "12 requested across three policies, 5 granted, which is host_capacity"
595 );
596 }
597
598 #[test]
599 fn zero_headroom_starts_nothing_even_at_maximum_demand() {
600 let host = host(2);
601 let policy = active_policy(1, "home", 2);
602
603 let full = vec![
604 attempt_in(AttemptState::Busy, 1, 1),
605 attempt_in(AttemptState::Busy, 2, 1),
606 ];
607 let mut alloc = HostAllocator::from_attempts(&host, &full);
608 assert_eq!(alloc.headroom(), 0);
609
610 let got = alloc.allocate(&policy, u32::from(u16::MAX));
611 assert_eq!(got.to_start, 0);
612 assert_eq!(got.headroom_before, 0);
613 assert_eq!(got.limiting_factor, LimitingFactor::MaxCapacity);
614 }
615
616 #[test]
617 fn headroom_smaller_than_the_per_policy_allowance_wins() {
618 let host = host(6);
621 let policy = active_policy(1, "home", 5);
622
623 let others = vec![
624 attempt_in(AttemptState::Busy, 1, 99),
625 attempt_in(AttemptState::Busy, 2, 99),
626 attempt_in(AttemptState::Idle, 3, 99),
627 attempt_in(AttemptState::Starting, 4, 99),
628 ];
629 let mut alloc = HostAllocator::from_attempts(&host, &others);
630 assert_eq!(alloc.headroom(), 2, "four slots are held by another policy");
631
632 let got = alloc.allocate(&policy, 5);
633 assert_eq!(got.desired, 5, "the policy's own ceiling would allow five");
634 assert_eq!(got.to_start, 2, "but the host has only two slots free");
635 assert_eq!(got.limiting_factor, LimitingFactor::HostCapacity);
636 assert_eq!(alloc.headroom(), 0);
637 }
638
639 #[test]
640 fn an_over_subscribed_host_reports_zero_headroom_rather_than_wrapping() {
641 let host = host(2);
646 let policy = active_policy(1, "home", 10);
647 let oversubscribed: Vec<RunnerAttempt> = (1..=9)
650 .map(|id| attempt_in(AttemptState::Busy, id, 99))
651 .collect();
652
653 let mut alloc = HostAllocator::from_attempts(&host, &oversubscribed);
654 assert_eq!(
655 alloc.active_total(),
656 9,
657 "the raw count is reported, over-subscription included"
658 );
659 assert_eq!(alloc.headroom(), 0);
660 assert_eq!(alloc.allocate(&policy, 10).to_start, 0);
661 }
662
663 #[test]
668 fn a_monitor_only_policy_under_maximum_demand_starts_nothing() {
669 let host = host(10);
673 let mut policy = ScalePolicy::new(
674 PolicyId::from_u128(1),
675 ScaleTarget::organization("acme").unwrap(),
676 1,
677 HOST,
678 PolicyMode::monitor_only(),
679 CachePolicy::default(),
680 );
681 policy.activate().unwrap();
682
683 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
684 let got = alloc.allocate(&policy, 1_000);
685 assert_eq!(got.to_start, 0);
686 assert_eq!(got.limiting_factor, LimitingFactor::MonitorOnly);
687 assert_eq!(
688 alloc.headroom(),
689 10,
690 "and it consumes no headroom, so an autoscale policy on the same host \
691 is unaffected"
692 );
693 }
694
695 #[test]
696 fn a_policy_that_is_not_active_and_enabled_starts_nothing() {
697 let host = host(10);
698
699 let pending = ScalePolicy::new(
701 PolicyId::from_u128(1),
702 ScaleTarget::repository("o/r").unwrap(),
703 1,
704 HOST,
705 PolicyMode::autoscale(labels("home"), 0, nz(5)).unwrap(),
706 CachePolicy::default(),
707 );
708 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
709 let got = alloc.allocate(&pending, 5);
710 assert_eq!(got.to_start, 0);
711 assert_eq!(got.limiting_factor, LimitingFactor::NotReconciling);
712
713 let mut draining = active_policy(2, "home", 5);
715 draining.request_disable().unwrap();
716 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
717 let got = alloc.allocate(&draining, 5);
718 assert_eq!(got.to_start, 0);
719 assert_eq!(got.limiting_factor, LimitingFactor::NotReconciling);
720 assert_eq!(alloc.headroom(), 10);
721 }
722
723 #[test]
724 fn a_policy_belonging_to_another_host_is_refused_before_any_headroom_is_spent() {
725 let host = host(4);
726 let mut theirs = ScalePolicy::new(
727 PolicyId::from_u128(1),
728 ScaleTarget::repository("o/r").unwrap(),
729 1,
730 HostId::from_u128(8),
731 PolicyMode::autoscale(labels("office"), 0, nz(4)).unwrap(),
732 CachePolicy::default(),
733 );
734 theirs.activate().unwrap();
735
736 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
737 let got = alloc.allocate(&theirs, 4);
738 assert_eq!(got.to_start, 0);
739 assert_eq!(got.limiting_factor, LimitingFactor::ForeignHost);
740 assert_eq!(alloc.headroom(), 4);
741 }
742
743 #[test]
748 fn max_capacity_beats_demand_and_host_capacity_beats_max_capacity() {
749 let host = host(2);
751 let policy = active_policy(1, "home", 4);
752
753 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
754 let got = alloc.allocate(&policy, 9);
755
756 assert_eq!(got.demand, 9);
757 assert_eq!(got.desired, 4, "max_capacity beats reported demand");
758 assert_eq!(got.to_start, 2, "host_capacity beats max_capacity");
759 assert_eq!(got.limiting_factor, LimitingFactor::HostCapacity);
760 }
761
762 #[test]
763 fn an_idle_host_with_no_demand_starts_no_runners() {
764 let host = host(8);
767 let policies = [active_policy(1, "home", 4), active_policy(2, "home", 4)];
768 let mut alloc = HostAllocator::from_attempts(&host, NO_ATTEMPTS);
769 for policy in &policies {
770 let got = alloc.allocate(policy, 0);
771 assert_eq!(got.to_start, 0);
772 assert_eq!(got.desired, 0);
773 }
774 assert_eq!(alloc.active_total(), 0);
775 assert_eq!(alloc.headroom(), 8);
776 }
777}