use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use runner_manager_domain::coordinator_protocol::AttemptUsage;
use runner_manager_domain::model::{AttemptId, HostId, ScaleTarget};
use runner_manager_domain::store::Store;
use runner_manager_platform::resource_usage::{TreeUsage, sample_trees};
use tokio::time::Instant;
pub const SAMPLE_EVERY: Duration = Duration::from_secs(2);
const ANNOUNCE_AGAIN: Duration = Duration::from_secs(30);
const KEEP_AFTER_LAST_SAMPLE: Duration = Duration::from_secs(10 * 60);
#[derive(Debug, Clone, Copy)]
struct Meter {
first_at: Instant,
last_at: Instant,
last_cpu: Duration,
used_cpu: Duration,
peak_rss: u64,
peak_millicores: u32,
samples: u32,
}
fn millicores(cpu: Duration, wall: Duration) -> Option<u32> {
if wall.is_zero() {
return None;
}
u32::try_from(cpu.as_micros().saturating_mul(1000) / wall.as_micros().max(1)).ok()
}
impl Meter {
fn usage(&self) -> AttemptUsage {
AttemptUsage {
peak_rss_bytes: self.peak_rss,
cpu_avg_millicores: millicores(
self.used_cpu,
self.last_at.saturating_duration_since(self.first_at),
)
.unwrap_or(0),
cpu_peak_millicores: self.peak_millicores,
samples: self.samples,
}
}
}
#[derive(Debug, Default)]
pub struct AttemptMeters {
meters: Mutex<HashMap<AttemptId, Meter>>,
exited: Mutex<HashMap<AttemptId, Instant>>,
}
impl AttemptMeters {
fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<AttemptId, Meter>> {
self.meters
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
pub fn record(&self, attempt: AttemptId, at: Instant, usage: TreeUsage) {
let mut meters = self.lock();
let meter = meters.entry(attempt).or_insert(Meter {
first_at: at,
last_at: at,
last_cpu: usage.cpu_time,
used_cpu: Duration::ZERO,
peak_rss: 0,
peak_millicores: 0,
samples: 0,
});
let increase = usage.cpu_time.saturating_sub(meter.last_cpu);
if let Some(busy) = millicores(increase, at.saturating_duration_since(meter.last_at)) {
meter.peak_millicores = meter.peak_millicores.max(busy);
}
meter.used_cpu = meter.used_cpu.saturating_add(increase);
meter.last_at = at;
meter.last_cpu = usage.cpu_time;
meter.peak_rss = meter.peak_rss.max(usage.rss_bytes);
meter.samples = meter.samples.saturating_add(1);
}
pub fn take(&self, attempt: AttemptId) -> Option<AttemptUsage> {
self.lock().remove(&attempt).map(|meter| meter.usage())
}
pub fn forget_stale(&self, live: &HashSet<AttemptId>, now: Instant) {
self.lock().retain(|attempt, meter| {
live.contains(attempt)
|| now.saturating_duration_since(meter.last_at) < KEEP_AFTER_LAST_SAMPLE
});
}
pub fn sample_once(&self, store: &dyn Store, host: HostId) -> Vec<ScaleTarget> {
let policies = store.policies().unwrap_or_default();
let mut live = HashSet::new();
let mut roots: Vec<(AttemptId, u32, &ScaleTarget)> = Vec::new();
for policy in policies.iter().filter(|policy| policy.is_owned_by(host)) {
let Ok(attempts) = store.active_attempts_for_policy(policy.id) else {
continue;
};
for attempt in attempts {
live.insert(attempt.id);
roots.extend(
attempt
.process_id()
.map(|pid| (attempt.id, pid, &policy.target)),
);
}
}
let pids: Vec<u32> = roots.iter().map(|(_, pid, _)| *pid).collect();
let trees = sample_trees(&pids);
let now = Instant::now();
let mut exited = Vec::new();
let mut seen = self
.exited
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
seen.retain(|attempt, _| live.contains(attempt));
for (attempt, pid, target) in roots {
let Some(usage) = trees.get(&pid) else {
let due = seen
.get(&attempt)
.is_none_or(|at| now.saturating_duration_since(*at) >= ANNOUNCE_AGAIN);
if due {
seen.insert(attempt, now);
exited.push(target.clone());
}
continue;
};
seen.remove(&attempt);
self.record(attempt, now, *usage);
}
drop(seen);
self.forget_stale(&live, now);
exited
}
}
pub async fn sample_attempts(
meters: Arc<AttemptMeters>,
store: Arc<dyn Store>,
host: HostId,
joined: impl Fn() -> bool + Send,
exited: impl Fn(&ScaleTarget) + Send,
) {
let mut ticks = tokio::time::interval(SAMPLE_EVERY);
ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
ticks.tick().await;
if !joined() {
continue;
}
let (meters, store) = (Arc::clone(&meters), Arc::clone(&store));
let gone = tokio::task::spawn_blocking(move || meters.sample_once(store.as_ref(), host))
.await
.unwrap_or_default();
for target in &gone {
exited(target);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tree(rss: u64, cpu_ms: u64) -> TreeUsage {
TreeUsage {
rss_bytes: rss,
cpu_time: Duration::from_millis(cpu_ms),
processes: 2,
}
}
#[test]
fn a_meter_keeps_the_peak_memory_and_the_average_and_busiest_cpu() {
let meters = AttemptMeters::default();
let attempt = AttemptId::new_random();
let start = Instant::now();
let at = |secs| start + Duration::from_secs(secs);
meters.record(attempt, at(0), tree(100, 0));
meters.record(attempt, at(2), tree(500, 4_000));
meters.record(attempt, at(4), tree(300, 3_000));
meters.record(attempt, at(8), tree(200, 6_000));
let usage = meters.take(attempt).expect("measured");
assert_eq!(
usage,
AttemptUsage {
peak_rss_bytes: 500,
cpu_avg_millicores: 875,
cpu_peak_millicores: 2000,
samples: 4,
}
);
assert_eq!(meters.take(attempt), None, "taken once");
}
#[test]
fn a_single_sample_has_memory_and_no_cpu_rate() {
let meters = AttemptMeters::default();
let attempt = AttemptId::new_random();
meters.record(attempt, Instant::now(), tree(42, 1_000));
let usage = meters.take(attempt).expect("measured");
assert_eq!(usage.peak_rss_bytes, 42);
assert_eq!(
(usage.cpu_avg_millicores, usage.cpu_peak_millicores),
(0, 0)
);
}
#[test]
fn a_stale_meter_is_forgotten_and_a_live_one_kept() {
let meters = AttemptMeters::default();
let (live, stale) = (AttemptId::new_random(), AttemptId::new_random());
let start = Instant::now();
meters.record(live, start, tree(1, 0));
meters.record(stale, start, tree(1, 0));
meters.forget_stale(
&HashSet::from([live]),
start + KEEP_AFTER_LAST_SAMPLE + Duration::from_secs(1),
);
assert!(meters.take(live).is_some());
assert!(meters.take(stale).is_none());
}
#[test]
fn a_runner_whose_process_exited_is_announced_and_again_while_it_stays_live() {
use runner_manager_domain::attempt::AttemptState;
use runner_manager_domain::store::SqliteStore;
use runner_manager_testkit::fixtures;
let host = fixtures::host().build();
let policy = fixtures::policy()
.repository("octo/app")
.host(host.id)
.autoscale("home", 2)
.active()
.build();
let store = SqliteStore::open_in_memory().expect("store");
store.put_host(&host).expect("host");
store.insert_policy(&policy).expect("policy");
let mut child = std::process::Command::new(std::env::current_exe().expect("exe"))
.args(["--exact", "no-such-test", "--quiet"])
.stdout(std::process::Stdio::null())
.spawn()
.expect("a short-lived child");
let exited_pid = child.id();
child.wait().expect("reaped");
let (running, exited) = (AttemptId::new_random(), AttemptId::new_random());
for (id, pid) in [(running, std::process::id()), (exited, exited_pid)] {
store
.record_attempt(
&fixtures::attempt()
.id(id)
.policy_id(policy.id)
.state(AttemptState::Busy)
.process_id(pid)
.build(),
)
.expect("recorded");
}
let meters = AttemptMeters::default();
assert_eq!(
meters.sample_once(&store, host.id),
vec![policy.target.clone()]
);
assert!(
meters.sample_once(&store, host.id).is_empty(),
"not announced again at once"
);
meters
.exited
.lock()
.expect("lock")
.values_mut()
.for_each(|at| *at = at.checked_sub(ANNOUNCE_AGAIN).expect("earlier"));
assert_eq!(
meters.sample_once(&store, host.id),
vec![policy.target.clone()]
);
assert!(
meters.take(running).is_some(),
"the live runner was sampled"
);
assert!(meters.take(exited).is_none());
}
}