qk_cache/
expectations.rs

1//! What tasks are expected to take, shared through the remote store, so that
2//! a machine without local history, such as a CI job on a fresh clone, can
3//! order tasks and size their shares of the cores.
4//!
5//! The store keeps one record per branch, `expectations/<name>`, holding for
6//! each task the last few runs that executed it: their time, the cores they
7//! held and the memory they used. A run reads the record of its branch, else
8//! that of the default branch, as warm state is read. A run that executed
9//! tasks adds their runs to the record the store holds at the time, the
10//! branch's or else the default branch's, so tasks it did not run keep
11//! theirs, and saves it for its branch; the newest save replaces the last.
12//! A task is expected to take what [`expect`] makes of its runs, as local
13//! history makes of its own. `QK_REMOTE_EXPECTATIONS` narrows what a run may
14//! do with them: `read` reads them without adding to them, and `off` leaves
15//! them alone.
16
17use std::collections::BTreeMap;
18use std::ffi::OsString;
19use std::sync::Arc;
20use std::sync::atomic::{AtomicBool, Ordering};
21use std::thread::JoinHandle;
22use std::time::{SystemTime, UNIX_EPOCH};
23
24use qk_config::Workspace;
25use serde::{Deserialize, Serialize};
26
27use crate::{paths, remote, warm};
28
29/// One run of a task, or what a task is expected to take.
30#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
31pub struct Expectation {
32    pub millis: u64,
33    /// The cores it held: one, unless its target has `qk:threads`.
34    pub cores: u64,
35    /// The most memory, in bytes, it used, when that was measured.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub memory: Option<u64>,
38}
39
40/// The runs kept of each task, as local history draws on its last five.
41const RUNS: usize = 5;
42/// How long a task no run has executed keeps its runs in the record.
43const KEPT_FOR: u64 = 30 * 24 * 60 * 60 * 1000;
44/// The most tasks a record keeps, those executed last.
45const TASKS: usize = 10_000;
46/// The most bytes a record takes, within the 4 MiB a store server accepts
47/// for one, whatever its tasks' ids: the tasks executed longest ago go first.
48const BYTES: usize = 4 * 1024 * 1024 - 64 * 1024;
49
50#[derive(Serialize, Deserialize)]
51struct Record {
52    version: u32,
53    /// The name it is stored under, which a store server checks.
54    name: String,
55    tasks: BTreeMap<String, Runs>,
56}
57
58#[derive(Serialize, Deserialize)]
59struct Runs {
60    /// When a run last added to it, in milliseconds since the epoch.
61    updated: u64,
62    /// Newest first.
63    runs: Vec<Expectation>,
64}
65
66/// What a task is expected to take from its recent `runs`: the median of
67/// their times, with the cores that run held, and the most memory any of
68/// them used, since expecting too little memory is what costs.
69pub fn expect(runs: &[Expectation]) -> Option<Expectation> {
70    let memory = runs.iter().filter_map(|run| run.memory).max();
71    let mut sorted = runs.to_vec();
72    sorted.sort_by_key(|run| run.millis);
73    let median = *sorted.get(sorted.len() / 2)?;
74    Some(Expectation { memory, ..median })
75}
76
77/// Whether `bytes` are a record saved for `branch`.
78pub(crate) fn usable(bytes: &[u8], branch: &str) -> bool {
79    parse(bytes, branch).is_some()
80}
81
82/// What each task in the record is expected to take; empty for a record that
83/// cannot be read or is not for `branch`. A run that held no core is no
84/// measure of its work, so a task expected on one is left out.
85pub(crate) fn read(bytes: &[u8], branch: &str) -> BTreeMap<String, Expectation> {
86    let Some(record) = parse(bytes, branch) else {
87        return BTreeMap::new();
88    };
89    record
90        .tasks
91        .into_iter()
92        .filter_map(|(id, runs)| {
93            let expected = expect(&runs.runs)?;
94            (expected.cores > 0).then_some((id, expected))
95        })
96        .collect()
97}
98
99fn parse(bytes: &[u8], branch: &str) -> Option<Record> {
100    let record: Record = serde_json::from_slice(bytes).ok()?;
101    (record.version == 1 && record.name == remote::expectations_name(branch)).then_some(record)
102}
103
104/// The record to save for `branch`: `base`, the record the store holds, with
105/// each task's run in `runs` added as its newest at `now`, and the tasks no
106/// run has added to for a while left out.
107pub(crate) fn merge(
108    base: Option<(&[u8], &str)>,
109    branch: &str,
110    runs: &BTreeMap<String, Expectation>,
111    now: u64,
112) -> Vec<u8> {
113    let mut tasks = base
114        .and_then(|(bytes, saved_for)| parse(bytes, saved_for))
115        .map(|record| record.tasks)
116        .unwrap_or_default();
117    for (id, run) in runs {
118        let kept = tasks.entry(id.clone()).or_insert(Runs {
119            updated: now,
120            runs: Vec::new(),
121        });
122        kept.updated = now;
123        kept.runs.insert(0, *run);
124        kept.runs.truncate(RUNS);
125    }
126    tasks.retain(|_, kept| kept.updated.saturating_add(KEPT_FOR) >= now);
127    // Newest first, so the tasks dropped are those executed longest ago.
128    let mut order: Vec<(u64, String)> = tasks
129        .iter()
130        .map(|(id, kept)| (kept.updated, id.clone()))
131        .collect();
132    order.sort_unstable_by(|a, b| b.0.cmp(&a.0).then_with(|| a.1.cmp(&b.1)));
133    for (_, id) in order.drain(TASKS.min(order.len())..) {
134        tasks.remove(&id);
135    }
136    let mut record = Record {
137        version: 1,
138        name: remote::expectations_name(branch),
139        tasks,
140    };
141    loop {
142        let bytes = serde_json::to_vec(&record).expect("expectations serialize");
143        if bytes.len() <= BYTES || record.tasks.is_empty() {
144            return bytes;
145        }
146        // A tenth at a time, so a record far over takes few passes.
147        let drop = order.len().div_ceil(10);
148        for (_, id) in order.drain(order.len() - drop..) {
149            record.tasks.remove(&id);
150        }
151    }
152}
153
154/// What `QK_REMOTE_EXPECTATIONS` lets a run do with the expectations the
155/// remote store holds.
156#[derive(Clone, Copy, Debug, PartialEq, Eq)]
157enum Sharing {
158    /// Read them, and add to them as the store's mode allows: the default.
159    ReadWrite,
160    /// Read them only.
161    Read,
162    /// Neither.
163    Off,
164}
165
166/// The variable's setting, or why it cannot be used.
167fn sharing(environment: &BTreeMap<OsString, OsString>) -> Result<Sharing, String> {
168    let value = environment
169        .get(std::ffi::OsStr::new("QK_REMOTE_EXPECTATIONS"))
170        .and_then(|value| value.to_str())
171        .filter(|value| !value.is_empty());
172    match value {
173        None | Some("read-write") => Ok(Sharing::ReadWrite),
174        Some("read" | "read-only") => Ok(Sharing::Read),
175        Some("off") => Ok(Sharing::Off),
176        Some(other) => Err(format!(
177            "unknown QK_REMOTE_EXPECTATIONS {other:?}; use read-write, read or off"
178        )),
179    }
180}
181
182/// The remote store the workspace configures, when there is one it may use;
183/// a store that cannot be configured is left out, as the cache reports it.
184fn remote(
185    workspace: &Workspace,
186    environment: &BTreeMap<OsString, OsString>,
187) -> Option<Arc<remote::Remote>> {
188    let (s3, server) = crate::remote_settings(workspace);
189    remote::configure(s3, server, environment)
190        .ok()
191        .flatten()
192        .map(Arc::new)
193}
194
195/// A fetch of the expectations the remote store holds, running beside
196/// whatever the run does until it needs them. Dropped unwaited, as when local
197/// history has every task, it is left to end on its own, and says nothing:
198/// the run does not wait for a fetch whose answer it does not need.
199pub struct Fetching(Option<(JoinHandle<Expected>, Arc<AtomicBool>)>);
200
201/// What each task is expected to take, by its id.
202type Expected = BTreeMap<String, Expectation>;
203
204impl Fetching {
205    /// What each task is expected to take, once the fetch has ended.
206    pub fn wait(mut self) -> BTreeMap<String, Expectation> {
207        self.0
208            .take()
209            .and_then(|(fetch, _)| fetch.join().ok())
210            .unwrap_or_default()
211    }
212}
213
214impl Drop for Fetching {
215    fn drop(&mut self) {
216        if let Some((_, wanted)) = self.0.take() {
217            wanted.store(false, Ordering::Relaxed);
218        }
219    }
220}
221
222/// Starts fetching what each task is expected to take from the record the
223/// remote store holds for the current branch, else for the default branch:
224/// nothing when no store is configured, `QK_REMOTE_EXPECTATIONS` is `off` or
225/// the store holds neither. A store that cannot be read is reported.
226pub fn fetch(workspace: &Workspace, environment: &BTreeMap<OsString, OsString>) -> Fetching {
227    match sharing(environment) {
228        Ok(Sharing::Off) => return Fetching(None),
229        Ok(_) => {}
230        Err(reason) => {
231            qk_executor::status!("qk: remote task expectations disabled: {reason}");
232            return Fetching(None);
233        }
234    }
235    let Some(remote) = remote(workspace, environment) else {
236        return Fetching(None);
237    };
238    let root = paths::cache_location(workspace);
239    let branches = warm::branches(workspace);
240    let wanted = Arc::new(AtomicBool::new(true));
241    let fetch = std::thread::spawn({
242        let wanted = wanted.clone();
243        move || match remote.fetch_usable_expectations(&root, &branches) {
244            Ok(Some((bytes, branch))) => read(&bytes, &branch),
245            Ok(None) => BTreeMap::new(),
246            Err(error) => {
247                if wanted.load(Ordering::Relaxed) {
248                    qk_executor::status!("qk: remote task expectations unavailable ({error:#})");
249                }
250                BTreeMap::new()
251            }
252        }
253    });
254    Fetching(Some((fetch, wanted)))
255}
256
257/// Adds each task's run in `runs` to the record the remote store holds, and
258/// saves it for the current branch, when the store's mode and
259/// `QK_REMOTE_EXPECTATIONS` let it be written and there is a branch to name.
260/// Failures are reported without failing the run.
261pub fn share(
262    workspace: &Workspace,
263    environment: &BTreeMap<OsString, OsString>,
264    runs: BTreeMap<String, Expectation>,
265) {
266    if runs.is_empty() || sharing(environment) != Ok(Sharing::ReadWrite) {
267        return;
268    }
269    let Some(remote) =
270        remote(workspace, environment).filter(|remote| remote.mode == remote::Mode::ReadWrite)
271    else {
272        return;
273    };
274    // A branch's save goes to its own name; without a branch there is none.
275    let branches = warm::branches(workspace);
276    if warm::current_branch(workspace).is_none() {
277        return;
278    }
279    let root = paths::cache_location(workspace);
280    if let Err(error) = std::fs::create_dir_all(root.join("tmp")) {
281        qk_executor::status!("qk: remote cache upload failed for task expectations: {error}");
282        return;
283    }
284    let now = SystemTime::now()
285        .duration_since(UNIX_EPOCH)
286        .map_or(0, |duration| duration.as_millis() as u64);
287    remote.upload_expectations(&root, branches, runs, now);
288    for failure in remote.finish() {
289        qk_executor::status!("qk: remote cache upload failed for {failure}");
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    fn run(millis: u64, memory: Option<u64>) -> Expectation {
298        Expectation {
299            millis,
300            cores: 2,
301            memory,
302        }
303    }
304
305    #[test]
306    fn a_task_expected_on_no_core_is_left_out() {
307        let run = |cores| Expectation {
308            millis: 1_000,
309            cores,
310            memory: None,
311        };
312        let runs: BTreeMap<String, Expectation> = [
313            ("app:none".to_owned(), run(0)),
314            ("app:one".to_owned(), run(1)),
315        ]
316        .into();
317        let record = merge(None, "main", &runs, 1);
318        let expected = read(&record, "main");
319        assert!(!expected.contains_key("app:none"));
320        assert_eq!(expected["app:one"].cores, 1);
321    }
322
323    #[test]
324    fn a_record_stays_within_what_a_store_server_accepts() {
325        let run = Expectation {
326            millis: 1_000,
327            cores: 1,
328            memory: Some(1 << 30),
329        };
330        let long = "x".repeat(600);
331        let runs: BTreeMap<String, Expectation> = (0..TASKS)
332            .map(|index| (format!("project-{index:05}-{long}:test"), run))
333            .collect();
334        let first = merge(None, "main", &runs, 1);
335        assert!(first.len() <= BYTES, "{}", first.len());
336        // The tasks run last are the ones kept.
337        let newest: BTreeMap<String, Expectation> = [("newest:test".to_owned(), run)].into();
338        let second = merge(Some((&first, "main")), "main", &newest, 2);
339        assert!(second.len() <= BYTES);
340        let record = parse(&second, "main").unwrap();
341        assert!(record.tasks.contains_key("newest:test"));
342        assert!(record.tasks.len() < TASKS);
343    }
344
345    #[test]
346    fn the_variable_narrows_what_a_run_does_with_them() {
347        let with = |value: &str| {
348            let environment = [("QK_REMOTE_EXPECTATIONS".into(), value.into())]
349                .into_iter()
350                .collect();
351            sharing(&environment)
352        };
353        assert_eq!(sharing(&BTreeMap::new()), Ok(Sharing::ReadWrite));
354        assert_eq!(with(""), Ok(Sharing::ReadWrite));
355        assert_eq!(with("read-write"), Ok(Sharing::ReadWrite));
356        assert_eq!(with("read"), Ok(Sharing::Read));
357        assert_eq!(with("read-only"), Ok(Sharing::Read));
358        assert_eq!(with("off"), Ok(Sharing::Off));
359        assert!(with("no").is_err());
360    }
361
362    #[test]
363    fn expects_the_median_time_and_the_most_memory() {
364        let runs = [run(30, None), run(10, Some(5)), run(20, Some(9))];
365        assert_eq!(
366            expect(&runs),
367            Some(Expectation {
368                millis: 20,
369                cores: 2,
370                memory: Some(9)
371            })
372        );
373        assert_eq!(expect(&[]), None);
374    }
375
376    #[test]
377    fn a_merge_keeps_the_tasks_a_run_did_not_run() {
378        let first = merge(
379            None,
380            "main",
381            &[
382                ("app:a".to_owned(), run(10, None)),
383                ("app:b".to_owned(), run(5, None)),
384            ]
385            .into_iter()
386            .collect(),
387            1_000,
388        );
389        // Saved for another branch from the default branch's record.
390        let second = merge(
391            Some((&first, "main")),
392            "feature",
393            &[("app:a".to_owned(), run(40, None))].into_iter().collect(),
394            2_000,
395        );
396        let read = read(&second, "feature");
397        assert_eq!(read["app:a"].millis, 40);
398        assert_eq!(read["app:b"].millis, 5);
399        // A record is read only under the name it was saved for.
400        assert!(super::read(&second, "main").is_empty());
401        // Runs are kept, the newest first, up to five.
402        let mut record = second;
403        for millis in [1, 2, 3, 4, 5] {
404            record = merge(
405                Some((&record, "feature")),
406                "feature",
407                &[("app:a".to_owned(), run(millis, None))]
408                    .into_iter()
409                    .collect(),
410                3_000,
411            );
412        }
413        assert_eq!(super::read(&record, "feature")["app:a"].millis, 3);
414        // A task no run has executed for a month is dropped.
415        let later = merge(
416            Some((&record, "feature")),
417            "feature",
418            &[("app:a".to_owned(), run(1, None))].into_iter().collect(),
419            3_000 + KEPT_FOR + 1,
420        );
421        let kept = super::read(&later, "feature");
422        assert!(kept.contains_key("app:a") && !kept.contains_key("app:b"));
423    }
424}
qk_cache/expectations.rs

qk_cache/
expectations.rs
