1use std::collections::{BTreeMap, BTreeSet};
25use std::fs;
26use std::io::Write;
27use std::path::{Path, PathBuf};
28use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
29use std::time::{Duration, SystemTime, UNIX_EPOCH};
30
31use crate::ahead::Taken;
32use anyhow::{Context, Result, bail};
33use qk_config::Workspace;
34use qk_taskgraph::{Task, TaskGraph};
35use serde::{Deserialize, Serialize};
36use serde_json::Value;
37
38use crate::paths::{self, Outputs};
39use crate::record::{self, Artifact, WarmGroup, WarmRecord, validate_link};
40use crate::store::{copy_blob_contents, mode, set_mode, set_modified, symlink};
41use crate::{Cache, hash::digest_file, profile};
42
43pub const RESTORED_AT: std::time::SystemTime = std::time::UNIX_EPOCH;
48
49const GROUP: &str = "group ";
51
52pub const FIELDS: &[&str] = &[
56 "outputs",
57 "paths",
58 "env",
59 "maxSize",
60 "remote",
61 "mtimes",
62 "portable",
63 "key",
64 "restoreKeys",
65 "group",
66 "survive",
67 "save",
68 "restore",
69 "prefetch",
70];
71
72const KEPT: usize = 8;
74const RESTORERS: usize = 8;
76
77type PendingFile<'a> = (&'a String, &'a str, u32, PathBuf, SystemTime);
79type RestoredFile = (Vec<i64>, u64);
81
82#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
85pub enum Portable {
86 #[default]
88 Own,
89 Path,
94 Anywhere,
96}
97
98#[derive(Clone, Debug, Default)]
100pub struct Warm {
101 pub identity: String,
103 pub outputs: bool,
104 pub paths: Vec<String>,
106 pub env: BTreeMap<String, String>,
108 pub directory: bool,
110 pub max_size: Option<u64>,
111 pub remote: bool,
112 pub preserve_mtimes: bool,
115 pub portable: Portable,
117 pub key: Vec<KeyPart>,
119 pub restore_keys: Option<usize>,
122 pub survive: Vec<String>,
125 pub background: bool,
128 pub each_path: bool,
131 pub prefetch: bool,
134}
135
136pub(crate) type Prefetches = crate::ahead::Ahead<(String, Warm, Vec<String>), (WarmRecord, String)>;
139
140const PREFETCHERS: usize = 2;
142
143#[derive(Clone, Debug)]
145pub enum KeyPart {
146 File(String),
148 Env(String),
150}
151
152pub fn config(workspace: &Workspace, task: &Task) -> Result<Vec<Warm>> {
158 const SHAPE: &str = "qk:warm must be an object or an array of objects";
159 let objects = match task.definition.extra.get("qk:warm") {
160 None => return Ok(Vec::new()),
161 Some(Value::Object(object)) => vec![object],
162 Some(Value::Array(entries)) => entries
163 .iter()
164 .map(|entry| entry.as_object().context(SHAPE))
165 .collect::<Result<_>>()?,
166 Some(_) => bail!(SHAPE),
167 };
168 let mut entries: Vec<Warm> = Vec::new();
169 for object in objects {
170 let warm = entry(workspace, task, object)?;
171 if entries.iter().any(|other| other.identity == warm.identity) {
172 bail!(match warm.group() {
173 Some(group) => format!("qk:warm names group {group:?} twice"),
174 None => "only one qk:warm entry may go without a group".to_owned(),
175 });
176 }
177 for name in warm.env.keys() {
178 if entries.iter().any(|other| other.env.contains_key(name)) {
179 bail!("qk:warm sets {name} in more than one entry");
180 }
181 }
182 entries.push(warm);
183 }
184 Ok(entries)
185}
186
187fn entry(
189 workspace: &Workspace,
190 task: &Task,
191 object: &serde_json::Map<String, Value>,
192) -> Result<Warm> {
193 let flag = |name: &str, default: bool| -> Result<bool> {
194 object.get(name).map_or(Ok(default), |value| {
195 value
196 .as_bool()
197 .with_context(|| format!("qk:warm.{name} must be true or false"))
198 })
199 };
200 let group = object
201 .get("group")
202 .map(|value| {
203 value
204 .as_str()
205 .filter(|name| !name.is_empty())
206 .context("qk:warm.group must be a name")
207 })
208 .transpose()?;
209 if group.is_some()
211 && ["outputs", "paths", "survive"]
212 .iter()
213 .any(|key| object.contains_key(*key))
214 {
215 bail!(
216 "a qk:warm.group shares {{warm}} alone; outputs, paths and survive belong to one task"
217 );
218 }
219 let identity = group.map_or_else(|| task.id.clone(), |group| format!("{GROUP}{group}"));
220 const SURVIVE_SHAPE: &str = "qk:warm.survive must be an array of target names";
221 let survive = object
222 .get("survive")
223 .map(|value| {
224 value
225 .as_array()
226 .context(SURVIVE_SHAPE)?
227 .iter()
228 .map(|target| {
229 target
230 .as_str()
231 .filter(|target| !target.is_empty())
232 .map(str::to_owned)
233 .context(SURVIVE_SHAPE)
234 })
235 .collect::<Result<Vec<_>>>()
236 })
237 .transpose()?
238 .unwrap_or_default();
239 let directory = directory(workspace, &identity);
240 let directory_text = directory.to_str().context("warm directory must be UTF-8")?;
241 let workspace_text = workspace
242 .root
243 .to_str()
244 .context("workspace root must be UTF-8")?;
245 let project_root = &workspace.projects[&task.project].root;
246 let project_directory = if project_root == "." {
247 workspace.root.clone()
248 } else {
249 workspace.root.join(project_root)
250 };
251 let project_text = project_directory
252 .to_str()
253 .context("project root must be UTF-8")?;
254 let expand = |text: &str| -> Result<String> {
255 if text.contains("{options.") || text.contains("{args.") {
256 bail!("dynamic warm environment values are not supported yet");
257 }
258 Ok(text
259 .replace("{workspaceRoot}", workspace_text)
260 .replace("{projectRoot}", project_text)
261 .replace("{warm}", directory_text))
262 };
263 let mut paths = Vec::new();
264 for path in object
265 .get("paths")
266 .map(|value| {
267 value
268 .as_array()
269 .context("qk:warm.paths must be an array of paths")
270 })
271 .transpose()?
272 .into_iter()
273 .flatten()
274 {
275 let path = path
276 .as_str()
277 .context("qk:warm.paths must be an array of paths")?;
278 if path.contains("{warm}") {
279 bail!("qk:warm.paths are workspace paths; {{warm}} is kept already");
280 }
281 let (excluded, path) = match path.strip_prefix('!') {
282 Some(path) => (true, path),
283 None => (false, path),
284 };
285 dynamic(path)?;
286 let path = paths::expand(workspace, &task.project, path)?
287 .trim_end_matches('/')
288 .to_owned();
289 paths::validate_path(&path)?;
290 paths.push(if excluded { format!("!{path}") } else { path });
291 }
292 let mut env = BTreeMap::new();
293 let mut uses_directory = false;
294 if let Some(values) = object.get("env") {
295 for (name, value) in values
296 .as_object()
297 .context("qk:warm.env must be an object")?
298 {
299 let value = value
300 .as_str()
301 .context("qk:warm.env values must be strings")?;
302 uses_directory |= value.contains("{warm}");
303 env.insert(name.clone(), expand(value)?);
304 }
305 }
306 let max_size = match object.get("maxSize") {
307 None => None,
308 Some(Value::String(text)) => Some(crate::parse_size(text)?),
309 Some(Value::Number(number)) => number.as_u64(),
310 Some(_) => bail!("qk:warm.maxSize must be a size such as \"2GB\""),
311 };
312 let preserve_mtimes = match object.get("mtimes").map(|value| value.as_str()) {
313 None | Some(Some("epoch")) => false,
314 Some(Some("preserve")) => true,
315 Some(_) => bail!("qk:warm.mtimes must be \"epoch\" or \"preserve\""),
316 };
317 const KEY_SHAPE: &str = "qk:warm.key must be an array of paths and {\"env\": name} objects";
318 let mut key = Vec::new();
319 for part in object
320 .get("key")
321 .map(|value| value.as_array().context(KEY_SHAPE))
322 .transpose()?
323 .into_iter()
324 .flatten()
325 {
326 key.push(match part {
327 Value::String(path) => {
328 dynamic(path)?;
329 let path = paths::expand(workspace, &task.project, path)?;
330 paths::validate_path(&path)?;
331 KeyPart::File(path)
332 }
333 Value::Object(object) if object.len() == 1 => KeyPart::Env(
334 object
335 .get("env")
336 .and_then(Value::as_str)
337 .context(KEY_SHAPE)?
338 .to_owned(),
339 ),
340 _ => bail!(KEY_SHAPE),
341 });
342 }
343 let restore_keys = object
344 .get("restoreKeys")
345 .map(|value| {
346 value
347 .as_u64()
348 .and_then(|count| usize::try_from(count).ok())
349 .filter(|count| *count <= key.len())
350 .context("qk:warm.restoreKeys must be a count no larger than qk:warm.key")
351 })
352 .transpose()?;
353 let each_path = match object.get("restore").map(Value::as_str) {
354 None | Some(Some("whole")) => false,
355 Some(Some("each-path")) if !paths.is_empty() => true,
356 Some(Some("each-path")) => bail!("qk:warm.restore \"each-path\" needs qk:warm.paths"),
357 Some(_) => bail!("qk:warm.restore must be \"whole\" or \"each-path\""),
358 };
359 let background = match object.get("save").map(Value::as_str) {
360 None | Some(Some("wait")) => false,
361 Some(Some("background")) => true,
362 Some(_) => bail!("qk:warm.save must be \"wait\" or \"background\""),
363 };
364 Ok(Warm {
365 identity,
366 outputs: flag("outputs", false)?,
367 paths,
368 env,
369 directory: uses_directory,
370 max_size,
371 remote: flag("remote", true)?,
372 preserve_mtimes,
373 portable: match object.get("portable") {
374 None | Some(Value::Bool(false)) => Portable::Own,
375 Some(Value::Bool(true)) => Portable::Anywhere,
376 Some(Value::String(path)) if path == "path" => Portable::Path,
377 Some(_) => bail!("qk:warm.portable must be true, false or \"path\""),
378 },
379 key,
380 restore_keys,
381 survive,
382 background,
383 each_path,
384 prefetch: flag("prefetch", false)?,
385 })
386}
387
388fn dynamic(path: &str) -> Result<()> {
392 if path.contains("{options.") || path.contains("{args.") {
393 bail!("dynamic cache paths are not supported yet");
394 }
395 Ok(())
396}
397
398fn directory(workspace: &Workspace, identity: &str) -> PathBuf {
399 paths::worktree_state(&workspace.root)
400 .join("warm")
401 .join(&blake3::hash(identity.as_bytes()).to_hex()[..32])
402}
403
404impl Warm {
405 pub fn group(&self) -> Option<&str> {
407 self.identity.strip_prefix(GROUP)
408 }
409
410 pub fn kept_paths(&self) -> Vec<String> {
413 self.paths
414 .iter()
415 .filter(|path| !path.starts_with('!'))
416 .cloned()
417 .collect()
418 }
419}
420
421pub fn kept_paths(entries: &[Warm]) -> Vec<String> {
423 entries.iter().flat_map(Warm::kept_paths).collect()
424}
425
426#[derive(Default)]
429pub struct Kept {
430 roots: Vec<PathBuf>,
432}
433
434impl Drop for Kept {
435 fn drop(&mut self) {
436 for root in self.roots.iter().rev() {
437 release(root);
438 }
439 }
440}
441
442#[derive(Default, Serialize, Deserialize)]
444struct Stash {
445 #[serde(skip)]
446 root: PathBuf,
447 #[serde(skip)]
448 workspace: PathBuf,
449 moved: Vec<String>,
451}
452
453struct Active {
456 users: usize,
457 stash: Stash,
458 _lock: fs::File,
459}
460
461static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
464
465impl Stash {
466 fn manifest(root: &std::path::Path) -> PathBuf {
468 root.join("moved.json")
469 }
470
471 fn save_manifest(&self) -> Result<()> {
474 let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
475 file.write_all(&serde_json::to_vec(self)?)?;
476 file.as_file().sync_all()?;
477 file.persist(Self::manifest(&self.root))?;
478 Ok(())
479 }
480
481 fn put_back(&self) -> bool {
487 let mut complete = true;
488 for (index, path) in self.moved.iter().enumerate().rev() {
489 let kept = self.root.join(index.to_string());
490 match fs::symlink_metadata(&kept) {
491 Ok(_) => {}
492 Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
493 Err(error) => {
494 complete = false;
495 qk_executor::status!(
496 "qk: could not inspect saved {path}, kept for the next run ({error})"
497 );
498 continue;
499 }
500 }
501 let original = self.workspace.join(path);
502 let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
503 remove(&original)?;
504 if let Some(parent) = original.parent() {
505 fs::create_dir_all(parent)?;
506 }
507 fs::rename(&kept, &original)?;
508 Ok(())
509 });
510 if let Err(error) = result {
511 complete = false;
512 qk_executor::status!(
513 "qk: could not put {path} back, kept for the next run ({error:#})"
514 );
515 }
516 }
517 if complete {
518 let _ = fs::remove_dir_all(&self.root);
519 }
520 complete
521 }
522}
523
524fn remove(path: &std::path::Path) -> std::io::Result<()> {
526 match fs::symlink_metadata(path) {
527 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
528 Err(error) => Err(error),
529 Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
530 Ok(_) => fs::remove_file(path),
531 }
532}
533
534fn release(root: &std::path::Path) {
537 let mut active = ACTIVE
538 .lock()
539 .unwrap_or_else(std::sync::PoisonError::into_inner);
540 let Some(entry) = active.get_mut(root) else {
541 return;
542 };
543 entry.users -= 1;
544 if entry.users > 0 {
545 return;
546 }
547 let entry = active.remove(root).expect("present above");
548 drop(active);
551 entry.stash.put_back();
552}
553
554fn names(owner: &Task, entry: &str, task: &Task) -> bool {
556 match entry.split_once(':') {
557 Some((project, target)) => task.project == project && task.target == target,
558 None => task.project == owner.project && task.target == entry,
559 }
560}
561
562fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
564 let mut seen = std::collections::BTreeSet::new();
565 let mut next = vec![from];
566 while let Some(id) = next.pop() {
567 for dependency in graph
568 .tasks
569 .get(id)
570 .into_iter()
571 .flat_map(|task| &task.dependencies)
572 {
573 if dependency == on {
574 return true;
575 }
576 if seen.insert(dependency.as_str()) {
577 next.push(dependency);
578 }
579 }
580 }
581 false
582}
583
584pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
587 let mut kept = Kept::default();
588 for owner in graph.tasks.values() {
589 let Ok(entries) = config(workspace, owner) else {
590 continue;
591 };
592 let Some(warm) = entries.into_iter().find(|warm| warm.group().is_none()) else {
594 continue;
595 };
596 if !warm.survive.iter().any(|entry| names(owner, entry, task))
597 || !depends_on(graph, &owner.id, &task.id)
598 {
599 continue;
600 }
601 match acquire(workspace, owner, &warm) {
602 Ok(root) => kept.roots.push(root),
603 Err(error) => qk_executor::status!(
604 "qk: {}: warm paths not kept across {} ({error:#})",
605 owner.id,
606 task.id
607 ),
608 }
609 }
610 kept
611}
612
613fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
616 let state = paths::worktree_state(&workspace.root).join("warm-kept");
617 let root = state.join(hash32(&owner.id));
618 fs::create_dir_all(&state)?;
619 let lock = fs::OpenOptions::new()
620 .create(true)
621 .truncate(false)
622 .read(true)
623 .write(true)
624 .open(state.join(format!("{}.lock", hash32(&owner.id))))?;
625 let mut active = loop {
628 let mut active = ACTIVE
629 .lock()
630 .unwrap_or_else(std::sync::PoisonError::into_inner);
631 if let Some(entry) = active.get_mut(&root) {
632 entry.users += 1;
633 return Ok(root);
634 }
635 match fs2::FileExt::try_lock_exclusive(&lock) {
636 Ok(()) => break active,
637 Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
638 drop(active);
639 std::thread::sleep(Duration::from_millis(20));
640 }
641 Err(error) => return Err(error.into()),
642 }
643 };
644 let left = match fs::read(Stash::manifest(&root)) {
651 Ok(bytes) => {
652 let mut left: Stash = serde_json::from_slice(&bytes)
653 .context("invalid warm stash manifest; saved paths retained for recovery")?;
654 left.root = root.clone();
655 left.workspace = workspace.root.clone();
656 Some(left)
657 }
658 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
659 if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
660 bail!("missing warm stash manifest; saved paths retained for recovery");
661 }
662 None
663 }
664 Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
665 };
666 let retained = left.filter(|left| !left.put_back());
667 let stash = stash(workspace, &root, warm, retained)?;
668 active.insert(
669 root.clone(),
670 Active {
671 users: 1,
672 stash,
673 _lock: lock,
674 },
675 );
676 Ok(root)
677}
678
679fn stash(
683 workspace: &Workspace,
684 root: &std::path::Path,
685 warm: &Warm,
686 retained: Option<Stash>,
687) -> Result<Stash> {
688 let mut stash = match retained {
689 Some(retained) => retained,
690 None => {
691 let _ = fs::remove_dir_all(root);
692 Stash {
693 root: root.to_owned(),
694 workspace: workspace.root.clone(),
695 moved: Vec::new(),
696 }
697 }
698 };
699 let held: Vec<String> = stash
700 .moved
701 .iter()
702 .enumerate()
703 .filter_map(
704 |(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
705 Ok(_) => Some(Ok(path.clone())),
706 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
707 Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
708 },
709 )
710 .collect::<Result<_>>()?;
711 let overlaps = |path: &str| {
712 held.iter().any(|held| {
713 held == path
714 || path.starts_with(&format!("{held}/"))
715 || held.starts_with(&format!("{path}/"))
716 })
717 };
718 let mut kept = std::collections::BTreeSet::new();
721 for path in warm.kept_paths() {
722 if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
723 .and_then(|outputs| outputs.paths(&workspace.root))
724 {
725 kept.extend(found);
726 }
727 }
728 let topmost: Vec<String> = kept
730 .iter()
731 .filter(|path| {
732 !std::path::Path::new(path.as_str())
733 .ancestors()
734 .skip(1)
735 .any(|ancestor| {
736 ancestor
737 .to_str()
738 .is_some_and(|ancestor| kept.contains(ancestor))
739 })
740 })
741 .filter(|path| !overlaps(path))
742 .cloned()
743 .collect();
744 if topmost.is_empty() {
745 return Ok(stash);
746 }
747 fs::create_dir_all(root)?;
748 for path in topmost {
749 let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
750 stash.moved.push(path.clone());
751 stash.save_manifest()?;
752 fs::rename(
755 workspace.root.join(&path),
756 root.join((stash.moved.len() - 1).to_string()),
757 )
758 .inspect_err(|_| {
759 stash.moved.pop();
760 })?;
761 Ok(())
762 });
763 if let Err(error) = moved {
764 stash.save_manifest()?;
767 stash.put_back();
768 return Err(error);
769 }
770 }
771 Ok(stash)
772}
773
774pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
777 let mut claimed: Vec<(String, &str)> = Vec::new();
778 for (id, task) in &graph.tasks {
779 for path in kept_paths(&config(workspace, task)?) {
780 if let Some((_, other)) = claimed.iter().find(|(known, _)| {
781 known == &path
782 || known.starts_with(&format!("{path}/"))
783 || path.starts_with(&format!("{known}/"))
784 }) {
785 bail!("{id} and {other} both keep {path} as warm state");
786 }
787 claimed.push((path, id));
788 }
789 }
790 Ok(())
791}
792
793impl WarmGroup {
794 fn size(&self, cache: &Cache) -> u64 {
796 self.artifacts
797 .values()
798 .filter_map(|artifact| match artifact {
799 Artifact::File { blob, .. } => {
800 fs::metadata(cache.root.join("blobs").join(blob)).ok()
801 }
802 _ => None,
803 })
804 .map(|metadata| metadata.len())
805 .sum()
806 }
807}
808
809#[derive(Default, Serialize, Deserialize)]
812struct RestoredFiles {
813 groups: BTreeMap<String, GroupFiles>,
814}
815
816type GroupFiles = BTreeMap<String, (Vec<i64>, Artifact)>;
818
819#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
821pub struct Restored {
822 #[serde(default, skip_serializing_if = "Option::is_none")]
824 pub group: Option<String>,
825 pub source: String,
828 pub groups: Vec<String>,
829 pub files: usize,
830 pub bytes: u64,
831}
832
833#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
835#[serde(rename_all = "camelCase")]
836pub struct WarmReport {
837 #[serde(default, skip_serializing_if = "Vec::is_empty")]
839 pub restored: Vec<Restored>,
840 #[serde(default, skip_serializing_if = "Vec::is_empty")]
844 pub present: Vec<String>,
845 pub save_ms: Option<u64>,
848 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
850 pub background: bool,
851}
852
853#[derive(Clone)]
856struct Location {
857 name: &'static str,
858 base: PathBuf,
859 outputs: Option<Outputs>,
861 each: Vec<Outputs>,
864}
865
866impl Location {
867 fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
869 if !self.base.exists() {
870 return Ok(Default::default());
871 }
872 match &self.outputs {
873 Some(outputs) => outputs.paths(&self.base),
874 None => {
875 let mut paths = std::collections::BTreeSet::new();
876 for entry in walkdir::WalkDir::new(&self.base)
877 .min_depth(1)
878 .follow_links(false)
879 {
880 let entry = entry?;
881 paths.insert(paths::relative(&self.base, entry.path())?);
882 }
883 Ok(paths)
884 }
885 }
886 }
887
888 fn matches(&self, path: &str) -> bool {
890 self.outputs
891 .as_ref()
892 .is_none_or(|outputs| outputs.matches(path))
893 }
894
895 fn present(&self) -> Result<Option<Vec<&Outputs>>> {
898 let mut present = Vec::new();
899 for outputs in &self.each {
900 if !outputs.paths(&self.base)?.is_empty() {
901 present.push(outputs);
902 }
903 }
904 Ok((!present.is_empty() && present.len() < self.each.len()).then_some(present))
905 }
906}
907
908fn remove_restored(location: &Location, kept: &[&Outputs]) {
912 let Ok(paths) = location.paths() else {
913 return;
914 };
915 for path in paths.iter().rev() {
916 if kept.iter().any(|outputs| outputs.matches(path)) {
917 continue;
918 }
919 let absolute = location.base.join(path);
920 let _ = match fs::symlink_metadata(&absolute) {
921 Ok(metadata) if metadata.is_dir() => fs::remove_dir(&absolute),
922 Ok(_) => fs::remove_file(&absolute),
923 Err(_) => continue,
924 };
925 }
926}
927
928fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
930 let mut locations = Vec::new();
931 if warm.outputs {
932 locations.push(Location {
933 name: "outputs",
934 base: workspace.root.clone(),
935 outputs: Some(Outputs::new(workspace, task)?),
936 each: Vec::new(),
937 });
938 }
939 if !warm.paths.is_empty() {
940 let mut each = Vec::new();
941 if warm.each_path {
942 let (excluded, included): (Vec<_>, Vec<_>) =
943 warm.paths.iter().partition(|path| path.starts_with('!'));
944 for path in included {
945 let mut patterns = vec![path.clone()];
946 patterns.extend(excluded.iter().map(|path| (*path).clone()));
947 each.push(Outputs::from_paths(&patterns)?);
948 }
949 }
950 locations.push(Location {
951 name: "paths",
952 base: workspace.root.clone(),
953 outputs: Some(Outputs::from_paths(&warm.paths)?),
954 each,
955 });
956 }
957 if warm.directory {
958 locations.push(Location {
959 name: "directory",
960 base: directory(workspace, &warm.identity),
961 outputs: None,
962 each: Vec::new(),
963 });
964 }
965 Ok(locations)
966}
967
968pub(crate) fn current_branch(workspace: &Workspace) -> Option<String> {
971 for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
972 if let Ok(value) = std::env::var(name)
973 && !value.is_empty()
974 {
975 return Some(value);
976 }
977 }
978 let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
979 let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
980 (output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
981}
982
983pub(crate) fn branches(workspace: &Workspace) -> Vec<String> {
986 let default = workspace
987 .config
988 .default_base
989 .clone()
990 .unwrap_or_else(|| "main".into());
991 let default = default
992 .strip_prefix("origin/")
993 .unwrap_or(&default)
994 .to_owned();
995 let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
996 if !branches.contains(&default) {
997 branches.push(default);
998 }
999 branches
1000}
1001
1002fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
1004 let modified = metadata
1005 .modified()
1006 .ok()
1007 .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
1008 .map_or(-1, |duration| duration.as_nanos() as i64);
1009 #[allow(unused_mut)]
1010 let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
1011 #[cfg(unix)]
1012 {
1013 use std::os::unix::fs::MetadataExt;
1014 stamp.extend([
1015 metadata.ctime(),
1016 metadata.ctime_nsec(),
1017 metadata.ino() as i64,
1018 ]);
1019 }
1020 stamp
1021}
1022
1023fn remote_identity(warm: &Warm, worktree: &str) -> String {
1028 match warm.portable {
1029 Portable::Path => format!("{}\0path\0{worktree}", warm.identity),
1030 _ => warm.identity.clone(),
1031 }
1032}
1033
1034fn worktree(workspace: &Workspace) -> String {
1035 fs::canonicalize(&workspace.root)
1036 .unwrap_or_else(|_| workspace.root.clone())
1037 .to_string_lossy()
1038 .into_owned()
1039}
1040
1041fn current_commit(workspace: &Workspace) -> Option<String> {
1043 let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
1044 let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
1045 (output.status.success() && !commit.is_empty()).then_some(commit)
1046}
1047
1048fn key_digests(
1051 workspace: &Workspace,
1052 warm: &Warm,
1053 prepared: &qk_executor::PreparedTask,
1054) -> Result<Vec<String>> {
1055 warm.key
1056 .iter()
1057 .map(|part| {
1058 Ok(match part {
1059 KeyPart::File(path) => {
1060 paths::safe_parents(&workspace.root, path)?;
1061 let absolute = workspace.root.join(path);
1062 if absolute.is_file() {
1063 format!("file:{}", digest_file(&absolute)?)
1064 } else {
1065 "file:absent".to_owned()
1066 }
1067 }
1068 KeyPart::Env(name) => match prepared
1071 .execution
1072 .get(std::ffi::OsStr::new(name))
1073 .or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
1074 {
1075 Some(value) => format!(
1076 "env:{}",
1077 blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
1078 ),
1079 None => "env:unset".to_owned(),
1080 },
1081 })
1082 })
1083 .collect()
1084}
1085
1086fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
1090 if saved == current {
1091 return Some(0);
1092 }
1093 let count = warm.restore_keys?;
1094 (saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
1095}
1096
1097fn descends_from(workspace: &Workspace, commit: &str) -> bool {
1099 paths::git_history(
1100 &workspace.root,
1101 &["merge-base", "--is-ancestor", commit, "HEAD"],
1102 )
1103 .is_ok_and(|output| output.status.success())
1104}
1105
1106fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
1109 let ancestor = paths::git_history(
1110 &workspace.root,
1111 &["merge-base", "--is-ancestor", commit, "HEAD"],
1112 )
1113 .ok()?;
1114 if !ancestor.status.success() {
1115 return None;
1116 }
1117 let output = paths::git_history(
1118 &workspace.root,
1119 &["rev-list", "--count", &format!("{commit}..HEAD")],
1120 )
1121 .ok()?;
1122 String::from_utf8(output.stdout).ok()?.trim().parse().ok()
1123}
1124
1125fn depth(graph: &TaskGraph, id: &str, known: &mut BTreeMap<String, usize>) -> usize {
1127 if let Some(depth) = known.get(id) {
1128 return *depth;
1129 }
1130 let depth = graph.tasks.get(id).map_or(0, |task| {
1131 task.dependencies
1132 .iter()
1133 .map(|dependency| depth(graph, dependency, known) + 1)
1134 .max()
1135 .unwrap_or(0)
1136 });
1137 known.insert(id.to_owned(), depth);
1138 depth
1139}
1140
1141fn warm_prefix(identity: &str) -> String {
1143 hash32(&format!(
1144 "{identity}\0{}\0{}",
1145 std::env::consts::OS,
1146 std::env::consts::ARCH
1147 ))
1148}
1149
1150fn warm_record(root: &Path, record: &WarmRecord) -> PathBuf {
1151 root.join("warm").join(format!(
1152 "{}-{}.json",
1153 warm_prefix(&record.task),
1154 hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
1155 ))
1156}
1157
1158fn store_warm(root: &Path, record: &WarmRecord) -> Result<()> {
1159 let path = warm_record(root, record);
1160 let directory = path.parent().context("warm records have a directory")?;
1161 fs::create_dir_all(directory)?;
1162 record::publish(root, &path, record)
1163}
1164
1165fn remote_warm(
1168 remote: &crate::remote::Remote,
1169 root: &Path,
1170 task: &str,
1171 warm: &Warm,
1172 own: &str,
1173 branches: &[String],
1174 current: &[String],
1175) -> Option<(WarmRecord, String)> {
1176 if !warm.remote || warm.portable == Portable::Own {
1177 return None;
1178 }
1179 let placed = |saved: Option<&str>| warm.portable != Portable::Path || saved == Some(own);
1181 for branch in branches {
1182 let suits = |saved: &[String], worktree: Option<&str>| {
1183 key_match(warm, current, saved).is_some() && placed(worktree)
1184 };
1185 let name = remote_identity(warm, own);
1186 match remote.fetch_warm(root, &name, task, branch, suits) {
1187 Ok(Some((bytes, _pin))) => {
1188 let Some(record) = WarmRecord::read(&bytes, &warm.identity).ok() else {
1189 continue;
1190 };
1191 if key_match(warm, current, &record.key).is_none()
1194 || !placed(Some(&record.worktree))
1195 {
1196 continue;
1197 }
1198 if store_warm(root, &record).is_err() {
1199 return None;
1200 }
1201 return Some((record, branch.clone()));
1202 }
1203 Ok(None) => {}
1204 Err(error) => {
1205 qk_executor::status!("qk: {task}: remote warm state unavailable ({error:#})");
1206 continue;
1209 }
1210 }
1211 }
1212 None
1213}
1214
1215fn hash32(text: &str) -> String {
1217 blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
1218}
1219
1220fn nanos(time: SystemTime) -> Option<i64> {
1222 time.duration_since(UNIX_EPOCH)
1223 .ok()
1224 .and_then(|duration| i64::try_from(duration.as_nanos()).ok())
1225}
1226
1227fn now_ms() -> u64 {
1229 SystemTime::now()
1230 .duration_since(UNIX_EPOCH)
1231 .map_or(0, |duration| duration.as_millis() as u64)
1232}
1233
1234fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
1236 paths::worktree_state(&workspace.root)
1237 .join("warm-restored")
1238 .join(format!("{}.json", hash32(&warm.identity)))
1239}
1240
1241fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
1243 fs::read(restored_files_path(workspace, warm))
1244 .ok()
1245 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1246 .unwrap_or_default()
1247}
1248
1249impl Cache {
1250 fn warm_prefix(&self, identity: &str) -> String {
1252 warm_prefix(identity)
1253 }
1254
1255 fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
1258 let directory = self.root.join("locks");
1259 fs::create_dir_all(&directory)?;
1260 let file = fs::OpenOptions::new()
1261 .create(true)
1262 .truncate(false)
1263 .read(true)
1264 .write(true)
1265 .open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
1266 fs2::FileExt::lock_exclusive(&file)?;
1267 Ok(file)
1268 }
1269
1270 fn warm_record(&self, record: &WarmRecord) -> PathBuf {
1271 warm_record(&self.root, record)
1272 }
1273
1274 fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, WarmRecord)> {
1276 let prefix = format!("{}-", self.warm_prefix(&warm.identity));
1277 let mut records: Vec<(PathBuf, WarmRecord)> = fs::read_dir(self.root.join("warm"))
1278 .into_iter()
1279 .flatten()
1280 .flatten()
1281 .filter(|entry| {
1282 entry
1283 .file_name()
1284 .to_str()
1285 .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
1286 })
1287 .filter_map(|entry| {
1288 let record =
1289 WarmRecord::read(&fs::read(entry.path()).ok()?, &warm.identity).ok()?;
1290 Some((entry.path(), record))
1291 })
1292 .collect();
1293 records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
1294 records
1295 }
1296
1297 fn choose_warm(
1304 &self,
1305 workspace: &Workspace,
1306 warm: &Warm,
1307 current: &[String],
1308 related: bool,
1309 ) -> Option<(WarmRecord, String)> {
1310 let own = worktree(workspace);
1311 let mut candidates: Vec<(u8, bool, WarmRecord)> = self
1312 .warm_records(warm)
1313 .into_iter()
1314 .filter_map(|(_, record)| {
1315 let quality = key_match(warm, current, &record.key)?;
1316 let other = record.worktree != own;
1317 (!other || warm.portable == Portable::Anywhere).then_some((quality, other, record))
1318 })
1319 .filter(|(_, _, record)| {
1320 let descends = record
1321 .commit
1322 .as_deref()
1323 .is_some_and(|commit| descends_from(workspace, commit));
1324 descends == related
1325 })
1326 .collect();
1327 candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
1329 let rank = candidates
1330 .first()
1331 .map(|(quality, other, _)| (*quality, *other))?;
1332 let tied = candidates
1333 .iter()
1334 .take_while(|(quality, other, _)| (*quality, *other) == rank)
1335 .count();
1336 let index = if tied > 1 {
1337 (0..tied)
1338 .min_by_key(|&index| {
1339 let behind = candidates[index]
1340 .2
1341 .commit
1342 .as_deref()
1343 .and_then(|commit| behind(workspace, commit));
1344 (behind.is_none(), behind, index)
1345 })
1346 .unwrap_or(0)
1347 } else {
1348 0
1349 };
1350 let (_, other, record) = candidates.swap_remove(index);
1351 let source = if other {
1352 format!("worktree {}", record.worktree)
1353 } else {
1354 "local".to_owned()
1355 };
1356 Some((record, source))
1357 }
1358
1359 fn store_warm(&self, record: &WarmRecord) -> Result<()> {
1360 store_warm(&self.root, record)
1361 }
1362
1363 pub(crate) fn clear_warm(&self, workspace: &Workspace, task: &Task, warm: &Warm) -> Result<()> {
1368 let _lock = self.lock_warm(warm)?;
1369 let identity = self.warm_prefix(&warm.identity);
1370 if self.cleared.lock().unwrap().contains(&identity) {
1371 return Ok(());
1372 }
1373 for location in locations(workspace, task, warm)? {
1374 for path in location.paths()?.iter().rev() {
1375 let absolute = location.base.join(path);
1376 let removed = match fs::symlink_metadata(&absolute) {
1377 Ok(metadata) if metadata.is_dir() => match fs::remove_dir(&absolute) {
1378 Err(error) if error.kind() == std::io::ErrorKind::DirectoryNotEmpty => {
1379 Ok(())
1380 }
1381 removed => removed,
1382 },
1383 Ok(_) => fs::remove_file(&absolute),
1384 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
1385 Err(error) => Err(error),
1386 };
1387 removed.with_context(|| format!("removing {}", absolute.display()))?;
1388 }
1389 }
1390 let notes = restored_files_path(workspace, warm);
1391 match fs::remove_file(¬es) {
1392 Err(error) if error.kind() != std::io::ErrorKind::NotFound => return Err(error.into()),
1393 _ => {}
1394 }
1395 self.cleared.lock().unwrap().insert(identity);
1396 Ok(())
1397 }
1398
1399 pub(crate) fn restore_warm(
1405 &self,
1406 workspace: &Workspace,
1407 task: &Task,
1408 warm: &Warm,
1409 prepared: &qk_executor::PreparedTask,
1410 ) -> Result<(Restored, Vec<String>)> {
1411 let mut restored = Restored {
1412 group: warm.group().map(str::to_owned),
1413 ..Restored::default()
1414 };
1415 let locking = profile::span(&task.id, "warm_restore_lock");
1416 let _lock = self.lock_warm(warm)?;
1417 drop(locking);
1418 let mut present = Vec::new();
1420 let mut missing = Vec::new();
1421 let looking = profile::span(&task.id, "warm_restore_present");
1422 let locations = locations(workspace, task, warm)?;
1423 for location in &locations {
1424 if location.paths()?.is_empty() {
1425 missing.push((location, Vec::new()));
1426 } else if let Some(kept) = location.present()? {
1427 missing.push((location, kept));
1430 } else {
1431 present.push(warm.group().map_or(location.name, |group| group).to_owned());
1432 }
1433 }
1434 drop(looking);
1435 if missing.is_empty() {
1436 return Ok((restored, present));
1437 }
1438 let choosing = profile::span(&task.id, "warm_restore_choose");
1439 let keying = profile::span(&task.id, "warm_restore_key");
1440 let current = key_digests(workspace, warm, prepared)?;
1441 drop(keying);
1442 let ahead = self.prefetches.get().map_or(Taken::Nothing, |prefetches| {
1443 let mut waiting = profile::span(&task.id, "warm_restore_prefetched");
1444 let taken = prefetches.take(&warm.identity, ¤t);
1445 waiting.note(
1446 "taken",
1447 match taken {
1448 Taken::Found(_) => "found",
1449 Taken::Missing => "missing",
1450 Taken::Nothing => "nothing",
1451 },
1452 );
1453 taken
1454 });
1455 let local = profile::span(&task.id, "warm_restore_local");
1456 let here = self.choose_warm(workspace, warm, ¤t, true);
1457 drop(local);
1458 let chosen = match (ahead, here) {
1459 (Taken::Found((record, _)), Some((mine, source))) if mine.saved > record.saved => {
1463 Some((mine, source))
1464 }
1465 (Taken::Found((record, branch)), _) => Some((record, format!("remote {branch}"))),
1466 (_, Some(here)) => Some(here),
1467 (Taken::Missing, None) => None,
1469 (Taken::Nothing, None) => {
1470 let mut fetching = profile::span(&task.id, "warm_restore_remote");
1471 let fetched = self
1472 .remote_warm(workspace, task, warm, ¤t)
1473 .map(|(record, branch)| (record, format!("remote {branch}")));
1474 fetching.note(
1475 "source",
1476 fetched
1477 .as_ref()
1478 .map_or("none", |(_, source)| source.as_str()),
1479 );
1480 fetched
1481 }
1482 }
1483 .or_else(|| self.choose_warm(workspace, warm, ¤t, false));
1485 let Some((record, source)) = chosen else {
1486 return Ok((restored, present));
1487 };
1488 drop(choosing);
1489 let _files = profile::span(&task.id, "warm_restore_files");
1490 let keep_mtimes = warm.preserve_mtimes && source == "local";
1491 restored.source = source;
1492 let mut noted = read_restored_files(workspace, warm);
1493 for (location, kept) in missing {
1494 let Some(group) = record.groups.get(location.name) else {
1495 continue;
1496 };
1497 let (files, bytes) =
1498 match self.restore_location(&task.id, location, &kept, group, keep_mtimes) {
1499 Ok(restored) => restored,
1500 Err(error) => {
1501 remove_restored(location, &kept);
1505 return Err(error);
1506 }
1507 };
1508 restored.files += files.len();
1509 restored.bytes += bytes;
1510 if kept.is_empty() {
1511 noted.groups.insert(location.name.to_owned(), files);
1512 } else {
1513 noted
1515 .groups
1516 .entry(location.name.to_owned())
1517 .or_default()
1518 .extend(files);
1519 }
1520 restored.groups.push(location.name.to_owned());
1521 }
1522 if !restored.groups.is_empty() {
1523 let path = restored_files_path(workspace, warm);
1524 fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
1525 fs::write(path, serde_json::to_vec(¬ed)?)?;
1526 }
1527 Ok((restored, present))
1528 }
1529
1530 fn restore_location(
1533 &self,
1534 task: &str,
1535 location: &Location,
1536 kept: &[&Outputs],
1537 group: &WarmGroup,
1538 keep_mtimes: bool,
1539 ) -> Result<(GroupFiles, u64)> {
1540 fs::create_dir_all(&location.base)?;
1541 let mut laying = profile::span(task, "warm_restore_layout");
1542 let mut created = BTreeSet::new();
1545 let mut directories = Vec::new();
1546 let mut pending = Vec::new();
1547 for (path, artifact) in &group.artifacts {
1548 paths::safe_parents(&location.base, path)?;
1549 if !location.matches(path) {
1550 bail!("warm state for {task} holds {path}, outside its paths");
1551 }
1552 if kept.iter().any(|outputs| outputs.matches(path)) {
1553 continue;
1554 }
1555 let destination = location.base.join(path);
1556 if let Some(parent) = destination.parent()
1557 && created.insert(parent.to_owned())
1558 {
1559 fs::create_dir_all(parent)?;
1560 }
1561 match artifact {
1562 Artifact::File { blob, mode } => {
1563 let modified = keep_mtimes
1564 .then(|| group.mtimes.get(path))
1565 .flatten()
1566 .and_then(|nanos| u64::try_from(*nanos).ok())
1567 .map_or(RESTORED_AT, |nanos| {
1568 UNIX_EPOCH + Duration::from_nanos(nanos)
1569 });
1570 pending.push((path, blob.as_str(), *mode, destination, modified));
1571 }
1572 Artifact::Directory { mode } => {
1573 fs::create_dir_all(&destination)?;
1574 created.insert(destination.clone());
1575 directories.push((destination, *mode));
1576 }
1577 Artifact::Symlink { target, directory } => {
1578 validate_link(path, target)?;
1579 symlink(target, &destination, *directory)?;
1580 }
1581 }
1582 }
1583 laying.count("directories", created.len() as u64);
1584 drop(laying);
1585 let mut copying = profile::span(task, "warm_restore_copy");
1586 copying.count("files", pending.len() as u64);
1587 let copied = self
1588 .restore_warm_files(&pending)
1589 .with_context(|| format!("warm state for {task} not restored"))?;
1590 copying.count("bytes", copied.iter().map(|(_, size)| size).sum());
1591 drop(copying);
1592 let _modes = profile::span(task, "warm_restore_modes");
1594 for (destination, mode) in directories {
1595 set_mode(&destination, mode)?;
1596 }
1597 let mut files = BTreeMap::new();
1598 let mut bytes = 0;
1599 for ((path, blob, mode, _, _), (metadata, size)) in pending.into_iter().zip(copied) {
1600 bytes += size;
1601 let artifact = Artifact::File {
1602 blob: blob.to_owned(),
1603 mode,
1604 };
1605 files.insert(path.clone(), (metadata, artifact));
1606 }
1607 Ok((files, bytes))
1608 }
1609
1610 fn restore_warm_files(&self, files: &[PendingFile]) -> Result<Vec<RestoredFile>> {
1613 let restore =
1614 |(_, blob, mode, destination, modified): &PendingFile| -> Result<RestoredFile> {
1615 let source = self.root.join("blobs").join(blob);
1616 if copy_blob_contents(&source, destination)? != *blob {
1617 let _ = fs::remove_file(destination);
1619 bail!("corrupt blob {blob}");
1620 }
1621 set_modified(destination, *modified)?;
1623 set_mode(destination, *mode)?;
1624 let metadata = fs::symlink_metadata(destination)?;
1625 Ok((stamp(&metadata), metadata.len()))
1626 };
1627 let workers = files.len().div_ceil(64).clamp(1, RESTORERS);
1628 if workers == 1 {
1629 return files.iter().map(restore).collect();
1630 }
1631 let next = AtomicUsize::new(0);
1634 let failed = AtomicBool::new(false);
1635 let mut results: Vec<(usize, Result<RestoredFile>)> = std::thread::scope(|scope| {
1636 let handles: Vec<_> = (0..workers)
1637 .map(|_| {
1638 scope.spawn(|| {
1639 let mut done = Vec::new();
1640 while !failed.load(Ordering::Relaxed) {
1641 let index = next.fetch_add(1, Ordering::Relaxed);
1642 let Some(file) = files.get(index) else { break };
1643 let result = restore(file);
1644 if result.is_err() {
1645 failed.store(true, Ordering::Relaxed);
1646 }
1647 done.push((index, result));
1648 }
1649 done
1650 })
1651 })
1652 .collect();
1653 handles
1654 .into_iter()
1655 .flat_map(|handle| handle.join().expect("warm restore worker panicked"))
1656 .collect()
1657 });
1658 results.sort_by_key(|(index, _)| *index);
1659 results.into_iter().map(|(_, result)| result).collect()
1660 }
1661
1662 fn remote_warm(
1665 &self,
1666 workspace: &Workspace,
1667 task: &Task,
1668 warm: &Warm,
1669 current: &[String],
1670 ) -> Option<(WarmRecord, String)> {
1671 let remote = self.remote.as_ref()?;
1672 let own = worktree(workspace);
1673 remote_warm(
1674 remote,
1675 &self.root,
1676 &task.id,
1677 warm,
1678 &own,
1679 &branches(workspace),
1680 current,
1681 )
1682 }
1683
1684 pub fn prefetch_warm(
1693 &self,
1694 workspace: &Workspace,
1695 graph: &TaskGraph,
1696 prepared: &BTreeMap<String, qk_executor::PreparedTask>,
1697 restores: impl Fn(&str) -> bool,
1698 ) {
1699 let Some(remote) = self.remote.as_ref().filter(|_| self.reuse && !self.cold) else {
1701 return;
1702 };
1703 if self.initialize().is_err() {
1704 return;
1705 }
1706 let own = worktree(workspace);
1707 let mut items = Vec::new();
1708 for (id, task) in &graph.tasks {
1709 let Some(prepared) = prepared.get(id).filter(|_| restores(id)) else {
1710 continue;
1711 };
1712 let Ok(entries) = config(workspace, task) else {
1713 continue;
1714 };
1715 if !entries.iter().any(|warm| warm.prefetch) {
1716 continue;
1717 }
1718 let mut running = prepared.clone();
1722 for warm in &entries {
1723 for (name, value) in &warm.env {
1724 running.execution.insert(name.into(), value.into());
1725 }
1726 }
1727 for warm in entries {
1728 if !warm.prefetch || !warm.remote || warm.portable == Portable::Own {
1729 continue;
1730 }
1731 let missing = locations(workspace, task, &warm).is_ok_and(|locations| {
1732 locations.iter().any(|location| {
1733 location.paths().map_or(true, |paths| paths.is_empty())
1734 || location.present().is_ok_and(|kept| kept.is_some())
1735 })
1736 });
1737 let Ok(current) = key_digests(workspace, &warm, &running) else {
1738 continue;
1739 };
1740 if !missing || self.choose_warm(workspace, &warm, ¤t, true).is_some() {
1741 continue;
1742 }
1743 let depth = depth(graph, id, &mut BTreeMap::new());
1744 items.push((depth, warm.identity.clone(), (id.clone(), warm, current)));
1745 }
1746 }
1747 if items.is_empty() {
1748 return;
1749 }
1750 items.sort_by(|a, b| (a.0, &(a.2).0).cmp(&(b.0, &(b.2).0)));
1751 let items = items
1752 .into_iter()
1753 .map(|(_, identity, job)| {
1754 let key = job.2.clone();
1755 (identity, key, job)
1756 })
1757 .collect();
1758 let (remote, root) = (remote.clone(), self.root.clone());
1759 let branches = branches(workspace);
1760 let prefetches = Prefetches::start(items, PREFETCHERS, move |(id, warm, current)| {
1761 let mut fetching = profile::span(&id, "warm_prefetch");
1762 let fetched = remote_warm(&remote, &root, &id, &warm, &own, &branches, ¤t);
1763 fetching.note(
1764 "branch",
1765 fetched
1766 .as_ref()
1767 .map_or("none", |(_, branch)| branch.as_str()),
1768 );
1769 fetched
1770 });
1771 let _ = self.prefetches.set(prefetches);
1772 }
1773
1774 pub fn abandon_prefetches(&self) {
1777 if let Some(prefetches) = self.prefetches.get() {
1778 prefetches.stop(false);
1779 }
1780 }
1781
1782 pub(crate) fn finish_prefetches(&self) {
1785 if let Some(prefetches) = self.prefetches.get() {
1786 prefetches.stop(true);
1787 }
1788 }
1789
1790 pub(crate) fn save_warm(
1793 &self,
1794 workspace: &Workspace,
1795 task: &Task,
1796 warm: &Warm,
1797 prepared: &qk_executor::PreparedTask,
1798 ) -> Result<()> {
1799 self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
1800 }
1801
1802 pub(crate) fn save_warm_in_background(
1805 &self,
1806 workspace: &Workspace,
1807 task: &Task,
1808 warm: &Warm,
1809 prepared: &qk_executor::PreparedTask,
1810 ) -> Result<()> {
1811 let save = Save::prepare(workspace, task, warm, prepared)?;
1812 let cache = Cache {
1813 remote: self.remote.clone(),
1814 ..Cache::new(self.root.clone())
1815 };
1816 let handle = std::thread::spawn(move || {
1817 if let Err(error) = cache.save_prepared(&save) {
1818 qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
1819 }
1820 });
1821 self.saves.lock().unwrap().push(handle);
1822 Ok(())
1823 }
1824
1825 fn save_prepared(&self, save: &Save) -> Result<()> {
1830 let warm = &save.warm;
1831 let _lock = self.lock_warm(warm)?;
1832 let _save = profile::span(&save.task, "warm_save");
1833 let records = self.warm_records(warm);
1834 let previous = records
1836 .iter()
1837 .filter(|(_, record)| record.worktree == save.worktree)
1838 .min_by_key(|(_, record)| record.key != save.key)
1839 .map(|(_, record)| record);
1840 let noted: RestoredFiles = fs::read(&save.notes)
1841 .ok()
1842 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1843 .unwrap_or_default();
1844 let mut record = WarmRecord::new(
1845 warm.identity.clone(),
1846 save.worktree.clone(),
1847 save.commit.clone(),
1848 now_ms(),
1849 save.key.clone(),
1850 );
1851 let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
1852 for location in &save.locations {
1853 if !location.base.exists() {
1854 continue;
1855 }
1856 let known = previous.and_then(|previous| previous.groups.get(location.name));
1857 let restored_files = noted.groups.get(location.name);
1858 let mut group = WarmGroup::default();
1859 let mut pending = Vec::new();
1860 for path in location.paths()? {
1861 paths::safe_parents(&location.base, &path)?;
1862 let absolute = location.base.join(&path);
1863 let metadata = match fs::symlink_metadata(&absolute) {
1864 Err(error) if gone(&error) => continue,
1865 result => result?,
1866 };
1867 let artifact = if metadata.file_type().is_symlink() {
1868 let target = match fs::read_link(&absolute) {
1869 Err(error) if gone(&error) => continue,
1870 result => result?,
1871 };
1872 let target = target
1873 .to_str()
1874 .context("symlink target must be UTF-8")?
1875 .to_owned();
1876 validate_link(&path, &target)?;
1877 Artifact::Symlink {
1878 target,
1879 directory: absolute.is_dir(),
1880 }
1881 } else if metadata.is_file() {
1882 let now = stamp(&metadata);
1883 let saved = known.and_then(|known| {
1884 (known.stamps.get(&path) == Some(&now))
1885 .then(|| known.artifacts.get(&path))
1886 .flatten()
1887 });
1888 let restored = restored_files
1889 .and_then(|files| files.get(&path))
1890 .filter(|(stamp, _)| *stamp == now)
1891 .map(|(_, artifact)| artifact);
1892 let reused = saved.or(restored).filter(|artifact| match artifact {
1893 Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
1894 _ => false,
1895 });
1896 let Some(artifact) = reused.cloned() else {
1897 pending.push((path, absolute, metadata));
1898 continue;
1899 };
1900 group.stamps.insert(path.clone(), now);
1901 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1902 group.mtimes.insert(path.clone(), modified);
1903 }
1904 artifact
1905 } else if metadata.is_dir() {
1906 Artifact::Directory {
1907 mode: mode(&metadata),
1908 }
1909 } else {
1910 continue;
1911 };
1912 group.artifacts.insert(path, artifact);
1913 }
1914 let storing = profile::span(&save.task, "warm_save_blobs");
1915 let sources: Vec<_> = pending
1916 .iter()
1917 .map(|(_, absolute, _)| absolute.clone())
1918 .collect();
1919 for ((path, _, metadata), result) in pending.into_iter().zip(self.put_blobs(&sources)) {
1920 let blob = match result {
1921 Ok(blob) => blob,
1922 Err(error) if error.downcast_ref::<std::io::Error>().is_some_and(gone) => {
1923 continue;
1924 }
1925 Err(error) => return Err(error),
1926 };
1927 group.stamps.insert(path.clone(), stamp(&metadata));
1928 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1929 group.mtimes.insert(path.clone(), modified);
1930 }
1931 group.artifacts.insert(
1932 path,
1933 Artifact::File {
1934 blob,
1935 mode: mode(&metadata),
1936 },
1937 );
1938 }
1939 drop(storing);
1940 if let Some(limit) = warm.max_size
1941 && group.size(self) > limit
1942 {
1943 qk_executor::status!(
1944 "qk: {}: warm {} is over its maxSize and was not saved",
1945 save.task,
1946 location.name
1947 );
1948 continue;
1949 }
1950 record.groups.insert(location.name.to_owned(), group);
1951 }
1952 self.store_warm(&record)?;
1953 let _ = fs::remove_file(&save.notes);
1954 let saved = self.warm_record(&record);
1957 for (path, _) in records
1958 .into_iter()
1959 .filter(|(path, _)| *path != saved)
1960 .skip(KEPT - 1)
1961 {
1962 let _ = fs::remove_file(path);
1963 }
1964 if warm.remote
1966 && warm.portable != Portable::Own
1967 && let Some(remote) = &self.remote
1968 && let Some(branch) = &save.branch
1969 {
1970 for group in record.groups.values_mut() {
1971 group.stamps.clear();
1972 }
1973 remote.upload_warm(
1974 &self.root,
1975 &warm.identity,
1976 &remote_identity(warm, &save.worktree),
1977 branch,
1978 serde_json::to_vec(&record)?,
1979 );
1980 }
1981 Ok(())
1982 }
1983}
1984
1985struct Save {
1988 task: String,
1989 warm: Warm,
1990 locations: Vec<Location>,
1991 worktree: String,
1992 commit: Option<String>,
1993 branch: Option<String>,
1994 key: Vec<String>,
1995 notes: PathBuf,
1996}
1997
1998impl Save {
1999 fn prepare(
2001 workspace: &Workspace,
2002 task: &Task,
2003 warm: &Warm,
2004 prepared: &qk_executor::PreparedTask,
2005 ) -> Result<Self> {
2006 Ok(Self {
2007 task: task.id.clone(),
2008 warm: warm.clone(),
2009 locations: locations(workspace, task, warm)?,
2010 worktree: worktree(workspace),
2011 commit: current_commit(workspace),
2012 branch: current_branch(workspace),
2013 key: key_digests(workspace, warm, prepared)?,
2014 notes: restored_files_path(workspace, warm),
2015 })
2016 }
2017}
2018
2019#[cfg(test)]
2020mod tests {
2021 use super::*;
2022 use qk_taskgraph::Request;
2023 use serde_json::json;
2024
2025 fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
2027 let temp = tempfile::tempdir().unwrap();
2028 fs::write(temp.path().join("nx.json"), "{}").unwrap();
2029 fs::write(
2030 temp.path().join("project.json"),
2031 json!({
2032 "name": "app", "targets": {
2033 "one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
2034 "two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
2035 }
2036 })
2037 .to_string(),
2038 )
2039 .unwrap();
2040 for path in ["scratch-one", "scratch-two"] {
2041 fs::create_dir(temp.path().join(path)).unwrap();
2042 fs::write(temp.path().join(path).join("state"), "saved").unwrap();
2043 }
2044 let workspace = Workspace::load(temp.path()).unwrap();
2045 let graph = TaskGraph::build(
2046 &workspace,
2047 &[
2048 Request::parse("app:one").unwrap(),
2049 Request::parse("app:two").unwrap(),
2050 ],
2051 )
2052 .unwrap();
2053 (temp, workspace, graph)
2054 }
2055
2056 #[test]
2059 fn asks_a_partial_clone_about_ancestors_without_fetching() {
2060 let temp = tempfile::tempdir().unwrap();
2061 let git = |root: &std::path::Path, args: &[&str]| {
2062 let output = paths::git(root, args).unwrap();
2063 assert!(output.status.success(), "git {args:?}: {output:?}");
2064 String::from_utf8(output.stdout).unwrap().trim().to_owned()
2065 };
2066 let source = temp.path().join("source");
2067 fs::create_dir(&source).unwrap();
2068 git(&source, &["init", "--quiet"]);
2069 git(&source, &["config", "uploadpack.allowFilter", "true"]);
2070 let commit = |message: &str| {
2071 fs::write(source.join("file"), message).unwrap();
2072 git(&source, &["add", "file"]);
2073 git(
2074 &source,
2075 &[
2076 "-c",
2077 "user.name=qk",
2078 "-c",
2079 "user.email=qk@example.com",
2080 "commit",
2081 "--quiet",
2082 "-m",
2083 message,
2084 ],
2085 );
2086 git(&source, &["rev-parse", "HEAD"])
2087 };
2088 let older = commit("older");
2089 let newer = commit("newer");
2090 let clone = temp.path().join("clone");
2091 let url = format!("file://{}", source.display());
2092 git(
2093 temp.path(),
2094 &[
2095 "clone",
2096 "--quiet",
2097 "--depth",
2098 "1",
2099 "--filter=blob:none",
2100 &url,
2101 clone.to_str().unwrap(),
2102 ],
2103 );
2104 fs::write(clone.join("nx.json"), "{}").unwrap();
2105 let workspace = Workspace::load(&clone).unwrap();
2106 assert!(descends_from(&workspace, &newer));
2107 assert!(!descends_from(&workspace, &older));
2108 assert_eq!(behind(&workspace, &older), None);
2109 let version = git(temp.path(), &["--version"]);
2111 let honoured = version
2112 .split_whitespace()
2113 .nth(2)
2114 .and_then(|version| {
2115 let mut parts = version.split('.').map(|part| part.parse::<u32>().ok());
2116 Some((parts.next()??, parts.next()??))
2117 })
2118 .is_some_and(|version| version >= (2, 44));
2119 let held = paths::git_history(&clone, &["cat-file", "-e", &format!("{older}^{{commit}}")])
2120 .unwrap();
2121 assert!(!honoured || !held.status.success(), "fetched {older}");
2122 }
2123
2124 fn entries(warm: Value) -> Result<Vec<Warm>> {
2126 let temp = tempfile::tempdir().unwrap();
2127 fs::write(temp.path().join("nx.json"), "{}").unwrap();
2128 fs::write(
2129 temp.path().join("project.json"),
2130 json!({"name": "app", "targets": {"build": {"command": "true", "qk:warm": warm}}})
2131 .to_string(),
2132 )
2133 .unwrap();
2134 let workspace = Workspace::load(temp.path()).unwrap();
2135 let graph = TaskGraph::build(&workspace, &[Request::parse("app:build").unwrap()]).unwrap();
2136 config(&workspace, &graph.tasks["app:build"])
2137 }
2138
2139 #[test]
2140 fn an_array_holds_the_task_entry_beside_groups() {
2141 let entries = entries(json!([
2142 {"group": "metro", "portable": true, "env": {"METRO_CACHE_ROOT": "{warm}/metro"}},
2143 {"paths": ["build"], "mtimes": "preserve"}
2144 ]))
2145 .unwrap();
2146 assert_eq!(entries.len(), 2);
2147 assert_eq!(entries[0].group(), Some("metro"));
2148 assert!(entries[0].portable == Portable::Anywhere && entries[0].directory);
2149 assert_eq!(entries[1].group(), None);
2150 assert!(entries[1].portable == Portable::Own && entries[1].preserve_mtimes);
2151 assert_eq!(kept_paths(&entries), ["build"]);
2152 }
2153
2154 #[test]
2155 fn portable_is_a_boolean_or_the_same_path() {
2156 let portable = |value| {
2157 entries(json!({"paths": ["build"], "portable": value}))
2158 .map(|entries| entries[0].portable)
2159 };
2160 assert_eq!(portable(json!("path")).unwrap(), Portable::Path);
2161 assert_eq!(portable(json!(false)).unwrap(), Portable::Own);
2162 assert_eq!(
2163 format!("{:#}", portable(json!("anywhere")).unwrap_err()),
2164 "qk:warm.portable must be true, false or \"path\""
2165 );
2166 }
2167
2168 #[test]
2169 fn entries_cannot_share_saves_or_variables() {
2170 let error = |warm| format!("{:#}", entries(warm).unwrap_err());
2171 assert_eq!(
2172 error(json!([{"paths": ["one"]}, {"paths": ["two"]}])),
2173 "only one qk:warm entry may go without a group"
2174 );
2175 assert_eq!(
2176 error(json!([{"group": "metro"}, {"group": "metro"}])),
2177 "qk:warm names group \"metro\" twice"
2178 );
2179 assert_eq!(
2180 error(json!([
2181 {"group": "one", "env": {"CACHE": "{warm}/one"}},
2182 {"group": "two", "env": {"CACHE": "{warm}/two"}}
2183 ])),
2184 "qk:warm sets CACHE in more than one entry"
2185 );
2186 assert_eq!(
2187 error(json!(["build"])),
2188 "qk:warm must be an object or an array of objects"
2189 );
2190 }
2191
2192 #[test]
2194 fn waiting_for_an_owner_does_not_block_other_releases() {
2195 let (_temp, workspace, graph) = fixture();
2196 let one = &graph.tasks["app:one"];
2197 let two = &graph.tasks["app:two"];
2198 let first = acquire(&workspace, one, &config(&workspace, one).unwrap()[0]).unwrap();
2199 let state = paths::worktree_state(&workspace.root).join("warm-kept");
2200 let lock = fs::OpenOptions::new()
2201 .create(true)
2202 .truncate(false)
2203 .read(true)
2204 .write(true)
2205 .open(state.join(format!("{}.lock", hash32(&two.id))))
2206 .unwrap();
2207 fs2::FileExt::lock_exclusive(&lock).unwrap();
2208 let (released, received) = std::sync::mpsc::channel();
2209 let completed = std::thread::scope(|scope| {
2210 scope.spawn(|| {
2211 let second =
2212 acquire(&workspace, two, &config(&workspace, two).unwrap()[0]).unwrap();
2213 release(&second);
2214 });
2215 std::thread::sleep(Duration::from_millis(100));
2216 scope.spawn(|| {
2217 release(&first);
2218 released.send(()).unwrap();
2219 });
2220 let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
2221 fs2::FileExt::unlock(&lock).unwrap();
2222 completed
2223 });
2224 assert!(
2225 completed,
2226 "a blocked owner lock prevented an unrelated release"
2227 );
2228 assert_eq!(
2229 fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
2230 "saved"
2231 );
2232 }
2233
2234 #[test]
2236 fn unreadable_ownership_preserves_saved_paths() {
2237 let (_temp, workspace, graph) = fixture();
2238 let owner = &graph.tasks["app:one"];
2239 let warm = config(&workspace, owner).unwrap().remove(0);
2240 let root = acquire(&workspace, owner, &warm).unwrap();
2241 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
2242 drop(entry); fs::write(Stash::manifest(&root), "{").unwrap();
2244 assert!(acquire(&workspace, owner, &warm).is_err());
2245 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2246 fs::remove_file(Stash::manifest(&root)).unwrap();
2247 assert!(acquire(&workspace, owner, &warm).is_err());
2248 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2249 fs::create_dir(Stash::manifest(&root)).unwrap();
2250 assert!(acquire(&workspace, owner, &warm).is_err());
2251 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2252 }
2253
2254 #[cfg(unix)]
2256 #[test]
2257 fn inspection_errors_keep_recovery_incomplete() {
2258 let (_temp, workspace, graph) = fixture();
2259 let owner = &graph.tasks["app:one"];
2260 let warm = config(&workspace, owner).unwrap().remove(0);
2261 let root = acquire(&workspace, owner, &warm).unwrap();
2262 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
2263 let saved = root.with_extension("saved");
2264 fs::rename(&root, &saved).unwrap();
2265 fs::write(&root, "blocked").unwrap(); assert!(!entry.stash.put_back());
2267 let retained = Stash {
2268 root: root.clone(),
2269 workspace: workspace.root.clone(),
2270 moved: entry.stash.moved.clone(),
2271 };
2272 assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
2273 assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
2274 assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
2275 }
2276}