1use std::collections::BTreeMap;
17use std::fs;
18use std::io::Write;
19use std::path::PathBuf;
20use std::time::{Duration, SystemTime, UNIX_EPOCH};
21
22use anyhow::{Context, Result, bail};
23use qk_config::Workspace;
24use qk_taskgraph::{Task, TaskGraph};
25use serde::{Deserialize, Serialize};
26use serde_json::Value;
27
28use crate::paths::{self, Outputs};
29use crate::store::{Artifact, mode, set_mode, symlink, validate_link};
30use crate::{Cache, hash::digest_file};
31
32pub const RESTORED_AT: std::time::SystemTime = std::time::UNIX_EPOCH;
37
38const KEPT: usize = 8;
40
41#[derive(Clone, Debug, Default)]
43pub struct Warm {
44 pub identity: String,
46 pub outputs: bool,
47 pub paths: Vec<String>,
49 pub env: BTreeMap<String, String>,
51 pub directory: bool,
53 pub max_size: Option<u64>,
54 pub remote: bool,
55 pub preserve_mtimes: bool,
58 pub portable: bool,
60 pub key: Vec<KeyPart>,
62 pub restore_keys: Option<usize>,
65 pub survive: Vec<String>,
68 pub background: bool,
71}
72
73#[derive(Clone, Debug)]
75pub enum KeyPart {
76 File(String),
78 Env(String),
80}
81
82pub fn config(workspace: &Workspace, task: &Task) -> Result<Option<Warm>> {
84 let Some(value) = task.definition.extra.get("qk:warm") else {
85 return Ok(None);
86 };
87 let object = value.as_object().context("qk:warm must be an object")?;
88 for key in object.keys() {
89 if !matches!(
90 key.as_str(),
91 "outputs"
92 | "paths"
93 | "env"
94 | "maxSize"
95 | "remote"
96 | "mtimes"
97 | "portable"
98 | "key"
99 | "restoreKeys"
100 | "group"
101 | "survive"
102 | "save"
103 ) {
104 bail!("unknown qk:warm field {key:?}");
105 }
106 }
107 let flag = |name: &str, default: bool| -> Result<bool> {
108 object.get(name).map_or(Ok(default), |value| {
109 value
110 .as_bool()
111 .with_context(|| format!("qk:warm.{name} must be true or false"))
112 })
113 };
114 let group = object
115 .get("group")
116 .map(|value| {
117 value
118 .as_str()
119 .filter(|name| !name.is_empty())
120 .context("qk:warm.group must be a name")
121 })
122 .transpose()?;
123 if group.is_some() && (object.contains_key("outputs") || object.contains_key("paths")) {
124 bail!("a qk:warm.group shares {{warm}} alone; outputs and paths belong to one task");
125 }
126 let identity = group.map_or_else(|| task.id.clone(), |group| format!("group {group}"));
127 const SURVIVE_SHAPE: &str = "qk:warm.survive must be an array of target names";
128 let survive = object
129 .get("survive")
130 .map(|value| {
131 value
132 .as_array()
133 .context(SURVIVE_SHAPE)?
134 .iter()
135 .map(|target| {
136 target
137 .as_str()
138 .filter(|target| !target.is_empty())
139 .map(str::to_owned)
140 .context(SURVIVE_SHAPE)
141 })
142 .collect::<Result<Vec<_>>>()
143 })
144 .transpose()?
145 .unwrap_or_default();
146 let directory = directory(workspace, &identity);
147 let directory_text = directory.to_str().context("warm directory must be UTF-8")?;
148 let workspace_text = workspace
149 .root
150 .to_str()
151 .context("workspace root must be UTF-8")?;
152 let project_root = &workspace.projects[&task.project].root;
153 let project_directory = if project_root == "." {
154 workspace.root.clone()
155 } else {
156 workspace.root.join(project_root)
157 };
158 let project_text = project_directory
159 .to_str()
160 .context("project root must be UTF-8")?;
161 let expand = |text: &str| -> Result<String> {
162 if text.contains("{options.") || text.contains("{args.") {
163 bail!("dynamic warm environment values are not supported yet");
164 }
165 Ok(text
166 .replace("{workspaceRoot}", workspace_text)
167 .replace("{projectRoot}", project_text)
168 .replace("{warm}", directory_text))
169 };
170 let mut paths = Vec::new();
171 for path in object
172 .get("paths")
173 .map(|value| {
174 value
175 .as_array()
176 .context("qk:warm.paths must be an array of paths")
177 })
178 .transpose()?
179 .into_iter()
180 .flatten()
181 {
182 let path = path
183 .as_str()
184 .context("qk:warm.paths must be an array of paths")?;
185 if path.contains("{warm}") {
186 bail!("qk:warm.paths are workspace paths; {{warm}} is kept already");
187 }
188 let (excluded, path) = match path.strip_prefix('!') {
189 Some(path) => (true, path),
190 None => (false, path),
191 };
192 let path = paths::expand(workspace, &task.project, path)?
193 .trim_end_matches('/')
194 .to_owned();
195 paths::validate_path(&path)?;
196 paths.push(if excluded { format!("!{path}") } else { path });
197 }
198 let mut env = BTreeMap::new();
199 let mut uses_directory = false;
200 if let Some(values) = object.get("env") {
201 for (name, value) in values
202 .as_object()
203 .context("qk:warm.env must be an object")?
204 {
205 let value = value
206 .as_str()
207 .context("qk:warm.env values must be strings")?;
208 uses_directory |= value.contains("{warm}");
209 env.insert(name.clone(), expand(value)?);
210 }
211 }
212 let max_size = match object.get("maxSize") {
213 None => None,
214 Some(Value::String(text)) => Some(crate::parse_size(text)?),
215 Some(Value::Number(number)) => number.as_u64(),
216 Some(_) => bail!("qk:warm.maxSize must be a size such as \"2GB\""),
217 };
218 let preserve_mtimes = match object.get("mtimes").map(|value| value.as_str()) {
219 None | Some(Some("epoch")) => false,
220 Some(Some("preserve")) => true,
221 Some(_) => bail!("qk:warm.mtimes must be \"epoch\" or \"preserve\""),
222 };
223 const KEY_SHAPE: &str = "qk:warm.key must be an array of paths and {\"env\": name} objects";
224 let mut key = Vec::new();
225 for part in object
226 .get("key")
227 .map(|value| value.as_array().context(KEY_SHAPE))
228 .transpose()?
229 .into_iter()
230 .flatten()
231 {
232 key.push(match part {
233 Value::String(path) => {
234 let path = paths::expand(workspace, &task.project, path)?;
235 paths::validate_path(&path)?;
236 KeyPart::File(path)
237 }
238 Value::Object(object) if object.len() == 1 => KeyPart::Env(
239 object
240 .get("env")
241 .and_then(Value::as_str)
242 .context(KEY_SHAPE)?
243 .to_owned(),
244 ),
245 _ => bail!(KEY_SHAPE),
246 });
247 }
248 let restore_keys = object
249 .get("restoreKeys")
250 .map(|value| {
251 value
252 .as_u64()
253 .and_then(|count| usize::try_from(count).ok())
254 .filter(|count| *count <= key.len())
255 .context("qk:warm.restoreKeys must be a count no larger than qk:warm.key")
256 })
257 .transpose()?;
258 let background = match object.get("save").map(Value::as_str) {
259 None | Some(Some("wait")) => false,
260 Some(Some("background")) => true,
261 Some(_) => bail!("qk:warm.save must be \"wait\" or \"background\""),
262 };
263 Ok(Some(Warm {
264 identity,
265 outputs: flag("outputs", false)?,
266 paths,
267 env,
268 directory: uses_directory,
269 max_size,
270 remote: flag("remote", true)?,
271 preserve_mtimes,
272 portable: flag("portable", false)?,
273 key,
274 restore_keys,
275 survive,
276 background,
277 }))
278}
279
280fn directory(workspace: &Workspace, identity: &str) -> PathBuf {
282 paths::worktree_state(&workspace.root)
283 .join("warm")
284 .join(&blake3::hash(identity.as_bytes()).to_hex()[..32])
285}
286
287impl Warm {
288 pub fn kept_paths(&self) -> Vec<String> {
291 self.paths
292 .iter()
293 .filter(|path| !path.starts_with('!'))
294 .cloned()
295 .collect()
296 }
297}
298
299#[derive(Default)]
302pub struct Kept {
303 roots: Vec<PathBuf>,
305}
306
307impl Drop for Kept {
308 fn drop(&mut self) {
309 for root in self.roots.iter().rev() {
310 release(root);
311 }
312 }
313}
314
315#[derive(Default, Serialize, Deserialize)]
317struct Stash {
318 #[serde(skip)]
319 root: PathBuf,
320 #[serde(skip)]
321 workspace: PathBuf,
322 moved: Vec<String>,
324}
325
326struct Active {
329 users: usize,
330 stash: Stash,
331 _lock: fs::File,
332}
333
334static ACTIVE: std::sync::Mutex<BTreeMap<PathBuf, Active>> = std::sync::Mutex::new(BTreeMap::new());
337
338impl Stash {
339 fn manifest(root: &std::path::Path) -> PathBuf {
341 root.join("moved.json")
342 }
343
344 fn save_manifest(&self) -> Result<()> {
347 use std::io::Write;
348 let mut file = tempfile::NamedTempFile::new_in(&self.root)?;
349 file.write_all(&serde_json::to_vec(self)?)?;
350 file.as_file().sync_all()?;
351 file.persist(Self::manifest(&self.root))?;
352 Ok(())
353 }
354
355 fn put_back(&self) -> bool {
361 let mut complete = true;
362 for (index, path) in self.moved.iter().enumerate().rev() {
363 let kept = self.root.join(index.to_string());
364 match fs::symlink_metadata(&kept) {
365 Ok(_) => {}
366 Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
367 Err(error) => {
368 complete = false;
369 qk_executor::status!(
370 "qk: could not inspect saved {path}, kept for the next run ({error})"
371 );
372 continue;
373 }
374 }
375 let original = self.workspace.join(path);
376 let result = paths::safe_parents(&self.workspace, path).and_then(|()| {
377 remove(&original)?;
378 if let Some(parent) = original.parent() {
379 fs::create_dir_all(parent)?;
380 }
381 fs::rename(&kept, &original)?;
382 Ok(())
383 });
384 if let Err(error) = result {
385 complete = false;
386 qk_executor::status!(
387 "qk: could not put {path} back, kept for the next run ({error:#})"
388 );
389 }
390 }
391 if complete {
392 let _ = fs::remove_dir_all(&self.root);
393 }
394 complete
395 }
396}
397
398fn remove(path: &std::path::Path) -> std::io::Result<()> {
400 match fs::symlink_metadata(path) {
401 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
402 Err(error) => Err(error),
403 Ok(metadata) if metadata.is_dir() => fs::remove_dir_all(path),
404 Ok(_) => fs::remove_file(path),
405 }
406}
407
408fn release(root: &std::path::Path) {
411 let mut active = ACTIVE
412 .lock()
413 .unwrap_or_else(std::sync::PoisonError::into_inner);
414 let Some(entry) = active.get_mut(root) else {
415 return;
416 };
417 entry.users -= 1;
418 if entry.users > 0 {
419 return;
420 }
421 let entry = active.remove(root).expect("present above");
422 drop(active);
425 entry.stash.put_back();
426}
427
428fn names(owner: &Task, entry: &str, task: &Task) -> bool {
430 match entry.split_once(':') {
431 Some((project, target)) => task.project == project && task.target == target,
432 None => task.project == owner.project && task.target == entry,
433 }
434}
435
436fn depends_on(graph: &TaskGraph, from: &str, on: &str) -> bool {
438 let mut seen = std::collections::BTreeSet::new();
439 let mut next = vec![from];
440 while let Some(id) = next.pop() {
441 for dependency in graph
442 .tasks
443 .get(id)
444 .into_iter()
445 .flat_map(|task| &task.dependencies)
446 {
447 if dependency == on {
448 return true;
449 }
450 if seen.insert(dependency.as_str()) {
451 next.push(dependency);
452 }
453 }
454 }
455 false
456}
457
458pub fn keep_across(workspace: &Workspace, graph: &TaskGraph, task: &Task) -> Kept {
461 let mut kept = Kept::default();
462 for owner in graph.tasks.values() {
463 let Ok(Some(warm)) = config(workspace, owner) else {
464 continue;
465 };
466 if !warm.survive.iter().any(|entry| names(owner, entry, task))
467 || !depends_on(graph, &owner.id, &task.id)
468 {
469 continue;
470 }
471 match acquire(workspace, owner, &warm) {
472 Ok(root) => kept.roots.push(root),
473 Err(error) => qk_executor::status!(
474 "qk: {}: warm paths not kept across {} ({error:#})",
475 owner.id,
476 task.id
477 ),
478 }
479 }
480 kept
481}
482
483fn acquire(workspace: &Workspace, owner: &Task, warm: &Warm) -> Result<PathBuf> {
486 let state = paths::worktree_state(&workspace.root).join("warm-kept");
487 let root = state.join(hash32(&owner.id));
488 fs::create_dir_all(&state)?;
489 let lock = fs::OpenOptions::new()
490 .create(true)
491 .truncate(false)
492 .read(true)
493 .write(true)
494 .open(state.join(format!("{}.lock", hash32(&owner.id))))?;
495 let mut active = loop {
498 let mut active = ACTIVE
499 .lock()
500 .unwrap_or_else(std::sync::PoisonError::into_inner);
501 if let Some(entry) = active.get_mut(&root) {
502 entry.users += 1;
503 return Ok(root);
504 }
505 match fs2::FileExt::try_lock_exclusive(&lock) {
506 Ok(()) => break active,
507 Err(error) if error.raw_os_error() == fs2::lock_contended_error().raw_os_error() => {
508 drop(active);
509 std::thread::sleep(Duration::from_millis(20));
510 }
511 Err(error) => return Err(error.into()),
512 }
513 };
514 let left = match fs::read(Stash::manifest(&root)) {
521 Ok(bytes) => {
522 let mut left: Stash = serde_json::from_slice(&bytes)
523 .context("invalid warm stash manifest; saved paths retained for recovery")?;
524 left.root = root.clone();
525 left.workspace = workspace.root.clone();
526 Some(left)
527 }
528 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
529 if root.is_dir() && fs::read_dir(&root)?.next().is_some() {
530 bail!("missing warm stash manifest; saved paths retained for recovery");
531 }
532 None
533 }
534 Err(error) => return Err(error).context("cannot read warm stash; saved paths retained"),
535 };
536 let retained = left.filter(|left| !left.put_back());
537 let stash = stash(workspace, &root, warm, retained)?;
538 active.insert(
539 root.clone(),
540 Active {
541 users: 1,
542 stash,
543 _lock: lock,
544 },
545 );
546 Ok(root)
547}
548
549fn stash(
553 workspace: &Workspace,
554 root: &std::path::Path,
555 warm: &Warm,
556 retained: Option<Stash>,
557) -> Result<Stash> {
558 let mut stash = match retained {
559 Some(retained) => retained,
560 None => {
561 let _ = fs::remove_dir_all(root);
562 Stash {
563 root: root.to_owned(),
564 workspace: workspace.root.clone(),
565 moved: Vec::new(),
566 }
567 }
568 };
569 let held: Vec<String> = stash
570 .moved
571 .iter()
572 .enumerate()
573 .filter_map(
574 |(index, path)| match fs::symlink_metadata(root.join(index.to_string())) {
575 Ok(_) => Some(Ok(path.clone())),
576 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
577 Err(error) => Some(Err(error).context("cannot inspect retained warm path")),
578 },
579 )
580 .collect::<Result<_>>()?;
581 let overlaps = |path: &str| {
582 held.iter().any(|held| {
583 held == path
584 || path.starts_with(&format!("{held}/"))
585 || held.starts_with(&format!("{path}/"))
586 })
587 };
588 let mut kept = std::collections::BTreeSet::new();
591 for path in warm.kept_paths() {
592 if let Ok(found) = Outputs::from_paths(std::slice::from_ref(&path))
593 .and_then(|outputs| outputs.paths(&workspace.root))
594 {
595 kept.extend(found);
596 }
597 }
598 let topmost: Vec<String> = kept
600 .iter()
601 .filter(|path| {
602 !std::path::Path::new(path.as_str())
603 .ancestors()
604 .skip(1)
605 .any(|ancestor| {
606 ancestor
607 .to_str()
608 .is_some_and(|ancestor| kept.contains(ancestor))
609 })
610 })
611 .filter(|path| !overlaps(path))
612 .cloned()
613 .collect();
614 if topmost.is_empty() {
615 return Ok(stash);
616 }
617 fs::create_dir_all(root)?;
618 for path in topmost {
619 let moved = paths::safe_parents(&workspace.root, &path).and_then(|()| {
620 stash.moved.push(path.clone());
621 stash.save_manifest()?;
622 fs::rename(
625 workspace.root.join(&path),
626 root.join((stash.moved.len() - 1).to_string()),
627 )
628 .inspect_err(|_| {
629 stash.moved.pop();
630 })?;
631 Ok(())
632 });
633 if let Err(error) = moved {
634 stash.save_manifest()?;
637 stash.put_back();
638 return Err(error);
639 }
640 }
641 Ok(stash)
642}
643
644pub fn check_overlaps(workspace: &Workspace, graph: &TaskGraph) -> Result<()> {
647 let mut claimed: Vec<(String, &str)> = Vec::new();
648 for (id, task) in &graph.tasks {
649 let Some(warm) = config(workspace, task)? else {
650 continue;
651 };
652 for path in warm.kept_paths() {
653 if let Some((_, other)) = claimed.iter().find(|(known, _)| {
654 known == &path
655 || known.starts_with(&format!("{path}/"))
656 || path.starts_with(&format!("{known}/"))
657 }) {
658 bail!("{id} and {other} both keep {path} as warm state");
659 }
660 claimed.push((path, id));
661 }
662 }
663 Ok(())
664}
665
666#[derive(Clone, Default, Serialize, Deserialize)]
668pub(crate) struct Group {
669 artifacts: BTreeMap<String, Artifact>,
670 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
673 stamps: BTreeMap<String, Vec<i64>>,
674 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
677 mtimes: BTreeMap<String, i64>,
678}
679
680impl Group {
681 fn size(&self, cache: &Cache) -> u64 {
683 self.artifacts
684 .values()
685 .filter_map(|artifact| match artifact {
686 Artifact::File { blob, .. } => {
687 fs::metadata(cache.root.join("blobs").join(blob)).ok()
688 }
689 _ => None,
690 })
691 .map(|metadata| metadata.len())
692 .sum()
693 }
694}
695
696#[derive(Clone, Default, Serialize, Deserialize)]
698pub(crate) struct Record {
699 version: u32,
700 task: String,
701 worktree: String,
703 #[serde(default, skip_serializing_if = "Option::is_none")]
705 commit: Option<String>,
706 saved: u64,
708 #[serde(default, skip_serializing_if = "Vec::is_empty")]
710 key: Vec<String>,
711 groups: BTreeMap<String, Group>,
712}
713
714const RECORD_VERSION: u32 = 2;
715
716#[derive(Default, Serialize, Deserialize)]
719struct RestoredFiles {
720 groups: BTreeMap<String, BTreeMap<String, (Vec<i64>, Artifact)>>,
721}
722
723#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
725pub struct Restored {
726 pub source: String,
729 pub groups: Vec<String>,
730 pub files: usize,
731 pub bytes: u64,
732}
733
734#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
736#[serde(rename_all = "camelCase")]
737pub struct WarmReport {
738 pub restored: Option<Restored>,
740 #[serde(default, skip_serializing_if = "Vec::is_empty")]
742 pub present: Vec<String>,
743 pub save_ms: Option<u64>,
746 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
748 pub background: bool,
749}
750
751#[derive(Clone)]
754struct Location {
755 name: &'static str,
756 base: PathBuf,
757 outputs: Option<Outputs>,
759}
760
761impl Location {
762 fn paths(&self) -> Result<std::collections::BTreeSet<String>> {
764 if !self.base.exists() {
765 return Ok(Default::default());
766 }
767 match &self.outputs {
768 Some(outputs) => outputs.paths(&self.base),
769 None => {
770 let mut paths = std::collections::BTreeSet::new();
771 for entry in walkdir::WalkDir::new(&self.base)
772 .min_depth(1)
773 .follow_links(false)
774 {
775 let entry = entry?;
776 paths.insert(paths::relative(&self.base, entry.path())?);
777 }
778 Ok(paths)
779 }
780 }
781 }
782
783 fn matches(&self, path: &str) -> bool {
785 self.outputs
786 .as_ref()
787 .is_none_or(|outputs| outputs.matches(path))
788 }
789}
790
791fn locations(workspace: &Workspace, task: &Task, warm: &Warm) -> Result<Vec<Location>> {
793 let mut locations = Vec::new();
794 if warm.outputs {
795 locations.push(Location {
796 name: "outputs",
797 base: workspace.root.clone(),
798 outputs: Some(Outputs::new(workspace, task)?),
799 });
800 }
801 if !warm.paths.is_empty() {
802 locations.push(Location {
803 name: "paths",
804 base: workspace.root.clone(),
805 outputs: Some(Outputs::from_paths(&warm.paths)?),
806 });
807 }
808 if warm.directory {
809 locations.push(Location {
810 name: "directory",
811 base: directory(workspace, &warm.identity),
812 outputs: None,
813 });
814 }
815 Ok(locations)
816}
817
818fn current_branch(workspace: &Workspace) -> Option<String> {
821 for name in ["GITHUB_HEAD_REF", "GITHUB_REF_NAME"] {
822 if let Ok(value) = std::env::var(name)
823 && !value.is_empty()
824 {
825 return Some(value);
826 }
827 }
828 let output = paths::git(&workspace.root, &["rev-parse", "--abbrev-ref", "HEAD"]).ok()?;
829 let branch = String::from_utf8(output.stdout).ok()?.trim().to_owned();
830 (output.status.success() && !branch.is_empty() && branch != "HEAD").then_some(branch)
831}
832
833fn branches(workspace: &Workspace) -> Vec<String> {
836 let default = workspace
837 .config
838 .default_base
839 .clone()
840 .unwrap_or_else(|| "main".into());
841 let default = default
842 .strip_prefix("origin/")
843 .unwrap_or(&default)
844 .to_owned();
845 let mut branches: Vec<String> = current_branch(workspace).into_iter().collect();
846 if !branches.contains(&default) {
847 branches.push(default);
848 }
849 branches
850}
851
852fn stamp(metadata: &fs::Metadata) -> Vec<i64> {
854 let modified = metadata
855 .modified()
856 .ok()
857 .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
858 .map_or(-1, |duration| duration.as_nanos() as i64);
859 #[allow(unused_mut)]
860 let mut stamp = vec![metadata.len() as i64, modified, i64::from(mode(metadata))];
861 #[cfg(unix)]
862 {
863 use std::os::unix::fs::MetadataExt;
864 stamp.extend([
865 metadata.ctime(),
866 metadata.ctime_nsec(),
867 metadata.ino() as i64,
868 ]);
869 }
870 stamp
871}
872
873fn worktree(workspace: &Workspace) -> String {
875 fs::canonicalize(&workspace.root)
876 .unwrap_or_else(|_| workspace.root.clone())
877 .to_string_lossy()
878 .into_owned()
879}
880
881fn current_commit(workspace: &Workspace) -> Option<String> {
883 let output = paths::git(&workspace.root, &["rev-parse", "HEAD"]).ok()?;
884 let commit = String::from_utf8(output.stdout).ok()?.trim().to_owned();
885 (output.status.success() && !commit.is_empty()).then_some(commit)
886}
887
888fn key_digests(
891 workspace: &Workspace,
892 warm: &Warm,
893 prepared: &qk_executor::PreparedTask,
894) -> Result<Vec<String>> {
895 warm.key
896 .iter()
897 .map(|part| {
898 Ok(match part {
899 KeyPart::File(path) => {
900 paths::safe_parents(&workspace.root, path)?;
901 let absolute = workspace.root.join(path);
902 if absolute.is_file() {
903 format!("file:{}", digest_file(&absolute)?)
904 } else {
905 "file:absent".to_owned()
906 }
907 }
908 KeyPart::Env(name) => match prepared
911 .execution
912 .get(std::ffi::OsStr::new(name))
913 .or_else(|| prepared.env.get(std::ffi::OsStr::new(name)))
914 {
915 Some(value) => format!(
916 "env:{}",
917 blake3::hash(value.as_encoded_bytes()).to_hex()[..32].to_owned()
918 ),
919 None => "env:unset".to_owned(),
920 },
921 })
922 })
923 .collect()
924}
925
926fn key_match(warm: &Warm, current: &[String], saved: &[String]) -> Option<u8> {
930 if saved == current {
931 return Some(0);
932 }
933 let count = warm.restore_keys?;
934 (saved.len() == current.len() && saved[..count] == current[..count]).then_some(1)
935}
936
937fn behind(workspace: &Workspace, commit: &str) -> Option<u64> {
940 let ancestor = paths::git(
941 &workspace.root,
942 &["merge-base", "--is-ancestor", commit, "HEAD"],
943 )
944 .ok()?;
945 if !ancestor.status.success() {
946 return None;
947 }
948 let output = paths::git(
949 &workspace.root,
950 &["rev-list", "--count", &format!("{commit}..HEAD")],
951 )
952 .ok()?;
953 String::from_utf8(output.stdout).ok()?.trim().parse().ok()
954}
955
956fn hash32(text: &str) -> String {
958 blake3::hash(text.as_bytes()).to_hex()[..32].to_owned()
959}
960
961fn nanos(time: SystemTime) -> Option<i64> {
963 time.duration_since(UNIX_EPOCH)
964 .ok()
965 .and_then(|duration| i64::try_from(duration.as_nanos()).ok())
966}
967
968fn now_ms() -> u64 {
970 SystemTime::now()
971 .duration_since(UNIX_EPOCH)
972 .map_or(0, |duration| duration.as_millis() as u64)
973}
974
975fn restored_files_path(workspace: &Workspace, warm: &Warm) -> PathBuf {
977 paths::worktree_state(&workspace.root)
978 .join("warm-restored")
979 .join(format!("{}.json", hash32(&warm.identity)))
980}
981
982fn read_restored_files(workspace: &Workspace, warm: &Warm) -> RestoredFiles {
984 fs::read(restored_files_path(workspace, warm))
985 .ok()
986 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
987 .unwrap_or_default()
988}
989
990impl Cache {
991 fn warm_prefix(&self, identity: &str) -> String {
993 hash32(&format!(
994 "{identity}\0{}\0{}",
995 std::env::consts::OS,
996 std::env::consts::ARCH
997 ))
998 }
999
1000 fn lock_warm(&self, warm: &Warm) -> Result<fs::File> {
1003 let directory = self.root.join("locks");
1004 fs::create_dir_all(&directory)?;
1005 let file = fs::OpenOptions::new()
1006 .create(true)
1007 .truncate(false)
1008 .read(true)
1009 .write(true)
1010 .open(directory.join(format!("warm-{}", self.warm_prefix(&warm.identity))))?;
1011 fs2::FileExt::lock_exclusive(&file)?;
1012 Ok(file)
1013 }
1014
1015 fn warm_record(&self, record: &Record) -> PathBuf {
1016 self.root.join("warm").join(format!(
1017 "{}-{}.json",
1018 self.warm_prefix(&record.task),
1019 hash32(&format!("{}\0{}", record.worktree, record.key.join("\0")))
1020 ))
1021 }
1022
1023 fn warm_records(&self, warm: &Warm) -> Vec<(PathBuf, Record)> {
1025 let prefix = format!("{}-", self.warm_prefix(&warm.identity));
1026 let mut records: Vec<(PathBuf, Record)> = fs::read_dir(self.root.join("warm"))
1027 .into_iter()
1028 .flatten()
1029 .flatten()
1030 .filter(|entry| {
1031 entry
1032 .file_name()
1033 .to_str()
1034 .is_some_and(|name| name.starts_with(&prefix) && name.ends_with(".json"))
1035 })
1036 .filter_map(|entry| {
1037 let record =
1038 serde_json::from_slice::<Record>(&fs::read(entry.path()).ok()?).ok()?;
1039 (record.version == RECORD_VERSION && record.task == warm.identity)
1040 .then(|| (entry.path(), record))
1041 })
1042 .collect();
1043 records.sort_by_key(|(_, record)| std::cmp::Reverse(record.saved));
1044 records
1045 }
1046
1047 fn choose_warm(
1053 &self,
1054 workspace: &Workspace,
1055 warm: &Warm,
1056 current: &[String],
1057 ) -> Option<(Record, String)> {
1058 let own = worktree(workspace);
1059 let mut candidates: Vec<(u8, bool, Record)> = self
1060 .warm_records(warm)
1061 .into_iter()
1062 .filter_map(|(_, record)| {
1063 let quality = key_match(warm, current, &record.key)?;
1064 let other = record.worktree != own;
1065 (!other || warm.portable).then_some((quality, other, record))
1066 })
1067 .collect();
1068 candidates.sort_by_key(|(quality, other, _)| (*quality, *other));
1070 let rank = candidates
1071 .first()
1072 .map(|(quality, other, _)| (*quality, *other))?;
1073 let tied = candidates
1074 .iter()
1075 .take_while(|(quality, other, _)| (*quality, *other) == rank)
1076 .count();
1077 let index = if tied > 1 {
1078 (0..tied)
1079 .min_by_key(|&index| {
1080 let behind = candidates[index]
1081 .2
1082 .commit
1083 .as_deref()
1084 .and_then(|commit| behind(workspace, commit));
1085 (behind.is_none(), behind, index)
1086 })
1087 .unwrap_or(0)
1088 } else {
1089 0
1090 };
1091 let (_, other, record) = candidates.swap_remove(index);
1092 let source = if other {
1093 format!("worktree {}", record.worktree)
1094 } else {
1095 "local".to_owned()
1096 };
1097 Some((record, source))
1098 }
1099
1100 fn store_warm(&self, record: &Record) -> Result<()> {
1101 let path = self.warm_record(record);
1102 let directory = path.parent().context("warm records have a directory")?;
1103 fs::create_dir_all(directory)?;
1104 let mut file = tempfile::NamedTempFile::new_in(self.root.join("tmp"))?;
1105 {
1106 let mut writer = std::io::BufWriter::new(file.as_file_mut());
1107 serde_json::to_writer(&mut writer, record)?;
1108 writer.flush()?;
1109 }
1110 file.persist(path)?;
1111 Ok(())
1112 }
1113
1114 pub(crate) fn restore_warm(
1120 &self,
1121 workspace: &Workspace,
1122 task: &Task,
1123 warm: &Warm,
1124 prepared: &qk_executor::PreparedTask,
1125 ) -> Result<(Restored, Vec<String>)> {
1126 let mut restored = Restored::default();
1127 let _lock = self.lock_warm(warm)?;
1128 let mut present = Vec::new();
1130 let mut missing = Vec::new();
1131 for location in locations(workspace, task, warm)? {
1132 if location.paths()?.is_empty() {
1133 missing.push(location);
1134 } else {
1135 present.push(location.name.to_owned());
1136 }
1137 }
1138 if missing.is_empty() {
1139 return Ok((restored, present));
1140 }
1141 let current = key_digests(workspace, warm, prepared)?;
1142 let (record, source) = match self.choose_warm(workspace, warm, ¤t) {
1143 Some(chosen) => chosen,
1144 None => match self.remote_warm(workspace, task, warm, ¤t) {
1145 Some((record, branch)) => (record, format!("remote {branch}")),
1146 None => return Ok((restored, present)),
1147 },
1148 };
1149 let keep_mtimes = warm.preserve_mtimes && source == "local";
1150 restored.source = source;
1151 let mut noted = read_restored_files(workspace, warm);
1152 for location in missing {
1153 let Some(group) = record.groups.get(location.name) else {
1154 continue;
1155 };
1156 fs::create_dir_all(&location.base)?;
1157 let mut files = BTreeMap::new();
1158 for (path, artifact) in &group.artifacts {
1159 paths::safe_parents(&location.base, path)?;
1160 if !location.matches(path) {
1161 bail!("warm state for {} holds {path}, outside its paths", task.id);
1162 }
1163 let destination = location.base.join(path);
1164 if let Some(parent) = destination.parent() {
1165 fs::create_dir_all(parent)?;
1166 }
1167 match artifact {
1168 Artifact::File { blob, mode } => {
1169 let source = self.root.join("blobs").join(blob);
1170 if digest_file(&source)? != *blob {
1171 bail!("warm state for {} has a corrupt blob", task.id);
1172 }
1173 fs::copy(&source, &destination)?;
1174 let modified = keep_mtimes
1175 .then(|| group.mtimes.get(path))
1176 .flatten()
1177 .and_then(|nanos| u64::try_from(*nanos).ok())
1178 .map_or(RESTORED_AT, |nanos| {
1179 UNIX_EPOCH + Duration::from_nanos(nanos)
1180 });
1181 fs::File::options()
1182 .write(true)
1183 .open(&destination)?
1184 .set_modified(modified)?;
1185 set_mode(&destination, *mode)?;
1186 restored.files += 1;
1187 restored.bytes += fs::metadata(&destination)?.len();
1188 files.insert(
1189 path.clone(),
1190 (
1191 stamp(&fs::symlink_metadata(&destination)?),
1192 artifact.clone(),
1193 ),
1194 );
1195 }
1196 Artifact::Directory { mode } => {
1197 fs::create_dir_all(&destination)?;
1198 set_mode(&destination, *mode)?;
1199 }
1200 Artifact::Symlink { target, directory } => {
1201 validate_link(path, target)?;
1202 symlink(target, &destination, *directory)?;
1203 }
1204 }
1205 }
1206 noted.groups.insert(location.name.to_owned(), files);
1207 restored.groups.push(location.name.to_owned());
1208 }
1209 if !restored.groups.is_empty() {
1210 let path = restored_files_path(workspace, warm);
1211 fs::create_dir_all(path.parent().context("restore notes have a directory")?)?;
1212 fs::write(path, serde_json::to_vec(¬ed)?)?;
1213 }
1214 Ok((restored, present))
1215 }
1216
1217 fn remote_warm(
1220 &self,
1221 workspace: &Workspace,
1222 task: &Task,
1223 warm: &Warm,
1224 current: &[String],
1225 ) -> Option<(Record, String)> {
1226 let remote = self
1227 .remote
1228 .as_ref()
1229 .filter(|_| warm.remote && warm.portable)?;
1230 for branch in branches(workspace) {
1231 match remote.fetch_warm(&self.root, &warm.identity, &branch) {
1232 Ok(Some(bytes)) => {
1233 let Some(record) =
1234 serde_json::from_slice::<Record>(&bytes)
1235 .ok()
1236 .filter(|record| {
1237 record.version == RECORD_VERSION && record.task == warm.identity
1238 })
1239 else {
1240 continue;
1241 };
1242 if key_match(warm, current, &record.key).is_none() {
1245 continue;
1246 }
1247 if self.store_warm(&record).is_err() {
1248 return None;
1249 }
1250 return Some((record, branch));
1251 }
1252 Ok(None) => {}
1253 Err(error) => {
1254 qk_executor::status!(
1255 "qk: {}: remote warm state unavailable ({error:#})",
1256 task.id
1257 );
1258 continue;
1261 }
1262 }
1263 }
1264 None
1265 }
1266
1267 pub(crate) fn save_warm(
1270 &self,
1271 workspace: &Workspace,
1272 task: &Task,
1273 warm: &Warm,
1274 prepared: &qk_executor::PreparedTask,
1275 ) -> Result<()> {
1276 self.save_prepared(&Save::prepare(workspace, task, warm, prepared)?)
1277 }
1278
1279 pub(crate) fn save_warm_in_background(
1282 &self,
1283 workspace: &Workspace,
1284 task: &Task,
1285 warm: &Warm,
1286 prepared: &qk_executor::PreparedTask,
1287 ) -> Result<()> {
1288 let save = Save::prepare(workspace, task, warm, prepared)?;
1289 let cache = Cache {
1290 remote: self.remote.clone(),
1291 ..Cache::new(self.root.clone())
1292 };
1293 let handle = std::thread::spawn(move || {
1294 if let Err(error) = cache.save_prepared(&save) {
1295 qk_executor::status!("qk: {}: warm state not saved ({error:#})", save.task);
1296 }
1297 });
1298 self.saves.lock().unwrap().push(handle);
1299 Ok(())
1300 }
1301
1302 fn save_prepared(&self, save: &Save) -> Result<()> {
1307 let warm = &save.warm;
1308 let _lock = self.lock_warm(warm)?;
1309 let records = self.warm_records(warm);
1310 let previous = records
1312 .iter()
1313 .filter(|(_, record)| record.worktree == save.worktree)
1314 .min_by_key(|(_, record)| record.key != save.key)
1315 .map(|(_, record)| record);
1316 let noted: RestoredFiles = fs::read(&save.notes)
1317 .ok()
1318 .and_then(|bytes| serde_json::from_slice(&bytes).ok())
1319 .unwrap_or_default();
1320 let mut record = Record {
1321 version: RECORD_VERSION,
1322 task: warm.identity.clone(),
1323 worktree: save.worktree.clone(),
1324 commit: save.commit.clone(),
1325 saved: now_ms(),
1326 key: save.key.clone(),
1327 groups: BTreeMap::new(),
1328 };
1329 let gone = |error: &std::io::Error| error.kind() == std::io::ErrorKind::NotFound;
1330 for location in &save.locations {
1331 if !location.base.exists() {
1332 continue;
1333 }
1334 let known = previous.and_then(|previous| previous.groups.get(location.name));
1335 let restored_files = noted.groups.get(location.name);
1336 let mut group = Group::default();
1337 let mut pending = Vec::new();
1338 for path in location.paths()? {
1339 paths::safe_parents(&location.base, &path)?;
1340 let absolute = location.base.join(&path);
1341 let metadata = match fs::symlink_metadata(&absolute) {
1342 Err(error) if gone(&error) => continue,
1343 result => result?,
1344 };
1345 let artifact = if metadata.file_type().is_symlink() {
1346 let target = match fs::read_link(&absolute) {
1347 Err(error) if gone(&error) => continue,
1348 result => result?,
1349 };
1350 let target = target
1351 .to_str()
1352 .context("symlink target must be UTF-8")?
1353 .to_owned();
1354 validate_link(&path, &target)?;
1355 Artifact::Symlink {
1356 target,
1357 directory: absolute.is_dir(),
1358 }
1359 } else if metadata.is_file() {
1360 let now = stamp(&metadata);
1361 let saved = known.and_then(|known| {
1362 (known.stamps.get(&path) == Some(&now))
1363 .then(|| known.artifacts.get(&path))
1364 .flatten()
1365 });
1366 let restored = restored_files
1367 .and_then(|files| files.get(&path))
1368 .filter(|(stamp, _)| *stamp == now)
1369 .map(|(_, artifact)| artifact);
1370 let reused = saved.or(restored).filter(|artifact| match artifact {
1371 Artifact::File { blob, .. } => self.root.join("blobs").join(blob).is_file(),
1372 _ => false,
1373 });
1374 let Some(artifact) = reused.cloned() else {
1375 pending.push((path, absolute, metadata));
1376 continue;
1377 };
1378 group.stamps.insert(path.clone(), now);
1379 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1380 group.mtimes.insert(path.clone(), modified);
1381 }
1382 artifact
1383 } else if metadata.is_dir() {
1384 Artifact::Directory {
1385 mode: mode(&metadata),
1386 }
1387 } else {
1388 continue;
1389 };
1390 group.artifacts.insert(path, artifact);
1391 }
1392 let sources: Vec<_> = pending
1393 .iter()
1394 .map(|(_, absolute, _)| absolute.clone())
1395 .collect();
1396 for ((path, _, metadata), result) in pending.into_iter().zip(self.put_blobs(&sources)) {
1397 let blob = match result {
1398 Ok(blob) => blob,
1399 Err(error) if error.downcast_ref::<std::io::Error>().is_some_and(gone) => {
1400 continue;
1401 }
1402 Err(error) => return Err(error),
1403 };
1404 group.stamps.insert(path.clone(), stamp(&metadata));
1405 if let Some(modified) = metadata.modified().ok().and_then(nanos) {
1406 group.mtimes.insert(path.clone(), modified);
1407 }
1408 group.artifacts.insert(
1409 path,
1410 Artifact::File {
1411 blob,
1412 mode: mode(&metadata),
1413 },
1414 );
1415 }
1416 if let Some(limit) = warm.max_size
1417 && group.size(self) > limit
1418 {
1419 qk_executor::status!(
1420 "qk: {}: warm {} is over its maxSize and was not saved",
1421 save.task,
1422 location.name
1423 );
1424 continue;
1425 }
1426 record.groups.insert(location.name.to_owned(), group);
1427 }
1428 self.store_warm(&record)?;
1429 let _ = fs::remove_file(&save.notes);
1430 let saved = self.warm_record(&record);
1433 for (path, _) in records
1434 .into_iter()
1435 .filter(|(path, _)| *path != saved)
1436 .skip(KEPT - 1)
1437 {
1438 let _ = fs::remove_file(path);
1439 }
1440 if warm.remote
1442 && warm.portable
1443 && let Some(remote) = &self.remote
1444 && let Some(branch) = &save.branch
1445 {
1446 for group in record.groups.values_mut() {
1447 group.stamps.clear();
1448 }
1449 remote.upload_warm(
1450 &self.root,
1451 &warm.identity,
1452 branch,
1453 serde_json::to_vec(&record)?,
1454 );
1455 }
1456 Ok(())
1457 }
1458}
1459
1460struct Save {
1463 task: String,
1464 warm: Warm,
1465 locations: Vec<Location>,
1466 worktree: String,
1467 commit: Option<String>,
1468 branch: Option<String>,
1469 key: Vec<String>,
1470 notes: PathBuf,
1471}
1472
1473impl Save {
1474 fn prepare(
1476 workspace: &Workspace,
1477 task: &Task,
1478 warm: &Warm,
1479 prepared: &qk_executor::PreparedTask,
1480 ) -> Result<Self> {
1481 Ok(Self {
1482 task: task.id.clone(),
1483 warm: warm.clone(),
1484 locations: locations(workspace, task, warm)?,
1485 worktree: worktree(workspace),
1486 commit: current_commit(workspace),
1487 branch: current_branch(workspace),
1488 key: key_digests(workspace, warm, prepared)?,
1489 notes: restored_files_path(workspace, warm),
1490 })
1491 }
1492}
1493
1494#[cfg(test)]
1495mod tests {
1496 use super::*;
1497 use qk_taskgraph::Request;
1498 use serde_json::json;
1499
1500 fn fixture() -> (tempfile::TempDir, Workspace, TaskGraph) {
1502 let temp = tempfile::tempdir().unwrap();
1503 fs::write(temp.path().join("nx.json"), "{}").unwrap();
1504 fs::write(
1505 temp.path().join("project.json"),
1506 json!({
1507 "name": "app", "targets": {
1508 "one": {"command": "true", "qk:warm": {"paths": ["scratch-one"]}},
1509 "two": {"command": "true", "qk:warm": {"paths": ["scratch-two"]}}
1510 }
1511 })
1512 .to_string(),
1513 )
1514 .unwrap();
1515 for path in ["scratch-one", "scratch-two"] {
1516 fs::create_dir(temp.path().join(path)).unwrap();
1517 fs::write(temp.path().join(path).join("state"), "saved").unwrap();
1518 }
1519 let workspace = Workspace::load(temp.path()).unwrap();
1520 let graph = TaskGraph::build(
1521 &workspace,
1522 &[
1523 Request::parse("app:one").unwrap(),
1524 Request::parse("app:two").unwrap(),
1525 ],
1526 )
1527 .unwrap();
1528 (temp, workspace, graph)
1529 }
1530
1531 #[test]
1533 fn waiting_for_an_owner_does_not_block_other_releases() {
1534 let (_temp, workspace, graph) = fixture();
1535 let one = &graph.tasks["app:one"];
1536 let two = &graph.tasks["app:two"];
1537 let first = acquire(&workspace, one, &config(&workspace, one).unwrap().unwrap()).unwrap();
1538 let state = paths::worktree_state(&workspace.root).join("warm-kept");
1539 let lock = fs::OpenOptions::new()
1540 .create(true)
1541 .truncate(false)
1542 .read(true)
1543 .write(true)
1544 .open(state.join(format!("{}.lock", hash32(&two.id))))
1545 .unwrap();
1546 fs2::FileExt::lock_exclusive(&lock).unwrap();
1547 let (released, received) = std::sync::mpsc::channel();
1548 let completed = std::thread::scope(|scope| {
1549 scope.spawn(|| {
1550 let second =
1551 acquire(&workspace, two, &config(&workspace, two).unwrap().unwrap()).unwrap();
1552 release(&second);
1553 });
1554 std::thread::sleep(Duration::from_millis(100));
1555 scope.spawn(|| {
1556 release(&first);
1557 released.send(()).unwrap();
1558 });
1559 let completed = received.recv_timeout(Duration::from_secs(2)).is_ok();
1560 fs2::FileExt::unlock(&lock).unwrap();
1561 completed
1562 });
1563 assert!(
1564 completed,
1565 "a blocked owner lock prevented an unrelated release"
1566 );
1567 assert_eq!(
1568 fs::read_to_string(workspace.root.join("scratch-one/state")).unwrap(),
1569 "saved"
1570 );
1571 }
1572
1573 #[test]
1575 fn unreadable_ownership_preserves_saved_paths() {
1576 let (_temp, workspace, graph) = fixture();
1577 let owner = &graph.tasks["app:one"];
1578 let warm = config(&workspace, owner).unwrap().unwrap();
1579 let root = acquire(&workspace, owner, &warm).unwrap();
1580 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
1581 drop(entry); fs::write(Stash::manifest(&root), "{").unwrap();
1583 assert!(acquire(&workspace, owner, &warm).is_err());
1584 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1585 fs::remove_file(Stash::manifest(&root)).unwrap();
1586 assert!(acquire(&workspace, owner, &warm).is_err());
1587 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1588 fs::create_dir(Stash::manifest(&root)).unwrap();
1589 assert!(acquire(&workspace, owner, &warm).is_err());
1590 assert_eq!(fs::read_to_string(root.join("0/state")).unwrap(), "saved");
1591 }
1592
1593 #[cfg(unix)]
1595 #[test]
1596 fn inspection_errors_keep_recovery_incomplete() {
1597 let (_temp, workspace, graph) = fixture();
1598 let owner = &graph.tasks["app:one"];
1599 let warm = config(&workspace, owner).unwrap().unwrap();
1600 let root = acquire(&workspace, owner, &warm).unwrap();
1601 let entry = ACTIVE.lock().unwrap().remove(&root).unwrap();
1602 let saved = root.with_extension("saved");
1603 fs::rename(&root, &saved).unwrap();
1604 fs::write(&root, "blocked").unwrap(); assert!(!entry.stash.put_back());
1606 let retained = Stash {
1607 root: root.clone(),
1608 workspace: workspace.root.clone(),
1609 moved: entry.stash.moved.clone(),
1610 };
1611 assert!(stash(&workspace, &root, &warm, Some(retained)).is_err());
1612 assert_eq!(fs::read_to_string(saved.join("0/state")).unwrap(), "saved");
1613 assert_eq!(fs::read_to_string(&root).unwrap(), "blocked");
1614 }
1615}