use std::collections::{BTreeMap, BTreeSet};
use std::time::Duration;
use super::{IDLE_TIMEOUT, Pool};
pub const ACTIVE_AFTER: Duration = Duration::from_secs(3600);
pub const LOOKBACK: u64 = 14 * DAY;
const USUAL_DAYS: usize = 3;
const MIN_RUNS: usize = 3;
const FORGET_AFTER: Duration = Duration::from_hours(24);
const UNKNOWN_SPAN: Duration = Duration::from_secs(3600);
const HOUR: u64 = 3600 * 1000;
type RunKey<'seen> = (&'seen str, &'seen str, &'seen str, u64, u32);
type WorkflowKey<'seen> = (&'seen str, &'seen str, &'seen str);
const DAY: u64 = 24 * HOUR;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Seen {
pub repository: String,
pub labels: Vec<String>,
pub workflow: Option<String>,
pub run_id: u64,
pub run_attempt: u32,
pub queued_at: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Burst {
pub jobs: usize,
pub span: Duration,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Learned {
usual: BTreeMap<String, BTreeSet<u64>>,
last: BTreeMap<String, u64>,
bursts: BTreeMap<(Pool, String), Burst>,
}
impl Learned {
pub fn new(seen: &[Seen], now: u64, profile: impl Fn(&Seen) -> Option<String>) -> Learned {
let since = now.saturating_sub(LOOKBACK);
let mut learned = Learned::default();
let mut days: BTreeMap<(&str, u64), BTreeSet<u64>> = BTreeMap::new();
let profiles: Vec<Option<String>> = seen.iter().map(profile).collect();
let mut runs: BTreeMap<RunKey<'_>, (usize, u64, u64)> = BTreeMap::new();
let mut attempts: BTreeMap<(&str, u64, u32), u64> = BTreeMap::new();
for (job, profile) in seen
.iter()
.zip(&profiles)
.filter(|(job, _)| job.queued_at >= since)
{
learned.note(&job.repository, job.queued_at);
days.entry((&job.repository, job.queued_at % DAY / HOUR))
.or_default()
.insert(job.queued_at / DAY);
if let (Some(workflow), Some(profile)) = (&job.workflow, profile) {
let run = runs
.entry((
&job.repository,
profile,
workflow,
job.run_id,
job.run_attempt,
))
.or_insert((0, job.queued_at, job.queued_at));
run.0 += 1;
run.1 = run.1.min(job.queued_at);
run.2 = run.2.max(job.queued_at);
let first = attempts
.entry((&job.repository, job.run_id, job.run_attempt))
.or_insert(job.queued_at);
*first = (*first).min(job.queued_at);
}
}
for ((repository, hour), days) in days {
if days.len() >= USUAL_DAYS {
learned
.usual
.entry(repository.to_owned())
.or_default()
.insert(hour);
}
}
let mut workflows: BTreeMap<WorkflowKey<'_>, Vec<(usize, u64)>> = BTreeMap::new();
for ((repository, profile, workflow, run_id, run_attempt), (jobs, first, last)) in runs {
let first = attempts
.get(&(repository, run_id, run_attempt))
.copied()
.unwrap_or(first);
workflows
.entry((repository, profile, workflow))
.or_default()
.push((jobs, last - first));
}
for ((repository, profile, workflow), runs) in workflows {
if runs.len() < MIN_RUNS {
continue;
}
let mut jobs: Vec<usize> = runs.iter().map(|(jobs, _)| *jobs).collect();
let mut spans: Vec<u64> = runs.iter().map(|(_, span)| *span).collect();
jobs.sort_unstable();
spans.sort_unstable();
let pool = Pool {
repository: repository.to_owned(),
profile: profile.to_owned(),
};
learned.bursts.insert(
(pool, workflow.to_owned()),
Burst {
jobs: jobs[(jobs.len() - 1) / 2],
span: Duration::from_millis(spans[(spans.len() * 90).div_ceil(100) - 1]),
},
);
}
learned
}
pub fn absorb(&mut self, older: &Learned) {
for (repository, at) in &older.last {
self.note(repository, *at);
}
}
pub fn note(&mut self, repository: &str, at: u64) {
let last = self.last.entry(repository.to_owned()).or_default();
*last = (*last).max(at);
}
pub fn active(&self, repository: &str, now: u64) -> bool {
let after = u64::try_from(ACTIVE_AFTER.as_millis()).unwrap_or(u64::MAX);
self.last
.get(repository)
.is_some_and(|last| now.saturating_sub(*last) < after)
|| self
.usual
.get(repository)
.is_some_and(|hours| hours.contains(&(now % DAY / HOUR)))
}
pub fn burst(&self, pool: &Pool, workflow: &str) -> Option<Burst> {
self.bursts
.get(&(pool.clone(), workflow.to_owned()))
.copied()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct Run {
repository: String,
workflow: Option<String>,
jobs: BTreeMap<u64, String>,
first: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct Attempt {
pub run_id: u64,
pub run_attempt: u32,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Runs(BTreeMap<Attempt, Run>);
impl Runs {
pub fn note(
&mut self,
pool: &Pool,
attempt: Attempt,
job: u64,
workflow: Option<&str>,
at: u64,
) {
let run = self.0.entry(attempt).or_insert_with(|| Run {
repository: pool.repository.clone(),
workflow: None,
jobs: BTreeMap::new(),
first: at,
});
run.jobs.insert(job, pool.profile.clone());
run.first = run.first.min(at);
if run.workflow.is_none() {
run.workflow = workflow.map(str::to_owned);
}
}
pub fn place(
&mut self,
pool: &Pool,
attempt: Attempt,
job: u64,
workflow: Option<&str>,
at: u64,
) {
let mut seen: Option<u64> = None;
let mut emptied = Vec::new();
for (other, run) in self.attempts_mut(attempt.run_id) {
if *other != attempt && run.jobs.remove(&job).is_some() {
if run.jobs.is_empty() {
emptied.push(*other);
seen = Some(run.first.min(seen.unwrap_or(at)));
} else {
seen = Some(seen.unwrap_or(at));
}
}
}
for other in emptied {
self.0.remove(&other);
}
if let Some(seen) = seen {
self.note(pool, attempt, job, workflow, seen);
}
}
pub fn name(&mut self, run_id: u64, workflow: &str) {
for (_, run) in self.attempts_mut(run_id) {
if run.workflow.is_none() {
run.workflow = Some(workflow.to_owned());
}
}
}
pub fn latest(&self, run_id: u64) -> Attempt {
self.0
.range(
Attempt {
run_id,
run_attempt: 0,
}..=Attempt {
run_id,
run_attempt: u32::MAX,
},
)
.next_back()
.map_or(
Attempt {
run_id,
run_attempt: 1,
},
|(attempt, _)| *attempt,
)
}
fn attempts_mut(&mut self, run_id: u64) -> impl Iterator<Item = (&Attempt, &mut Run)> {
self.0.range_mut(
Attempt {
run_id,
run_attempt: 0,
}..=Attempt {
run_id,
run_attempt: u32::MAX,
},
)
}
pub fn expire(&mut self, now: u64) {
let after = u64::try_from(FORGET_AFTER.as_millis()).unwrap_or(u64::MAX);
self.0
.retain(|_, run| now.saturating_sub(run.first) < after);
}
pub fn remaining(&self, pool: &Pool, learned: &Learned, now: u64) -> usize {
self.0
.values()
.filter(|run| run.repository == pool.repository && run.counts(pool, learned, now))
.filter_map(|run| {
let burst = learned.burst(pool, run.workflow.as_deref()?)?;
let seen = run
.jobs
.values()
.filter(|profile| **profile == pool.profile)
.count();
Some(burst.jobs.saturating_sub(seen))
})
.sum()
}
}
impl Run {
fn counts(&self, pool: &Pool, learned: &Learned, now: u64) -> bool {
let span = self
.workflow
.as_deref()
.and_then(|workflow| learned.burst(pool, workflow))
.map_or(UNKNOWN_SPAN, |burst| burst.span + IDLE_TIMEOUT);
now.saturating_sub(self.first) < u64::try_from(span.as_millis()).unwrap_or(u64::MAX)
}
}
pub fn target(warm: u32, active: bool, remaining: usize) -> usize {
if warm == 0 {
return 0;
}
let base = if active {
usize::try_from(warm).unwrap_or(usize::MAX)
} else {
0
};
base.saturating_add(remaining)
}
#[cfg(test)]
mod tests;