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(crate) fn shared_remotely(&self) -> bool {
434 self.remote && self.portable != Portable::Own
435 }
436
437 pub fn group(&self) -> Option<&str> {
439 self.identity.strip_prefix(GROUP)
440 }
441
442 pub fn kept_paths(&self) -> Vec<String> {
445 self.paths
446 .iter()
447 .filter(|path| !path.starts_with('!'))
448 .cloned()
449 .collect()
450 }
451}
452
453pub fn kept_paths(entries: &[Warm]) -> Vec<String> {
455 entries.iter().flat_map(Warm::kept_paths).collect()
456}
457
458#[derive(Default)]
461pub struct Kept {
462 roots: Vec<PathBuf>,
464}
465
466impl Drop for Kept {
467 fn drop(&mut self) {
468 for root in self.roots.iter().rev() {
469 release(root);
470 }
471 }
472}
473
474#[derive(Default, Serialize, Deserialize)]
476struct Stash {
477 #[serde(skip)]
478 root: PathBuf,
479 #[serde(skip)]
480 workspace: PathBuf,
481 moved: Vec<String>,
483}
484
485struct Active {
488 users: usize,
489 stash: Stash,
490 _lock: fs::File,
491}
492
493static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
496
497impl Stash {
498 fn manifest(root: &std::path::Path) -> PathBuf {
500 root.join("moved.json")
501 }
502
503 fn save_manifest(&self) -> Result<()> {
506 let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
507 file.write_all(&serde_json::to_vec(self)?)?;
508 file.as_file().sync_all()?;
509 file.persist(Self::manifest(&self.root))?;
510 Ok(())
511 }
512
513 fn put_back(&self) -> bool {
519 let mut complete = true;
520 for (index, path) in self.moved.iter().enumerate().rev() {
521 let kept = self.root.join(index.to_string());
522 match fs::symlink_metadata(&kept) {
523 Ok(_) => {}
524 Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
525 Err(error) => {
526 complete = false;
527 qk_executor::status!(
528 "qk: could not inspect saved {path}, kept for the next run ({error})"
529 );
530 continue;
531 }
532 }
533 let original = self.workspace.join(path);
534 let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
535 remove(&original)?;
536 if let Some(parent) = original.parent() {
537 fs::create_dir_all(parent)?;
538 }
539 fs::rename(&kept, &original)?;
540 Ok(())
541 });
542 if let Err(error) = result {
543 complete = false;
544 qk_executor::status!(
545 "qk: could not put {path} back, kept for the next run ({error:#})"
546 );
547 }
548 }
549 if complete {
550 let _ = fs::remove_dir_all(&self.root);
551 }
552 complete
553 }
554}
555
556fn remove(path: &std::path::Path) -> std::io::Result<()> {
558 match fs::symlink_metadata(path) {
559 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
560 Err(error) => Err(error),
561 Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
562 Ok(_) => fs::remove_file(path),
563 }
564}
565
566fn release(root: &std::path::Path) {
569 let mut active = ACTIVE
570 .lock()
571 .unwrap_or_else(std::sync::PoisonError::into_inner);
572 let Some(entry) = active.get_mut(root) else {
573 return;
574 };
575 entry.users -= 1;
576 if entry.users > 0 {
577 return;
578 }
579 let entry = active.remove(root).expect("present above");
580 drop(active);
583 entry.stash.put_back();
584}
585
586fn names(owner: &Task, entry: &str, task: &Task) -> bool {
588 match entry.split_once(':') {
589 Some((project, target)) => task.project == project && task.target == target,
590 None => task.project == owner.project && task.target == entry,
591 }
592}
593
594fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
596 let mut seen = std::collections::BTreeSet::new();
597 let mut next = vec![from];
598 while let Some(id) = next.pop() {
599 for dependency in graph
600 .tasks
601 .get(id)
602 .into_iter()
603 .flat_map(|task| &task.dependencies)
604 {
605 if dependency == on {
606 return true;
607 }
608 if seen.insert(dependency.as_str()) {
609 next.push(dependency);
610 }
611 }
612 }
613 false
614}
615
616pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
619 let mut kept = Kept::default();
620 for owner in graph.tasks.values() {
621 let Ok(entries) = config(workspace, owner) else {
622 continue;
623 };
624 let Some(warm) = entries.into_iter().find(|warm| warm.group().is_none()) else {
626 continue;
627 };
628 if !warm.survive.iter().any(|entry| names(owner, entry, task))
629 || !depends_on(graph, &owner.id, &task.id)
630 {
631 continue;
632 }
633 match acquire(workspace, owner, &warm) {
634 Ok(root) => kept.roots.push(root),
635 Err(error) => qk_executor::status!(
636 "qk: {}: warm paths not kept across {} ({error:#})",
637 owner.id,
638 task.id
639 ),
640 }
641 }
642 kept
643}
644
645fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
648 let state = paths::worktree_state(&workspace.root).join("warm-kept");
649 let root = state.join(hash32(&owner.id));
650 fs::create_dir_all(&state)?;
651 let lock = fs::OpenOptions::new()
652 .create(true)
653 .truncate(false)
654 .read(true)
655 .write(true)
656 .open(state.join(format!("{}.lock", hash32(&owner.id))))?;
657 let mut active = loop {
660 let mut active = ACTIVE
661 .lock()
662 .unwrap_or_else(std::sync::PoisonError::into_inner);
663 if let Some(entry) = active.get_mut(&root) {
664 entry.users += 1;
665 return Ok(root);
666 }
667 match fs2::FileExt::try_lock_exclusive(&lock) {
668 Ok(()) => break active,
669 Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
670 drop(active);
671 std::thread::sleep(Duration::from_millis(20));
672 }
673 Err(error) => return Err(error.into()),
674 }
675 };
676 let left = match fs::read(Stash::manifest(&root)) {
683 Ok(bytes) => {
684 let mut left: Stash = serde_json::from_slice(&bytes)
685 .context("invalid warm stash manifest; saved paths retained for recovery")?;
686 left.root = root.clone();
687 left.workspace = workspace.root.clone();
688 Some(left)
689 }
690 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
691 if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
692 bail!("missing warm stash manifest; saved paths retained for recovery");
693 }
694 None
695 }
696 Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
697 };
698 let retained = left.filter(|left| !left.put_back());
699 let stash = stash(workspace, &root, warm, retained)?;
700 active.insert(
701 root.clone(),
702 Active {
703 users: 1,
704 stash,
705 _lock: lock,
706 },
707 );
708 Ok(root)
709}
710
711fn stash(
715 workspace: &Workspace,
716 root: &std::path::Path,
717 warm: &Warm,
718 retained: Option<Stash>,
719) -> Result<Stash> {
720 let mut stash = match retained {
721 Some(retained) => retained,
722 None => {
723 let _ = fs::remove_dir_all(root);
724 Stash {
725 root: root.to_owned(),
726 workspace: workspace.root.clone(),
727 moved: Vec::new(),
728 }
729 }
730 };
731 let held: Vec<String> = stash
732 .moved
733 .iter()
734 .enumerate()
735 .filter_map(
736 |(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
737 Ok(_) => Some(Ok(path.clone())),
738 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
739 Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
740 },
741 )
742 .collect::<Result<_>>()?;
743 let overlaps = |path: &str| {
744 held.iter().any(|held| {
745 held == path
746 || path.starts_with(&format!("{held}/"))
747 || held.starts_with(&format!("{path}/"))
748 })
749 };
750 let mut kept = std::collections::BTreeSet::new();
753 for path in warm.kept_paths() {
754 if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
755 .and_then(|outputs| outputs.paths(&workspace.root))
756 {
757 kept.extend(found);
758 }
759 }
760 let topmost: Vec<String> = kept
762 .iter()
763 .filter(|path| {
764 !std::path::Path::new(path.as_str())
765 .ancestors()
766 .skip(1)
767 .any(|ancestor| {
768 ancestor
769 .to_str()
770 .is_some_and(|ancestor| kept.contains(ancestor))
771 })
772 })
773 .filter(|path| !overlaps(path))
774 .cloned()
775 .collect();
776 if topmost.is_empty() {
777 return Ok(stash);
778 }
779 fs::create_dir_all(root)?;
780 for path in topmost {
781 let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
782 stash.moved.push(path.clone());
783 stash.save_manifest()?;
784 fs::rename(
787 workspace.root.join(&path),
788 root.join((stash.moved.len() - 1).to_string()),
789 )
790 .inspect_err(|_| {
791 stash.moved.pop();
792 })?;
793 Ok(())
794 });
795 if let Err(error) = moved {
796 stash.save_manifest()?;
799 stash.put_back();
800 return Err(error);
801 }
802 }
803 Ok(stash)
804}
805
806pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
809 let mut claimed: Vec<(String, &str)> = Vec::new();
810 for (id, task) in &graph.tasks {
811 for path in kept_paths(&config(workspace, task)?) {
812 if let Some((_, other)) = claimed.iter().find(|(known, _)| {
813 known == &path
814 || known.starts_with(&format!("{path}/"))
815 || path.starts_with(&format!("{known}/"))
816 }) {
817 bail!("{id} and {other} both keep {path} as warm state");
818 }
819 claimed.push((path, id));
820 }
821 }
822 Ok(())
823}
824
825impl WarmGroup {
826 fn size(&self, cache: &Cache) -> u64 {
828 self.artifacts
829 .values()
830 .filter_map(|artifact| match artifact {
831 Artifact::File { blob, .. } => {
832 fs::metadata(cache.root.join("blobs").join(blob)).ok()
833 }
834 _ => None,
835 })
836 .map(|metadata| metadata.len())
837 .sum()
838 }
839}
840
841#[derive(Default, Serialize, Deserialize)]
844struct RestoredFiles {
845 groups: BTreeMap<String, GroupFiles>,
846}
847
848type GroupFiles = BTreeMap<String, (Vec<i64>, Artifact)>;
850
851#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
853#[serde(rename_all = "camelCase")]
854pub struct Restored {
855 #[serde(default, skip_serializing_if = "Option::is_none")]
857 pub group: Option<String>,
858 pub source: String,
861 pub groups: Vec<String>,
862 pub files: usize,
863 pub bytes: u64,
864 #[serde(default)]
867 pub ms: u64,
868 #[serde(default, skip_serializing_if = "Option::is_none")]
872 pub fetch_ms: Option<u64>,
873}
874
875#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
877#[serde(rename_all = "camelCase")]
878pub struct WarmReport {
879 #[serde(default, skip_serializing_if = "Vec::is_empty")]
881 pub restored: Vec<Restored>,
882 #[serde(default, skip_serializing_if = "Vec::is_empty")]
886 pub present: Vec<String>,
887 pub save_ms: Option<u64>,
890 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
892 pub background: bool,
893}
894
895#[derive(Clone)]
898struct Location {
899 name: &'static str,
900 base: PathBuf,
901 outputs: Option<Outputs>,
903 each: Vec<Outputs>,
906}
907
908impl Location {
909 fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
911 if !self.base.exists() {
912 return Ok(Default::default());
913 }
914 match &self.outputs {
915 Some(outputs) => outputs.paths(&self.base),
916 None => {
917 let mut paths = std::collections::BTreeSet::new();
918 for entry in walkdir::WalkDir::new(&self.base)
919 .min_depth(1)
920 .follow_links(false)
921 {
922 let entry = entry?;
923 paths.insert(paths::relative(&self.base, entry.path())?);
924 }
925 Ok(paths)
926 }
927 }
928 }
929
930 fn matches(&self, path: &str) -> bool {
932 self.outputs
933 .as_ref()
934 .is_none_or(|outputs| outputs.matches(path))
935 }
936
937 fn present(&self) -> Result<Option<Vec<&Outputs>>> {
940 let mut present = Vec::new();
941 for outputs in &self.each {
942 if !outputs.paths(&self.base)?.is_empty() {
943 present.push(outputs);
944 }
945 }
946 Ok((!present.is_empty() && present.len() < self.each.len()).then_some(present))
947 }
948}
949
950fn remove_restored(location: &Location, kept: &[&Outputs]) {
954 let Ok(paths) = location.paths() else {
955 return;
956 };
957 for path in paths.iter().rev() {
958 if kept.iter().any(|outputs| outputs.matches(path)) {
959 continue;
960 }
961 let absolute = location.base.join(path);
962 let _ = match fs::symlink_metadata(&absolute) {
963 Ok(metadata) if metadata.is_dir() => fs::remove_dir(&absolute),
964 Ok(_) => fs::remove_file(&absolute),
965 Err(_) => continue,
966 };
967 }
968}
969
970fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
972 let mut locations = Vec::new();
973 if warm.outputs {
974 locations.push(Location {
975 name: "outputs",
976 base: workspace.root.clone(),
977 outputs: Some(Outputs::new(workspace, task)?),
978 each: Vec::new(),
979 });
980 }
981 if !warm.paths.is_empty() {
982 let mut each = Vec::new();
983 if warm.each_path {
984 let (excluded, included): (Vec<_>, Vec<_>) =
985 warm.paths.iter().partition(|path| path.starts_with('!'));
986 for path in included {
987 let mut patterns = vec![path.clone()];
988 patterns.extend(excluded.iter().map(|path| (*path).clone()));
989 each.push(Outputs::from_paths(&patterns)?);
990 }
991 }
992 locations.push(Location {
993 name: "paths",
994 base: workspace.root.clone(),
995 outputs: Some(Outputs::from_paths(&warm.paths)?),
996 each,
997 });
998 }
999 if warm.directory {
1000 locations.push(Location {
1001 name: "directory",
1002 base: directory(workspace, &warm.identity),
1003 outputs: None,
1004 each: Vec::new(),
1005 });
1006 }
1007 Ok(locations)
1008}
1009
1010pub(crate) fn current_branch(workspace: &Workspace) -> Option<String> {
1013 for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
1014 if let Ok(value) = std::env::var(name)
1015 && !value.is_empty()
1016 {
1017 return Some(value);
1018 }
1019 }
1020 let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
1021 let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
1022 (output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
1023}
1024
1025pub(crate) fn branches(workspace: &Workspace) -> Vec<String> {
1028 let default = workspace
1029 .config
1030 .default_base
1031 .clone()
1032 .unwrap_or_else(|| "main".into());
1033 let default = default
1034 .strip_prefix("origin/")
1035 .unwrap_or(&default)
1036 .to_owned();
1037 let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
1038 if !branches.contains(&default) {
1039 branches.push(default);
1040 }
1041 branches
1042}
1043
1044fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
1046 let modified = metadata
1047 .modified()
1048 .ok()
1049 .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
1050 .map_or(-1, |duration| duration.as_nanos() as i64);
1051 #[allow(unused_mut)]
1052 let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
1053 #[cfg(unix)]
1054 {
1055 use std::os::unix::fs::MetadataExt;
1056 stamp.extend([
1057 metadata.ctime(),
1058 metadata.ctime_nsec(),
1059 metadata.ino() as i64,
1060 ]);
1061 }
1062 stamp
1063}
1064
1065fn remote_identity(warm: &Warm, worktree: &str) -> String {
1070 match warm.portable {
1071 Portable::Path => format!("{}\0path\0{worktree}", warm.identity),
1072 _ => warm.identity.clone(),
1073 }
1074}
1075
1076fn worktree(workspace: &Workspace) -> String {
1077 fs::canonicalize(&workspace.root)
1078 .unwrap_or_else(|_| workspace.root.clone())
1079 .to_string_lossy()
1080 .into_owned()
1081}
1082
1083fn current_commit(workspace: &Workspace) -> Option<String> {
1085 let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
1086 let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
1087 (output.status.success() && !commit.is_empty()).then_some(commit)
1088}
1089
1090fn key_digests(
1093 workspace: &Workspace,
1094 warm: &Warm,
1095 prepared: &qk_executor::PreparedTask,
1096) -> Result<Vec<String>> {
1097 warm.key
1098 .iter()
1099 .map(|part| {
1100 Ok(match part {
1101 KeyPart::File(path) => {
1102 paths::safe_parents(&workspace.root, path)?;
1103 let absolute = workspace.root.join(path);
1104 if absolute.is_file() {
1105 format!("file:{}", digest_file(&absolute)?)
1106 } else {
1107 "file:absent".to_owned()
1108 }
1109 }
1110 KeyPart::Env(name) => match prepared
1113 .execution
1114 .get(std::ffi::OsStr::new(name))
1115 .or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
1116 {
1117 Some(value) => format!(
1118 "env:{}",
1119 blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
1120 ),
1121 None => "env:unset".to_owned(),
1122 },
1123 })
1124 })
1125 .collect()
1126}
1127
1128fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
1132 if saved == current {
1133 return Some(0);
1134 }
1135 let count = warm.restore_keys?;
1136 (saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
1137}
1138
1139fn descends_from(workspace: &Workspace, commit: &str) -> bool {
1141 paths::git_history(
1142 &workspace.root,
1143 &["merge-base", "--is-ancestor", commit, "HEAD"],
1144 )
1145 .is_ok_and(|output| output.status.success())
1146}
1147
1148fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
1151 let ancestor = paths::git_history(
1152 &workspace.root,
1153 &["merge-base", "--is-ancestor", commit, "HEAD"],
1154 )
1155 .ok()?;
1156 if !ancestor.status.success() {
1157 return None;
1158 }
1159 let output = paths::git_history(
1160 &workspace.root,
1161 &["rev-list", "--count", &format!("{commit}..HEAD")],
1162 )
1163 .ok()?;
1164 String::from_utf8(output.stdout).ok()?.trim().parse().ok()
1165}
1166
1167fn depth(graph: &TaskGraph, id: &str, known: &mut BTreeMap<String, usize>) -> usize {
1169 if let Some(depth) = known.get(id) {
1170 return *depth;
1171 }
1172 let depth = graph.tasks.get(id).map_or(0, |task| {
1173 task.dependencies
1174 .iter()
1175 .map(|dependency| depth(graph, dependency, known) + 1)
1176 .max()
1177 .unwrap_or(0)
1178 });
1179 known.insert(id.to_owned(), depth);
1180 depth
1181}
1182
1183fn warm_prefix(identity: &str) -> String {
1185 hash32(&format!(
1186 "{identity}\0{}\0{}",
1187 std::env::consts::OS,
1188 std::env::consts::ARCH
1189 ))
1190}
1191
1192fn warm_record(root: &Path, record: &WarmRecord) -> PathBuf {
1193 root.join("warm").join(format!(
1194 "{}-{}.json",
1195 warm_prefix(&record.task),
1196 hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
1197 ))
1198}
1199
1200fn store_warm(root: &Path, record: &WarmRecord) -> Result<()> {
1201 let path = warm_record(root, record);
1202 let directory = path.parent().context("warm records have a directory")?;
1203 fs::create_dir_all(directory)?;
1204 record::publish(root, &path, record)
1205}
1206
1207fn remote_warm(
1210 remote: &crate::remote::Remote,
1211 root: &Path,
1212 task: &str,
1213 warm: &Warm,
1214 own: &str,
1215 branches: &[String],
1216 current: &[String],
1217) -> Option<(WarmRecord, String)> {
1218 if !warm.shared_remotely() {
1219 return None;
1220 }
1221 let placed = |saved: Option<&str>| warm.portable != Portable::Path || saved == Some(own);
1223 for branch in branches {
1224 let suits = |saved: &[String], worktree: Option<&str>| {
1225 key_match(warm, current, saved).is_some() && placed(worktree)
1226 };
1227 let name = remote_identity(warm, own);
1228 match remote.fetch_warm(root, &name, task, branch, suits) {
1229 Ok(Some((bytes, _pin))) => {
1230 let Some(record) = WarmRecord::read(&bytes, &warm.identity).ok() else {
1231 continue;
1232 };
1233 if key_match(warm, current, &record.key).is_none()
1236 || !placed(Some(&record.worktree))
1237 {
1238 continue;
1239 }
1240 if store_warm(root, &record).is_err() {
1241 return None;
1242 }
1243 return Some((record, branch.clone()));
1244 }
1245 Ok(None) => {}
1246 Err(error) => {
1247 qk_executor::status!("qk: {task}: remote warm state unavailable ({error:#})");
1248 continue;
1251 }
1252 }
1253 }
1254 None
1255}
1256
1257fn hash32(text: &str) -> String {
1259 blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
1260}
1261
1262fn nanos(time: SystemTime) -> Option<i64> {
1264 time.duration_since(UNIX_EPOCH)
1265 .ok()
1266 .and_then(|duration| i64::try_from(duration.as_nanos()).ok())
1267}
1268
1269fn now_ms() -> u64 {
1271 SystemTime::now()
1272 .duration_since(UNIX_EPOCH)
1273 .map_or(0, |duration| duration.as_millis() as u64)
1274}
1275
1276fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
1278 paths::worktree_state(&workspace.root)
1279 .join("warm-restored")
1280 .join(format!("{}.json", hash32(&warm.identity)))
1281}
1282
1283fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
1285 fs::read(restored_files_path(workspace, warm))
1286 .ok()
1287 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1288 .unwrap_or_default()
1289}
1290
1291impl Cache {
1292 fn warm_prefix(&self, identity: &str) -> String {
1294 warm_prefix(identity)
1295 }
1296
1297 fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
1300 let directory = self.root.join("locks");
1301 fs::create_dir_all(&directory)?;
1302 let file = fs::OpenOptions::new()
1303 .create(true)
1304 .truncate(false)
1305 .read(true)
1306 .write(true)
1307 .open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
1308 fs2::FileExt::lock_exclusive(&file)?;
1309 Ok(file)
1310 }
1311
1312 fn warm_record(&self, record: &WarmRecord) -> PathBuf {
1313 warm_record(&self.root, record)
1314 }
1315
1316 fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, WarmRecord)> {
1318 let prefix = format!("{}-", self.warm_prefix(&warm.identity));
1319 let mut records: Vec<(PathBuf, WarmRecord)> = fs::read_dir(self.root.join("warm"))
1320 .into_iter()
1321 .flatten()
1322 .flatten()
1323 .filter(|entry| {
1324 entry
1325 .file_name()
1326 .to_str()
1327 .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
1328 })
1329 .filter_map(|entry| {
1330 let record =
1331 WarmRecord::read(&fs::read(entry.path()).ok()?, &warm.identity).ok()?;
1332 Some((entry.path(), record))
1333 })
1334 .collect();
1335 records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
1336 records
1337 }
1338
1339 fn choose_warm(
1346 &self,
1347 workspace: &Workspace,
1348 warm: &Warm,
1349 current: &[String],
1350 related: bool,
1351 ) -> Option<(WarmRecord, String)> {
1352 let own = worktree(workspace);
1353 let mut candidates: Vec<(u8, bool, WarmRecord)> = self
1354 .warm_records(warm)
1355 .into_iter()
1356 .filter_map(|(_, record)| {
1357 let quality = key_match(warm, current, &record.key)?;
1358 let other = record.worktree != own;
1359 (!other || warm.portable == Portable::Anywhere).then_some((quality, other, record))
1360 })
1361 .filter(|(_, _, record)| {
1362 let descends = record
1363 .commit
1364 .as_deref()
1365 .is_some_and(|commit| descends_from(workspace, commit));
1366 descends == related
1367 })
1368 .collect();
1369 candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
1371 let rank = candidates
1372 .first()
1373 .map(|(quality, other, _)| (*quality, *other))?;
1374 let tied = candidates
1375 .iter()
1376 .take_while(|(quality, other, _)| (*quality, *other) == rank)
1377 .count();
1378 let index = if tied > 1 {
1379 (0..tied)
1380 .min_by_key(|&index| {
1381 let behind = candidates[index]
1382 .2
1383 .commit
1384 .as_deref()
1385 .and_then(|commit| behind(workspace, commit));
1386 (behind.is_none(), behind, index)
1387 })
1388 .unwrap_or(0)
1389 } else {
1390 0
1391 };
1392 let (_, other, record) = candidates.swap_remove(index);
1393 let source = if other {
1394 format!("worktree {}", record.worktree)
1395 } else {
1396 "local".to_owned()
1397 };
1398 Some((record, source))
1399 }
1400
1401 fn store_warm(&self, record: &WarmRecord) -> Result<()> {
1402 store_warm(&self.root, record)
1403 }
1404
1405 pub(crate) fn clear_warm(&self, workspace: &Workspace, task: &Task, warm: &Warm) -> Result<()> {
1410 let _lock = self.lock_warm(warm)?;
1411 let identity = self.warm_prefix(&warm.identity);
1412 if self.cleared.lock().unwrap().contains(&identity) {
1413 return Ok(());
1414 }
1415 for location in locations(workspace, task, warm)? {
1416 for path in location.paths()?.iter().rev() {
1417 let absolute = location.base.join(path);
1418 let removed = match fs::symlink_metadata(&absolute) {
1419 Ok(metadata) if metadata.is_dir() => match fs::remove_dir(&absolute) {
1420 Err(error) if error.kind() == std::io::ErrorKind::DirectoryNotEmpty => {
1421 Ok(())
1422 }
1423 removed => removed,
1424 },
1425 Ok(_) => fs::remove_file(&absolute),
1426 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
1427 Err(error) => Err(error),
1428 };
1429 removed.with_context(|| format!("removing {}", absolute.display()))?;
1430 }
1431 }
1432 let notes = restored_files_path(workspace, warm);
1433 match fs::remove_file(¬es) {
1434 Err(error) if error.kind() != std::io::ErrorKind::NotFound => return Err(error.into()),
1435 _ => {}
1436 }
1437 self.cleared.lock().unwrap().insert(identity);
1438 Ok(())
1439 }
1440
1441 pub(crate) fn restore_warm(
1446 &self,
1447 workspace: &Workspace,
1448 task: &Task,
1449 warm: &Warm,
1450 prepared: &qk_executor::PreparedTask,
1451 ) -> Result<(Restored, Vec<String>)> {
1452 let started = std::time::Instant::now();
1453 let mut restored = Restored {
1454 group: warm.group().map(str::to_owned),
1455 ..Restored::default()
1456 };
1457 let locking = profile::span(&task.id, "warm_restore_lock");
1458 let _lock = self.lock_warm(warm)?;
1459 drop(locking);
1460 let mut present = Vec::new();
1462 let mut missing = Vec::new();
1463 let looking = profile::span(&task.id, "warm_restore_present");
1464 let locations = locations(workspace, task, warm)?;
1465 for location in &locations {
1466 if location.paths()?.is_empty() {
1467 missing.push((location, Vec::new()));
1468 } else if let Some(kept) = location.present()? {
1469 missing.push((location, kept));
1472 } else {
1473 present.push(warm.group().map_or(location.name, |group| group).to_owned());
1474 }
1475 }
1476 drop(looking);
1477 if missing.is_empty() {
1478 return Ok((restored, present));
1479 }
1480 let choosing = profile::span(&task.id, "warm_restore_choose");
1481 let keying = profile::span(&task.id, "warm_restore_key");
1482 let current = key_digests(workspace, warm, prepared)?;
1483 drop(keying);
1484 let mut fetch_ms = None;
1485 let ahead = self.prefetches.get().map_or(Taken::Nothing, |prefetches| {
1486 let mut waiting = profile::span(&task.id, "warm_restore_prefetched");
1487 let taking = std::time::Instant::now();
1488 let taken = prefetches.take(&warm.identity, ¤t);
1489 if !matches!(taken, Taken::Nothing) {
1490 fetch_ms = Some(taking.elapsed().as_millis() as u64);
1491 }
1492 waiting.note(
1493 "taken",
1494 match taken {
1495 Taken::Found(_) => "found",
1496 Taken::Missing => "missing",
1497 Taken::Nothing => "nothing",
1498 },
1499 );
1500 taken
1501 });
1502 let local = profile::span(&task.id, "warm_restore_local");
1503 let here = self.choose_warm(workspace, warm, ¤t, true);
1504 drop(local);
1505 let chosen = match (ahead, here) {
1506 (Taken::Found((record, _)), Some((mine, source))) if mine.saved > record.saved => {
1510 Some((mine, source))
1511 }
1512 (Taken::Found((record, branch)), _) => Some((record, format!("remote {branch}"))),
1513 (_, Some(here)) => Some(here),
1514 (Taken::Missing, None) => None,
1516 (Taken::Nothing, None) => {
1517 let mut fetching = profile::span(&task.id, "warm_restore_remote");
1518 let asking = std::time::Instant::now();
1519 let fetched = self
1520 .remote_warm(workspace, task, warm, ¤t)
1521 .map(|(record, branch)| (record, format!("remote {branch}")));
1522 if self.remote.is_some() && warm.shared_remotely() {
1523 fetch_ms = Some(asking.elapsed().as_millis() as u64);
1524 }
1525 fetching.note(
1526 "source",
1527 fetched
1528 .as_ref()
1529 .map_or("none", |(_, source)| source.as_str()),
1530 );
1531 fetched
1532 }
1533 }
1534 .or_else(|| self.choose_warm(workspace, warm, ¤t, false));
1536 let Some((record, source)) = chosen else {
1537 return Ok((restored, present));
1538 };
1539 drop(choosing);
1540 let _files = profile::span(&task.id, "warm_restore_files");
1541 let keep_mtimes = warm.preserve_mtimes;
1545 restored.source = source;
1546 let mut noted = read_restored_files(workspace, warm);
1547 for (location, kept) in missing {
1548 let Some(group) = record.groups.get(location.name) else {
1549 continue;
1550 };
1551 let (files, bytes) =
1552 match self.restore_location(&task.id, location, &kept, group, keep_mtimes) {
1553 Ok(restored) => restored,
1554 Err(error) => {
1555 remove_restored(location, &kept);
1559 return Err(error);
1560 }
1561 };
1562 restored.files += files.len();
1563 restored.bytes += bytes;
1564 if kept.is_empty() {
1565 noted.groups.insert(location.name.to_owned(), files);
1566 } else {
1567 noted
1569 .groups
1570 .entry(location.name.to_owned())
1571 .or_default()
1572 .extend(files);
1573 }
1574 restored.groups.push(location.name.to_owned());
1575 }
1576 if !restored.groups.is_empty() {
1577 let path = restored_files_path(workspace, warm);
1578 fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
1579 fs::write(path, serde_json::to_vec(¬ed)?)?;
1580 }
1581 restored.fetch_ms = fetch_ms;
1582 restored.ms = started.elapsed().as_millis() as u64;
1583 Ok((restored, present))
1584 }
1585
1586 fn restore_location(
1589 &self,
1590 task: &str,
1591 location: &Location,
1592 kept: &[&Outputs],
1593 group: &WarmGroup,
1594 keep_mtimes: bool,
1595 ) -> Result<(GroupFiles, u64)> {
1596 fs::create_dir_all(&location.base)?;
1597 let mut laying = profile::span(task, "warm_restore_layout");
1598 let mut created = BTreeSet::new();
1601 let mut directories = Vec::new();
1602 let mut pending = Vec::new();
1603 for (path, artifact) in &group.artifacts {
1604 paths::safe_parents(&location.base, path)?;
1605 if !location.matches(path) {
1606 bail!("warm state for {task} holds {path}, outside its paths");
1607 }
1608 if kept.iter().any(|outputs| outputs.matches(path)) {
1609 continue;
1610 }
1611 let destination = location.base.join(path);
1612 if let Some(parent) = destination.parent()
1613 && created.insert(parent.to_owned())
1614 {
1615 fs::create_dir_all(parent)?;
1616 }
1617 match artifact {
1618 Artifact::File { blob, mode } => {
1619 let modified = keep_mtimes
1620 .then(|| group.mtimes.get(path))
1621 .flatten()
1622 .and_then(|nanos| u64::try_from(*nanos).ok())
1623 .map_or(RESTORED_AT, |nanos| {
1624 UNIX_EPOCH + Duration::from_nanos(nanos)
1625 });
1626 pending.push((path, blob.as_str(), *mode, destination, modified));
1627 }
1628 Artifact::Directory { mode } => {
1629 fs::create_dir_all(&destination)?;
1630 created.insert(destination.clone());
1631 directories.push((destination, *mode));
1632 }
1633 Artifact::Symlink { target, directory } => {
1634 validate_link(path, target)?;
1635 symlink(target, &destination, *directory)?;
1636 }
1637 }
1638 }
1639 laying.count("directories", created.len() as u64);
1640 drop(laying);
1641 let mut copying = profile::span(task, "warm_restore_copy");
1642 copying.count("files", pending.len() as u64);
1643 let copied = self
1644 .restore_warm_files(&pending)
1645 .with_context(|| format!("warm state for {task} not restored"))?;
1646 copying.count("bytes", copied.iter().map(|(_, size)| size).sum());
1647 drop(copying);
1648 let _modes = profile::span(task, "warm_restore_modes");
1650 for (destination, mode) in directories {
1651 set_mode(&destination, mode)?;
1652 }
1653 let mut files = BTreeMap::new();
1654 let mut bytes = 0;
1655 for ((path, blob, mode, _, _), (metadata, size)) in pending.into_iter().zip(copied) {
1656 bytes += size;
1657 let artifact = Artifact::File {
1658 blob: blob.to_owned(),
1659 mode,
1660 };
1661 files.insert(path.clone(), (metadata, artifact));
1662 }
1663 Ok((files, bytes))
1664 }
1665
1666 fn restore_warm_files(&self, files: &[PendingFile]) -> Result<Vec<RestoredFile>> {
1669 let restore =
1670 |(_, blob, mode, destination, modified): &PendingFile| -> Result<RestoredFile> {
1671 let source = self.root.join("blobs").join(blob);
1672 if !restore_blob(&source, destination, blob, self.verified())? {
1673 let _ = fs::remove_file(destination);
1675 bail!("corrupt blob {blob}");
1676 }
1677 set_modified(destination, *modified)?;
1679 set_mode(destination, *mode)?;
1680 let metadata = fs::symlink_metadata(destination)?;
1681 Ok((stamp(&metadata), metadata.len()))
1682 };
1683 let workers = files.len().div_ceil(64).clamp(1, RESTORERS);
1684 if workers == 1 {
1685 return files.iter().map(restore).collect();
1686 }
1687 let next = AtomicUsize::new(0);
1690 let failed = AtomicBool::new(false);
1691 let mut results: Vec<(usize, Result<RestoredFile>)> = std::thread::scope(|scope| {
1692 let handles: Vec<_> = (0..workers)
1693 .map(|_| {
1694 scope.spawn(|| {
1695 let mut done = Vec::new();
1696 while !failed.load(Ordering::Relaxed) {
1697 let index = next.fetch_add(1, Ordering::Relaxed);
1698 let Some(file) = files.get(index) else { break };
1699 let result = restore(file);
1700 if result.is_err() {
1701 failed.store(true, Ordering::Relaxed);
1702 }
1703 done.push((index, result));
1704 }
1705 done
1706 })
1707 })
1708 .collect();
1709 handles
1710 .into_iter()
1711 .flat_map(|handle| handle.join().expect("warm restore worker panicked"))
1712 .collect()
1713 });
1714 results.sort_by_key(|(index, _)| *index);
1715 results.into_iter().map(|(_, result)| result).collect()
1716 }
1717
1718 fn remote_warm(
1721 &self,
1722 workspace: &Workspace,
1723 task: &Task,
1724 warm: &Warm,
1725 current: &[String],
1726 ) -> Option<(WarmRecord, String)> {
1727 let remote = self.remote.as_ref()?;
1728 let own = worktree(workspace);
1729 remote_warm(
1730 remote,
1731 &self.root,
1732 &task.id,
1733 warm,
1734 &own,
1735 &branches(workspace),
1736 current,
1737 )
1738 }
1739
1740 pub fn prefetch_warm(
1749 &self,
1750 workspace: &Workspace,
1751 graph: &TaskGraph,
1752 prepared: &BTreeMap<String, qk_executor::PreparedTask>,
1753 restores: impl Fn(&str) -> bool,
1754 ) {
1755 let Some(remote) = self.remote.as_ref().filter(|_| self.reuse && !self.cold) else {
1757 return;
1758 };
1759 if self.initialize().is_err() {
1760 return;
1761 }
1762 let own = worktree(workspace);
1763 let mut items = Vec::new();
1764 for (id, task) in &graph.tasks {
1765 let Some(prepared) = prepared.get(id).filter(|_| restores(id)) else {
1766 continue;
1767 };
1768 let Ok(entries) = config(workspace, task) else {
1769 continue;
1770 };
1771 if !entries.iter().any(|warm| warm.prefetch) {
1772 continue;
1773 }
1774 let mut running = prepared.clone();
1778 for warm in &entries {
1779 for (name, value) in &warm.env {
1780 running.execution.insert(name.into(), value.into());
1781 }
1782 }
1783 for warm in entries {
1784 if !warm.prefetch || !warm.shared_remotely() {
1785 continue;
1786 }
1787 let missing = locations(workspace, task, &warm).is_ok_and(|locations| {
1788 locations.iter().any(|location| {
1789 location.paths().map_or(true, |paths| paths.is_empty())
1790 || location.present().is_ok_and(|kept| kept.is_some())
1791 })
1792 });
1793 let Ok(current) = key_digests(workspace, &warm, &running) else {
1794 continue;
1795 };
1796 if !missing || self.choose_warm(workspace, &warm, ¤t, true).is_some() {
1797 continue;
1798 }
1799 let depth = depth(graph, id, &mut BTreeMap::new());
1800 items.push((depth, warm.identity.clone(), (id.clone(), warm, current)));
1801 }
1802 }
1803 if items.is_empty() {
1804 return;
1805 }
1806 items.sort_by(|a, b| (a.0, &(a.2).0).cmp(&(b.0, &(b.2).0)));
1807 let items = items
1808 .into_iter()
1809 .map(|(_, identity, job)| {
1810 let key = job.2.clone();
1811 (identity, key, job)
1812 })
1813 .collect();
1814 let (remote, root) = (remote.clone(), self.root.clone());
1815 let branches = branches(workspace);
1816 let prefetches = Prefetches::start(items, PREFETCHERS, move |(id, warm, current)| {
1817 let mut fetching = profile::span(&id, "warm_prefetch");
1818 let fetched = remote_warm(&remote, &root, &id, &warm, &own, &branches, ¤t);
1819 fetching.note(
1820 "branch",
1821 fetched
1822 .as_ref()
1823 .map_or("none", |(_, branch)| branch.as_str()),
1824 );
1825 fetched
1826 });
1827 let _ = self.prefetches.set(prefetches);
1828 }
1829
1830 pub fn abandon_prefetches(&self) {
1833 if let Some(prefetches) = self.prefetches.get() {
1834 prefetches.stop(false);
1835 }
1836 }
1837
1838 pub(crate) fn finish_prefetches(&self) {
1841 if let Some(prefetches) = self.prefetches.get() {
1842 prefetches.stop(true);
1843 }
1844 }
1845
1846 pub(crate) fn save_warm(
1849 &self,
1850 workspace: &Workspace,
1851 task: &Task,
1852 warm: &Warm,
1853 prepared: &qk_executor::PreparedTask,
1854 ) -> Result<()> {
1855 self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
1856 }
1857
1858 pub(crate) fn save_warm_in_background(
1861 &self,
1862 workspace: &Workspace,
1863 task: &Task,
1864 warm: &Warm,
1865 prepared: &qk_executor::PreparedTask,
1866 ) -> Result<()> {
1867 let save = Save::prepare(workspace, task, warm, prepared)?;
1868 let cache = Cache {
1869 remote: self.remote.clone(),
1870 ..Cache::new(self.root.clone())
1871 };
1872 let handle = std::thread::spawn(move || {
1873 if let Err(error) = cache.save_prepared(&save) {
1874 qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
1875 }
1876 });
1877 self.saves.lock().unwrap().push(handle);
1878 Ok(())
1879 }
1880
1881 fn save_prepared(&self, save: &Save) -> Result<()> {
1886 let warm = &save.warm;
1887 let _lock = self.lock_warm(warm)?;
1888 let _save = profile::span(&save.task, "warm_save");
1889 let records = self.warm_records(warm);
1890 let previous = records
1892 .iter()
1893 .filter(|(_, record)| record.worktree == save.worktree)
1894 .min_by_key(|(_, record)| record.key != save.key)
1895 .map(|(_, record)| record);
1896 let noted: RestoredFiles = fs::read(&save.notes)
1897 .ok()
1898 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1899 .unwrap_or_default();
1900 let mut record = WarmRecord::new(
1901 warm.identity.clone(),
1902 save.worktree.clone(),
1903 save.commit.clone(),
1904 now_ms(),
1905 save.key.clone(),
1906 );
1907 let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
1908 for location in &save.locations {
1909 if !location.base.exists() {
1910 continue;
1911 }
1912 let known = previous.and_then(|previous| previous.groups.get(location.name));
1913 let restored_files = noted.groups.get(location.name);
1914 let mut group = WarmGroup::default();
1915 let mut pending = Vec::new();
1916 for path in location.paths()? {
1917 paths::safe_parents(&location.base, &path)?;
1918 let absolute = location.base.join(&path);
1919 let metadata = match fs::symlink_metadata(&absolute) {
1920 Err(error) if gone(&error) => continue,
1921 result => result?,
1922 };
1923 let artifact = if metadata.file_type().is_symlink() {
1924 let target = match fs::read_link(&absolute) {
1925 Err(error) if gone(&error) => continue,
1926 result => result?,
1927 };
1928 let target = target
1929 .to_str()
1930 .context("symlink target must be UTF-8")?
1931 .to_owned();
1932 validate_link(&path, &target)?;
1933 Artifact::Symlink {
1934 target,
1935 directory: absolute.is_dir(),
1936 }
1937 } else if metadata.is_file() {
1938 let now = stamp(&metadata);
1939 let saved = known.and_then(|known| {
1940 (known.stamps.get(&path) == Some(&now))
1941 .then(|| known.artifacts.get(&path))
1942 .flatten()
1943 });
1944 let restored = restored_files
1945 .and_then(|files| files.get(&path))
1946 .filter(|(stamp, _)| *stamp == now)
1947 .map(|(_, artifact)| artifact);
1948 let reused = saved.or(restored).filter(|artifact| match artifact {
1949 Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
1950 _ => false,
1951 });
1952 let Some(artifact) = reused.cloned() else {
1953 pending.push((path, absolute, metadata));
1954 continue;
1955 };
1956 group.stamps.insert(path.clone(), now);
1957 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1958 group.mtimes.insert(path.clone(), modified);
1959 }
1960 artifact
1961 } else if metadata.is_dir() {
1962 Artifact::Directory {
1963 mode: mode(&metadata),
1964 }
1965 } else {
1966 continue;
1967 };
1968 group.artifacts.insert(path, artifact);
1969 }
1970 let storing = profile::span(&save.task, "warm_save_blobs");
1971 let sources: Vec<_> = pending
1972 .iter()
1973 .map(|(_, absolute, _)| absolute.clone())
1974 .collect();
1975 for ((path, _, metadata), result) in pending.into_iter().zip(self.put_blobs(&sources)) {
1976 let blob = match result {
1977 Ok(blob) => blob,
1978 Err(error) if error.downcast_ref::<std::io::Error>().is_some_and(gone) => {
1979 continue;
1980 }
1981 Err(error) => return Err(error),
1982 };
1983 group.stamps.insert(path.clone(), stamp(&metadata));
1984 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1985 group.mtimes.insert(path.clone(), modified);
1986 }
1987 group.artifacts.insert(
1988 path,
1989 Artifact::File {
1990 blob,
1991 mode: mode(&metadata),
1992 },
1993 );
1994 }
1995 drop(storing);
1996 if let Some(limit) = warm.max_size
1997 && group.size(self) > limit
1998 {
1999 qk_executor::status!(
2000 "qk: {}: warm {} is over its maxSize and was not saved",
2001 save.task,
2002 location.name
2003 );
2004 continue;
2005 }
2006 record.groups.insert(location.name.to_owned(), group);
2007 }
2008 self.store_warm(&record)?;
2009 let _ = fs::remove_file(&save.notes);
2010 let saved = self.warm_record(&record);
2013 for (path, _) in records
2014 .into_iter()
2015 .filter(|(path, _)| *path != saved)
2016 .skip(KEPT - 1)
2017 {
2018 let _ = fs::remove_file(path);
2019 }
2020 if warm.shared_remotely()
2022 && let Some(remote) = &self.remote
2023 && let Some(branch) = &save.branch
2024 {
2025 for group in record.groups.values_mut() {
2026 group.stamps.clear();
2027 }
2028 remote.upload_warm(
2029 &self.root,
2030 &warm.identity,
2031 &remote_identity(warm, &save.worktree),
2032 branch,
2033 serde_json::to_vec(&record)?,
2034 );
2035 }
2036 Ok(())
2037 }
2038}
2039
2040struct Save {
2043 task: String,
2044 warm: Warm,
2045 locations: Vec<Location>,
2046 worktree: String,
2047 commit: Option<String>,
2048 branch: Option<String>,
2049 key: Vec<String>,
2050 notes: PathBuf,
2051}
2052
2053impl Save {
2054 fn prepare(
2056 workspace: &Workspace,
2057 task: &Task,
2058 warm: &Warm,
2059 prepared: &qk_executor::PreparedTask,
2060 ) -> Result<Self> {
2061 Ok(Self {
2062 task: task.id.clone(),
2063 warm: warm.clone(),
2064 locations: locations(workspace, task, warm)?,
2065 worktree: worktree(workspace),
2066 commit: current_commit(workspace),
2067 branch: current_branch(workspace),
2068 key: key_digests(workspace, warm, prepared)?,
2069 notes: restored_files_path(workspace, warm),
2070 })
2071 }
2072}
2073
2074#[cfg(test)]
2075mod tests {
2076 use super::*;
2077 use qk_taskgraph::Request;
2078 use serde_json::json;
2079
2080 fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
2082 let temp = tempfile::tempdir().unwrap();
2083 fs::write(temp.path().join("nx.json"), "{}").unwrap();
2084 fs::write(
2085 temp.path().join("project.json"),
2086 json!({
2087 "name": "app", "targets": {
2088 "one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
2089 "two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
2090 }
2091 })
2092 .to_string(),
2093 )
2094 .unwrap();
2095 for path in ["scratch-one", "scratch-two"] {
2096 fs::create_dir(temp.path().join(path)).unwrap();
2097 fs::write(temp.path().join(path).join("state"), "saved").unwrap();
2098 }
2099 let workspace = Workspace::load(temp.path()).unwrap();
2100 let graph = TaskGraph::build(
2101 &workspace,
2102 &[
2103 Request::parse("app:one").unwrap(),
2104 Request::parse("app:two").unwrap(),
2105 ],
2106 )
2107 .unwrap();
2108 (temp, workspace, graph)
2109 }
2110
2111 #[test]
2114 fn asks_a_partial_clone_about_ancestors_without_fetching() {
2115 let temp = tempfile::tempdir().unwrap();
2116 let git = |root: &std::path::Path, args: &[&str]| {
2117 let output = paths::git(root, args).unwrap();
2118 assert!(output.status.success(), "git {args:?}: {output:?}");
2119 String::from_utf8(output.stdout).unwrap().trim().to_owned()
2120 };
2121 let source = temp.path().join("source");
2122 fs::create_dir(&source).unwrap();
2123 git(&source, &["init", "--quiet"]);
2124 git(&source, &["config", "uploadpack.allowFilter", "true"]);
2125 let commit = |message: &str| {
2126 fs::write(source.join("file"), message).unwrap();
2127 git(&source, &["add", "file"]);
2128 git(
2129 &source,
2130 &[
2131 "-c",
2132 "user.name=qk",
2133 "-c",
2134 "user.email=qk@example.com",
2135 "commit",
2136 "--quiet",
2137 "-m",
2138 message,
2139 ],
2140 );
2141 git(&source, &["rev-parse", "HEAD"])
2142 };
2143 let older = commit("older");
2144 let newer = commit("newer");
2145 let clone = temp.path().join("clone");
2146 let url = format!("file://{}", source.display());
2147 git(
2148 temp.path(),
2149 &[
2150 "clone",
2151 "--quiet",
2152 "--depth",
2153 "1",
2154 "--filter=blob:none",
2155 &url,
2156 clone.to_str().unwrap(),
2157 ],
2158 );
2159 fs::write(clone.join("nx.json"), "{}").unwrap();
2160 let workspace = Workspace::load(&clone).unwrap();
2161 assert!(descends_from(&workspace, &newer));
2162 assert!(!descends_from(&workspace, &older));
2163 assert_eq!(behind(&workspace, &older), None);
2164 let version = git(temp.path(), &["--version"]);
2166 let honoured = version
2167 .split_whitespace()
2168 .nth(2)
2169 .and_then(|version| {
2170 let mut parts = version.split('.').map(|part| part.parse::<u32>().ok());
2171 Some((parts.next()??, parts.next()??))
2172 })
2173 .is_some_and(|version| version >= (2, 44));
2174 let held = paths::git_history(&clone, &["cat-file", "-e", &format!("{older}^{{commit}}")])
2175 .unwrap();
2176 assert!(!honoured || !held.status.success(), "fetched {older}");
2177 }
2178
2179 fn entries(warm: Value) -> Result<Vec<Warm>> {
2181 entries_in(warm, None)
2182 }
2183
2184 fn entries_in(warm: Value, configuration: Option<&str>) -> Result<Vec<Warm>> {
2187 let temp = tempfile::tempdir().unwrap();
2188 fs::write(temp.path().join("nx.json"), "{}").unwrap();
2189 fs::write(
2190 temp.path().join("project.json"),
2191 json!({"name": "app", "targets": {"build": {
2192 "command": "true",
2193 "qk:warm": warm,
2194 "configurations": {"kept": {}, "skipped": {}}
2195 }}})
2196 .to_string(),
2197 )
2198 .unwrap();
2199 let workspace = Workspace::load(temp.path()).unwrap();
2200 let id = configuration.map_or_else(
2201 || "app:build".to_owned(),
2202 |name| format!("app:build:{name}"),
2203 );
2204 let graph = TaskGraph::build(&workspace, &[Request::parse(&id).unwrap()]).unwrap();
2205 config(&workspace, &graph.tasks[&id])
2206 }
2207
2208 #[test]
2209 fn entries_conflict_only_where_a_configuration_selects_both() {
2210 let apart = json!([
2212 {"paths": ["one"], "env": {"CACHE": "{warm}/one"}, "configurations": ["kept"]},
2213 {"paths": ["two"], "env": {"CACHE": "{warm}/two"}, "configurations": ["skipped"]}
2214 ]);
2215 let kept = entries_in(apart.clone(), Some("kept")).unwrap();
2216 assert_eq!(kept.len(), 1);
2217 assert!(kept[0].env["CACHE"].ends_with("/one"));
2218 let skipped = entries_in(apart.clone(), Some("skipped")).unwrap();
2219 assert_eq!(skipped.len(), 1);
2220 assert!(skipped[0].env["CACHE"].ends_with("/two"));
2221 assert!(entries_in(apart, None).unwrap().is_empty());
2222 let both = json!([
2224 {"paths": ["one"], "configurations": ["kept"]},
2225 {"paths": ["two"], "configurations": ["kept", "skipped"]}
2226 ]);
2227 assert_eq!(entries_in(both.clone(), Some("skipped")).unwrap().len(), 1);
2228 assert_eq!(
2229 format!("{:#}", entries_in(both, Some("kept")).unwrap_err()),
2230 "only one qk:warm entry may go without a group"
2231 );
2232 let shared = json!([
2234 {"group": "one", "env": {"CACHE": "{warm}"}, "configurations": ["kept"]},
2235 {"group": "two", "env": {"CACHE": "{warm}"}, "configurations": ["kept", "skipped"]}
2236 ]);
2237 assert_eq!(
2238 entries_in(shared.clone(), Some("skipped")).unwrap().len(),
2239 1
2240 );
2241 assert_eq!(
2242 format!("{:#}", entries_in(shared, Some("kept")).unwrap_err()),
2243 "qk:warm sets CACHE in more than one entry"
2244 );
2245 }
2246
2247 #[test]
2248 fn an_array_holds_the_task_entry_beside_groups() {
2249 let entries = entries(json!([
2250 {"group": "metro", "portable": true, "env": {"METRO_CACHE_ROOT": "{warm}/metro"}},
2251 {"paths": ["build"], "mtimes": "preserve"}
2252 ]))
2253 .unwrap();
2254 assert_eq!(entries.len(), 2);
2255 assert_eq!(entries[0].group(), Some("metro"));
2256 assert!(entries[0].portable == Portable::Anywhere && entries[0].directory);
2257 assert_eq!(entries[1].group(), None);
2258 assert!(entries[1].portable == Portable::Own && entries[1].preserve_mtimes);
2259 assert_eq!(kept_paths(&entries), ["build"]);
2260 }
2261
2262 #[test]
2263 fn portable_is_a_boolean_or_the_same_path() {
2264 let portable = |value| {
2265 entries(json!({"paths": ["build"], "portable": value}))
2266 .map(|entries| entries[0].portable)
2267 };
2268 assert_eq!(portable(json!("path")).unwrap(), Portable::Path);
2269 assert_eq!(portable(json!(false)).unwrap(), Portable::Own);
2270 assert_eq!(
2271 format!("{:#}", portable(json!("anywhere")).unwrap_err()),
2272 "qk:warm.portable must be true, false or \"path\""
2273 );
2274 }
2275
2276 #[test]
2277 fn entries_cannot_share_saves_or_variables() {
2278 let error = |warm| format!("{:#}", entries(warm).unwrap_err());
2279 assert_eq!(
2280 error(json!([{"paths": ["one"]}, {"paths": ["two"]}])),
2281 "only one qk:warm entry may go without a group"
2282 );
2283 assert_eq!(
2284 error(json!([{"group": "metro"}, {"group": "metro"}])),
2285 "qk:warm names group \"metro\" twice"
2286 );
2287 assert_eq!(
2288 error(json!([
2289 {"group": "one", "env": {"CACHE": "{warm}/one"}},
2290 {"group": "two", "env": {"CACHE": "{warm}/two"}}
2291 ])),
2292 "qk:warm sets CACHE in more than one entry"
2293 );
2294 assert_eq!(
2295 error(json!(["build"])),
2296 "qk:warm must be an object or an array of objects"
2297 );
2298 }
2299
2300 #[test]
2302 fn waiting_for_an_owner_does_not_block_other_releases() {
2303 let (_temp, workspace, graph) = fixture();
2304 let one = &graph.tasks["app:one"];
2305 let two = &graph.tasks["app:two"];
2306 let first = acquire(&workspace, one, &config(&workspace, one).unwrap()[0]).unwrap();
2307 let state = paths::worktree_state(&workspace.root).join("warm-kept");
2308 let lock = fs::OpenOptions::new()
2309 .create(true)
2310 .truncate(false)
2311 .read(true)
2312 .write(true)
2313 .open(state.join(format!("{}.lock", hash32(&two.id))))
2314 .unwrap();
2315 fs2::FileExt::lock_exclusive(&lock).unwrap();
2316 let (released, received) = std::sync::mpsc::channel();
2317 let completed = std::thread::scope(|scope| {
2318 scope.spawn(|| {
2319 let second =
2320 acquire(&workspace, two, &config(&workspace, two).unwrap()[0]).unwrap();
2321 release(&second);
2322 });
2323 std::thread::sleep(Duration::from_millis(100));
2324 scope.spawn(|| {
2325 release(&first);
2326 released.send(()).unwrap();
2327 });
2328 let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
2329 fs2::FileExt::unlock(&lock).unwrap();
2330 completed
2331 });
2332 assert!(
2333 completed,
2334 "a blocked owner lock prevented an unrelated release"
2335 );
2336 assert_eq!(
2337 fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
2338 "saved"
2339 );
2340 }
2341
2342 #[test]
2344 fn unreadable_ownership_preserves_saved_paths() {
2345 let (_temp, workspace, graph) = fixture();
2346 let owner = &graph.tasks["app:one"];
2347 let warm = config(&workspace, owner).unwrap().remove(0);
2348 let root = acquire(&workspace, owner, &warm).unwrap();
2349 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
2350 drop(entry); fs::write(Stash::manifest(&root), "{").unwrap();
2352 assert!(acquire(&workspace, owner, &warm).is_err());
2353 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2354 fs::remove_file(Stash::manifest(&root)).unwrap();
2355 assert!(acquire(&workspace, owner, &warm).is_err());
2356 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2357 fs::create_dir(Stash::manifest(&root)).unwrap();
2358 assert!(acquire(&workspace, owner, &warm).is_err());
2359 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
2360 }
2361
2362 #[cfg(unix)]
2364 #[test]
2365 fn inspection_errors_keep_recovery_incomplete() {
2366 let (_temp, workspace, graph) = fixture();
2367 let owner = &graph.tasks["app:one"];
2368 let warm = config(&workspace, owner).unwrap().remove(0);
2369 let root = acquire(&workspace, owner, &warm).unwrap();
2370 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
2371 let saved = root.with_extension("saved");
2372 fs::rename(&root, &saved).unwrap();
2373 fs::write(&root, "blocked").unwrap(); assert!(!entry.stash.put_back());
2375 let retained = Stash {
2376 root: root.clone(),
2377 workspace: workspace.root.clone(),
2378 moved: entry.stash.moved.clone(),
2379 };
2380 assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
2381 assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
2382 assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
2383 }
2384}