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