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