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