| 1 | use std::{ |
| 2 | collections::BTreeMap, |
| 3 | path::Path, |
| 4 | time::{Duration, SystemTime, UNIX_EPOCH}, |
| 5 | }; |
| 6 | |
| 7 | use serde::{Deserialize, Serialize}; |
| 8 | |
| 9 | use crate::{sched::Need, store, units::Bytes}; |
| 10 | |
| 11 | const KEEP: usize = 1000; |
| 12 | |
| 13 | #[derive(Clone, Copy, Debug, Serialize, Deserialize)] |
| 14 | pub struct Record { |
| 15 | pub runs: u32, |
| 16 | pub peak_rss: Bytes, |
| 17 | /// ewma, so one freak run doesn't define the command forever |
| 18 | pub wall: u64, |
| 19 | pub cpu: u64, |
| 20 | pub cores: u32, |
| 21 | pub last: u64, |
| 22 | } |
| 23 | |
| 24 | type History = BTreeMap<String, Record>; |
| 25 | |
| 26 | pub fn key(root: &Path, cmd: &[String]) -> String { |
| 27 | let mut k = format!("{}\t{}", root.display(), cmd.join(" ")); |
| 28 | k.truncate(k.floor_char_boundary(400)); |
| 29 | k |
| 30 | } |
| 31 | |
| 32 | fn now() -> u64 { |
| 33 | SystemTime::now() |
| 34 | .duration_since(UNIX_EPOCH) |
| 35 | .map_or(0, |d| d.as_secs()) |
| 36 | } |
| 37 | |
| 38 | pub fn record(key: String, peak_rss: Bytes, wall: Duration, cpu: Duration, cores: u32) { |
| 39 | let _ = store::update("history", |h: &mut History| { |
| 40 | let wall = wall.as_secs(); |
| 41 | h.entry(key) |
| 42 | .and_modify(|r| { |
| 43 | r.runs += 1; |
| 44 | r.peak_rss = r.peak_rss.max(peak_rss); |
| 45 | r.wall = (r.wall * 7 + wall * 3) / 10; |
| 46 | r.cpu = cpu.as_secs(); |
| 47 | r.cores = cores; |
| 48 | r.last = now(); |
| 49 | }) |
| 50 | .or_insert_with(|| Record { |
| 51 | runs: 1, |
| 52 | peak_rss, |
| 53 | wall, |
| 54 | cpu: cpu.as_secs(), |
| 55 | cores, |
| 56 | last: now(), |
| 57 | }); |
| 58 | if h.len() > KEEP { |
| 59 | let mut ages: Vec<_> = h.iter().map(|(k, r)| (r.last, k.clone())).collect(); |
| 60 | ages.sort(); |
| 61 | for (_, k) in ages.into_iter().take(h.len() - KEEP) { |
| 62 | h.remove(&k); |
| 63 | } |
| 64 | } |
| 65 | }); |
| 66 | } |
| 67 | |
| 68 | pub fn lookup(key: &str) -> Option<Record> { |
| 69 | store::load::<History>("history").get(key).copied() |
| 70 | } |
| 71 | |
| 72 | /// padded, since the agent that under-asks is the one that ooms the video. |
| 73 | /// nix builds happen in nix-daemon, so their rss never shows up in ours and |
| 74 | /// the base estimate stands |
| 75 | pub fn suggest(rec: &Record, base: Need, cmd: &[String]) -> Need { |
| 76 | let via_daemon = cmd |
| 77 | .iter() |
| 78 | .any(|t| t.contains("nix") || t.contains("home-manager")); |
| 79 | Need { |
| 80 | ram: if via_daemon { |
| 81 | base.ram |
| 82 | } else { |
| 83 | Bytes(rec.peak_rss.0 + rec.peak_rss.0 / 4).max(Bytes::mib(256)) |
| 84 | }, |
| 85 | cores: rec.cores.max(1), |
| 86 | time: Duration::from_secs((rec.wall * 3 / 2).max(30)), |
| 87 | ..base |
| 88 | } |
| 89 | } |