Skip to main content

qk_cache/
warm.rs

1//! Warm state: scratch state a task runs faster with, kept between runs of
2//! the same task and restored before it runs, but never part of a result.
3//! A target opts in with `qk:warm`:
4//!
5//! - `outputs: true`: the task's previous outputs, so an incremental tool
6//!   such as `tsc --build` starts from its last build;
7//! - `paths`: scratch directories in the workspace;
8//! - `env`: variables for the task, which may name `{warm}`, a directory qk
9//!   keeps for the task outside the working tree.
10//!
11//! A group already present on disk is left alone. Otherwise it is restored
12//! from this worktree's own save. `portable: true` also permits other
13//! worktrees' and remote saves. `portable: "path"` permits remote saves made
14//! in a worktree at the same path, as CI runners checking out in one place
15//! share state that records where it was made. A local save made at a commit `HEAD` does not
16//! descend from, as on another branch, comes last, after the remote's: on a
17//! machine that builds many branches in one directory, as a CI runner does,
18//! another branch's state is the least like this checkout's. It is saved after
19//! successful runs, one record per worktree.
20//!
21//! `qk:warm` may also be an array of such objects, each restored and saved on
22//! its own, so state that relocates can be portable beside state that cannot.
23
24use std::collections::{BTreeMap, BTreeSet};
25use std::fs;
26use std::io::Write;
27use std::path::PathBuf;
28use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
29use std::time::{Duration, SystemTime, UNIX_EPOCH};
30
31use anyhow::{Context, Result, bail};
32use qk_config::Workspace;
33use qk_taskgraph::{Task, TaskGraph};
34use serde::{Deserialize, Serialize};
35use serde_json::Value;
36
37use crate::paths::{self, Outputs};
38use crate::record::{self, Artifact, WarmGroup, WarmRecord, validate_link};
39use crate::store::{copy_blob_contents, mode, set_mode, set_modified, symlink};
40use crate::{Cache, hash::digest_file, profile};
41
42/// When restored files say they were written. Warm state comes from another
43/// checkout, so it must look older than every file in this one: tools that
44/// trust timestamps, as `tsc --build` does, then check the sources against it
45/// instead of taking a restored build for an up-to-date one.
46pub const RESTORED_AT: std::time::SystemTime = std::time::UNIX_EPOCH;
47
48/// What a group's identity starts with, before its name.
49const GROUP: &str = "group ";
50
51/// The fields a `qk:warm` object can set. Others are ignored, as Nx ignores
52/// keys it does not know, so a configuration can use fields a newer qk added;
53/// `qk doctor` reports them.
54pub const FIELDS: &[&str] = &[
55    "outputs",
56    "paths",
57    "env",
58    "maxSize",
59    "remote",
60    "mtimes",
61    "portable",
62    "key",
63    "restoreKeys",
64    "group",
65    "survive",
66    "save",
67    "restore",
68];
69
70/// How many saves are kept for a task, across worktrees and keys.
71const KEPT: usize = 8;
72/// Files restored at once.
73const RESTORERS: usize = 8;
74
75/// A warm file to restore: its path, blob, mode, destination and modification time.
76type PendingFile<'a> = (&'a String, &'a str, u32, PathBuf, SystemTime);
77/// A restored file's stamp and size.
78type RestoredFile = (Vec<i64>, u64);
79
80/// Which saves of a target's warm state, besides its own worktree's, may be
81/// restored.
82#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
83pub enum Portable {
84    /// None: `portable: false`, the default.
85    #[default]
86    Own,
87    /// The remote's, when made in a worktree at the same path, and so this
88    /// worktree's own on another machine: `portable: "path"`. For state that
89    /// records where it was made, which CI runners checking out in one
90    /// place can share and other checkouts cannot use.
91    Path,
92    /// Any other worktree's and the remote's: `portable: true`.
93    Anywhere,
94}
95
96/// A target's `qk:warm`.
97#[derive(Clone, Debug, Default)]
98pub struct Warm {
99    /// What its saves belong to: the task, or the group it shares them with.
100    pub identity: String,
101    pub outputs: bool,
102    /// Workspace-relative scratch paths.
103    pub paths: Vec<String>,
104    /// Variables for the task, with tokens expanded.
105    pub env: BTreeMap<String, String>,
106    /// Whether any variable names `{warm}`, which makes it a group.
107    pub directory: bool,
108    pub max_size: Option<u64>,
109    pub remote: bool,
110    /// Whether a worktree's own save comes back with the modification times it
111    /// was saved with, rather than [`RESTORED_AT`].
112    pub preserve_mtimes: bool,
113    /// Which saves besides this worktree's own may be restored.
114    pub portable: Portable,
115    /// What decides whether a save suits this checkout, in order.
116    pub key: Vec<KeyPart>,
117    /// How many leading parts of the key a save must match when none matches
118    /// it whole, or `None` to restore only exact matches.
119    pub restore_keys: Option<usize>,
120    /// Dependencies that delete the paths, which are moved aside while one
121    /// runs: `target` in the task's project, or `project:target`.
122    pub survive: Vec<String>,
123    /// Whether saving happens after the task has reported, in the
124    /// background, rather than before.
125    pub background: bool,
126    /// Whether each of `paths` is restored when it alone is missing, rather
127    /// than all of them only when none is on disk.
128    pub each_path: bool,
129}
130
131/// One part of a warm key.
132#[derive(Clone, Debug)]
133pub enum KeyPart {
134    /// A workspace-relative file, by its content.
135    File(String),
136    /// An environment variable, by its value.
137    Env(String),
138}
139
140/// The target's `qk:warm` entries, none without one.
141///
142/// Entries of an array keep their saves apart, so at most one of them may
143/// belong to the task itself, and the rest name groups. No two may set the
144/// same variable.
145pub fn config(workspace: &Workspace, task: &Task) -> Result<Vec<Warm>> {
146    const SHAPE: &str = "qk:warm must be an object or an array of objects";
147    let objects = match task.definition.extra.get("qk:warm") {
148        None => return Ok(Vec::new()),
149        Some(Value::Object(object)) => vec![object],
150        Some(Value::Array(entries)) => entries
151            .iter()
152            .map(|entry| entry.as_object().context(SHAPE))
153            .collect::<Result<_>>()?,
154        Some(_) => bail!(SHAPE),
155    };
156    let mut entries: Vec<Warm> = Vec::new();
157    for object in objects {
158        let warm = entry(workspace, task, object)?;
159        if entries.iter().any(|other| other.identity == warm.identity) {
160            bail!(match warm.group() {
161                Some(group) => format!("qk:warm names group {group:?} twice"),
162                None => "only one qk:warm entry may go without a group".to_owned(),
163            });
164        }
165        for name in warm.env.keys() {
166            if entries.iter().any(|other| other.env.contains_key(name)) {
167                bail!("qk:warm sets {name} in more than one entry");
168            }
169        }
170        entries.push(warm);
171    }
172    Ok(entries)
173}
174
175/// One `qk:warm` object.
176fn entry(
177    workspace: &Workspace,
178    task: &Task,
179    object: &serde_json::Map<String, Value>,
180) -> Result<Warm> {
181    let flag = |name: &str, default: bool| -> Result<bool> {
182        object.get(name).map_or(Ok(default), |value| {
183            value
184                .as_bool()
185                .with_context(|| format!("qk:warm.{name} must be true or false"))
186        })
187    };
188    let group = object
189        .get("group")
190        .map(|value| {
191            value
192                .as_str()
193                .filter(|name| !name.is_empty())
194                .context("qk:warm.group must be a name")
195        })
196        .transpose()?;
197    // `survive` keeps `paths` across a dependency, which a group has none of.
198    if group.is_some()
199        && ["outputs", "paths", "survive"]
200            .iter()
201            .any(|key| object.contains_key(*key))
202    {
203        bail!(
204            "a qk:warm.group shares {{warm}} alone; outputs, paths and survive belong to one task"
205        );
206    }
207    let identity = group.map_or_else(|| task.id.clone(), |group| format!("{GROUP}{group}"));
208    const SURVIVE_SHAPE: &str = "qk:warm.survive must be an array of target names";
209    let survive = object
210        .get("survive")
211        .map(|value| {
212            value
213                .as_array()
214                .context(SURVIVE_SHAPE)?
215                .iter()
216                .map(|target| {
217                    target
218                        .as_str()
219                        .filter(|target| !target.is_empty())
220                        .map(str::to_owned)
221                        .context(SURVIVE_SHAPE)
222                })
223                .collect::<Result<Vec<_>>>()
224        })
225        .transpose()?
226        .unwrap_or_default();
227    let directory = directory(workspace, &identity);
228    let directory_text = directory.to_str().context("warm directory must be UTF-8")?;
229    let workspace_text = workspace
230        .root
231        .to_str()
232        .context("workspace root must be UTF-8")?;
233    let project_root = &workspace.projects[&task.project].root;
234    let project_directory = if project_root == "." {
235        workspace.root.clone()
236    } else {
237        workspace.root.join(project_root)
238    };
239    let project_text = project_directory
240        .to_str()
241        .context("project root must be UTF-8")?;
242    let expand = |text: &str| -> Result<String> {
243        if text.contains("{options.") || text.contains("{args.") {
244            bail!("dynamic warm environment values are not supported yet");
245        }
246        Ok(text
247            .replace("{workspaceRoot}", workspace_text)
248            .replace("{projectRoot}", project_text)
249            .replace("{warm}", directory_text))
250    };
251    let mut paths = Vec::new();
252    for path in object
253        .get("paths")
254        .map(|value| {
255            value
256                .as_array()
257                .context("qk:warm.paths must be an array of paths")
258        })
259        .transpose()?
260        .into_iter()
261        .flatten()
262    {
263        let path = path
264            .as_str()
265            .context("qk:warm.paths must be an array of paths")?;
266        if path.contains("{warm}") {
267            bail!("qk:warm.paths are workspace paths; {{warm}} is kept already");
268        }
269        let (excluded, path) = match path.strip_prefix('!') {
270            Some(path) => (true, path),
271            None => (false, path),
272        };
273        dynamic(path)?;
274        let path = paths::expand(workspace, &task.project, path)?
275            .trim_end_matches('/')
276            .to_owned();
277        paths::validate_path(&path)?;
278        paths.push(if excluded { format!("!{path}") } else { path });
279    }
280    let mut env = BTreeMap::new();
281    let mut uses_directory = false;
282    if let Some(values) = object.get("env") {
283        for (name, value) in values
284            .as_object()
285            .context("qk:warm.env must be an object")?
286        {
287            let value = value
288                .as_str()
289                .context("qk:warm.env values must be strings")?;
290            uses_directory |= value.contains("{warm}");
291            env.insert(name.clone(), expand(value)?);
292        }
293    }
294    let max_size = match object.get("maxSize") {
295        None => None,
296        Some(Value::String(text)) => Some(crate::parse_size(text)?),
297        Some(Value::Number(number)) => number.as_u64(),
298        Some(_) => bail!("qk:warm.maxSize must be a size such as \"2GB\""),
299    };
300    let preserve_mtimes = match object.get("mtimes").map(|value| value.as_str()) {
301        None | Some(Some("epoch")) => false,
302        Some(Some("preserve")) => true,
303        Some(_) => bail!("qk:warm.mtimes must be \"epoch\" or \"preserve\""),
304    };
305    const KEY_SHAPE: &str = "qk:warm.key must be an array of paths and {\"env\": name} objects";
306    let mut key = Vec::new();
307    for part in object
308        .get("key")
309        .map(|value| value.as_array().context(KEY_SHAPE))
310        .transpose()?
311        .into_iter()
312        .flatten()
313    {
314        key.push(match part {
315            Value::String(path) => {
316                dynamic(path)?;
317                let path = paths::expand(workspace, &task.project, path)?;
318                paths::validate_path(&path)?;
319                KeyPart::File(path)
320            }
321            Value::Object(object) if object.len() == 1 => KeyPart::Env(
322                object
323                    .get("env")
324                    .and_then(Value::as_str)
325                    .context(KEY_SHAPE)?
326                    .to_owned(),
327            ),
328            _ => bail!(KEY_SHAPE),
329        });
330    }
331    let restore_keys = object
332        .get("restoreKeys")
333        .map(|value| {
334            value
335                .as_u64()
336                .and_then(|count| usize::try_from(count).ok())
337                .filter(|count| *count <= key.len())
338                .context("qk:warm.restoreKeys must be a count no larger than qk:warm.key")
339        })
340        .transpose()?;
341    let each_path = match object.get("restore").map(Value::as_str) {
342        None | Some(Some("whole")) => false,
343        Some(Some("each-path")) if !paths.is_empty() => true,
344        Some(Some("each-path")) => bail!("qk:warm.restore \"each-path\" needs qk:warm.paths"),
345        Some(_) => bail!("qk:warm.restore must be \"whole\" or \"each-path\""),
346    };
347    let background = match object.get("save").map(Value::as_str) {
348        None | Some(Some("wait")) => false,
349        Some(Some("background")) => true,
350        Some(_) => bail!("qk:warm.save must be \"wait\" or \"background\""),
351    };
352    Ok(Warm {
353        identity,
354        outputs: flag("outputs", false)?,
355        paths,
356        env,
357        directory: uses_directory,
358        max_size,
359        remote: flag("remote", true)?,
360        preserve_mtimes,
361        portable: match object.get("portable") {
362            None | Some(Value::Bool(false)) => Portable::Own,
363            Some(Value::Bool(true)) => Portable::Anywhere,
364            Some(Value::String(path)) if path == "path" => Portable::Path,
365            Some(_) => bail!("qk:warm.portable must be true, false or \"path\""),
366        },
367        key,
368        restore_keys,
369        survive,
370        background,
371        each_path,
372    })
373}
374
375/// The directory `{warm}` names for a task or group, in the worktree's state.
376/// Rejects a warm path that would change with the task's options or
377/// arguments, which qk cannot resolve before the task runs.
378fn dynamic(path: &str) -> Result<()> {
379    if path.contains("{options.") || path.contains("{args.") {
380        bail!("dynamic cache paths are not supported yet");
381    }
382    Ok(())
383}
384
385fn directory(workspace: &Workspace, identity: &str) -> PathBuf {
386    paths::worktree_state(&workspace.root)
387        .join("warm")
388        .join(&blake3::hash(identity.as_bytes()).to_hex()[..32])
389}
390
391impl Warm {
392    /// The group the entry shares its saves with, if it names one.
393    pub fn group(&self) -> Option<&str> {
394        self.identity.strip_prefix(GROUP)
395    }
396
397    /// The scratch paths without their exclusions: what the task keeps,
398    /// whether or not it is saved, and so never an input.
399    pub fn kept_paths(&self) -> Vec<String> {
400        self.paths
401            .iter()
402            .filter(|path| !path.starts_with('!'))
403            .cloned()
404            .collect()
405    }
406}
407
408/// The scratch paths every entry keeps.
409pub fn kept_paths(entries: &[Warm]) -> Vec<String> {
410    entries.iter().flat_map(Warm::kept_paths).collect()
411}
412
413/// Paths moved aside while a dependency that deletes them runs, and moved
414/// back, over whatever it left there, when this is dropped.
415#[derive(Default)]
416pub struct Kept {
417    /// The stashes this dependency holds a share of.
418    roots: Vec<PathBuf>,
419}
420
421impl Drop for Kept {
422    fn drop(&mut self) {
423        for root in self.roots.iter().rev() {
424            release(root);
425        }
426    }
427}
428
429/// One task's paths, moved into its directory in the worktree's state.
430#[derive(Default, Serialize, Deserialize)]
431struct Stash {
432    #[serde(skip)]
433    root: PathBuf,
434    #[serde(skip)]
435    workspace: PathBuf,
436    /// Each moved path, relative to the workspace, by the order it was moved.
437    moved: Vec<String>,
438}
439
440/// A stash in use, and how many of the running dependencies share it. The
441/// lock keeps another qk process from taking it for abandoned.
442struct Active {
443    users: usize,
444    stash: Stash,
445    _lock: fs::File,
446}
447
448/// The stashes this process holds, by root. Dependencies that run at once
449/// share one: the first moves the paths, the last puts them back.
450static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
451
452impl Stash {
453    /// The ownership record for the numbered paths in a stash.
454    fn manifest(root: &std::path::Path) -> PathBuf {
455        root.join("moved.json")
456    }
457
458    /// Publishes ownership before moving a path, without truncating the last
459    /// valid record. Sync the file before the atomic replacement.
460    fn save_manifest(&self) -> Result<()> {
461        let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
462        file.write_all(&serde_json::to_vec(self)?)?;
463        file.as_file().sync_all()?;
464        file.persist(Self::manifest(&self.root))?;
465        Ok(())
466    }
467
468    /// Moves the paths back, replacing what is there now, and removes the
469    /// stash once every path is back. A path whose parents are not plain
470    /// directories inside the workspace, as when a dependency put a link or a
471    /// file in their place, is not touched and stays in the stash for the
472    /// next run. Whether every path went back.
473    fn put_back(&self) -> bool {
474        let mut complete = true;
475        for (index, path) in self.moved.iter().enumerate().rev() {
476            let kept = self.root.join(index.to_string());
477            match fs::symlink_metadata(&kept) {
478                Ok(_) => {}
479                Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
480                Err(error) => {
481                    complete = false;
482                    qk_executor::status!(
483                        "qk: could not inspect saved {path}, kept for the next run ({error})"
484                    );
485                    continue;
486                }
487            }
488            let original = self.workspace.join(path);
489            let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
490                remove(&original)?;
491                if let Some(parent) = original.parent() {
492                    fs::create_dir_all(parent)?;
493                }
494                fs::rename(&kept, &original)?;
495                Ok(())
496            });
497            if let Err(error) = result {
498                complete = false;
499                qk_executor::status!(
500                    "qk: could not put {path} back, kept for the next run ({error:#})"
501                );
502            }
503        }
504        if complete {
505            let _ = fs::remove_dir_all(&self.root);
506        }
507        complete
508    }
509}
510
511/// Removes a path itself, including a symlink, without following it.
512fn remove(path: &std::path::Path) -> std::io::Result<()> {
513    match fs::symlink_metadata(path) {
514        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
515        Err(error) => Err(error),
516        Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
517        Ok(_) => fs::remove_file(path),
518    }
519}
520
521/// Gives up one dependency's share of a stash, putting the paths back when
522/// it was the last.
523fn release(root: &std::path::Path) {
524    let mut active = ACTIVE
525        .lock()
526        .unwrap_or_else(std::sync::PoisonError::into_inner);
527    let Some(entry) = active.get_mut(root) else {
528        return;
529    };
530    entry.users -= 1;
531    if entry.users > 0 {
532        return;
533    }
534    let entry = active.remove(root).expect("present above");
535    // Another dependency taking the stash again waits on the lock, not on
536    // every other task's stash.
537    drop(active);
538    entry.stash.put_back();
539}
540
541/// Whether `entry`, from `owner`'s `survive`, names `task`.
542fn names(owner: &Task, entry: &str, task: &Task) -> bool {
543    match entry.split_once(':') {
544        Some((project, target)) => task.project == project && task.target == target,
545        None => task.project == owner.project && task.target == entry,
546    }
547}
548
549/// Whether `from` depends on `on`, directly or through other tasks.
550fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
551    let mut seen = std::collections::BTreeSet::new();
552    let mut next = vec![from];
553    while let Some(id) = next.pop() {
554        for dependency in graph
555            .tasks
556            .get(id)
557            .into_iter()
558            .flat_map(|task| &task.dependencies)
559        {
560            if dependency == on {
561                return true;
562            }
563            if seen.insert(dependency.as_str()) {
564                next.push(dependency);
565            }
566        }
567    }
568    false
569}
570
571/// Moves aside the warm paths of every task in the graph that depends on
572/// `task` and names it in `survive`, so that `task` cannot delete them.
573pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
574    let mut kept = Kept::default();
575    for owner in graph.tasks.values() {
576        let Ok(entries) = config(workspace, owner) else {
577            continue;
578        };
579        // Only the task's own entry has paths to keep.
580        let Some(warm) = entries.into_iter().find(|warm| warm.group().is_none()) else {
581            continue;
582        };
583        if !warm.survive.iter().any(|entry| names(owner, entry, task))
584            || !depends_on(graph, &owner.id, &task.id)
585        {
586            continue;
587        }
588        match acquire(workspace, owner, &warm) {
589            Ok(root) => kept.roots.push(root),
590            Err(error) => qk_executor::status!(
591                "qk: {}: warm paths not kept across {} ({error:#})",
592                owner.id,
593                task.id
594            ),
595        }
596    }
597    kept
598}
599
600/// Takes a share of `owner`'s stash, moving its paths aside if no running
601/// dependency has yet.
602fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
603    let state = paths::worktree_state(&workspace.root).join("warm-kept");
604    let root = state.join(hash32(&owner.id));
605    fs::create_dir_all(&state)?;
606    let lock = fs::OpenOptions::new()
607        .create(true)
608        .truncate(false)
609        .read(true)
610        .write(true)
611        .open(state.join(format!("{}.lock", hash32(&owner.id))))?;
612    // Only the nonblocking attempt holds ACTIVE. Waiting for another process
613    // must leave this process's other stashes free to acquire and release.
614    let mut active = loop {
615        let mut active = ACTIVE
616            .lock()
617            .unwrap_or_else(std::sync::PoisonError::into_inner);
618        if let Some(entry) = active.get_mut(&root) {
619            entry.users += 1;
620            return Ok(root);
621        }
622        match fs2::FileExt::try_lock_exclusive(&lock) {
623            Ok(()) => break active,
624            Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
625                drop(active);
626                std::thread::sleep(Duration::from_millis(20));
627            }
628            Err(error) => return Err(error.into()),
629        }
630    };
631    // Under the lock, a stash on disk is one no run holds: one an earlier run
632    // could not put back, or left when it was stopped. It goes back first.
633    // What cannot go back yet, as when its parent is not a directory, stays
634    // in the stash this dependency holds, beside the paths that did go back,
635    // and is put back after the dependency, which may well recreate the
636    // parent.
637    let left = match fs::read(Stash::manifest(&root)) {
638        Ok(bytes) => {
639            let mut left: Stash = serde_json::from_slice(&bytes)
640                .context("invalid warm stash manifest; saved paths retained for recovery")?;
641            left.root = root.clone();
642            left.workspace = workspace.root.clone();
643            Some(left)
644        }
645        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
646            if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
647                bail!("missing warm stash manifest; saved paths retained for recovery");
648            }
649            None
650        }
651        Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
652    };
653    let retained = left.filter(|left| !left.put_back());
654    let stash = stash(workspace, &root, warm, retained)?;
655    active.insert(
656        root.clone(),
657        Active {
658            users: 1,
659            stash,
660            _lock: lock,
661        },
662    );
663    Ok(root)
664}
665
666/// Moves the warm paths into a stash at `root`: a new one, or `retained`,
667/// what an earlier run could not put back, which keeps its entries. A path
668/// that is one of those, or inside or around one, stays where it is.
669fn stash(
670    workspace: &Workspace,
671    root: &std::path::Path,
672    warm: &Warm,
673    retained: Option<Stash>,
674) -> Result<Stash> {
675    let mut stash = match retained {
676        Some(retained) => retained,
677        None => {
678            let _ = fs::remove_dir_all(root);
679            Stash {
680                root: root.to_owned(),
681                workspace: workspace.root.clone(),
682                moved: Vec::new(),
683            }
684        }
685    };
686    let held: Vec<String> = stash
687        .moved
688        .iter()
689        .enumerate()
690        .filter_map(
691            |(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
692                Ok(_) => Some(Ok(path.clone())),
693                Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
694                Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
695            },
696        )
697        .collect::<Result<_>>()?;
698    let overlaps = |path: &str| {
699        held.iter().any(|held| {
700            held == path
701                || path.starts_with(&format!("{held}/"))
702                || held.starts_with(&format!("{path}/"))
703        })
704    };
705    // Each path on its own: one whose parents are not usable now, as one
706    // retained above may have, has nothing live there to move.
707    let mut kept = std::collections::BTreeSet::new();
708    for path in warm.kept_paths() {
709        if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
710            .and_then(|outputs| outputs.paths(&workspace.root))
711        {
712            kept.extend(found);
713        }
714    }
715    // The topmost paths only: moving one moves what is below it.
716    let topmost: Vec<String> = kept
717        .iter()
718        .filter(|path| {
719            !std::path::Path::new(path.as_str())
720                .ancestors()
721                .skip(1)
722                .any(|ancestor| {
723                    ancestor
724                        .to_str()
725                        .is_some_and(|ancestor| kept.contains(ancestor))
726                })
727        })
728        .filter(|path| !overlaps(path))
729        .cloned()
730        .collect();
731    if topmost.is_empty() {
732        return Ok(stash);
733    }
734    fs::create_dir_all(root)?;
735    for path in topmost {
736        let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
737            stash.moved.push(path.clone());
738            stash.save_manifest()?;
739            // An entry put back earlier left its slot empty, and a retained
740            // one's slot is taken; a new one goes at the end either way.
741            fs::rename(
742                workspace.root.join(&path),
743                root.join((stash.moved.len() - 1).to_string()),
744            )
745            .inspect_err(|_| {
746                stash.moved.pop();
747            })?;
748            Ok(())
749        });
750        if let Err(error) = moved {
751            // What was moved before goes back, rather than waiting on a run
752            // that holds no share of it.
753            stash.save_manifest()?;
754            stash.put_back();
755            return Err(error);
756        }
757    }
758    Ok(stash)
759}
760
761/// Two tasks in a run cannot keep the same scratch path, since both would
762/// write it.
763pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
764    let mut claimed: Vec<(String, &str)> = Vec::new();
765    for (id, task) in &graph.tasks {
766        for path in kept_paths(&config(workspace, task)?) {
767            if let Some((_, other)) = claimed.iter().find(|(known, _)| {
768                known == &path
769                    || known.starts_with(&format!("{path}/"))
770                    || path.starts_with(&format!("{known}/"))
771            }) {
772                bail!("{id} and {other} both keep {path} as warm state");
773            }
774            claimed.push((path, id));
775        }
776    }
777    Ok(())
778}
779
780impl WarmGroup {
781    /// Total saved file bytes, excluding directories and links.
782    fn size(&self, cache: &Cache) -> u64 {
783        self.artifacts
784            .values()
785            .filter_map(|artifact| match artifact {
786                Artifact::File { blob, .. } => {
787                    fs::metadata(cache.root.join("blobs").join(blob)).ok()
788                }
789                _ => None,
790            })
791            .map(|metadata| metadata.len())
792            .sum()
793    }
794}
795
796/// What this worktree last restored, per group: each file's metadata once
797/// restored and the artifact it came from, so a save need not read it again.
798#[derive(Default, Serialize, Deserialize)]
799struct RestoredFiles {
800    groups: BTreeMap<String, GroupFiles>,
801}
802
803/// A group's restored files: each one's stamp and artifact, by path.
804type GroupFiles = BTreeMap<String, (Vec<i64>, Artifact)>;
805
806/// What a restore brought back.
807#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
808pub struct Restored {
809    /// The group the entry shares its saves with, absent for the task's own.
810    #[serde(default, skip_serializing_if = "Option::is_none")]
811    pub group: Option<String>,
812    /// `local` for this worktree's own save, `worktree <root>` for another
813    /// worktree's, or `remote <branch>`.
814    pub source: String,
815    pub groups: Vec<String>,
816    pub files: usize,
817    pub bytes: u64,
818}
819
820/// What warm state did for one run of a task.
821#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
822#[serde(rename_all = "camelCase")]
823pub struct WarmReport {
824    /// What each entry restored before the run, for those that restored any.
825    #[serde(default, skip_serializing_if = "Vec::is_empty")]
826    pub restored: Vec<Restored>,
827    /// The groups already on disk before the run, which were left as they
828    /// were: `outputs`, `paths` and `directory` for the task's own entry, and
829    /// a named group's name for its directory.
830    #[serde(default, skip_serializing_if = "Vec::is_empty")]
831    pub present: Vec<String>,
832    /// How long saving the entries after the run took, when any was saved
833    /// before the task reported.
834    pub save_ms: Option<u64>,
835    /// Whether any entry was saved in the background, after the task reported.
836    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
837    pub background: bool,
838}
839
840/// A warm group: its name, the directory its paths are relative to, and
841/// which paths below that directory belong to it.
842#[derive(Clone)]
843struct Location {
844    name: &'static str,
845    base: PathBuf,
846    /// `None` for everything below `base`.
847    outputs: Option<Outputs>,
848    /// With `restore: "each-path"`, each configured path on its own, with the
849    /// exclusions, so each can be restored when it alone is missing.
850    each: Vec<Outputs>,
851}
852
853impl Location {
854    /// The group's paths on disk, relative to `base`.
855    fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
856        if !self.base.exists() {
857            return Ok(Default::default());
858        }
859        match &self.outputs {
860            Some(outputs) => outputs.paths(&self.base),
861            None => {
862                let mut paths = std::collections::BTreeSet::new();
863                for entry in walkdir::WalkDir::new(&self.base)
864                    .min_depth(1)
865                    .follow_links(false)
866                {
867                    let entry = entry?;
868                    paths.insert(paths::relative(&self.base, entry.path())?);
869                }
870                Ok(paths)
871            }
872        }
873    }
874
875    /// Whether a relative path belongs to this configured group.
876    fn matches(&self, path: &str) -> bool {
877        self.outputs
878            .as_ref()
879            .is_none_or(|outputs| outputs.matches(path))
880    }
881
882    /// The configured paths that are on disk, when restoring each path on its
883    /// own and some of them are missing; `None` restores the group whole.
884    fn present(&self) -> Result<Option<Vec<&Outputs>>> {
885        let mut present = Vec::new();
886        for outputs in &self.each {
887            if !outputs.paths(&self.base)?.is_empty() {
888                present.push(outputs);
889            }
890        }
891        Ok((!present.is_empty() && present.len() < self.each.len()).then_some(present))
892    }
893}
894
895/// Removes what a failed restore put in a location, children before their
896/// directories, leaving the paths it kept. Nothing else of the location was on
897/// disk, so all of it came from the attempt.
898fn remove_restored(location: &Location, kept: &[&Outputs]) {
899    let Ok(paths) = location.paths() else {
900        return;
901    };
902    for path in paths.iter().rev() {
903        if kept.iter().any(|outputs| outputs.matches(path)) {
904            continue;
905        }
906        let absolute = location.base.join(path);
907        let _ = match fs::symlink_metadata(&absolute) {
908            Ok(metadata) if metadata.is_dir() => fs::remove_dir(&absolute),
909            Ok(_) => fs::remove_file(&absolute),
910            Err(_) => continue,
911        };
912    }
913}
914
915/// Resolves output, scratch and external environment groups for this task.
916fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
917    let mut locations = Vec::new();
918    if warm.outputs {
919        locations.push(Location {
920            name: "outputs",
921            base: workspace.root.clone(),
922            outputs: Some(Outputs::new(workspace, task)?),
923            each: Vec::new(),
924        });
925    }
926    if !warm.paths.is_empty() {
927        let mut each = Vec::new();
928        if warm.each_path {
929            let (excluded, included): (Vec<_>, Vec<_>) =
930                warm.paths.iter().partition(|path| path.starts_with('!'));
931            for path in included {
932                let mut patterns = vec![path.clone()];
933                patterns.extend(excluded.iter().map(|path| (*path).clone()));
934                each.push(Outputs::from_paths(&patterns)?);
935            }
936        }
937        locations.push(Location {
938            name: "paths",
939            base: workspace.root.clone(),
940            outputs: Some(Outputs::from_paths(&warm.paths)?),
941            each,
942        });
943    }
944    if warm.directory {
945        locations.push(Location {
946            name: "directory",
947            base: directory(workspace, &warm.identity),
948            outputs: None,
949            each: Vec::new(),
950        });
951    }
952    Ok(locations)
953}
954
955/// The branch being built: from CI when it checks out a detached head, else
956/// from git. `None` when there is no branch to name.
957pub(crate) fn current_branch(workspace: &Workspace) -> Option<String> {
958    for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
959        if let Ok(value) = std::env::var(name)
960            && !value.is_empty()
961        {
962            return Some(value);
963        }
964    }
965    let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
966    let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
967    (output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
968}
969
970/// The branches to take warm state from, in order: the current one, then the
971/// workspace's default branch.
972pub(crate) fn branches(workspace: &Workspace) -> Vec<String> {
973    let default = workspace
974        .config
975        .default_base
976        .clone()
977        .unwrap_or_else(|| "main".into());
978    let default = default
979        .strip_prefix("origin/")
980        .unwrap_or(&default)
981        .to_owned();
982    let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
983    if !branches.contains(&default) {
984        branches.push(default);
985    }
986    branches
987}
988
989/// Metadata used to detect files unchanged since their save or restore.
990fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
991    let modified = metadata
992        .modified()
993        .ok()
994        .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
995        .map_or(-1, |duration| duration.as_nanos() as i64);
996    #[allow(unused_mut)]
997    let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
998    #[cfg(unix)]
999    {
1000        use std::os::unix::fs::MetadataExt;
1001        stamp.extend([
1002            metadata.ctime(),
1003            metadata.ctime_nsec(),
1004            metadata.ino() as i64,
1005        ]);
1006    }
1007    stamp
1008}
1009
1010/// This worktree's root, as records name it.
1011/// The name a remote save of `warm` from a checkout at `worktree` is stored
1012/// under: path-portable saves are kept apart by path, so a checkout at
1013/// another path never replaces the save the checkouts at this one restore.
1014fn remote_identity(warm: &Warm, worktree: &str) -> String {
1015    match warm.portable {
1016        Portable::Path => format!("{}\0path\0{worktree}", warm.identity),
1017        _ => warm.identity.clone(),
1018    }
1019}
1020
1021fn worktree(workspace: &Workspace) -> String {
1022    fs::canonicalize(&workspace.root)
1023        .unwrap_or_else(|_| workspace.root.clone())
1024        .to_string_lossy()
1025        .into_owned()
1026}
1027
1028/// The checkout commit, absent outside Git workspaces.
1029fn current_commit(workspace: &Workspace) -> Option<String> {
1030    let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
1031    let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
1032    (output.status.success() && !commit.is_empty()).then_some(commit)
1033}
1034
1035/// A digest of each part of the warm key, as this checkout and the task's
1036/// environment have them.
1037fn key_digests(
1038    workspace: &Workspace,
1039    warm: &Warm,
1040    prepared: &qk_executor::PreparedTask,
1041) -> Result<Vec<String>> {
1042    warm.key
1043        .iter()
1044        .map(|part| {
1045            Ok(match part {
1046                KeyPart::File(path) => {
1047                    paths::safe_parents(&workspace.root, path)?;
1048                    let absolute = workspace.root.join(path);
1049                    if absolute.is_file() {
1050                        format!("file:{}", digest_file(&absolute)?)
1051                    } else {
1052                        "file:absent".to_owned()
1053                    }
1054                }
1055                // As the task's process gets them: its execution variables
1056                // over its environment, and nothing else.
1057                KeyPart::Env(name) => match prepared
1058                    .execution
1059                    .get(std::ffi::OsStr::new(name))
1060                    .or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
1061                {
1062                    Some(value) => format!(
1063                        "env:{}",
1064                        blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
1065                    ),
1066                    None => "env:unset".to_owned(),
1067                },
1068            })
1069        })
1070        .collect()
1071}
1072
1073/// How well a save's key suits this checkout: `Some(0)` when it matches
1074/// whole, `Some(1)` when it matches only in the leading parts `restoreKeys`
1075/// allows, `None` when it does not suit.
1076fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
1077    if saved == current {
1078        return Some(0);
1079    }
1080    let count = warm.restore_keys?;
1081    (saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
1082}
1083
1084/// Whether `HEAD` descends from `commit`, or is it.
1085fn descends_from(workspace: &Workspace, commit: &str) -> bool {
1086    paths::git(
1087        &workspace.root,
1088        &["merge-base", "--is-ancestor", commit, "HEAD"],
1089    )
1090    .is_ok_and(|output| output.status.success())
1091}
1092
1093/// How many commits `commit` is behind the checkout's `HEAD`, or `None` when
1094/// it is not an ancestor of it.
1095fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
1096    let ancestor = paths::git(
1097        &workspace.root,
1098        &["merge-base", "--is-ancestor", commit, "HEAD"],
1099    )
1100    .ok()?;
1101    if !ancestor.status.success() {
1102        return None;
1103    }
1104    let output = paths::git(
1105        &workspace.root,
1106        &["rev-list", "--count", &format!("{commit}..HEAD")],
1107    )
1108    .ok()?;
1109    String::from_utf8(output.stdout).ok()?.trim().parse().ok()
1110}
1111
1112/// A compact stable filename component for task and group identities.
1113fn hash32(text: &str) -> String {
1114    blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
1115}
1116
1117/// A representable timestamp for restoring saved modification times.
1118fn nanos(time: SystemTime) -> Option<i64> {
1119    time.duration_since(UNIX_EPOCH)
1120        .ok()
1121        .and_then(|duration| i64::try_from(duration.as_nanos()).ok())
1122}
1123
1124/// The save time used to order records across worktrees.
1125fn now_ms() -> u64 {
1126    SystemTime::now()
1127        .duration_since(UNIX_EPOCH)
1128        .map_or(0, |duration| duration.as_millis() as u64)
1129}
1130
1131/// Where a worktree notes what it last restored for a task or group.
1132fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
1133    paths::worktree_state(&workspace.root)
1134        .join("warm-restored")
1135        .join(format!("{}.json", hash32(&warm.identity)))
1136}
1137
1138/// Reads optional restore metadata; a missing note forces fresh file reads.
1139fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
1140    fs::read(restored_files_path(workspace, warm))
1141        .ok()
1142        .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1143        .unwrap_or_default()
1144}
1145
1146impl Cache {
1147    /// The file-name prefix every save of a task or group shares.
1148    fn warm_prefix(&self, identity: &str) -> String {
1149        hash32(&format!(
1150            "{identity}\0{}\0{}",
1151            std::env::consts::OS,
1152            std::env::consts::ARCH
1153        ))
1154    }
1155
1156    /// Serializes restoring and saving a task's or group's state, which the
1157    /// tasks of a group may do at once.
1158    fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
1159        let directory = self.root.join("locks");
1160        fs::create_dir_all(&directory)?;
1161        let file = fs::OpenOptions::new()
1162            .create(true)
1163            .truncate(false)
1164            .read(true)
1165            .write(true)
1166            .open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
1167        fs2::FileExt::lock_exclusive(&file)?;
1168        Ok(file)
1169    }
1170
1171    fn warm_record(&self, record: &WarmRecord) -> PathBuf {
1172        self.root.join("warm").join(format!(
1173            "{}-{}.json",
1174            self.warm_prefix(&record.task),
1175            hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
1176        ))
1177    }
1178
1179    /// Every save of the task or group, newest first.
1180    fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, WarmRecord)> {
1181        let prefix = format!("{}-", self.warm_prefix(&warm.identity));
1182        let mut records: Vec<(PathBuf, WarmRecord)> = fs::read_dir(self.root.join("warm"))
1183            .into_iter()
1184            .flatten()
1185            .flatten()
1186            .filter(|entry| {
1187                entry
1188                    .file_name()
1189                    .to_str()
1190                    .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
1191            })
1192            .filter_map(|entry| {
1193                let record =
1194                    WarmRecord::read(&fs::read(entry.path()).ok()?, &warm.identity).ok()?;
1195                Some((entry.path(), record))
1196            })
1197            .collect();
1198        records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
1199        records
1200    }
1201
1202    /// The save to restore from, and how to name where it came from. A save
1203    /// whose key matches whole comes before one that matches in part; then
1204    /// this worktree's own before another's, which only a portable target
1205    /// restores; then the one made at the commit nearest behind `HEAD`; then
1206    /// the newest. With `related`, only saves made at a commit `HEAD`
1207    /// descends from are considered, and without, only the others.
1208    fn choose_warm(
1209        &self,
1210        workspace: &Workspace,
1211        warm: &Warm,
1212        current: &[String],
1213        related: bool,
1214    ) -> Option<(WarmRecord, String)> {
1215        let own = worktree(workspace);
1216        let mut candidates: Vec<(u8, bool, WarmRecord)> = self
1217            .warm_records(warm)
1218            .into_iter()
1219            .filter_map(|(_, record)| {
1220                let quality = key_match(warm, current, &record.key)?;
1221                let other = record.worktree != own;
1222                (!other || warm.portable == Portable::Anywhere).then_some((quality, other, record))
1223            })
1224            .filter(|(_, _, record)| {
1225                let descends = record
1226                    .commit
1227                    .as_deref()
1228                    .is_some_and(|commit| descends_from(workspace, commit));
1229                descends == related
1230            })
1231            .collect();
1232        // Newest first within each rank, as warm_records returned them.
1233        candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
1234        let rank = candidates
1235            .first()
1236            .map(|(quality, other, _)| (*quality, *other))?;
1237        let tied = candidates
1238            .iter()
1239            .take_while(|(quality, other, _)| (*quality, *other) == rank)
1240            .count();
1241        let index = if tied > 1 {
1242            (0..tied)
1243                .min_by_key(|&index| {
1244                    let behind = candidates[index]
1245                        .2
1246                        .commit
1247                        .as_deref()
1248                        .and_then(|commit| behind(workspace, commit));
1249                    (behind.is_none(), behind, index)
1250                })
1251                .unwrap_or(0)
1252        } else {
1253            0
1254        };
1255        let (_, other, record) = candidates.swap_remove(index);
1256        let source = if other {
1257            format!("worktree {}", record.worktree)
1258        } else {
1259            "local".to_owned()
1260        };
1261        Some((record, source))
1262    }
1263
1264    fn store_warm(&self, record: &WarmRecord) -> Result<()> {
1265        let path = self.warm_record(record);
1266        let directory = path.parent().context("warm records have a directory")?;
1267        fs::create_dir_all(directory)?;
1268        record::publish(&self.root, &path, record)
1269    }
1270
1271    /// Restores each warm group not already on disk from the save
1272    /// [`Cache::choose_warm`] picks, or the remote's, and notes what it
1273    /// restored so the next save can tell which files are unchanged. Files are
1274    /// dated [`RESTORED_AT`], unless the target keeps modification times and
1275    /// the save is this worktree's own.
1276    pub(crate) fn restore_warm(
1277        &self,
1278        workspace: &Workspace,
1279        task: &Task,
1280        warm: &Warm,
1281        prepared: &qk_executor::PreparedTask,
1282    ) -> Result<(Restored, Vec<String>)> {
1283        let mut restored = Restored {
1284            group: warm.group().map(str::to_owned),
1285            ..Restored::default()
1286        };
1287        let _lock = self.lock_warm(warm)?;
1288        // Local state wins: it is the newest for this checkout.
1289        let mut present = Vec::new();
1290        let mut missing = Vec::new();
1291        let locations = locations(workspace, task, warm)?;
1292        for location in &locations {
1293            if location.paths()?.is_empty() {
1294                missing.push((location, Vec::new()));
1295            } else if let Some(kept) = location.present()? {
1296                // Only what no path on disk covers comes back, so a file the
1297                // tool removed from a kept path is not restored into it.
1298                missing.push((location, kept));
1299            } else {
1300                present.push(warm.group().map_or(location.name, |group| group).to_owned());
1301            }
1302        }
1303        if missing.is_empty() {
1304            return Ok((restored, present));
1305        }
1306        let choosing = profile::span(&task.id, "warm_restore_choose");
1307        let current = key_digests(workspace, warm, prepared)?;
1308        // Another branch's save, here, only when the remote has none.
1309        let chosen = self
1310            .choose_warm(workspace, warm, &current, true)
1311            .or_else(|| {
1312                self.remote_warm(workspace, task, warm, &current)
1313                    .map(|(record, branch)| (record, format!("remote {branch}")))
1314            })
1315            .or_else(|| self.choose_warm(workspace, warm, &current, false));
1316        let Some((record, source)) = chosen else {
1317            return Ok((restored, present));
1318        };
1319        drop(choosing);
1320        let _files = profile::span(&task.id, "warm_restore_files");
1321        let keep_mtimes = warm.preserve_mtimes && source == "local";
1322        restored.source = source;
1323        let mut noted = read_restored_files(workspace, warm);
1324        for (location, kept) in missing {
1325            let Some(group) = record.groups.get(location.name) else {
1326                continue;
1327            };
1328            let (files, bytes) =
1329                match self.restore_location(&task.id, location, &kept, group, keep_mtimes) {
1330                    Ok(restored) => restored,
1331                    Err(error) => {
1332                        // What no kept path covers was not on disk, so all of it
1333                        // there now came from this attempt. A fragment left behind
1334                        // would count as present and keep the rest from coming back.
1335                        remove_restored(location, &kept);
1336                        return Err(error);
1337                    }
1338                };
1339            restored.files += files.len();
1340            restored.bytes += bytes;
1341            if kept.is_empty() {
1342                noted.groups.insert(location.name.to_owned(), files);
1343            } else {
1344                // The kept paths' notes from an earlier restore still hold.
1345                noted
1346                    .groups
1347                    .entry(location.name.to_owned())
1348                    .or_default()
1349                    .extend(files);
1350            }
1351            restored.groups.push(location.name.to_owned());
1352        }
1353        if !restored.groups.is_empty() {
1354            let path = restored_files_path(workspace, warm);
1355            fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
1356            fs::write(path, serde_json::to_vec(&noted)?)?;
1357        }
1358        Ok((restored, present))
1359    }
1360
1361    /// Restores one missing location from its saved group: the restored files'
1362    /// stamps and artifacts, and how many bytes they hold.
1363    fn restore_location(
1364        &self,
1365        task: &str,
1366        location: &Location,
1367        kept: &[&Outputs],
1368        group: &WarmGroup,
1369        keep_mtimes: bool,
1370    ) -> Result<(GroupFiles, u64)> {
1371        fs::create_dir_all(&location.base)?;
1372        // Validate every path and lay out directories and links first, so the
1373        // files, which are most of the work, can be copied in parallel.
1374        let mut created = BTreeSet::new();
1375        let mut directories = Vec::new();
1376        let mut pending = Vec::new();
1377        for (path, artifact) in &group.artifacts {
1378            paths::safe_parents(&location.base, path)?;
1379            if !location.matches(path) {
1380                bail!("warm state for {task} holds {path}, outside its paths");
1381            }
1382            if kept.iter().any(|outputs| outputs.matches(path)) {
1383                continue;
1384            }
1385            let destination = location.base.join(path);
1386            if let Some(parent) = destination.parent()
1387                && created.insert(parent.to_owned())
1388            {
1389                fs::create_dir_all(parent)?;
1390            }
1391            match artifact {
1392                Artifact::File { blob, mode } => {
1393                    let modified = keep_mtimes
1394                        .then(|| group.mtimes.get(path))
1395                        .flatten()
1396                        .and_then(|nanos| u64::try_from(*nanos).ok())
1397                        .map_or(RESTORED_AT, |nanos| {
1398                            UNIX_EPOCH + Duration::from_nanos(nanos)
1399                        });
1400                    pending.push((path, blob.as_str(), *mode, destination, modified));
1401                }
1402                Artifact::Directory { mode } => {
1403                    fs::create_dir_all(&destination)?;
1404                    created.insert(destination.clone());
1405                    directories.push((destination, *mode));
1406                }
1407                Artifact::Symlink { target, directory } => {
1408                    validate_link(path, target)?;
1409                    symlink(target, &destination, *directory)?;
1410                }
1411            }
1412        }
1413        let copied = self
1414            .restore_warm_files(&pending)
1415            .with_context(|| format!("warm state for {task} not restored"))?;
1416        // Directory modes come last, since one may forbid writing into it.
1417        for (destination, mode) in directories {
1418            set_mode(&destination, mode)?;
1419        }
1420        let mut files = BTreeMap::new();
1421        let mut bytes = 0;
1422        for ((path, blob, mode, _, _), (metadata, size)) in pending.into_iter().zip(copied) {
1423            bytes += size;
1424            let artifact = Artifact::File {
1425                blob: blob.to_owned(),
1426                mode,
1427            };
1428            files.insert(path.clone(), (metadata, artifact));
1429        }
1430        Ok((files, bytes))
1431    }
1432
1433    /// Copies warm files out of their blobs, verifying each copy, with bounded
1434    /// parallelism. Returns each file's stamp and size, in the given order.
1435    fn restore_warm_files(&self, files: &[PendingFile]) -> Result<Vec<RestoredFile>> {
1436        let restore =
1437            |(_, blob, mode, destination, modified): &PendingFile| -> Result<RestoredFile> {
1438                let source = self.root.join("blobs").join(blob);
1439                if copy_blob_contents(&source, destination)? != *blob {
1440                    // A corrupt copy must not stay where the tool would use it.
1441                    let _ = fs::remove_file(destination);
1442                    bail!("corrupt blob {blob}");
1443                }
1444                // Before its mode, which may prevent opening the file.
1445                set_modified(destination, *modified)?;
1446                set_mode(destination, *mode)?;
1447                let metadata = fs::symlink_metadata(destination)?;
1448                Ok((stamp(&metadata), metadata.len()))
1449            };
1450        let workers = files.len().div_ceil(64).clamp(1, RESTORERS);
1451        if workers == 1 {
1452            return files.iter().map(restore).collect();
1453        }
1454        // Sizes vary widely, so each worker takes the next file rather than a
1455        // fixed share, and none starts another after a failure.
1456        let next = AtomicUsize::new(0);
1457        let failed = AtomicBool::new(false);
1458        let mut results: Vec<(usize, Result<RestoredFile>)> = std::thread::scope(|scope| {
1459            let handles: Vec<_> = (0..workers)
1460                .map(|_| {
1461                    scope.spawn(|| {
1462                        let mut done = Vec::new();
1463                        while !failed.load(Ordering::Relaxed) {
1464                            let index = next.fetch_add(1, Ordering::Relaxed);
1465                            let Some(file) = files.get(index) else { break };
1466                            let result = restore(file);
1467                            if result.is_err() {
1468                                failed.store(true, Ordering::Relaxed);
1469                            }
1470                            done.push((index, result));
1471                        }
1472                        done
1473                    })
1474                })
1475                .collect();
1476            handles
1477                .into_iter()
1478                .flat_map(|handle| handle.join().expect("warm restore worker panicked"))
1479                .collect()
1480        });
1481        results.sort_by_key(|(index, _)| *index);
1482        results.into_iter().map(|(_, result)| result).collect()
1483    }
1484
1485    /// The task's warm state from the remote store: the current branch's, else
1486    /// the default branch's. It is kept locally from then on.
1487    fn remote_warm(
1488        &self,
1489        workspace: &Workspace,
1490        task: &Task,
1491        warm: &Warm,
1492        current: &[String],
1493    ) -> Option<(WarmRecord, String)> {
1494        let remote = self
1495            .remote
1496            .as_ref()
1497            .filter(|_| warm.remote && warm.portable != Portable::Own)?;
1498        let own = worktree(workspace);
1499        // Saved in a worktree at another path, it is no save for this one.
1500        let placed =
1501            |saved: Option<&str>| warm.portable != Portable::Path || saved == Some(own.as_str());
1502        for branch in branches(workspace) {
1503            let suits = |saved: &[String], worktree: Option<&str>| {
1504                key_match(warm, current, saved).is_some() && placed(worktree)
1505            };
1506            let name = remote_identity(warm, &own);
1507            match remote.fetch_warm(&self.root, &name, &task.id, &branch, suits) {
1508                Ok(Some((bytes, _pin))) => {
1509                    let Some(record) = WarmRecord::read(&bytes, &warm.identity).ok() else {
1510                        continue;
1511                    };
1512                    // A branch's save that does not suit this checkout gives
1513                    // way to the next branch's.
1514                    if key_match(warm, current, &record.key).is_none()
1515                        || !placed(Some(&record.worktree))
1516                    {
1517                        continue;
1518                    }
1519                    if self.store_warm(&record).is_err() {
1520                        return None;
1521                    }
1522                    return Some((record, branch));
1523                }
1524                Ok(None) => {}
1525                Err(error) => {
1526                    qk_executor::status!(
1527                        "qk: {}: remote warm state unavailable ({error:#})",
1528                        task.id
1529                    );
1530                    // An invalid record or unavailable blob on this branch
1531                    // does not rule out a usable default-branch save.
1532                    continue;
1533                }
1534            }
1535        }
1536        None
1537    }
1538
1539    /// Saves the task's warm groups after a successful run, as this worktree's
1540    /// save.
1541    pub(crate) fn save_warm(
1542        &self,
1543        workspace: &Workspace,
1544        task: &Task,
1545        warm: &Warm,
1546        prepared: &qk_executor::PreparedTask,
1547    ) -> Result<()> {
1548        self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
1549    }
1550
1551    /// Saves the task's warm groups on a thread of its own, which
1552    /// [`Cache::finish`] waits for.
1553    pub(crate) fn save_warm_in_background(
1554        &self,
1555        workspace: &Workspace,
1556        task: &Task,
1557        warm: &Warm,
1558        prepared: &qk_executor::PreparedTask,
1559    ) -> Result<()> {
1560        let save = Save::prepare(workspace, task, warm, prepared)?;
1561        let cache = Cache {
1562            remote: self.remote.clone(),
1563            ..Cache::new(self.root.clone())
1564        };
1565        let handle = std::thread::spawn(move || {
1566            if let Err(error) = cache.save_prepared(&save) {
1567                qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
1568            }
1569        });
1570        self.saves.lock().unwrap().push(handle);
1571        Ok(())
1572    }
1573
1574    /// Saves what [`Save::prepare`] gathered. A file whose metadata matches
1575    /// this worktree's last save or restore keeps its blob unread; one that
1576    /// is gone by the time it is read, as another task may remove it, is left
1577    /// out.
1578    fn save_prepared(&self, save: &Save) -> Result<()> {
1579        let warm = &save.warm;
1580        let _lock = self.lock_warm(warm)?;
1581        let _save = profile::span(&save.task, "warm_save");
1582        let records = self.warm_records(warm);
1583        // This worktree's save under the same key, else under any.
1584        let previous = records
1585            .iter()
1586            .filter(|(_, record)| record.worktree == save.worktree)
1587            .min_by_key(|(_, record)| record.key != save.key)
1588            .map(|(_, record)| record);
1589        let noted: RestoredFiles = fs::read(&save.notes)
1590            .ok()
1591            .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1592            .unwrap_or_default();
1593        let mut record = WarmRecord::new(
1594            warm.identity.clone(),
1595            save.worktree.clone(),
1596            save.commit.clone(),
1597            now_ms(),
1598            save.key.clone(),
1599        );
1600        let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
1601        for location in &save.locations {
1602            if !location.base.exists() {
1603                continue;
1604            }
1605            let known = previous.and_then(|previous| previous.groups.get(location.name));
1606            let restored_files = noted.groups.get(location.name);
1607            let mut group = WarmGroup::default();
1608            let mut pending = Vec::new();
1609            for path in location.paths()? {
1610                paths::safe_parents(&location.base, &path)?;
1611                let absolute = location.base.join(&path);
1612                let metadata = match fs::symlink_metadata(&absolute) {
1613                    Err(error) if gone(&error) => continue,
1614                    result => result?,
1615                };
1616                let artifact = if metadata.file_type().is_symlink() {
1617                    let target = match fs::read_link(&absolute) {
1618                        Err(error) if gone(&error) => continue,
1619                        result => result?,
1620                    };
1621                    let target = target
1622                        .to_str()
1623                        .context("symlink target must be UTF-8")?
1624                        .to_owned();
1625                    validate_link(&path, &target)?;
1626                    Artifact::Symlink {
1627                        target,
1628                        directory: absolute.is_dir(),
1629                    }
1630                } else if metadata.is_file() {
1631                    let now = stamp(&metadata);
1632                    let saved = known.and_then(|known| {
1633                        (known.stamps.get(&path) == Some(&now))
1634                            .then(|| known.artifacts.get(&path))
1635                            .flatten()
1636                    });
1637                    let restored = restored_files
1638                        .and_then(|files| files.get(&path))
1639                        .filter(|(stamp, _)| *stamp == now)
1640                        .map(|(_, artifact)| artifact);
1641                    let reused = saved.or(restored).filter(|artifact| match artifact {
1642                        Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
1643                        _ => false,
1644                    });
1645                    let Some(artifact) = reused.cloned() else {
1646                        pending.push((path, absolute, metadata));
1647                        continue;
1648                    };
1649                    group.stamps.insert(path.clone(), now);
1650                    if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1651                        group.mtimes.insert(path.clone(), modified);
1652                    }
1653                    artifact
1654                } else if metadata.is_dir() {
1655                    Artifact::Directory {
1656                        mode: mode(&metadata),
1657                    }
1658                } else {
1659                    continue;
1660                };
1661                group.artifacts.insert(path, artifact);
1662            }
1663            let storing = profile::span(&save.task, "warm_save_blobs");
1664            let sources: Vec<_> = pending
1665                .iter()
1666                .map(|(_, absolute, _)| absolute.clone())
1667                .collect();
1668            for ((path, _, metadata), result) in pending.into_iter().zip(self.put_blobs(&sources)) {
1669                let blob = match result {
1670                    Ok(blob) => blob,
1671                    Err(error) if error.downcast_ref::<std::io::Error>().is_some_and(gone) => {
1672                        continue;
1673                    }
1674                    Err(error) => return Err(error),
1675                };
1676                group.stamps.insert(path.clone(), stamp(&metadata));
1677                if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1678                    group.mtimes.insert(path.clone(), modified);
1679                }
1680                group.artifacts.insert(
1681                    path,
1682                    Artifact::File {
1683                        blob,
1684                        mode: mode(&metadata),
1685                    },
1686                );
1687            }
1688            drop(storing);
1689            if let Some(limit) = warm.max_size
1690                && group.size(self) > limit
1691            {
1692                qk_executor::status!(
1693                    "qk: {}: warm {} is over its maxSize and was not saved",
1694                    save.task,
1695                    location.name
1696                );
1697                continue;
1698            }
1699            record.groups.insert(location.name.to_owned(), group);
1700        }
1701        self.store_warm(&record)?;
1702        let _ = fs::remove_file(&save.notes);
1703        // The oldest saves beyond the kept number go; the blobs they alone
1704        // cite are evicted with the cache.
1705        let saved = self.warm_record(&record);
1706        for (path, _) in records
1707            .into_iter()
1708            .filter(|(path, _)| *path != saved)
1709            .skip(KEPT - 1)
1710        {
1711            let _ = fs::remove_file(path);
1712        }
1713        // Stamps describe this machine's files; the remote record goes without.
1714        if warm.remote
1715            && warm.portable != Portable::Own
1716            && let Some(remote) = &self.remote
1717            && let Some(branch) = &save.branch
1718        {
1719            for group in record.groups.values_mut() {
1720                group.stamps.clear();
1721            }
1722            remote.upload_warm(
1723                &self.root,
1724                &warm.identity,
1725                &remote_identity(warm, &save.worktree),
1726                branch,
1727                serde_json::to_vec(&record)?,
1728            );
1729        }
1730        Ok(())
1731    }
1732}
1733
1734/// What a save needs from the workspace, gathered where the task ran so that
1735/// the save can run elsewhere.
1736struct Save {
1737    task: String,
1738    warm: Warm,
1739    locations: Vec<Location>,
1740    worktree: String,
1741    commit: Option<String>,
1742    branch: Option<String>,
1743    key: Vec<String>,
1744    notes: PathBuf,
1745}
1746
1747impl Save {
1748    /// Captures the effective key and locations before a background save.
1749    fn prepare(
1750        workspace: &Workspace,
1751        task: &Task,
1752        warm: &Warm,
1753        prepared: &qk_executor::PreparedTask,
1754    ) -> Result<Self> {
1755        Ok(Self {
1756            task: task.id.clone(),
1757            warm: warm.clone(),
1758            locations: locations(workspace, task, warm)?,
1759            worktree: worktree(workspace),
1760            commit: current_commit(workspace),
1761            branch: current_branch(workspace),
1762            key: key_digests(workspace, warm, prepared)?,
1763            notes: restored_files_path(workspace, warm),
1764        })
1765    }
1766}
1767
1768#[cfg(test)]
1769mod tests {
1770    use super::*;
1771    use qk_taskgraph::Request;
1772    use serde_json::json;
1773
1774    /// Two independent owners let a blocked file lock contend with a release.
1775    fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
1776        let temp = tempfile::tempdir().unwrap();
1777        fs::write(temp.path().join("nx.json"), "{}").unwrap();
1778        fs::write(
1779            temp.path().join("project.json"),
1780            json!({
1781                "name": "app", "targets": {
1782                    "one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
1783                    "two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
1784                }
1785            })
1786            .to_string(),
1787        )
1788        .unwrap();
1789        for path in ["scratch-one", "scratch-two"] {
1790            fs::create_dir(temp.path().join(path)).unwrap();
1791            fs::write(temp.path().join(path).join("state"), "saved").unwrap();
1792        }
1793        let workspace = Workspace::load(temp.path()).unwrap();
1794        let graph = TaskGraph::build(
1795            &workspace,
1796            &[
1797                Request::parse("app:one").unwrap(),
1798                Request::parse("app:two").unwrap(),
1799            ],
1800        )
1801        .unwrap();
1802        (temp, workspace, graph)
1803    }
1804
1805    /// The `qk:warm` entries of an `app:build` whose target declares `warm`.
1806    fn entries(warm: Value) -> Result<Vec<Warm>> {
1807        let temp = tempfile::tempdir().unwrap();
1808        fs::write(temp.path().join("nx.json"), "{}").unwrap();
1809        fs::write(
1810            temp.path().join("project.json"),
1811            json!({"name": "app", "targets": {"build": {"command": "true", "qk:warm": warm}}})
1812                .to_string(),
1813        )
1814        .unwrap();
1815        let workspace = Workspace::load(temp.path()).unwrap();
1816        let graph = TaskGraph::build(&workspace, &[Request::parse("app:build").unwrap()]).unwrap();
1817        config(&workspace, &graph.tasks["app:build"])
1818    }
1819
1820    #[test]
1821    fn an_array_holds_the_task_entry_beside_groups() {
1822        let entries = entries(json!([
1823            {"group": "metro", "portable": true, "env": {"METRO_CACHE_ROOT": "{warm}/metro"}},
1824            {"paths": ["build"], "mtimes": "preserve"}
1825        ]))
1826        .unwrap();
1827        assert_eq!(entries.len(), 2);
1828        assert_eq!(entries[0].group(), Some("metro"));
1829        assert!(entries[0].portable == Portable::Anywhere && entries[0].directory);
1830        assert_eq!(entries[1].group(), None);
1831        assert!(entries[1].portable == Portable::Own && entries[1].preserve_mtimes);
1832        assert_eq!(kept_paths(&entries), ["build"]);
1833    }
1834
1835    #[test]
1836    fn portable_is_a_boolean_or_the_same_path() {
1837        let portable = |value| {
1838            entries(json!({"paths": ["build"], "portable": value}))
1839                .map(|entries| entries[0].portable)
1840        };
1841        assert_eq!(portable(json!("path")).unwrap(), Portable::Path);
1842        assert_eq!(portable(json!(false)).unwrap(), Portable::Own);
1843        assert_eq!(
1844            format!("{:#}", portable(json!("anywhere")).unwrap_err()),
1845            "qk:warm.portable must be true, false or \"path\""
1846        );
1847    }
1848
1849    #[test]
1850    fn entries_cannot_share_saves_or_variables() {
1851        let error = |warm| format!("{:#}", entries(warm).unwrap_err());
1852        assert_eq!(
1853            error(json!([{"paths": ["one"]}, {"paths": ["two"]}])),
1854            "only one qk:warm entry may go without a group"
1855        );
1856        assert_eq!(
1857            error(json!([{"group": "metro"}, {"group": "metro"}])),
1858            "qk:warm names group \"metro\" twice"
1859        );
1860        assert_eq!(
1861            error(json!([
1862                {"group": "one", "env": {"CACHE": "{warm}/one"}},
1863                {"group": "two", "env": {"CACHE": "{warm}/two"}}
1864            ])),
1865            "qk:warm sets CACHE in more than one entry"
1866        );
1867        assert_eq!(
1868            error(json!(["build"])),
1869            "qk:warm must be an object or an array of objects"
1870        );
1871    }
1872
1873    /// An external owner lock must not prevent another stash's final release.
1874    #[test]
1875    fn waiting_for_an_owner_does_not_block_other_releases() {
1876        let (_temp, workspace, graph) = fixture();
1877        let one = &graph.tasks["app:one"];
1878        let two = &graph.tasks["app:two"];
1879        let first = acquire(&workspace, one, &config(&workspace, one).unwrap()[0]).unwrap();
1880        let state = paths::worktree_state(&workspace.root).join("warm-kept");
1881        let lock = fs::OpenOptions::new()
1882            .create(true)
1883            .truncate(false)
1884            .read(true)
1885            .write(true)
1886            .open(state.join(format!("{}.lock", hash32(&two.id))))
1887            .unwrap();
1888        fs2::FileExt::lock_exclusive(&lock).unwrap();
1889        let (released, received) = std::sync::mpsc::channel();
1890        let completed = std::thread::scope(|scope| {
1891            scope.spawn(|| {
1892                let second =
1893                    acquire(&workspace, two, &config(&workspace, two).unwrap()[0]).unwrap();
1894                release(&second);
1895            });
1896            std::thread::sleep(Duration::from_millis(100));
1897            scope.spawn(|| {
1898                release(&first);
1899                released.send(()).unwrap();
1900            });
1901            let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
1902            fs2::FileExt::unlock(&lock).unwrap();
1903            completed
1904        });
1905        assert!(
1906            completed,
1907            "a blocked owner lock prevented an unrelated release"
1908        );
1909        assert_eq!(
1910            fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
1911            "saved"
1912        );
1913    }
1914
1915    /// Damaged or absent ownership metadata never makes saved contents disposable.
1916    #[test]
1917    fn unreadable_ownership_preserves_saved_paths() {
1918        let (_temp, workspace, graph) = fixture();
1919        let owner = &graph.tasks["app:one"];
1920        let warm = config(&workspace, owner).unwrap().remove(0);
1921        let root = acquire(&workspace, owner, &warm).unwrap();
1922        let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
1923        drop(entry); // Simulate a stopped process, without returning its paths.
1924        fs::write(Stash::manifest(&root), "{").unwrap();
1925        assert!(acquire(&workspace, owner, &warm).is_err());
1926        assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1927        fs::remove_file(Stash::manifest(&root)).unwrap();
1928        assert!(acquire(&workspace, owner, &warm).is_err());
1929        assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1930        fs::create_dir(Stash::manifest(&root)).unwrap();
1931        assert!(acquire(&workspace, owner, &warm).is_err());
1932        assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1933    }
1934
1935    /// Inspection errors are incomplete recovery, not evidence of an empty slot.
1936    #[cfg(unix)]
1937    #[test]
1938    fn inspection_errors_keep_recovery_incomplete() {
1939        let (_temp, workspace, graph) = fixture();
1940        let owner = &graph.tasks["app:one"];
1941        let warm = config(&workspace, owner).unwrap().remove(0);
1942        let root = acquire(&workspace, owner, &warm).unwrap();
1943        let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
1944        let saved = root.with_extension("saved");
1945        fs::rename(&root, &saved).unwrap();
1946        fs::write(&root, "blocked").unwrap(); // Numbered slot inspection returns ENOTDIR.
1947        assert!(!entry.stash.put_back());
1948        let retained = Stash {
1949            root: root.clone(),
1950            workspace: workspace.root.clone(),
1951            moved: entry.stash.moved.clone(),
1952        };
1953        assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
1954        assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
1955        assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
1956    }
1957}