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_pattern(&path)?;
317 paths.push(if excluded { format!("!{path}") } else { path });
318 }
319 let mut env = BTreeMap::new();
320 let mut uses_directory = false;
321 if let Some(values) = object.get("env") {
322 for (name, value) in values
323 .as_object()
324 .context("qk:warm.env must be an object")?
325 {
326 let value = value
327 .as_str()
328 .context("qk:warm.env values must be strings")?;
329 uses_directory |= value.contains("{warm}");
330 env.insert(name.clone(), expand(value)?);
331 }
332 }
333 let max_size = match object.get("maxSize") {
334 None => None,
335 Some(Value::String(text)) => Some(crate::parse_size(text)?),
336 Some(Value::Number(number)) => number.as_u64(),
337 Some(_) => bail!("qk:warm.maxSize must be a size such as \"2GB\""),
338 };
339 let preserve_mtimes = match object.get("mtimes").map(|value| value.as_str()) {
340 None | Some(Some("epoch")) => false,
341 Some(Some("preserve")) => true,
342 Some(_) => bail!("qk:warm.mtimes must be \"epoch\" or \"preserve\""),
343 };
344 const KEY_SHAPE: &str = "qk:warm.key must be an array of paths and {\"env\": name} objects";
345 let mut key = Vec::new();
346 for part in object
347 .get("key")
348 .map(|value| value.as_array().context(KEY_SHAPE))
349 .transpose()?
350 .into_iter()
351 .flatten()
352 {
353 key.push(match part {
354 Value::String(path) => {
355 dynamic(path)?;
356 let path = paths::expand(workspace, &task.project, path)?;
357 paths::validate_path(&path)?;
358 KeyPart::File(path)
359 }
360 Value::Object(object) if object.len() == 1 => KeyPart::Env(
361 object
362 .get("env")
363 .and_then(Value::as_str)
364 .context(KEY_SHAPE)?
365 .to_owned(),
366 ),
367 _ => bail!(KEY_SHAPE),
368 });
369 }
370 let restore_keys = object
371 .get("restoreKeys")
372 .map(|value| {
373 value
374 .as_u64()
375 .and_then(|count| usize::try_from(count).ok())
376 .filter(|count| *count <= key.len())
377 .context("qk:warm.restoreKeys must be a count no larger than qk:warm.key")
378 })
379 .transpose()?;
380 let each_path = match object.get("restore").map(Value::as_str) {
381 None | Some(Some("whole")) => false,
382 Some(Some("each-path")) if !paths.is_empty() => true,
383 Some(Some("each-path")) => bail!("qk:warm.restore \"each-path\" needs qk:warm.paths"),
384 Some(_) => bail!("qk:warm.restore must be \"whole\" or \"each-path\""),
385 };
386 let background = match object.get("save").map(Value::as_str) {
387 None | Some(Some("wait")) => false,
388 Some(Some("background")) => true,
389 Some(_) => bail!("qk:warm.save must be \"wait\" or \"background\""),
390 };
391 Ok(Warm {
392 identity,
393 outputs: flag("outputs", false)?,
394 paths,
395 env,
396 directory: uses_directory,
397 max_size,
398 remote: flag("remote", true)?,
399 preserve_mtimes,
400 portable: match object.get("portable") {
401 None | Some(Value::Bool(false)) => Portable::Own,
402 Some(Value::Bool(true)) => Portable::Anywhere,
403 Some(Value::String(path)) if path == "path" => Portable::Path,
404 Some(_) => bail!("qk:warm.portable must be true, false or \"path\""),
405 },
406 key,
407 restore_keys,
408 survive,
409 background,
410 each_path,
411 prefetch: flag("prefetch", false)?,
412 })
413}
414
415fn dynamic(path: &str) -> Result<()> {
419 if path.contains("{options.") || path.contains("{args.") {
420 bail!("dynamic cache paths are not supported yet");
421 }
422 Ok(())
423}
424
425fn directory(workspace: &Workspace, identity: &str) -> PathBuf {
426 paths::worktree_state(&workspace.root)
427 .join("warm")
428 .join(&blake3::hash(identity.as_bytes()).to_hex()[..32])
429}
430
431impl Warm {
432 pub fn group(&self) -> Option<&str> {
434 self.identity.strip_prefix(GROUP)
435 }
436
437 pub fn kept_paths(&self) -> Vec<String> {
440 self.paths
441 .iter()
442 .filter(|path| !path.starts_with('!'))
443 .cloned()
444 .collect()
445 }
446}
447
448pub fn kept_paths(entries: &[Warm]) -> Vec<String> {
450 entries.iter().flat_map(Warm::kept_paths).collect()
451}
452
453#[derive(Default)]
456pub struct Kept {
457 roots: Vec<PathBuf>,
459}
460
461impl Drop for Kept {
462 fn drop(&mut self) {
463 for root in self.roots.iter().rev() {
464 release(root);
465 }
466 }
467}
468
469#[derive(Default, Serialize, Deserialize)]
471struct Stash {
472 #[serde(skip)]
473 root: PathBuf,
474 #[serde(skip)]
475 workspace: PathBuf,
476 moved: Vec<String>,
478}
479
480struct Active {
483 users: usize,
484 stash: Stash,
485 _lock: fs::File,
486}
487
488static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
491
492impl Stash {
493 fn manifest(root: &std::path::Path) -> PathBuf {
495 root.join("moved.json")
496 }
497
498 fn save_manifest(&self) -> Result<()> {
501 let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
502 file.write_all(&serde_json::to_vec(self)?)?;
503 file.as_file().sync_all()?;
504 file.persist(Self::manifest(&self.root))?;
505 Ok(())
506 }
507
508 fn put_back(&self) -> bool {
514 let mut complete = true;
515 for (index, path) in self.moved.iter().enumerate().rev() {
516 let kept = self.root.join(index.to_string());
517 match fs::symlink_metadata(&kept) {
518 Ok(_) => {}
519 Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
520 Err(error) => {
521 complete = false;
522 qk_executor::status!(
523 "qk: could not inspect saved {path}, kept for the next run ({error})"
524 );
525 continue;
526 }
527 }
528 let original = self.workspace.join(path);
529 let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
530 remove(&original)?;
531 if let Some(parent) = original.parent() {
532 fs::create_dir_all(parent)?;
533 }
534 fs::rename(&kept, &original)?;
535 Ok(())
536 });
537 if let Err(error) = result {
538 complete = false;
539 qk_executor::status!(
540 "qk: could not put {path} back, kept for the next run ({error:#})"
541 );
542 }
543 }
544 if complete {
545 let _ = fs::remove_dir_all(&self.root);
546 }
547 complete
548 }
549}
550
551fn remove(path: &std::path::Path) -> std::io::Result<()> {
553 match fs::symlink_metadata(path) {
554 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
555 Err(error) => Err(error),
556 Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
557 Ok(_) => fs::remove_file(path),
558 }
559}
560
561fn release(root: &std::path::Path) {
564 let mut active = ACTIVE
565 .lock()
566 .unwrap_or_else(std::sync::PoisonError::into_inner);
567 let Some(entry) = active.get_mut(root) else {
568 return;
569 };
570 entry.users -= 1;
571 if entry.users > 0 {
572 return;
573 }
574 let entry = active.remove(root).expect("present above");
575 drop(active);
578 entry.stash.put_back();
579}
580
581fn names(owner: &Task, entry: &str, task: &Task) -> bool {
583 match entry.split_once(':') {
584 Some((project, target)) => task.project == project && task.target == target,
585 None => task.project == owner.project && task.target == entry,
586 }
587}
588
589fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
591 let mut seen = std::collections::BTreeSet::new();
592 let mut next = vec![from];
593 while let Some(id) = next.pop() {
594 for dependency in graph
595 .tasks
596 .get(id)
597 .into_iter()
598 .flat_map(|task| &task.dependencies)
599 {
600 if dependency == on {
601 return true;
602 }
603 if seen.insert(dependency.as_str()) {
604 next.push(dependency);
605 }
606 }
607 }
608 false
609}
610
611pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
614 let mut kept = Kept::default();
615 for owner in graph.tasks.values() {
616 let Ok(entries) = config(workspace, owner) else {
617 continue;
618 };
619 let Some(warm) = entries.into_iter().find(|warm| warm.group().is_none()) else {
621 continue;
622 };
623 if !warm.survive.iter().any(|entry| names(owner, entry, task))
624 || !depends_on(graph, &owner.id, &task.id)
625 {
626 continue;
627 }
628 match acquire(workspace, owner, &warm) {
629 Ok(root) => kept.roots.push(root),
630 Err(error) => qk_executor::status!(
631 "qk: {}: warm paths not kept across {} ({error:#})",
632 owner.id,
633 task.id
634 ),
635 }
636 }
637 kept
638}
639
640fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
643 let state = paths::worktree_state(&workspace.root).join("warm-kept");
644 let root = state.join(hash32(&owner.id));
645 fs::create_dir_all(&state)?;
646 let lock = fs::OpenOptions::new()
647 .create(true)
648 .truncate(false)
649 .read(true)
650 .write(true)
651 .open(state.join(format!("{}.lock", hash32(&owner.id))))?;
652 let mut active = loop {
655 let mut active = ACTIVE
656 .lock()
657 .unwrap_or_else(std::sync::PoisonError::into_inner);
658 if let Some(entry) = active.get_mut(&root) {
659 entry.users += 1;
660 return Ok(root);
661 }
662 match fs2::FileExt::try_lock_exclusive(&lock) {
663 Ok(()) => break active,
664 Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
665 drop(active);
666 std::thread::sleep(Duration::from_millis(20));
667 }
668 Err(error) => return Err(error.into()),
669 }
670 };
671 let left = match fs::read(Stash::manifest(&root)) {
678 Ok(bytes) => {
679 let mut left: Stash = serde_json::from_slice(&bytes)
680 .context("invalid warm stash manifest; saved paths retained for recovery")?;
681 left.root = root.clone();
682 left.workspace = workspace.root.clone();
683 Some(left)
684 }
685 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
686 if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
687 bail!("missing warm stash manifest; saved paths retained for recovery");
688 }
689 None
690 }
691 Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
692 };
693 let retained = left.filter(|left| !left.put_back());
694 let stash = stash(workspace, &root, warm, retained)?;
695 active.insert(
696 root.clone(),
697 Active {
698 users: 1,
699 stash,
700 _lock: lock,
701 },
702 );
703 Ok(root)
704}
705
706fn stash(
710 workspace: &Workspace,
711 root: &std::path::Path,
712 warm: &Warm,
713 retained: Option<Stash>,
714) -> Result<Stash> {
715 let mut stash = match retained {
716 Some(retained) => retained,
717 None => {
718 let _ = fs::remove_dir_all(root);
719 Stash {
720 root: root.to_owned(),
721 workspace: workspace.root.clone(),
722 moved: Vec::new(),
723 }
724 }
725 };
726 let held: Vec<String> = stash
727 .moved
728 .iter()
729 .enumerate()
730 .filter_map(
731 |(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
732 Ok(_) => Some(Ok(path.clone())),
733 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
734 Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
735 },
736 )
737 .collect::<Result<_>>()?;
738 let overlaps = |path: &str| {
739 held.iter().any(|held| {
740 held == path
741 || path.starts_with(&format!("{held}/"))
742 || held.starts_with(&format!("{path}/"))
743 })
744 };
745 let mut kept = std::collections::BTreeSet::new();
748 for path in warm.kept_paths() {
749 if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
750 .and_then(|outputs| outputs.paths(&workspace.root))
751 {
752 kept.extend(found);
753 }
754 }
755 let topmost: Vec<String> = kept
757 .iter()
758 .filter(|path| {
759 !std::path::Path::new(path.as_str())
760 .ancestors()
761 .skip(1)
762 .any(|ancestor| {
763 ancestor
764 .to_str()
765 .is_some_and(|ancestor| kept.contains(ancestor))
766 })
767 })
768 .filter(|path| !overlaps(path))
769 .cloned()
770 .collect();
771 if topmost.is_empty() {
772 return Ok(stash);
773 }
774 fs::create_dir_all(root)?;
775 for path in topmost {
776 let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
777 stash.moved.push(path.clone());
778 stash.save_manifest()?;
779 fs::rename(
782 workspace.root.join(&path),
783 root.join((stash.moved.len() - 1).to_string()),
784 )
785 .inspect_err(|_| {
786 stash.moved.pop();
787 })?;
788 Ok(())
789 });
790 if let Err(error) = moved {
791 stash.save_manifest()?;
794 stash.put_back();
795 return Err(error);
796 }
797 }
798 Ok(stash)
799}
800
801pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
804 let mut claimed: Vec<(String, &str)> = Vec::new();
805 for (id, task) in &graph.tasks {
806 for path in kept_paths(&config(workspace, task)?) {
807 if let Some((_, other)) = claimed.iter().find(|(known, _)| {
808 known == &path
809 || known.starts_with(&format!("{path}/"))
810 || path.starts_with(&format!("{known}/"))
811 }) {
812 bail!("{id} and {other} both keep {path} as warm state");
813 }
814 claimed.push((path, id));
815 }
816 }
817 Ok(())
818}
819
820impl WarmGroup {
821 fn size(&self, cache: &Cache) -> u64 {
823 self.artifacts
824 .values()
825 .filter_map(|artifact| match artifact {
826 Artifact::File { blob, .. } => {
827 fs::metadata(cache.root.join("blobs").join(blob)).ok()
828 }
829 _ => None,
830 })
831 .map(|metadata| metadata.len())
832 .sum()
833 }
834}
835
836#[derive(Default, Serialize, Deserialize)]
839struct RestoredFiles {
840 groups: BTreeMap<String, GroupFiles>,
841}
842
843type GroupFiles = BTreeMap<String, (Vec<i64>, Artifact)>;
845
846#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
848pub struct Restored {
849 #[serde(default, skip_serializing_if = "Option::is_none")]
851 pub group: Option<String>,
852 pub source: String,
855 pub groups: Vec<String>,
856 pub files: usize,
857 pub bytes: u64,
858}
859
860#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
862#[serde(rename_all = "camelCase")]
863pub struct WarmReport {
864 #[serde(default, skip_serializing_if = "Vec::is_empty")]
866 pub restored: Vec<Restored>,
867 #[serde(default, skip_serializing_if = "Vec::is_empty")]
871 pub present: Vec<String>,
872 pub save_ms: Option<u64>,
875 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
877 pub background: bool,
878}
879
880#[derive(Clone)]
883struct Location {
884 name: &'static str,
885 base: PathBuf,
886 outputs: Option<Outputs>,
888 each: Vec<Outputs>,
891}
892
893impl Location {
894 fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
896 if !self.base.exists() {
897 return Ok(Default::default());
898 }
899 match &self.outputs {
900 Some(outputs) => outputs.paths(&self.base),
901 None => {
902 let mut paths = std::collections::BTreeSet::new();
903 for entry in walkdir::WalkDir::new(&self.base)
904 .min_depth(1)
905 .follow_links(false)
906 {
907 let entry = entry?;
908 paths.insert(paths::relative(&self.base, entry.path())?);
909 }
910 Ok(paths)
911 }
912 }
913 }
914
915 fn matches(&self, path: &str) -> bool {
917 self.outputs
918 .as_ref()
919 .is_none_or(|outputs| outputs.matches(path))
920 }
921
922 fn present(&self) -> Result<Option<Vec<&Outputs>>> {
925 let mut present = Vec::new();
926 for outputs in &self.each {
927 if !outputs.paths(&self.base)?.is_empty() {
928 present.push(outputs);
929 }
930 }
931 Ok((!present.is_empty() && present.len() < self.each.len()).then_some(present))
932 }
933}
934
935fn remove_restored(location: &Location, kept: &[&Outputs]) {
939 let Ok(paths) = location.paths() else {
940 return;
941 };
942 for path in paths.iter().rev() {
943 if kept.iter().any(|outputs| outputs.matches(path)) {
944 continue;
945 }
946 let absolute = location.base.join(path);
947 let _ = match fs::symlink_metadata(&absolute) {
948 Ok(metadata) if metadata.is_dir() => fs::remove_dir(&absolute),
949 Ok(_) => fs::remove_file(&absolute),
950 Err(_) => continue,
951 };
952 }
953}
954
955fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
957 let mut locations = Vec::new();
958 if warm.outputs {
959 locations.push(Location {
960 name: "outputs",
961 base: workspace.root.clone(),
962 outputs: Some(Outputs::new(workspace, task)?),
963 each: Vec::new(),
964 });
965 }
966 if !warm.paths.is_empty() {
967 let mut each = Vec::new();
968 if warm.each_path {
969 let (excluded, included): (Vec<_>, Vec<_>) =
970 warm.paths.iter().partition(|path| path.starts_with('!'));
971 for path in included {
972 let mut patterns = vec![path.clone()];
973 patterns.extend(excluded.iter().map(|path| (*path).clone()));
974 each.push(Outputs::from_paths(&patterns)?);
975 }
976 }
977 locations.push(Location {
978 name: "paths",
979 base: workspace.root.clone(),
980 outputs: Some(Outputs::from_paths(&warm.paths)?),
981 each,
982 });
983 }
984 if warm.directory {
985 locations.push(Location {
986 name: "directory",
987 base: directory(workspace, &warm.identity),
988 outputs: None,
989 each: Vec::new(),
990 });
991 }
992 Ok(locations)
993}
994
995pub(crate) fn current_branch(workspace: &Workspace) -> Option<String> {
998 for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
999 if let Ok(value) = std::env::var(name)
1000 && !value.is_empty()
1001 {
1002 return Some(value);
1003 }
1004 }
1005 let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
1006 let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
1007 (output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
1008}
1009
1010pub(crate) fn branches(workspace: &Workspace) -> Vec<String> {
1013 let default = workspace
1014 .config
1015 .default_base
1016 .clone()
1017 .unwrap_or_else(|| "main".into());
1018 let default = default
1019 .strip_prefix("origin/")
1020 .unwrap_or(&default)
1021 .to_owned();
1022 let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
1023 if !branches.contains(&default) {
1024 branches.push(default);
1025 }
1026 branches
1027}
1028
1029fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
1031 let modified = metadata
1032 .modified()
1033 .ok()
1034 .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
1035 .map_or(-1, |duration| duration.as_nanos() as i64);
1036 #[allow(unused_mut)]
1037 let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
1038 #[cfg(unix)]
1039 {
1040 use std::os::unix::fs::MetadataExt;
1041 stamp.extend([
1042 metadata.ctime(),
1043 metadata.ctime_nsec(),
1044 metadata.ino() as i64,
1045 ]);
1046 }
1047 stamp
1048}
1049
1050fn remote_identity(warm: &Warm, worktree: &str) -> String {
1055 match warm.portable {
1056 Portable::Path => format!("{}\0path\0{worktree}", warm.identity),
1057 _ => warm.identity.clone(),
1058 }
1059}
1060
1061fn worktree(workspace: &Workspace) -> String {
1062 fs::canonicalize(&workspace.root)
1063 .unwrap_or_else(|_| workspace.root.clone())
1064 .to_string_lossy()
1065 .into_owned()
1066}
1067
1068fn current_commit(workspace: &Workspace) -> Option<String> {
1070 let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
1071 let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
1072 (output.status.success() && !commit.is_empty()).then_some(commit)
1073}
1074
1075fn key_digests(
1078 workspace: &Workspace,
1079 warm: &Warm,
1080 prepared: &qk_executor::PreparedTask,
1081) -> Result<Vec<String>> {
1082 warm.key
1083 .iter()
1084 .map(|part| {
1085 Ok(match part {
1086 KeyPart::File(path) => {
1087 paths::safe_parents(&workspace.root, path)?;
1088 let absolute = workspace.root.join(path);
1089 if absolute.is_file() {
1090 format!("file:{}", digest_file(&absolute)?)
1091 } else {
1092 "file:absent".to_owned()
1093 }
1094 }
1095 KeyPart::Env(name) => match prepared
1098 .execution
1099 .get(std::ffi::OsStr::new(name))
1100 .or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
1101 {
1102 Some(value) => format!(
1103 "env:{}",
1104 blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
1105 ),
1106 None => "env:unset".to_owned(),
1107 },
1108 })
1109 })
1110 .collect()
1111}
1112
1113fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
1117 if saved == current {
1118 return Some(0);
1119 }
1120 let count = warm.restore_keys?;
1121 (saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
1122}
1123
1124fn descends_from(workspace: &Workspace, commit: &str) -> bool {
1126 paths::git_history(
1127 &workspace.root,
1128 &["merge-base", "--is-ancestor", commit, "HEAD"],
1129 )
1130 .is_ok_and(|output| output.status.success())
1131}
1132
1133fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
1136 let ancestor = paths::git_history(
1137 &workspace.root,
1138 &["merge-base", "--is-ancestor", commit, "HEAD"],
1139 )
1140 .ok()?;
1141 if !ancestor.status.success() {
1142 return None;
1143 }
1144 let output = paths::git_history(
1145 &workspace.root,
1146 &["rev-list", "--count", &format!("{commit}..HEAD")],
1147 )
1148 .ok()?;
1149 String::from_utf8(output.stdout).ok()?.trim().parse().ok()
1150}
1151
1152fn depth(graph: &TaskGraph, id: &str, known: &mut BTreeMap<String, usize>) -> usize {
1154 if let Some(depth) = known.get(id) {
1155 return *depth;
1156 }
1157 let depth = graph.tasks.get(id).map_or(0, |task| {
1158 task.dependencies
1159 .iter()
1160 .map(|dependency| depth(graph, dependency, known) + 1)
1161 .max()
1162 .unwrap_or(0)
1163 });
1164 known.insert(id.to_owned(), depth);
1165 depth
1166}
1167
1168fn warm_prefix(identity: &str) -> String {
1170 hash32(&format!(
1171 "{identity}\0{}\0{}",
1172 std::env::consts::OS,
1173 std::env::consts::ARCH
1174 ))
1175}
1176
1177fn warm_record(root: &Path, record: &WarmRecord) -> PathBuf {
1178 root.join("warm").join(format!(
1179 "{}-{}.json",
1180 warm_prefix(&record.task),
1181 hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
1182 ))
1183}
1184
1185fn store_warm(root: &Path, record: &WarmRecord) -> Result<()> {
1186 let path = warm_record(root, record);
1187 let directory = path.parent().context("warm records have a directory")?;
1188 fs::create_dir_all(directory)?;
1189 record::publish(root, &path, record)
1190}
1191
1192fn remote_warm(
1195 remote: &crate::remote::Remote,
1196 root: &Path,
1197 task: &str,
1198 warm: &Warm,
1199 own: &str,
1200 branches: &[String],
1201 current: &[String],
1202) -> Option<(WarmRecord, String)> {
1203 if !warm.remote || warm.portable == Portable::Own {
1204 return None;
1205 }
1206 let placed = |saved: Option<&str>| warm.portable != Portable::Path || saved == Some(own);
1208 for branch in branches {
1209 let suits = |saved: &[String], worktree: Option<&str>| {
1210 key_match(warm, current, saved).is_some() && placed(worktree)
1211 };
1212 let name = remote_identity(warm, own);
1213 match remote.fetch_warm(root, &name, task, branch, suits) {
1214 Ok(Some((bytes, _pin))) => {
1215 let Some(record) = WarmRecord::read(&bytes, &warm.identity).ok() else {
1216 continue;
1217 };
1218 if key_match(warm, current, &record.key).is_none()
1221 || !placed(Some(&record.worktree))
1222 {
1223 continue;
1224 }
1225 if store_warm(root, &record).is_err() {
1226 return None;
1227 }
1228 return Some((record, branch.clone()));
1229 }
1230 Ok(None) => {}
1231 Err(error) => {
1232 qk_executor::status!("qk: {task}: remote warm state unavailable ({error:#})");
1233 continue;
1236 }
1237 }
1238 }
1239 None
1240}
1241
1242fn hash32(text: &str) -> String {
1244 blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
1245}
1246
1247fn nanos(time: SystemTime) -> Option<i64> {
1249 time.duration_since(UNIX_EPOCH)
1250 .ok()
1251 .and_then(|duration| i64::try_from(duration.as_nanos()).ok())
1252}
1253
1254fn now_ms() -> u64 {
1256 SystemTime::now()
1257 .duration_since(UNIX_EPOCH)
1258 .map_or(0, |duration| duration.as_millis() as u64)
1259}
1260
1261fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
1263 paths::worktree_state(&workspace.root)
1264 .join("warm-restored")
1265 .join(format!("{}.json", hash32(&warm.identity)))
1266}
1267
1268fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
1270 fs::read(restored_files_path(workspace, warm))
1271 .ok()
1272 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1273 .unwrap_or_default()
1274}
1275
1276impl Cache {
1277 fn warm_prefix(&self, identity: &str) -> String {
1279 warm_prefix(identity)
1280 }
1281
1282 fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
1285 let directory = self.root.join("locks");
1286 fs::create_dir_all(&directory)?;
1287 let file = fs::OpenOptions::new()
1288 .create(true)
1289 .truncate(false)
1290 .read(true)
1291 .write(true)
1292 .open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
1293 fs2::FileExt::lock_exclusive(&file)?;
1294 Ok(file)
1295 }
1296
1297 fn warm_record(&self, record: &WarmRecord) -> PathBuf {
1298 warm_record(&self.root, record)
1299 }
1300
1301 fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, WarmRecord)> {
1303 let prefix = format!("{}-", self.warm_prefix(&warm.identity));
1304 let mut records: Vec<(PathBuf, WarmRecord)> = fs::read_dir(self.root.join("warm"))
1305 .into_iter()
1306 .flatten()
1307 .flatten()
1308 .filter(|entry| {
1309 entry
1310 .file_name()
1311 .to_str()
1312 .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
1313 })
1314 .filter_map(|entry| {
1315 let record =
1316 WarmRecord::read(&fs::read(entry.path()).ok()?, &warm.identity).ok()?;
1317 Some((entry.path(), record))
1318 })
1319 .collect();
1320 records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
1321 records
1322 }
1323
1324 fn choose_warm(
1331 &self,
1332 workspace: &Workspace,
1333 warm: &Warm,
1334 current: &[String],
1335 related: bool,
1336 ) -> Option<(WarmRecord, String)> {
1337 let own = worktree(workspace);
1338 let mut candidates: Vec<(u8, bool, WarmRecord)> = self
1339 .warm_records(warm)
1340 .into_iter()
1341 .filter_map(|(_, record)| {
1342 let quality = key_match(warm, current, &record.key)?;
1343 let other = record.worktree != own;
1344 (!other || warm.portable == Portable::Anywhere).then_some((quality, other, record))
1345 })
1346 .filter(|(_, _, record)| {
1347 let descends = record
1348 .commit
1349 .as_deref()
1350 .is_some_and(|commit| descends_from(workspace, commit));
1351 descends == related
1352 })
1353 .collect();
1354 candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
1356 let rank = candidates
1357 .first()
1358 .map(|(quality, other, _)| (*quality, *other))?;
1359 let tied = candidates
1360 .iter()
1361 .take_while(|(quality, other, _)| (*quality, *other) == rank)
1362 .count();
1363 let index = if tied > 1 {
1364 (0..tied)
1365 .min_by_key(|&index| {
1366 let behind = candidates[index]
1367 .2
1368 .commit
1369 .as_deref()
1370 .and_then(|commit| behind(workspace, commit));
1371 (behind.is_none(), behind, index)
1372 })
1373 .unwrap_or(0)
1374 } else {
1375 0
1376 };
1377 let (_, other, record) = candidates.swap_remove(index);
1378 let source = if other {
1379 format!("worktree {}", record.worktree)
1380 } else {
1381 "local".to_owned()
1382 };
1383 Some((record, source))
1384 }
1385
1386 fn store_warm(&self, record: &WarmRecord) -> Result<()> {
1387 store_warm(&self.root, record)
1388 }
1389
1390 pub(crate) fn clear_warm(&self, workspace: &Workspace, task: &Task, warm: &Warm) -> Result<()> {
1395 let _lock = self.lock_warm(warm)?;
1396 let identity = self.warm_prefix(&warm.identity);
1397 if self.cleared.lock().unwrap().contains(&identity) {
1398 return Ok(());
1399 }
1400 for location in locations(workspace, task, warm)? {
1401 for path in location.paths()?.iter().rev() {
1402 let absolute = location.base.join(path);
1403 let removed = match fs::symlink_metadata(&absolute) {
1404 Ok(metadata) if metadata.is_dir() => match fs::remove_dir(&absolute) {
1405 Err(error) if error.kind() == std::io::ErrorKind::DirectoryNotEmpty => {
1406 Ok(())
1407 }
1408 removed => removed,
1409 },
1410 Ok(_) => fs::remove_file(&absolute),
1411 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
1412 Err(error) => Err(error),
1413 };
1414 removed.with_context(|| format!("removing {}", absolute.display()))?;
1415 }
1416 }
1417 let notes = restored_files_path(workspace, warm);
1418 match fs::remove_file(¬es) {
1419 Err(error) if error.kind() != std::io::ErrorKind::NotFound => return Err(error.into()),
1420 _ => {}
1421 }
1422 self.cleared.lock().unwrap().insert(identity);
1423 Ok(())
1424 }
1425
1426 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;
1520 restored.source = source;
1521 let mut noted = read_restored_files(workspace, warm);
1522 for (location, kept) in missing {
1523 let Some(group) = record.groups.get(location.name) else {
1524 continue;
1525 };
1526 let (files, bytes) =
1527 match self.restore_location(&task.id, location, &kept, group, keep_mtimes) {
1528 Ok(restored) => restored,
1529 Err(error) => {
1530 remove_restored(location, &kept);
1534 return Err(error);
1535 }
1536 };
1537 restored.files += files.len();
1538 restored.bytes += bytes;
1539 if kept.is_empty() {
1540 noted.groups.insert(location.name.to_owned(), files);
1541 } else {
1542 noted
1544 .groups
1545 .entry(location.name.to_owned())
1546 .or_default()
1547 .extend(files);
1548 }
1549 restored.groups.push(location.name.to_owned());
1550 }
1551 if !restored.groups.is_empty() {
1552 let path = restored_files_path(workspace, warm);
1553 fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
1554 fs::write(path, serde_json::to_vec(¬ed)?)?;
1555 }
1556 Ok((restored, present))
1557 }
1558
1559 fn restore_location(
1562 &self,
1563 task: &str,
1564 location: &Location,
1565 kept: &[&Outputs],
1566 group: &WarmGroup,
1567 keep_mtimes: bool,
1568 ) -> Result<(GroupFiles, u64)> {
1569 fs::create_dir_all(&location.base)?;
1570 let mut laying = profile::span(task, "warm_restore_layout");
1571 let mut created = BTreeSet::new();
1574 let mut directories = Vec::new();
1575 let mut pending = Vec::new();
1576 for (path, artifact) in &group.artifacts {
1577 paths::safe_parents(&location.base, path)?;
1578 if !location.matches(path) {
1579 bail!("warm state for {task} holds {path}, outside its paths");
1580 }
1581 if kept.iter().any(|outputs| outputs.matches(path)) {
1582 continue;
1583 }
1584 let destination = location.base.join(path);
1585 if let Some(parent) = destination.parent()
1586 && created.insert(parent.to_owned())
1587 {
1588 fs::create_dir_all(parent)?;
1589 }
1590 match artifact {
1591 Artifact::File { blob, mode } => {
1592 let modified = keep_mtimes
1593 .then(|| group.mtimes.get(path))
1594 .flatten()
1595 .and_then(|nanos| u64::try_from(*nanos).ok())
1596 .map_or(RESTORED_AT, |nanos| {
1597 UNIX_EPOCH + Duration::from_nanos(nanos)
1598 });
1599 pending.push((path, blob.as_str(), *mode, destination, modified));
1600 }
1601 Artifact::Directory { mode } => {
1602 fs::create_dir_all(&destination)?;
1603 created.insert(destination.clone());
1604 directories.push((destination, *mode));
1605 }
1606 Artifact::Symlink { target, directory } => {
1607 validate_link(path, target)?;
1608 symlink(target, &destination, *directory)?;
1609 }
1610 }
1611 }
1612 laying.count("directories", created.len() as u64);
1613 drop(laying);
1614 let mut copying = profile::span(task, "warm_restore_copy");
1615 copying.count("files", pending.len() as u64);
1616 let copied = self
1617 .restore_warm_files(&pending)
1618 .with_context(|| format!("warm state for {task} not restored"))?;
1619 copying.count("bytes", copied.iter().map(|(_, size)| size).sum());
1620 drop(copying);
1621 let _modes = profile::span(task, "warm_restore_modes");
1623 for (destination, mode) in directories {
1624 set_mode(&destination, mode)?;
1625 }
1626 let mut files = BTreeMap::new();
1627 let mut bytes = 0;
1628 for ((path, blob, mode, _, _), (metadata, size)) in pending.into_iter().zip(copied) {
1629 bytes += size;
1630 let artifact = Artifact::File {
1631 blob: blob.to_owned(),
1632 mode,
1633 };
1634 files.insert(path.clone(), (metadata, artifact));
1635 }
1636 Ok((files, bytes))
1637 }
1638
1639 fn restore_warm_files(&self, files: &[PendingFile]) -> Result<Vec<RestoredFile>> {
1642 let restore =
1643 |(_, blob, mode, destination, modified): &PendingFile| -> Result<RestoredFile> {
1644 let source = self.root.join("blobs").join(blob);
1645 if !restore_blob(&source, destination, blob, self.verified())? {
1646 let _ = fs::remove_file(destination);
1648 bail!("corrupt blob {blob}");
1649 }
1650 set_modified(destination, *modified)?;
1652 set_mode(destination, *mode)?;
1653 let metadata = fs::symlink_metadata(destination)?;
1654 Ok((stamp(&metadata), metadata.len()))
1655 };
1656 let workers = files.len().div_ceil(64).clamp(1, RESTORERS);
1657 if workers == 1 {
1658 return files.iter().map(restore).collect();
1659 }
1660 let next = AtomicUsize::new(0);
1663 let failed = AtomicBool::new(false);
1664 let mut results: Vec<(usize, Result<RestoredFile>)> = std::thread::scope(|scope| {
1665 let handles: Vec<_> = (0..workers)
1666 .map(|_| {
1667 scope.spawn(|| {
1668 let mut done = Vec::new();
1669 while !failed.load(Ordering::Relaxed) {
1670 let index = next.fetch_add(1, Ordering::Relaxed);
1671 let Some(file) = files.get(index) else { break };
1672 let result = restore(file);
1673 if result.is_err() {
1674 failed.store(true, Ordering::Relaxed);
1675 }
1676 done.push((index, result));
1677 }
1678 done
1679 })
1680 })
1681 .collect();
1682 handles
1683 .into_iter()
1684 .flat_map(|handle| handle.join().expect("warm restore worker panicked"))
1685 .collect()
1686 });
1687 results.sort_by_key(|(index, _)| *index);
1688 results.into_iter().map(|(_, result)| result).collect()
1689 }
1690
1691 fn remote_warm(
1694 &self,
1695 workspace: &Workspace,
1696 task: &Task,
1697 warm: &Warm,
1698 current: &[String],
1699 ) -> Option<(WarmRecord, String)> {
1700 let remote = self.remote.as_ref()?;
1701 let own = worktree(workspace);
1702 remote_warm(
1703 remote,
1704 &self.root,
1705 &task.id,
1706 warm,
1707 &own,
1708 &branches(workspace),
1709 current,
1710 )
1711 }
1712
1713 pub fn prefetch_warm(
1722 &self,
1723 workspace: &Workspace,
1724 graph: &TaskGraph,
1725 prepared: &BTreeMap<String, qk_executor::PreparedTask>,
1726 restores: impl Fn(&str) -> bool,
1727 ) {
1728 let Some(remote) = self.remote.as_ref().filter(|_| self.reuse && !self.cold) else {
1730 return;
1731 };
1732 if self.initialize().is_err() {
1733 return;
1734 }
1735 let own = worktree(workspace);
1736 let mut items = Vec::new();
1737 for (id, task) in &graph.tasks {
1738 let Some(prepared) = prepared.get(id).filter(|_| restores(id)) else {
1739 continue;
1740 };
1741 let Ok(entries) = config(workspace, task) else {
1742 continue;
1743 };
1744 if !entries.iter().any(|warm| warm.prefetch) {
1745 continue;
1746 }
1747 let mut running = prepared.clone();
1751 for warm in &entries {
1752 for (name, value) in &warm.env {
1753 running.execution.insert(name.into(), value.into());
1754 }
1755 }
1756 for warm in entries {
1757 if !warm.prefetch || !warm.remote || warm.portable == Portable::Own {
1758 continue;
1759 }
1760 let missing = locations(workspace, task, &warm).is_ok_and(|locations| {
1761 locations.iter().any(|location| {
1762 location.paths().map_or(true, |paths| paths.is_empty())
1763 || location.present().is_ok_and(|kept| kept.is_some())
1764 })
1765 });
1766 let Ok(current) = key_digests(workspace, &warm, &running) else {
1767 continue;
1768 };
1769 if !missing || self.choose_warm(workspace, &warm, ¤t, true).is_some() {
1770 continue;
1771 }
1772 let depth = depth(graph, id, &mut BTreeMap::new());
1773 items.push((depth, warm.identity.clone(), (id.clone(), warm, current)));
1774 }
1775 }
1776 if items.is_empty() {
1777 return;
1778 }
1779 items.sort_by(|a, b| (a.0, &(a.2).0).cmp(&(b.0, &(b.2).0)));
1780 let items = items
1781 .into_iter()
1782 .map(|(_, identity, job)| {
1783 let key = job.2.clone();
1784 (identity, key, job)
1785 })
1786 .collect();
1787 let (remote, root) = (remote.clone(), self.root.clone());
1788 let branches = branches(workspace);
1789 let prefetches = Prefetches::start(items, PREFETCHERS, move |(id, warm, current)| {
1790 let mut fetching = profile::span(&id, "warm_prefetch");
1791 let fetched = remote_warm(&remote, &root, &id, &warm, &own, &branches, ¤t);
1792 fetching.note(
1793 "branch",
1794 fetched
1795 .as_ref()
1796 .map_or("none", |(_, branch)| branch.as_str()),
1797 );
1798 fetched
1799 });
1800 let _ = self.prefetches.set(prefetches);
1801 }
1802
1803 pub fn abandon_prefetches(&self) {
1806 if let Some(prefetches) = self.prefetches.get() {
1807 prefetches.stop(false);
1808 }
1809 }
1810
1811 pub(crate) fn finish_prefetches(&self) {
1814 if let Some(prefetches) = self.prefetches.get() {
1815 prefetches.stop(true);
1816 }
1817 }
1818
1819 pub(crate) fn save_warm(
1822 &self,
1823 workspace: &Workspace,
1824 task: &Task,
1825 warm: &Warm,
1826 prepared: &qk_executor::PreparedTask,
1827 ) -> Result<()> {
1828 self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
1829 }
1830
1831 pub(crate) fn save_warm_in_background(
1834 &self,
1835 workspace: &Workspace,
1836 task: &Task,
1837 warm: &Warm,
1838 prepared: &qk_executor::PreparedTask,
1839 ) -> Result<()> {
1840 let save = Save::prepare(workspace, task, warm, prepared)?;
1841 let cache = Cache {
1842 remote: self.remote.clone(),
1843 ..Cache::new(self.root.clone())
1844 };
1845 let handle = std::thread::spawn(move || {
1846 if let Err(error) = cache.save_prepared(&save) {
1847 qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
1848 }
1849 });
1850 self.saves.lock().unwrap().push(handle);
1851 Ok(())
1852 }
1853
1854 fn save_prepared(&self, save: &Save) -> Result<()> {
1859 let warm = &save.warm;
1860 let _lock = self.lock_warm(warm)?;
1861 let _save = profile::span(&save.task, "warm_save");
1862 let records = self.warm_records(warm);
1863 let previous = records
1865 .iter()
1866 .filter(|(_, record)| record.worktree == save.worktree)
1867 .min_by_key(|(_, record)| record.key != save.key)
1868 .map(|(_, record)| record);
1869 let noted: RestoredFiles = fs::read(&save.notes)
1870 .ok()
1871 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1872 .unwrap_or_default();
1873 let mut record = WarmRecord::new(
1874 warm.identity.clone(),
1875 save.worktree.clone(),
1876 save.commit.clone(),
1877 now_ms(),
1878 save.key.clone(),
1879 );
1880 let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
1881 for location in &save.locations {
1882 if !location.base.exists() {
1883 continue;
1884 }
1885 let known = previous.and_then(|previous| previous.groups.get(location.name));
1886 let restored_files = noted.groups.get(location.name);
1887 let mut group = WarmGroup::default();
1888 let mut pending = Vec::new();
1889 for path in location.paths()? {
1890 paths::safe_parents(&location.base, &path)?;
1891 let absolute = location.base.join(&path);
1892 let metadata = match fs::symlink_metadata(&absolute) {
1893 Err(error) if gone(&error) => continue,
1894 result => result?,
1895 };
1896 let artifact = if metadata.file_type().is_symlink() {
1897 let target = match fs::read_link(&absolute) {
1898 Err(error) if gone(&error) => continue,
1899 result => result?,
1900 };
1901 let target = target
1902 .to_str()
1903 .context("symlink target must be UTF-8")?
1904 .to_owned();
1905 validate_link(&path, &target)?;
1906 Artifact::Symlink {
1907 target,
1908 directory: absolute.is_dir(),
1909 }
1910 } else if metadata.is_file() {
1911 let now = stamp(&metadata);
1912 let saved = known.and_then(|known| {
1913 (known.stamps.get(&path) == Some(&now))
1914 .then(|| known.artifacts.get(&path))
1915 .flatten()
1916 });
1917 let restored = restored_files
1918 .and_then(|files| files.get(&path))
1919 .filter(|(stamp, _)| *stamp == now)
1920 .map(|(_, artifact)| artifact);
1921 let reused = saved.or(restored).filter(|artifact| match artifact {
1922 Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
1923 _ => false,
1924 });
1925 let Some(artifact) = reused.cloned() else {
1926 pending.push((path, absolute, metadata));
1927 continue;
1928 };
1929 group.stamps.insert(path.clone(), now);
1930 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1931 group.mtimes.insert(path.clone(), modified);
1932 }
1933 artifact
1934 } else if metadata.is_dir() {
1935 Artifact::Directory {
1936 mode: mode(&metadata),
1937 }
1938 } else {
1939 continue;
1940 };
1941 group.artifacts.insert(path, artifact);
1942 }
1943 let storing = profile::span(&save.task, "warm_save_blobs");
1944 let sources: Vec<_> = pending
1945 .iter()
1946 .map(|(_, absolute, _)| absolute.clone())
1947 .collect();
1948 for ((path, _, metadata), result) in pending.into_iter().zip(self.put_blobs(&sources)) {
1949 let blob = match result {
1950 Ok(blob) => blob,
1951 Err(error) if error.downcast_ref::<std::io::Error>().is_some_and(gone) => {
1952 continue;
1953 }
1954 Err(error) => return Err(error),
1955 };
1956 group.stamps.insert(path.clone(), stamp(&metadata));
1957 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1958 group.mtimes.insert(path.clone(), modified);
1959 }
1960 group.artifacts.insert(
1961 path,
1962 Artifact::File {
1963 blob,
1964 mode: mode(&metadata),
1965 },
1966 );
1967 }
1968 drop(storing);
1969 if let Some(limit) = warm.max_size
1970 && group.size(self) > limit
1971 {
1972 qk_executor::status!(
1973 "qk: {}: warm {} is over its maxSize and was not saved",
1974 save.task,
1975 location.name
1976 );
1977 continue;
1978 }
1979 record.groups.insert(location.name.to_owned(), group);
1980 }
1981 self.store_warm(&record)?;
1982 let _ = fs::remove_file(&save.notes);
1983 let saved = self.warm_record(&record);
1986 for (path, _) in records
1987 .into_iter()
1988 .filter(|(path, _)| *path != saved)
1989 .skip(KEPT - 1)
1990 {
1991 let _ = fs::remove_file(path);
1992 }
1993 if warm.remote
1995 && warm.portable != Portable::Own
1996 && let Some(remote) = &self.remote
1997 && let Some(branch) = &save.branch
1998 {
1999 for group in record.groups.values_mut() {
2000 group.stamps.clear();
2001 }
2002 remote.upload_warm(
2003 &self.root,
2004 &warm.identity,
2005 &remote_identity(warm, &save.worktree),
2006 branch,
2007 serde_json::to_vec(&record)?,
2008 );
2009 }
2010 Ok(())
2011 }
2012}
2013
2014struct Save {
2017 task: String,
2018 warm: Warm,
2019 locations: Vec<Location>,
2020 worktree: String,
2021 commit: Option<String>,
2022 branch: Option<String>,
2023 key: Vec<String>,
2024 notes: PathBuf,
2025}
2026
2027impl Save {
2028 fn prepare(
2030 workspace: &Workspace,
2031 task: &Task,
2032 warm: &Warm,
2033 prepared: &qk_executor::PreparedTask,
2034 ) -> Result<Self> {
2035 Ok(Self {
2036 task: task.id.clone(),
2037 warm: warm.clone(),
2038 locations: locations(workspace, task, warm)?,
2039 worktree: worktree(workspace),
2040 commit: current_commit(workspace),
2041 branch: current_branch(workspace),
2042 key: key_digests(workspace, warm, prepared)?,
2043 notes: restored_files_path(workspace, warm),
2044 })
2045 }
2046}
2047
2048#[cfg(test)]
2049mod tests {
2050 use super::*;
2051 use qk_taskgraph::Request;
2052 use serde_json::json;
2053
2054 fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
2056 let temp = tempfile::tempdir().unwrap();
2057 fs::write(temp.path().join("nx.json"), "{}").unwrap();
2058 fs::write(
2059 temp.path().join("project.json"),
2060 json!({
2061 "name": "app", "targets": {
2062 "one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
2063 "two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
2064 }
2065 })
2066 .to_string(),
2067 )
2068 .unwrap();
2069 for path in ["scratch-one", "scratch-two"] {
2070 fs::create_dir(temp.path().join(path)).unwrap();
2071 fs::write(temp.path().join(path).join("state"), "saved").unwrap();
2072 }
2073 let workspace = Workspace::load(temp.path()).unwrap();
2074 let graph = TaskGraph::build(
2075 &workspace,
2076 &[
2077 Request::parse("app:one").unwrap(),
2078 Request::parse("app:two").unwrap(),
2079 ],
2080 )
2081 .unwrap();
2082 (temp, workspace, graph)
2083 }
2084
2085 #[test]
2088 fn asks_a_partial_clone_about_ancestors_without_fetching() {
2089 let temp = tempfile::tempdir().unwrap();
2090 let git = |root: &std::path::Path, args: &[&str]| {
2091 let output = paths::git(root, args).unwrap();
2092 assert!(output.status.success(), "git {args:?}: {output:?}");
2093 String::from_utf8(output.stdout).unwrap().trim().to_owned()
2094 };
2095 let source = temp.path().join("source");
2096 fs::create_dir(&source).unwrap();
2097 git(&source, &["init", "--quiet"]);
2098 git(&source, &["config", "uploadpack.allowFilter", "true"]);
2099 let commit = |message: &str| {
2100 fs::write(source.join("file"), message).unwrap();
2101 git(&source, &["add", "file"]);
2102 git(
2103 &source,
2104 &[
2105 "-c",
2106 "user.name=qk",
2107 "-c",
2108 "user.email=qk@example.com",
2109 "commit",
2110 "--quiet",
2111 "-m",
2112 message,
2113 ],
2114 );
2115 git(&source, &["rev-parse", "HEAD"])
2116 };
2117 let older = commit("older");
2118 let newer = commit("newer");
2119 let clone = temp.path().join("clone");
2120 let url = format!("file://{}", source.display());
2121 git(
2122 temp.path(),
2123 &[
2124 "clone",
2125 "--quiet",
2126 "--depth",
2127 "1",
2128 "--filter=blob:none",
2129 &url,
2130 clone.to_str().unwrap(),
2131 ],
2132 );
2133 fs::write(clone.join("nx.json"), "{}").unwrap();
2134 let workspace = Workspace::load(&clone).unwrap();
2135 assert!(descends_from(&workspace, &newer));
2136 assert!(!descends_from(&workspace, &older));
2137 assert_eq!(behind(&workspace, &older), None);
2138 let version = git(temp.path(), &["--version"]);
2140 let honoured = version
2141 .split_whitespace()
2142 .nth(2)
2143 .and_then(|version| {
2144 let mut parts = version.split('.').map(|part| part.parse::<u32>().ok());
2145 Some((parts.next()??, parts.next()??))
2146 })
2147 .is_some_and(|version| version >= (2, 44));
2148 let held = paths::git_history(&clone, &["cat-file", "-e", &format!("{older}^{{commit}}")])
2149 .unwrap();
2150 assert!(!honoured || !held.status.success(), "fetched {older}");
2151 }
2152
2153 fn entries(warm: Value) -> Result<Vec<Warm>> {
2155 entries_in(warm, None)
2156 }
2157
2158 fn entries_in(warm: Value, configuration: Option<&str>) -> Result<Vec<Warm>> {
2161 let temp = tempfile::tempdir().unwrap();
2162 fs::write(temp.path().join("nx.json"), "{}").unwrap();
2163 fs::write(
2164 temp.path().join("project.json"),
2165 json!({"name": "app", "targets": {"build": {
2166 "command": "true",
2167 "qk:warm": warm,
2168 "configurations": {"kept": {}, "skipped": {}}
2169 }}})
2170 .to_string(),
2171 )
2172 .unwrap();
2173 let workspace = Workspace::load(temp.path()).unwrap();
2174 let id = configuration.map_or_else(
2175 || "app:build".to_owned(),
2176 |name| format!("app:build:{name}"),
2177 );
2178 let graph = TaskGraph::build(&workspace, &[Request::parse(&id).unwrap()]).unwrap();
2179 config(&workspace, &graph.tasks[&id])
2180 }
2181
2182 #[test]
2183 fn entries_conflict_only_where_a_configuration_selects_both() {
2184 let apart = json!([
2186 {"paths": ["one"], "env": {"CACHE": "{warm}/one"}, "configurations": ["kept"]},
2187 {"paths": ["two"], "env": {"CACHE": "{warm}/two"}, "configurations": ["skipped"]}
2188 ]);
2189 let kept = entries_in(apart.clone(), Some("kept")).unwrap();
2190 assert_eq!(kept.len(), 1);
2191 assert!(kept[0].env["CACHE"].ends_with("/one"));
2192 let skipped = entries_in(apart.clone(), Some("skipped")).unwrap();
2193 assert_eq!(skipped.len(), 1);
2194 assert!(skipped[0].env["CACHE"].ends_with("/two"));
2195 assert!(entries_in(apart, None).unwrap().is_empty());
2196 let both = json!([
2198 {"paths": ["one"], "configurations": ["kept"]},
2199 {"paths": ["two"], "configurations": ["kept", "skipped"]}
2200 ]);
2201 assert_eq!(entries_in(both.clone(), Some("skipped")).unwrap().len(), 1);
2202 assert_eq!(
2203 format!("{:#}", entries_in(both, Some("kept")).unwrap_err()),
2204 "only one qk:warm entry may go without a group"
2205 );
2206 let shared = json!([
2208 {"group": "one", "env": {"CACHE": "{warm}"}, "configurations": ["kept"]},
2209 {"group": "two", "env": {"CACHE": "{warm}"}, "configurations": ["kept", "skipped"]}
2210 ]);
2211 assert_eq!(
2212 entries_in(shared.clone(), Some("skipped")).unwrap().len(),
2213 1
2214 );
2215 assert_eq!(
2216 format!("{:#}", entries_in(shared, Some("kept")).unwrap_err()),
2217 "qk:warm sets CACHE in more than one entry"
2218 );
2219 }
2220
2221 #[test]
2222 fn an_array_holds_the_task_entry_beside_groups() {
2223 let entries = entries(json!([
2224 {"group": "metro", "portable": true, "env": {"METRO_CACHE_ROOT": "{warm}/metro"}},
2225 {"paths": ["build"], "mtimes": "preserve"}
2226 ]))
2227 .unwrap();
2228 assert_eq!(entries.len(), 2);
2229 assert_eq!(entries[0].group(), Some("metro"));
2230 assert!(entries[0].portable == Portable::Anywhere && entries[0].directory);
2231 assert_eq!(entries[1].group(), None);
2232 assert!(entries[1].portable == Portable::Own && entries[1].preserve_mtimes);
2233 assert_eq!(kept_paths(&entries), ["build"]);
2234 }
2235
2236 #[test]
2237 fn portable_is_a_boolean_or_the_same_path() {
2238 let portable = |value| {
2239 entries(json!({"paths": ["build"], "portable": value}))
2240 .map(|entries| entries[0].portable)
2241 };
2242 assert_eq!(portable(json!("path")).unwrap(), Portable::Path);
2243 assert_eq!(portable(json!(false)).unwrap(), Portable::Own);
2244 assert_eq!(
2245 format!("{:#}", portable(json!("anywhere")).unwrap_err()),
2246 "qk:warm.portable must be true, false or \"path\""
2247 );
2248 }
2249
2250 #[test]
2251 fn entries_cannot_share_saves_or_variables() {
2252 let error = |warm| format!("{:#}", entries(warm).unwrap_err());
2253 assert_eq!(
2254 error(json!([{"paths": ["one"]}, {"paths": ["two"]}])),
2255 "only one qk:warm entry may go without a group"
2256 );
2257 assert_eq!(
2258 error(json!([{"group": "metro"}, {"group": "metro"}])),
2259 "qk:warm names group \"metro\" twice"
2260 );
2261 assert_eq!(
2262 error(json!([
2263 {"group": "one", "env": {"CACHE": "{warm}/one"}},
2264 {"group": "two", "env": {"CACHE": "{warm}/two"}}
2265 ])),
2266 "qk:warm sets CACHE in more than one entry"
2267 );
2268 assert_eq!(
2269 error(json!(["build"])),
2270 "qk:warm must be an object or an array of objects"
2271 );
2272 }
2273
2274 #[test]
2276 fn waiting_for_an_owner_does_not_block_other_releases() {
2277 let (_temp, workspace, graph) = fixture();
2278 let one = &graph.tasks["app:one"];
2279 let two = &graph.tasks["app:two"];
2280 let first = acquire(&workspace, one, &config(&workspace, one).unwrap()[0]).unwrap();
2281 let state = paths::worktree_state(&workspace.root).join("warm-kept");
2282 let lock = fs::OpenOptions::new()
2283 .create(true)
2284 .truncate(false)
2285 .read(true)
2286 .write(true)
2287 .open(state.join(format!("{}.lock", hash32(&two.id))))
2288 .unwrap();
2289 fs2::FileExt::lock_exclusive(&lock).unwrap();
2290 let (released, received) = std::sync::mpsc::channel();
2291 let completed = std::thread::scope(|scope| {
2292 scope.spawn(|| {
2293 let second =
2294 acquire(&workspace, two, &config(&workspace, two).unwrap()[0]).unwrap();
2295 release(&second);
2296 });
2297 std::thread::sleep(Duration::from_millis(100));
2298 scope.spawn(|| {
2299 release(&first);
2300 released.send(()).unwrap();
2301 });
2302 let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
2303 fs2::FileExt::unlock(&lock).unwrap();
2304 completed
2305 });
2306 assert!(
2307 completed,
2308 "a blocked owner lock prevented an unrelated release"
2309 );
2310 assert_eq!(
2311 fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
2312 "saved"
2313 );
2314 }
2315
2316 #[test]
2318 fn unreadable_ownership_preserves_saved_paths() {
2319 let (_temp, workspace, graph) = fixture();
2320 let owner = &graph.tasks["app:one"];
2321 let warm = config(&workspace, owner).unwrap().remove(0);
2322 let root = acquire(&workspace, owner, &warm).unwrap();
2323 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
2324 drop(entry); fs::write(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::remove_file(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 fs::create_dir(Stash::manifest(&root)).unwrap();
2332 assert!(acquire(&workspace, owner, &warm).is_err());
2333 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2334 }
2335
2336 #[cfg(unix)]
2338 #[test]
2339 fn inspection_errors_keep_recovery_incomplete() {
2340 let (_temp, workspace, graph) = fixture();
2341 let owner = &graph.tasks["app:one"];
2342 let warm = config(&workspace, owner).unwrap().remove(0);
2343 let root = acquire(&workspace, owner, &warm).unwrap();
2344 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
2345 let saved = root.with_extension("saved");
2346 fs::rename(&root, &saved).unwrap();
2347 fs::write(&root, "blocked").unwrap(); assert!(!entry.stash.put_back());
2349 let retained = Stash {
2350 root: root.clone(),
2351 workspace: workspace.root.clone(),
2352 moved: entry.stash.moved.clone(),
2353 };
2354 assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
2355 assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
2356 assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
2357 }
2358}