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